TileLang：一种面向人工智能系统的可组合式分块编程模型

摘要

        现代AI工作负载在训练和推理过程中高度依赖优化的计算内核。这些AI内核遵循明确的数据流模式，例如在DRAM与SRAM之间移动数据块，并对这些数据块执行一系列计算。尽管这些模式清晰可见，编写高性能内核依然复杂。要实现峰值性能，需要以硬件为中心进行细致优化，以充分利用现代加速器。虽然领域专用编译器试图减轻编写高性能内核的负担，但它们在易用性和表达性方面仍存在不足。

        本文提出TileLang，一种通用的分块编程模型，旨在实现更高效的AI内核编程。TileLang将调度空间（线程绑定、内存布局、张量化与流水线）与数据流解耦，并将其封装为一组可定制的标注和原语。这种方法允许用户专注于内核数据流本身，而将大多数其他优化交由编译器处理。我们在常用设备上进行了全面实验，评估结果表明，TileLang在关键内核中能够达到最先进的性能，证明其统一的块与线程范式以及透明的调度能力能够满足现代AI系统开发对性能和灵活性的需求。

1 引言

        近年来，对AI工作负载高性能的追求加速了专用内核的发展，这些内核驱动着训练和推理过程。矩阵乘法尤其支撑了广泛的神经网络架构，从简单的前馈层到基于 Transformer 的大规模模型。为应对这些网络巨大的计算负担，出现了如 FlashAttention 等定制内核以优化注意力机制，降低内存开销并提升处理吞吐量。然而，在不断演进的加速器硬件上实现高效率，依赖于对硬件感知设计与复杂调优的精细结合——这些挑战激发了对更具表达力的领域专用编译器的广泛兴趣。

        深度学习内核通常表现为数据流模式，涉及在DRAM与SRAM之间移动数据块，并在这些块上执行计算序列。尽管这些模式看似清晰，但构建高性能内核仍具挑战性，因为开发者必须手动处理多项关键优化：

• 线程绑定：指将块操作和数据映射到适当线程的过程。在现代加速器架构（如GPU）中，这涉及在线程块 block、线程束 warp 和单个线程 thread 间仔细分配任务，以最大化并行性并减少负载不均衡。最优的绑定策略可提升数据局部性，降低线程同步与分岔开销，从而提升计算吞吐量。

• 内存布局：内存布局优化涉及在物理内存中有组织地排列数据，以消除存储体冲突并确保高效访问模式。如近期研究所示，该过程通常需要将自然数据表示转换为与架构内存子系统对齐的分块或分块格式。此类重组有助于实现合并访问与高效缓存利用，从而降低内存延迟并提升整体系统性能。

• 内禀张量化：利用内禀函数指直接使用针对性能优化的目标专用指令。现代处理器和加速器提供专用操作（如Tensor Core和Matrix Core），可同时执行多个算术运算，并配备向量复制和异步复制等机制以更好地利用带宽。使用这些内禀指令需要精确管理数据类型、内存对齐和控制流，以充分发挥硬件计算能力，从而在关键内核操作中实现显著加速。

• 流水线：流水线技术通过重叠数据移动与计算来缓解内存访问延迟。通过并发调度数据传输与计算任务，流水线确保处理单元保持活跃，并最小化因内存延迟导致的空闲时间。在先进的Nvidia Hopper架构中，张量内存加速器（TMA）可通过为不同计算单元（如CUDA核心与Tensor Core）启用异步处理来促进这一过程，进一步提升并发性。

        尽管近期针对AI工作负载的领域专用编译器极大简化了高性能内核的创建，它们仍将大多数底层优化与内核实现交织在一起，即使数据流已明确暴露。例如，Triton提供了直观的块级原语，但将线程行为、内存布局和地址空间标注隐藏于自动生成的策略之后。这种抽象简化了编程，却阻碍了经验丰富的开发者寻求极致性能——例如在实现量化权重矩阵乘法时。此类内核通常需要内联汇编来执行向量化数据类型转换，以及精心设计与硬件特定内存缓冲区对齐的自定义数据布局。虽然Triton提供了如tl.dot等向量化操作，但将其扩展至定制用例（例如通过PTX注册手工编写的高性能块操作符）仍十分繁琐。此外，尽管Triton提供了用户友好的流水线调节参数（num_stage），它不允许用户定义完全自定义的流水线。因此，领域专家在开发需要显式控制内存层级及其他细粒度优化的内核时受到限制。

        为应对这些局限性，我们提出TileLang，一种既保留Triton的简洁性又提供更高灵活性的编程模型。TileLang 旨在为用户提供对调度空间的细粒度控制，以实现更高性能。我们认为实现这一点的关键是将数据流与调度解耦：用户仅专注于使用可组合的块操作符定义数据流，而编译器负责探索并应用调度策略。当编译器默认优化不足时，用户可在前端进行更精确的控制。我们引入了一种可组合的分块编程抽象，其中核心计算模式（如 GEMM、COPY、ATOMIC 和 REDUCE ）通过 tile 操作符表达。这些操作符独立于调度决策定义内核的数据流。同时，提供一组调度原语和标注以捕获进一步优化，使用户可选择依赖编译器生成的调度或手动微调内核的性能关键部分。

        为提升 TileLang 的易用性，我们在 Python 中实现了前端语言，支持灵活的编程风格且仅需最少的类型标注。此外，我们为 TileLang 设计了一个编译器，可将用户定义的程序转换为高度优化的底层代码，以在现代硬件上高效执行。该编译器自动化关键优化，减少了性能调优所需的手动工作。

