TileLang 入门教程，用领域特定语言描述矩阵分块策略

为什么我们需要 TileLang

在将大模型推理服务迁移到 AMD ROCm 平台的过程中，很多开发者会发现一个尴尬的现象：代码虽然通过HIPify成功转换了，框架也用SGLang跑通了，但最终的推理延迟和吞吐量却总是不如预期。这往往不是因为硬件不行，而是通用的算子实现无法完全吃透 AMD GPU 独特的架构特性。

AMD 的 CDNA 架构拥有特殊的矩阵核心（Matrix Cores）和复杂的内存层级（如 LDS 共享内存）。如果直接沿用从 CUDA 平移过来的逻辑，很容易导致计算单元闲置或者内存带宽成为瓶颈。这时候，我们就需要一种更精细的工具来描述数据如何在芯片内部流动，这就是TileLang登场的原因。它不是让你去写晦涩的汇编，而是用一种领域特定语言（DSL）清晰地定义“矩阵分块”策略，让编译器自动生成针对特定架构高度优化的内核代码。

理解矩阵分块的核心逻辑

要写好 TileLang 代码，首先得跳出“逐元素计算”的思维惯性，转而思考“数据块”的搬运与计算。在 GPU 上，全局显存（Global Memory）的访问速度远慢于片上共享内存（LDS）。高效的算子优化，本质上就是设计一套精密的流水线：先把大块数据切分成适合放入 LDS 的小_tile_，由多个线程协作将其从显存预取到共享内存，然后在片上完成密集计算，最后写回结果。

TileLang 的核心价值在于它将这个过程显式化了。你不需要手动管理线程索引的复杂偏移量，只需声明块的大小（Block Size）、循环的展开方式以及数据在层级间的映射关系。编译器会据此生成完美的指令序列，确保 Wavefront（AMD 的线程束）内的线程协同工作，避免分支发散，最大化利用向量指令集。

手把手实现一个矩阵乘法 Kernel

理论说得再多，不如看一段真实的代码。下面我们通过一个最经典的矩阵乘法（C=A×BC = A \times BC=A×B）示例，演示如何用 TileLang 描述这一过程。假设我们要计算两个M×KM \times KM×K和K×NK \times NK×N的矩阵相乘。

首先，我们需要定义程序的入口和迭代空间。在 TileLang 中，我们使用@tilelang.kernel装饰器来标记函数，并通过iter_vars声明逻辑上的循环维度。

importtilelangastl@tl.kerneldefmatmul_kernel(A:tl.Buffer["float16",[M,K]],B:tl.Buffer["float16",[K,N]],C:tl.Buffer["float16",[M,N]]):# 定义逻辑迭代变量m,n,k=tl.iter_vars()# 设定分块大小，这是优化的关键参数BLOCK_M=64BLOCK_N=64BLOCK_K=32# 将逻辑坐标映射到具体的 Block IDpid_m=m//BLOCK_M pid_n=n//BLOCK_N# 初始化共享内存缓冲区# LDS 是片上高速缓存，必须显式声明shared_A=tl.alloc_shared([BLOCK_M,BLOCK_K],dtype="float16")shared_B=tl.alloc_shared([BLOCK_K,BLOCK_N],dtype="float16")# 累加器，用于存放中间计算结果acc=tl.zeros([BLOCK_M,BLOCK_N],dtype="float32")# 主循环：沿着 K 维度进行分块迭代fork_iterintl.range(0,K,BLOCK_K):# 阶段一：数据加载 (Data Movement)# 将全局显存中的数据异步加载到共享内存# 这里隐含了线程协作的逻辑，每个线程负责搬运一部分tl.copy(A[pid_m*BLOCK_M:(pid_m+1)*BLOCK_M,k_iter:k_iter+BLOCK_K],shared_A)tl.copy(B[k_iter:k_iter+BLOCK_K,pid_n*BLOCK_N:(pid_n+1)*BLOCK_N],shared_B)# 等待数据加载完成，确保同步tl.sync()# 阶段二：矩阵计算 (Compute)# 在共享内存上进行小块矩阵乘法，并累加到 acc# 编译器会将此操作映射为 AMD Matrix Core 指令acc+=tl.matmul(shared_A,shared_B)# 再次同步，确保下一轮迭代不会覆盖正在使用的数据tl.sync()# 阶段三：写回结果# 将累加器中的高精度结果转换并写回全局显存tl.copy(acc,C[pid_m*BLOCK_M:(pid_m+1)*BLOCK_M,pid_n*BLOCK_N:(pid_n+1)*BLOCK_N])

这段代码看似简洁，但背后蕴含了完整的优化逻辑。注意看BLOCK_M、BLOCK_N和BLOCK_K的定义，这三个数值直接决定了寄存器压力和 LDS 的使用率。在 AMD CDNA 架构上，通常需要根据 Wavefront 的大小（通常是 64）来对齐这些块尺寸，以消除线程束内的空闲线程。

代码中的tl.copy并非简单的内存拷贝，在编译后的 HIP 代码中，它会被展开为高效的vector_load和vector_store指令，甚至利用 DMA 引擎进行异步搬运，从而掩盖内存访问延迟。而tl.matmul在共享内存上的操作，则会被直接 lowering 为mfma(Matrix Fused Multiply-Add) 指令，这是 AMD 矩阵核心的杀手锏，能在一个时钟周期内完成大量浮点运算。

从 DSL 到机器码的蜕变

当你运行这段 TileLang 代码时，编译器前端会解析你的分块策略，构建出中间表示（IR）。接着，后端会根据目标架构（例如 MI250 或 MI300 系列）的具体参数，进行指令调度和寄存器分配。

最关键的一步是循环展开与指令重排。编译器会自动分析依赖关系，将数据加载指令提前发起，使得计算单元在处理上一块数据时，下一块数据已经在传输路上。这种软件流水线（Software Pipelining）技术，如果手动用 C++/HIP 编写，不仅代码量巨大，而且极易出错。而在 TileLang 中，你只需要关注数据流动的拓扑结构，复杂的调度交给编译器即可。

此外，TileLang 还能自动处理边界条件。当矩阵尺寸不能被块大小整除时，生成的内核会自动插入掩码（Mask）逻辑，防止越界访问，无需开发者手动编写繁琐的if-else判断，这进一步保证了生成代码的整洁与高效。

实战中的调优心得

在实际项目中，不要指望一套参数打天下。不同的模型层（如 Attention 的 QKV 投影 vs MLP 层）对算力与带宽的需求比例不同。对于计算密集型层，可以尝试增大BLOCK_K以复用更多共享内存中的数据；对于访存密集型层，则可能需要调整BLOCK_M和BLOCK_N的比例来匹配带宽峰值。

建议在使用 TileLang 时，结合rocprof等性能分析工具，观察生成的内核在 L1/L2 缓存命中率以及 Matrix Core 利用率上的表现。很多时候，仅仅微调几个分块常数，就能带来 20% 以上的性能提升。这种细粒度的控制能力，正是我们在非 NVIDIA 环境下构建高性能推理服务的底气所在。

通过这种“描述即优化”的方式，我们不再是被动的代码搬运工，而是成为了硬件资源的调度者。TileLang 让算子优化变得可解释、可维护，也让 AMD GPU 的潜力得以真正释放。

200小时GPU算力已就位，快来领取：https://marketing.csdn.net/questions/Q2604140858304426315?utm_source=AIpaper