arXiv:2504.17577cs.LG2025-04被引 64

TileLang让AI核程序编写更简单高效,自动处理复杂优化。

TileLang: A Composable Tiled Programming Model for AI Systems

  • 将调度与数据流分离,用注解和原语简化编程
  • 在多种设备上实现顶尖性能,关键核函数加速显著
  • 适合需要高性能、可复用AI内核的开发者使用

现代AI工作负载严重依赖为训练和推理优化的计算内核。这些内核遵循明确的数据流模式,如在DRAM与SRAM间移动数据块(tile)并在其上执行一系列计算。然而,即使模式清晰,编写高性能内核仍极复杂,需深入硬件特性的精细优化。尽管领域专用编译器试图减轻负担,却常面临可用性与表达力不足的问题。本文提出TileLang,一种通用的分块编程模型,用于更高效的AI内核开发。TileLang将调度空间(线程绑定、布局、张量化、流水线)与数据流解耦,封装为可定制的注解与原语。该方法使用户聚焦于内核的数据流逻辑,其余优化交由编译器完成。我们在常见设备上开展全面实验,结果显示,TileLang在多个关键内核中达到当前最优性能,证明其统一的块-线程范式与透明调度能力,能同时满足现代AI系统开发对性能与灵活性的需求。

原文摘要 · Abstract (English)

Modern AI workloads rely heavily on optimized computing kernels for both training and inference. These AI kernels follow well-defined data-flow patterns, such as moving tiles between DRAM and SRAM and performing a sequence of computations on those tiles. However, writing high-performance kernels remains complex despite the clarity of these patterns. Achieving peak performance requires careful, hardware-centric optimizations to fully leverage modern accelerators. While domain-specific compilers attempt to reduce the burden of writing high-performance kernels, they often struggle with usability and expressiveness gaps. In this paper, we present TileLang, a generalized tiled programming model for more efficient AI Kernel programming. TileLang decouples scheduling space (thread binding, layout, tensorize and pipeline) from dataflow, and encapsulated them as a set of customization annotations and primitives. This approach allows users to focus on the kernel's data-flow itself, while leaving most other optimizations to compilers. We conduct comprehensive experiments on commonly-used devices, across numerous experiments, our evaluation shows that TileLang can achieve state-of-the-art performance in key kernels, demonstrating that its unified block-and-thread paradigm and transparent scheduling capabilities deliver both the power and flexibility demanded by modern AI system development.

AI内核编程模型性能优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。