本文贡献总结如下：
（1）块级编程语言：我们设计了一种块级编程语言，允许用户显式声明缓冲区在硬件内存层级中的放置。通过布局推断机制，系统抽象了高效并行化缓冲区操作的复杂性，同时暴露线程级控制接口，使专家能够精确管理每个线程与缓冲区的交互。
（2）具备自动化优化的编译器：我们为 TileLang 提供了配套编译器，包含一系列自动化编译流程，涵盖通过布局推断机制实现自动并行化、内核库动态参数简化、自动流水线推导以及动态形状的循环尾部拆分优化等功能。该编译器确保 TileLang 程序既高效又易于编写。
（3）最先进的性能：在实际AI内核上的实验评估表明，TileLang 在 NVIDIA 和 AMD GPU 上均实现了与专用厂商库及其他基于 DSL 的方法（如 Triton ）相当甚至更优的性能。

        在本文后续部分，我们将介绍 TileLang 的设计与实现。首先描述语言语法和底层编程模型，然后详细说明 TileLang JIT 编译器架构，涵盖硬件无关与硬件感知优化。最后，我们将 TileLang 与现有方案进行比较，总结研究发现，并概述这一统一高性能AI内核开发方法的未来方向。我们已开源 TileLang。

2 TileLang 示例

        现有的将调度与计算分离的机器学习编译器（如 TVM）要求用户明确区分计算和调度。此外，用户必须手动注册新的张量指令并指定缓冲区布局才能实现最佳性能。然而，编写和理解调度程序仍然具有挑战性。尽管像 Triton 这样的现代框架允许用户专注于分块级编程，但其数据流表示通常不够清晰，并且需要使用某些变通方法（例如掩码条件加载）或特定于硬件的功能（如张量内存加速器）。虽然像 ThunderKitten 这样的框架将程序抽象为加载、计算、存储和同步操作的分块粒度组合，但其数据流仍然不够透明，限制了用户应用进一步优化的能力。此外，随着基于 Python 的深度学习框架 [3, 22] 的广泛采用，手动将模型转换为 C++ 进行优化是不切实际的。

Fig. 1. An example TileLang program and the corresponding lowered ir and generated cuda c code. The code snippets are simplified for demonstration purposes.

        因此，在设计 TileLang 时，我们强调三个关键原则：
(1) Pythonic 设计，与 Python 生态系统无缝集成，提供熟悉的编码体验并降低学习曲线；
(2) 以数据流为中心，使用户能够主要关注数据流，同时抽象化底层调度的复杂性。它将线程绑定、内存布局、张量化和流水线等调度方面与数据流解耦，并将其封装为一组可自定义的注解和原语，以提高可编程性和可维护性；
(3) 可组合性，确保内核、原语和调度策略可以无缝组合以构建复杂的设计。

        接下来，我们在 TileLang 中实现一个通用矩阵乘法内核来展示其基本语法，并说明它如何提高生产力。如图 11(a) 所示，实现首先定义 GEMM 内核的输入和输出（第 8 行），指定其形状和数据类型。随后，我们初始化内核上下文（第 9-11 行），这决定了网格大小和总线程数，接着是内核主体（第 12-27 行），其中包括片上内存分配和数据流管理。

        由于 TileLang 是一种嵌入 Python 的编程语言，它支持 Python 的所有命令式构造（例如 if-else、for 和 while），关键区别在于用户必须为函数参数和变量声明提供显式类型注解。这一要求源于 Python 的动态类型特性，其本身可能不适用于设备代码生成（例如 CUDA/HIP），因为静态数据类型对于确定精确的数据位宽至关重要。在 TileLang 中，类型注解明确定义了元素类型和张量形状，确保了正确性和高效的代码生成。

        此外，TileLang 允许显式内存分配，从而对数据放置和访问模式提供更好的控制。在给定的实现中，TileLang 使用 T.alloc_shared 将矩阵 𝐴 和 𝐵 的子矩阵存储在共享内存中，而 T.alloc_fragments 用于在块级别的寄存器文件中分配累加器。此外，流水线执行的使用使得内存传输与计算重叠，有效地隐藏内存延迟并提高整体吞吐量。T.gemm 操作利用 NVIDIA CUTLASS 或手动编写的 HIP 代码来高效地执行分块级矩阵计算。通过自动化底层调度和同步，TileLang 使开发人员能够专注于算法设计而非特定于硬件的优化，从而在保持计算效率的同时提高生产力。

        最后，我们调用 tilelang.compile（第 31 行）将 tilelang 程序降级为中间表示（IR），如图 11(b) 所示。然后，该 IR 被进一步编译成可执行文件，生成最终的优化代码，如图 11(c) 所示。

 

3 TileLang 语言

        本节将介绍我们基于分块的编程模型的基础，阐述TileLang如何系统高效地管理AI内核开发，并概述TileLang将数据流与其他调度空间分离的设计理念。

Fig. 2. Stages of TileLang Compile Pipeline.

        图2展示了TileLang的五阶段编译流水线。首先，开发者使用TileLang编写高层程序来描述计算逻辑和数据访问模式。在解析阶段，TileLang程序被解析为Python AST，随后转换为TileLang AST。接着，IR构建器将AST转换为TVM中间表示，使我们能够利用TVM的语法树及相关基础设施。随后，优化阶段执行一系列图优化和调度转换以提升执行效率。最后，代码生成阶段将优化后的IR转换为后端代码，如LLVM IR、CUDA C/C++或HIP C/C++，以支持多种硬件平台。

Table 1. A partial list of the dataflow operators and scheduling primitives supported by TileLang.

        表1展示了TileLang所提供的数据流操作符和调度原语的一个代表性子集。TileLang采用以数据为中心的编程范式，核心计算语义通过分块级操作符（如T.copy、T.gemm和T.reduce）来表达。作为这些操作符的补充，TileLang提供了一组调度原语，使开发者能够微调并行性、流水线和内存布局等性能关键方面。我们将在后续章节详细阐述这两个组件的设计。

3.1 基于分块的编程模型

        图11提供了一个简洁的TileLang矩阵乘法（GEMM）示例，展示了开发者如何利用分块、内存放置、流水线和操作符调用等高层结构来精细控制数据移动与计算。具体而言，图11(a)中的代码片段展示了多级分块如何利用不同内存层次（全局内存、共享内存和寄存器）来优化带宽利用率并降低延迟。整体而言，图11(b)展示了TileLang类Python语法如何让开发者在友好的编程模型中推理性能关键优化。

Fig. 3. Optimizing GEMM with Multi-Level Tiling on GPUs via TileLang.

分块声明。我们方法的核心是将分块作为编程模型中的一等对象。一个分块代表一块具有特定形状的数据，可以由一个线程束、线程块或等效的并行单元拥有和操作。在矩阵乘法示例中，A和B缓冲区在内核循环中以分块形式（由block_M、block_N、block_K决定）读入。通过T.Kernel，TileLang定义了执行上下文，其中包括线程块索引（bx和by）以及线程数量。这些上下文有助于我们计算每个线程块的索引，并使TileLang更易于自动推断和优化内存访问与计算。此外，这些上下文允许用户在线程块内手动控制每个独立线程的行为。

显式硬件内存分配。TileLang的一大特点是能够将分块缓冲区显式放置在硬件内存层次结构中。它不依赖于编译器不透明的优化过程，而是提供直接映射到物理内存空间或加速器特定结构的用户层内禀函数。具体包括：

• T.alloc_shared：在快速的片上存储空间（对应NVIDIA GPU的共享内存）中分配内存。共享内存非常适合缓存计算过程中的中间数据，因为它比全局内存快得多，并且允许同一线程块内的线程高效共享数据。例如，在矩阵乘法中，矩阵的分块可加载到共享内存中，以减少对全局内存带宽的需求并提升性能。

• T.alloc_fragment：在片段内存（对应NVIDIA GPU的寄存器文件）中分配累加器。通过将输入和部分和保存在寄存器或硬件级缓存中，可进一步最小化延迟。需要注意的是，在这个分块程序中，每个分块分配了与共享内存相同的本地缓冲区，这似乎有违直觉，因为共享内存通常更快但容量更大，而寄存器文件有限。这是因为此处的分配指的是针对整个线程块的寄存器文件。TileLang在编译期间会使用布局推断通道来推导出布局对象T.Fragment，该对象决定了如何为每个线程分配相应的寄存器文件。这个过程将在后续章节详细讨论。

        使用T.copy可以管理全局内存与硬件特定内存之间的数据传输。此外，可以使用T.clear或T.fill初始化硬件特定的缓冲区。对于数据赋值操作，也可以使用T.Parallel并行执行，如第8行所示。

3.2 以数据流为中心的分块操作符

        TileLang抽象出一组分块操作符，使开发者能够专注于数据流逻辑，而无需管理每个分块操作的低层实现细节。图4展示了分块操作符的接口及其几个代表性示例，包括GEMM、Copy和Parallel。每个分块操作符都需要实现两个关键接口：Lower和InferLayout。Lower接口定义了如何将高层分块操作符降级为低层中间表示（IR），例如线程绑定或向量化内存访问。例如，Copy可以降级为具有显式线程绑定和向量化加载/存储的循环。InferLayout接口负责确定与分块操作符相关的内存和循环布局。这包括推断缓冲区布局（例如，交换内存）或循环级布局（例如，线程绑定）。例如，T.gemm对其共享内存输入应用交换布局，并使用矩阵特定的布局来回写MMA片段。类似地，T.Parallel中的并行循环结构可以使用线程级绑定和向量化访问模式来表达，这两者都是通过布局推断推导出来的。第4.1节将更详细地讨论布局组合及其在降级过程中的作用。

Fig. 4. Interface of a Tile-Operator, and example instances of TileOP.

        表1列出了TileLang操作符的一个子集，以简化基于分块的编程中的常见操作。这些内置操作符抽象了硬件内存访问和计算的低层细节，使开发者能够从数据流的角度专注于高层算法设计，同时保持对性能关键方面的细粒度控制。每个操作符的设计都与分块编程模型无缝集成，确保在硬件内存层次结构中进行高效的数据移动和计算。下面，我们将描述几个关键操作符及其在优化内存传输和算术计算中的作用。

• copy：copy操作符是带有内存复制功能的T.Parallel的语法糖。它允许在寄存器的作用域片段、静态共享内存的shared作用域、动态共享内存的shared.dyn作用域以及全局内存的global作用域之间进行数据复制。

• gemm：内置的T.gemm操作符是一个高度优化的通用矩阵乘法实现，支持各种内存访问模式（ss, sr, rs, rr），其中r表示寄存器内存，s表示共享内存。该操作符会根据内核配置自动选择最优实现。对于CUDA后端，T.gemm利用Nvidia的CUTLASS库来高效利用Tensor Cores或CUDA Cores；而对于AMD GPU，它则同时使用可组合内核和手写的HIP代码进行性能优化。用户还可以通过在Python中注册自定义原语来扩展T.gemm，使其能够灵活适应特定的使用场景。

• reduce：T.reduce操作符提供了一个灵活高效的归约机制，用于跨维度聚合数据。它支持多种归约操作，如求和、最小值、最大值和乘积等。归约可以在指定的轴上进行，从而实现对矩阵进行行方向或列方向的归约等操作。T.reduce的实现利用了线程束级和线程块级的并行性，以确保在CUDA和AMD后端上都能获得最佳性能。用户也可以通过定义自己的归约内核来自定义归约操作。

• atomic：T.atomic操作符提供了原子操作，用于在并行上下文中安全地更新共享内存或全局内存。它开箱即用地支持常见的原子操作，如加法、最小值和最大值。T.atomic确保在并发更新过程中的线程安全性，这使得它在直方图更新、使用共享内存的归约以及无锁计数器等操作中至关重要。它的设计目标是利用NVIDIA和AMD GPU上的原生硬件原子指令，在保证并行执行正确性的同时实现高性能。

3.3 调度标注与原语

        虽然数据流模式构成了计算组织的基础，但现代高性能计算需要对执行模式进行更细粒度的控制。为了满足这一需求，TileLang提供了一套全面的调度原语，使开发者能够精确调整其应用程序的性能关键方面，详见表1：

• Pipelined：T.Pipelined原语允许循环进行高效的流水线执行，通过重叠计算和内存操作来提高性能。在图11中，遍历k（归约维度）的循环被设置为num_stages=3的流水线，创建了一个3级流水线。这个流水线允许数据传输、计算和后续的数据准备工作重叠进行，从而有效减少内存瓶颈并提高计算吞吐量。将T.Pipelined降级为CUDA源代码的详细设计将在第4.4节讨论。

• Parallel：T.Parallel原语通过将循环迭代映射到线程来实现循环的自动并行化。在图8中，将数据复制到A_shared的操作使用T.Parallel(8, 32)在两个维度（8和32）上进行并行化。这不仅通过利用硬件并行性提高了性能，还自动将线程映射到迭代，并支持向量化以进行进一步优化。

• annotate_layout：T.annotate_layout原语允许您使用用户定义的内存布局来指定共享内存或全局内存的内存布局优化。默认情况下，TileLang采用经过优化的内存布局，旨在最小化NVIDIA和AMD GPU上的存储体冲突。

• use_swizzle：T.P.use_swizzle原语通过启用交换内存访问来提高L2缓存局部性，从而改善光栅化过程中的数据重用。该原语在处理并行线程块中的分块数据时特别有效。

4 调度设计与自动化

        本节将讨论除数据流外，TileLang中的四种调度空间及其自动化设计。其中一些相对独立（如流水线和张量化），而另一些则关联更紧密，例如线程绑定和内存布局设计。在接下来的小节中，我们将首先解释内存布局基础设施的设计，然后是线程绑定，接着讨论张量化的自动化设计，最后分享流水线的设计。

4.1 内存布局组合

        在TileLang中，我们支持使用诸如A[i, k]之类的高级接口对多维数组进行索引。这种高级索引最终会通过一系列软硬件抽象层转换为物理内存地址。为了对这一索引转换过程进行建模，我们引入了核心抽象概念——布局，它描述了数据在内存中是如何组织和映射的。

        在物理地址层面，布局可以表示为形式为∑𝑖 𝑦𝑖𝑠𝑖的线性化地址表达式，其中𝑦𝑖表示沿第𝑖维度的索引，𝑠𝑖是该维度对整个线性内存地址的跨度贡献。给定一个布局𝐿 = 𝑠 : 𝑑 = (𝑠0, 𝑠1, …, 𝑠𝑛−1) : (𝑑0, 𝑑1, …, 𝑑𝑛−1)，TileLang借鉴了TVM的设计思路，引入了一种基于IterVar的、可组合且可堆叠的布局函数抽象。由于IterVar可以封装跨度信息，布局表达式可以被简化为关于IterVar的代数形式。因此，布局函数可以形式化地表达为一个映射𝑓: K𝑛 → K𝑚，其中𝑓编码了从高级索引到内存地址的转换。

Fig. 5. Interface and example instances of Layout Function.

        图5(a)展示了TileLang中布局的定义。其核心组件包括iter_vars（可选地携带范围信息）以及一组基于这些迭代变量计算内存位置的forward_index表达式。这些表达式共同定义了一个代数函数𝑓: K𝑛 → K𝑚。如图5(b)所示，这允许表达一个从2D到1D的布局转换。给定缓冲区的形状，iter_vars被绑定到特定区域，生成的表达式被传递给算术分析器以确定符号或常量边界。这些边界用于推断转换后缓冲区的形状，并相应调整缓冲区访问索引。

        TileLang还支持非双射的布局转换。例如，图5(c)展示了如何使用布局为缓冲区访问应用填充。这些布局转换是可组合的，TileLang包含多种内置布局策略，例如常用于减轻GPU上共享内存存储体冲突的布局交换。

        此外，TileLang引入了布局抽象的一个扩展，称为片段。与标准布局相比，片段布局总是产生形式为𝑓: K𝑛 → K²的输出，其中两个输出维度分别代表线程在寄存器文件内的位置以及本地寄存器文件的索引。例如，在图11中，内核在块级别分配了一个寄存器文件𝐶local。然而，由于GPU寄存器文件必须在块内的线程间进行分区，片段布局准确地描述了这种分区方案。

        图6(a)展示了片段布局的定义。TileLang提供了四种原语操作来帮助用户扩展现有的片段布局。图6(b)展示了如何使用这些原语，从一个用于m16k16矩阵片段的mma_ldmatrix指令的基本布局，推导出完整的块级布局的示例。其中，base_layout表示单个线程束消费一个m16k16矩阵的布局。通过repeat原语扩展该布局，形成warp_layout，允许单个线程束消费一个m32k16矩阵。图6(c)可视化了这个转换。接着，warp_layout使用repeat_on_thread和replicate等原语进一步扩展，产生block_layout，代表四个线程束共同消费一个m128k16矩阵。

Fig. 6. Interface and example instances of Fragment Layout.

4.2 线程绑定

        基于片段布局的抽象，随之而来的一个关键挑战是如何在执行过程中将这些布局映射到线程上。这就引出了线程绑定问题，它涉及如何将块级寄存器文件分配到各个线程，以及如何推断合适的片段布局。此外，还需要确定如何正确地并行化循环以满足布局约束。

        尽管第4.1节引入了片段布局以帮助简化此过程，但对于任意的计算表达式，为所有缓冲区确定合适的片段布局仍然困难。我们通过两个关键观察来指导这一过程。首先，由于多个分块操作符经常共享相同的缓冲区，它们各自的布局和线程绑定策略是相互依赖的。其次，不同操作符对布局和线程绑定的严格程度要求不同。例如，在GPU上，GEMM操作符（利用Tensor Cores）对布局和线程绑定施加了严格的约束，而逐元素操作符通常允许更大的灵活性。

        基于这些观察，我们提出了一种基于布局和片段对象的推断方案，以优化缓冲区布局和线程绑定。为了系统化管理缓冲区布局，我们维护一个记录所有缓冲区布局信息的LayoutMap。我们为分块操作符布局定义了一个分层优先级系统，优先级越高表示布局要求越严格，对性能的影响也越大。TileLang以自顶向下的方式进行布局推断，按优先级从高到低依次推断布局。在每个优先级级别上，TileLang会尝试推断所有未确定缓冲区的布局，直到无法取得进一步进展，然后再进入下一个更低的优先级级别。

        如图7所示，考虑这样一个场景：矩阵C是GEMM操作的结果，对应一个片段对象，并且需要在GEMM计算后加上偏置D。鉴于GEMM在推断过程中具有最高优先级，其线程绑定配置是预先确定的，而D的线程绑定策略仍有待确定。输出矩阵C的维度是4×4，分布在8个线程上，每个线程负责2个元素。因此，偏置缓冲区D的布局必须与此配置对齐。由于张量C的每一行由2个线程处理，两个线程在执行加法操作时都需要访问D中相同的元素。因此，必须对D进行复制，以确保每个线程都能访问到相应的元素。可以使用相同的方法推断出D的布局。

Fig. 7. An example of thread binding inference for Fragments.

        图8展示了线程绑定推断过程的一个示例。具体来说，图8(a)展示了一个用于复制数据的简单代码片段，描述了一个子分块从全局内存传输到共享内存的数据流。适当的线程绑定和向量化访问可以充分利用GPU的并行性，并利用高性能内存访问指令。在图8(b)中，T.copy操作被展开为多个循环轴。在应用布局推断通道后，如图8(c)所示，程序经历了自动向量化和并行化。最后，在图8(d)所示的阶段，应用了布局交换。

Fig. 8. Multi-Stage Automatic Thread Binding Inference for Efficient Parallel Memory Access.

4.3 利用高性能硬件指令

        现代硬件架构通常支持多种指令路径来实现相同的计算操作。例如，在NVIDIA GPU上，8位乘加操作可以通过几种类型的指令实现。IMAD指令执行标量融合乘加操作，计算𝑑 = 𝑎 · 𝑏 + 𝑐，所有操作数在内部被提升为32位整数进行计算。DP4A指令支持向量化点积操作，计算𝑑 = ⟨a, b⟩ + 𝑐 = Σ𝑖₌₀³ 𝑎𝑖𝑏𝑖 + 𝑐，其中a和b是长度为4的8位整数向量，偏置𝑐和输出𝑑均以32位整数精度表示。对于更高吞吐量的矩阵计算，MMA指令利用Tensor Cores执行D = A · B + C，其中A ∈ R¹⁶ˣ³²，B ∈ R³²ˣ⁸，C, D ∈ R¹⁶ˣ⁸；在此情况下，A和B是8位整数矩阵，而C及累加结果D使用32位整数精度。在NVIDIA RTX 3090 GPU上，这些指令的吞吐量分别约为17.8 TOPS、71.2 TOPS和284 TOPS。此外，MMA指令在同一精度设置下支持多种形状。

        在TileLang中，如图10(a)和(b)所示，有两种调用硬件张量指令的方法。第一种方法（图10(a)）使用C++源码注入，其中像dp4a这样的指令通过C++模板手动封装，并通过T.import_source和T.call_extern注入到内核中。这使得用户能够在利用熟悉的C风格语法的同时进行底层控制。注入的函数定义在生成代码的开头，并在内核内部调用。或者，如图10(b)所示，TileLang提供了内置的T.ptx原语，允许直接在内核中生成内联PTX指令（例如mma.m16n8k32.row.col.s32.s8.s8.s32）。这提供了另一种利用专用指令的底层机制，尤其适用于线程束级操作。

Fig. 9. Different methods of using high performance hardware instructions in tilelang

        然而，根据输入形状和数据类型选择最合适的指令可能具有挑战性。为了简化此过程，TileLang还支持与分块库的集成，如图10(c)所示。分块库——如NVIDIA的CUTLASS（cute）或AMD的可组合内核（CK）——为GEMM等操作提供了高级、标准化的基于分块的API（例如tl::gemm_ss）。这些库抽象了硬件特定的细节，允许底层实现根据给定的输入配置自动选择最高效的指令。在TileLang中，开发者可以使用T.call_extern以一种直接且一致的方式调用这些库。

        总而言之，TileLang提供了两种互补的方法来利用高性能指令。第一种利用分块库，这简化了集成并受益于厂商优化的性能。然而，高级抽象可能会限制底层控制。例如，cute::gemm_ss接口对共享内存输入执行GEMM操作，但从共享内存到寄存器的数据流由cute模板内部管理。这使得无法从外部标注或覆盖内部布局，从而降低了灵活性。此外，由于大量使用模板，编译时间可能会显著增加。使用NVCC 12.8跟踪工具的分析显示，对于tilelang生成的CUDA代码，模板展开约占90%的编译时间。

Fig. 10. Different methods of using DP4A and mma in tilelang

        相比之下，TileLang允许通过T.gemm使用tilelang本身直接实现指令。这避免了布局标注的限制，并减少了编译时间。然而，它要求用户为每个目标硬件指令在tilelang内部实现完整的指令集。目前，TileLang同时支持这两种方法，默认使用基于分块库的方法，以便快速支持新的硬件指令。

4.4 软件定义流水线

        TileLang采用自动化的软件流水线推断机制来分析计算块（例如，本例中的Copy和GEMM）之间的依赖关系，并生成结构化的流水线调度，以在保持正确执行顺序的同时最大化并行性。具体而言，该机制将Copy任务与其他计算密集型操作交错排列以减少空闲时间，并且在检测到异步处理机会时，自动将这些任务映射到可用的硬件资源上并发执行。因此，TileLang只需向用户暴露一个简单的num_stages接口，极大地简化了流程。然而，我们也允许用户在需要时显式提供关于顺序和阶段的信息。

Fig. 11. Software pipeline scheduling in TileLang. This illustration demonstrates how TileLang interleaves Copy and GEMM.

        对于Ampere架构，TileLang支持使用cp.async进行异步内存复制操作。cp.async指令促进全局内存和共享内存之间的快速数据移动，使内存传输与计算重叠以提高性能。TileLang通过分析循环结构并为符合条件的传输自动插入cp.async指令来整合此能力。此外，TileLang确保正确使用cp.async.commit和cp.async.wait指令来处理同步，保证数据正确性。这种优化特别有效，因为它减轻了寄存器文件的压力，并实现了对硬件带宽的更有效利用。

        在Hopper架构中引入了两项新特性。首先，引入了新的TMA单元作为专用硬件单元，负责全局内存和共享内存之间的数据复制。其次，PTX指令集引入了新的wgmma指令，该指令允许由一个线程束组（由四个线程束组成）执行矩阵乘法（MMA）操作，以提高TensorCore利用率。此外，wgmma.mma_async指令是异步的。此外，针对Hopper架构的内核优化通常采用线程束专用化，即将线程划分为生产者和消费者。生产者线程使用TMA移动数据，而消费者线程负责计算。

        在TileLang中，我们在降级过程中自动执行线程束专用化优化。具体来说，TileLang分析所有语句的缓冲区使用情况并确定其角色（生产者或消费者）。基于此分析，生产者和消费者根据threadIdx被划分到不同的执行路径。为确保计算正确性，TileLang利用活跃变量分析来确定适当的同步点，并相应地插入内存屏障。

        异步复制指令和DMA支持在AMD CDNA架构中同样提供，TileLang通过HIP封装的Copy原语来支持。具体而言，TileLang利用诸如s_waitcnt lgkmcnt和buffer_load_dword lds等指令来高效管理内存传输。这种集成使系统能够充分利用硬件的重叠数据移动与计算的能力，进一步提升了流水线性能并减少了空闲时间。

5 数值实验

        本节通过一系列在不同硬件平台和工作负载上的综合数值实验评估TileLang的性能。我们的目标是展示TileLang在优化构成现代机器学习工作负载核心的关键算子内核方面的有效性、通用性和可扩展性。通过与最先进解决方案进行基准测试，我们旨在突出TileLang在处理混合精度计算方面的多功能性，以及其在多种GPU架构上实现显著性能提升的能力。

5.1 实验设置

硬件平台。我们在NVIDIA和AMD GPU上评估TileLang，因为它们是应用最广泛的加速器。我们的实验使用了三款前沿GPU：NVIDIA H100（80 GB）、NVIDIA A100（80 GB）和AMD Instinct MI300X（192 GB）。对于NVIDIA H100，我们使用CUDA 12.4；对于MI300X，我们使用ROCm 6.1.0。所有平台均在Ubuntu 20.04下运行。

算子工作负载。我们在一系列大规模深度学习流程中频繁出现的算子工作负载上评估TileLang。在NVIDIA H100上，我们重点关注多头注意力（MHA）、线性注意力和通用矩阵乘法（GEMM）。对于NVIDIA A100，我们测量反量化GEMM内核的性能。同时，在AMD Instinct MI300X上，我们对GEMM和MHA进行基准测试，以涵盖跨越不同GPU架构的代表性用例。这些工作负载构成了包括大语言模型在内的许多当代神经网络模型的基础构建模块。

基线方法。为了评估TileLang的性能，我们将其与机器学习和GPU编程中广泛使用的几种最先进基线进行比较。这些基线包括：

• FlashAttention-3：针对多头注意力的优化实现，利用了如tma和wgmma.mma_async等CUDA指令。

• Triton：一个用于高效GPU内核的开源框架，支持NVIDIA和AMD GPU，但需要手动优化。

• cuBLAS：NVIDIA的高性能密集线性代数库。

• rocBLAS：AMD的BLAS库。

• PyTorch：提供手写优化的内核（如GEMM和FlashAttention-2），但未完全优化。

• BitsandBytes：专为支持如𝑊NF4𝐴FP16等格式并提供高效内核而设计。

• Marlin：针对𝑊INT4𝐴FP16计算进行高度优化的内核。

这一选择为TileLang提供了跨多种优化策略和硬件兼容性的全面比较。

5.2 实验

Flash Attention性能。与FlashAttention-3、Triton和PyTorch相比，TileLang分别实现了1.36倍、1.41倍和1.70倍的加速。由于FlashAttention-3是手工实现的方法，无法高效适应不同大小的工作负载。具体而言，其固定的分块大小导致对于较短的序列长度性能不佳。对于较长的序列长度（例如8k），TileLang的性能仍接近FlashAttention-3。PyTorch使用了手写优化的FlashAttention-2内核，其性能低于FlashAttention-3。与这些基于手动模板的实现相比，TileLang可以自动利用诸如cp.async.bulk和wgmma.mma_async等指令，并自动应用线程束专用化等优化。值得注意的是，在H100 GPU上，TileLang能够表达与FlashAttention-3中使用的同样复杂的流水线调度方案。

Fig. 12. FlashAttention, LinearAtten Performance on Hopper Architecture.

线性注意力性能。在我们的线性注意力实验中，我们使用Mamba-2中的分块扫描和分块状态函数。与Triton相比，TileLang实现了平均1.77倍和2.10倍的加速。

多头潜在注意力性能。图14展示了MLA的性能以及在H100和MI300X GPU上相应内核实现的代码行数（LOC）。在H100上，TileLang相比Torch实现了1075.9倍的加速，显著优于Triton和FlashInfer，并达到了手写优化的FlashMLA实现性能的98%。此外，TileLang仅需约70行Python代码，与其他基线相比展现了显著更好的可用性。在MI300X上，TileLang相比Torch实现了129.2倍的加速，并在性能和代码紧凑性方面都超越了Triton。与手写库AITER相比，TileLang达到了其性能的95%。由于AITER的内核实现未开源，其代码行数未包含在图中。

矩阵乘法性能。图13展示了GEMM工作负载在NVIDIA和AMD GPU上的性能，将TileLang与Triton及厂商优化库进行了比较。在RTX 4090、A100、H100和MI300X上，TileLang相比厂商库分别实现了1.10倍、0.97倍、1.00倍和1.04倍的加速。与Triton相比，TileLang在相同GPU上分别实现了1.08倍、1.03倍、1.13倍和1.25倍的加速。对于矩阵乘法，TileLang使用简洁的语法达到了与厂商优化库相媲美的性能。此外，通过采用布局交换，TileLang确保了在所有测试设备上执行时无存储体冲突。

Fig. 13. GEMM performance on Nvidia and AMD GPUs. 

Fig. 14. Comparison of MLA performance and code lines on H100 and MI300X.

反量化矩阵乘法性能。BitBLAS是一个用于混合精度计算的高性能库，具有先进的定制类型系统以及对张量数值类型和属性的调度功能。其最初基于TensorIR构建，我们已将其底层后端替换为TileLang，从而能够直接与其他混合精度加速库进行比较。相比cuBLAS-𝑊FP16𝐴FP16，在BitBLAS-TileLang-𝑊INT2𝐴INT8配置的驱动下，TileLang实现了最高7.65倍的加速。此外，对于𝑊INT4𝐴FP16格式，我们的方法相比Marlin实现了平均1.04倍的加速；对于𝑊NF4𝐴FP16格式，相比BitsandBytes实现了平均1.62倍的加速。通过暴露线程级编程接口并允许控制数据布局和流水线配置，TileLang为开发者提供了更细粒度的优化能力。例如，开发者可以利用基于PTX的快速数值精度转换指令，并利用Ladder在分块内实现更平滑的内存访问。这些优化在Triton中实现起来具有挑战性，这使得TileLang具备独特的能力，能够提供Triton难以实现的优异性能。

Fig. 15. Dequantize Matmul Performance on A100 GPU.

6 结论与展望

        为应对为现代硬件加速器编写高性能内核的挑战，本文提出了TileLang，一种类Python的领域特定语言（DSL），允许用户在分块粒度上进行编程。与Triton不同，TileLang使用户能够在前端显式声明硬件内存层次结构中不同级别的缓冲区，并利用布局推断机制高效地并行化缓冲区操作。这意味着用户只需描述缓冲区的计算逻辑，而无需关心并行化如何实现。同时，TileLang为专家提供了在操作缓冲区时显式指定单个线程具体行为的灵活性。这种方法在易用性和细粒度控制之间取得了平衡，同时兼顾了灵活性与性能。

        与ThunderKittens相比，TileLang通过允许开发者完全使用Python进行编程，并默认抽象化流水线等优化细节，从而简化了编程过程。例如，在Flash Attention的实现中，TileLang会在Ampere GPU上自动使用异步复制进行数据移动，并在Hopper GPU上将流水线降级为TMA。尽管如此，TileLang仍然为用户提供了在必要时于前端显式实现流水线化的选项。此外，TileLang对动态参数、动态形状等高级功能提供了有力支持，使其特别适用于编写内核库。

        我们还希望探讨未来工作中扩展和增强TileLang的几个有前景的方向：首先，我们计划基于TileLang构建一个自承载的分块库，以消除当前对内置算子依赖CUTLASS及手动包装的CUDA/HIP代码的情况。其次，我们的目标是通过引入分块级通信原语和调度策略，将TileLang扩展至支持一系列分布式场景。这将使用户能够实现针对特定通信和计算资源配置的高性能内核。此外，我们计划研究为TileLang设计成本模型。鉴于其基于分块的编程范式显式暴露了线程映射细节，内存访问模式与计算行为被明确定义，这有助于进行硬件行为分析并开发更有效的成本模型。最后，我们计划探索动态形状调优的优化，特别是针对维度动态变化的程序，如何选择最合适的分块配置。TileLang设计中对内存层次结构的显式暴露，将进一步有助于支持多种硬件平台（如CPU、NPU等）的后端。我们将探索一种通用设计方法，以扩展对多后端的支持，使TileLang能够无缝适配多样化的硬件架构。

        我们的系统已开源，以支持未来的发展和社区贡献：https://github.com/tile-ai/tilelang。

References

[1] AMD CDNA Architecture. https://www.amd.com/en/technologies/cdna.
[2] NVIDIA Tensor Cores. https://www.nvidia.com/en-us/data-center/tensor-cores/.
[3] PyTorch. https://pytorch.org/.
[4] ThunderKittens. https://github.com/HazyResearch/ThunderKittens.
[5] Inc. Advanced Micro Devices. Amd cdna™ 3 architecture. Technical report, Advanced Micro Devices, Inc., 2023.
[6] Advanced Micro Devices (AMD). AMD Composable Kernel. https://github.com/ROCm/composable_kernel.
[7] Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang,
Yuwei Hu, Luis Ceze, et al. {TVM}: An automated {End-to-End} optimizing compiler for deep learning. In 13th
USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), pages 578–594, 2018.
[8] Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan
Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. TVM: An automated end-to-end optimizing
compiler for deep learning. In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18),
pages 578–594, Carlsbad, CA, 2018. USENIX Association.
[9] NVIDIA Corporation. Nvidia a100 tensor core gpu architecture. Technical report, NVIDIA Corporation, 2020.
[10] NVIDIA Corporation. Nvidia h100 tensor core gpu architecture. Technical report, NVIDIA Corporation, 2023.
[11] NVIDIA Corporation. Cutlass: Cuda templates for linear algebra subroutines. https://github.com/NVIDIA/cutlass,
2024.
[12] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact
attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344–16359, 2022.
[13] Google. Google assistant with bard: Generative ai. https://blog.google/products/assistant/google-assistant-bard-
generative-ai/, 2024.
[14] Bastian Hagedorn, Bin Fan, Hanfeng Chen, Cris Cecka, Michael Garland, and Vinod Grover. Graphene: An ir for
optimized tensor computations on gpus. In Proceedings of the 28th ACM International Conference on Architectural
Support for Programming Languages and Operating Systems, Volume 3, pages 302–313, 2023.
[15] Young Jin Kim, Rawn Henry, Raffy Fahim, and Hany Hassan Awadalla. Who says elephants can’t run: Bringing large
scale moe models into cloud scale production. arXiv preprint arXiv:2211.10017, 2022.
[16] Microsoft. The new bing. https://www.microsoft.com/en-us/edge/features/the-new-bing?form=MT00D8, 2024.
[17] OpenAI. Introducing chatgpt, 2022. Available: https://openai.com/blog/chatgpt.
[18] Phitchaya Mangpo Phothilimthana, Archibald Samuel Elliott, An Wang, Abhinav Jangda, Bastian Hagedorn, Henrik
Barthels, Samuel J Kaufman, Vinod Grover, Emina Torlak, and Rastislav Bodik. Swizzle inventor: data movement
synthesis for gpu kernels. In Proceedings of the Twenty-Fourth International Conference on Architectural Support for
Programming Languages and Operating Systems, pages 65–78, 2019.
[19] Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and
accurate attention with asynchrony and low-precision. arXiv preprint arXiv:2407.08608, 2024.
[20] Philippe Tillet, H. T. Kung, and David Cox. Triton: An Intermediate Language and Compiler for Tiled Neural Network
Computations, page 10–19. Association for Computing Machinery, New York, NY, USA, 2019.
[21] Lei Wang, Lingxiao Ma, Shijie Cao, Quanlu Zhang, Jilong Xue, Yining Shi, Ningxin Zheng, Ziming Miao, Fan Yang,
Ting Cao, et al. Ladder: Enabling efficient {Low-Precision} deep learning computing through hardware-aware tensor
transformation. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), pages 307–323,
2024.
[22] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim
Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface’s transformers: State-of-the-art natural language processing.
arXiv preprint arXiv:1910.03771, 2019.
[23] Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan
Yang. Diffusion models: A comprehensive survey of methods and applications. ACM Computing Surveys, 56(4):1–39,
2023.
[24] Lianmin Zheng, Chengfan Jia, Minmin Sun, Zhao Wu, Cody Hao Yu, Ameer Haj-Ali, Yida Wang, Jun Yang, Danyang
Zhuo, Koushik Sen, Joseph E. Gonzalez, and Ion Stoica. Ansor: Generating high-performance tensor programs for
deep learning. In 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20), pages 863–879.
USENIX Association, November 2020.
, Vol. 1, No. 1, Article . Publication date: April 2025.
TileLang: A Composable Tiled Programming Model for AI Systems 19
[25] Hongyu Zhu, Ruofan Wu, Yijia Diao, Shanbin Ke, Haoyu Li, Chen Zhang, Jilong Xue, Lingxiao Ma, Yuqing Xia, Wei
Cui, Fan Yang, Mao Yang, Lidong Zhou, Asaf Cidon, and Gennady Pekhimenko. ROLLER: Fast and efficient tensor
compilation for deep learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22),
pages 233–248, Carlsbad, CA, July 2022. USENIX Association.

A Operator shapes in our benchmark