TileLang 算子优化入门，让 AMD GPU 跑得更快

2600_96323227

309人浏览 · 2026-06-24 10:36:30

2600_96323227 · 2026-06-24 10:36:30 发布

为什么通用算子在 AMD GPU 上“跑不满”

很多开发者在将大模型推理服务迁移到 AMD Instinct 系列显卡（如 MI300X）时，都会遇到一个尴尬的局面：代码能跑通，但性能就是不如预期。尤其是在运行 Attention 或 MLP 层这类计算密集型算子时，即便使用了 SGLang 这样高效的框架，吞吐量依然难以达到理论峰值。

这背后的核心原因往往在于“水土不服”。NVIDIA 的 CUDA 生态经过多年打磨，其默认算子库（如 cuBLAS）针对 SM 架构做了极深的优化。而 AMD GPU 基于 CDNA 或 RDNA 架构，拥有独特的内存层级和计算单元组织方式——特别是 Wavefront（波前）机制，它与 NVIDIA 的 Warp 在执行逻辑和尺寸上存在显著差异。如果直接沿用从 CUDA 平移过来的通用实现，很容易导致线程束发散、共享内存（LDS）访问冲突或寄存器压力过大，从而让昂贵的算力闲置。

要解决这个问题，通用的框架配置往往不够用，我们需要更细粒度的控制手段。这时候，TileLang 这样的领域特定语言（DSL）就成了进阶开发者的利器。它允许我们用高层语言描述矩阵分块（Tiling）和数据流动策略，然后自动编译成高度适配特定 GPU 架构的内核代码。

用 TileLang 重写矩阵乘法：从理论到代码

为了直观展示优化过程，我们选取最基础的矩阵乘法（GEMM）作为案例。假设我们要计算 $\times B$ ，其中 $A$ 是 $\times K$ 矩阵， $B$ 是 $\times N$ 矩阵。在通用实现中，全局内存的频繁访问通常是性能瓶颈。优化的核心思路是将数据切分成小块（Tile），加载到速度更快的共享内存中进行计算。

在 TileLang 中，我们可以清晰地定义这种分块策略。以下是一个针对 AMD GPU 优化的简化代码示例：

import tilelang as tl

# 定义矩阵维度
M, N, K = 4096, 4096, 4096
# 关键：根据 AMD GPU 架构调整 Block Size
# 对于 MI300X (gfx942)，Wavefront 大小为 64
# 设置 Block 为 128x128 可以完美映射 2x2 个 Wavefront，减少调度开销
BLOCK_SIZE_M = 128
BLOCK_SIZE_N = 128
BLOCK_SIZE_K = 32

@tl.kernel
def gemm_kernel(
    A: tl.Tensor[M, K],
    B: tl.Tensor[K, N],
    C: tl.Tensor[M, N]
):
    # 分配共享内存 (LDS)
    shared_A = tl.alloc_shared([BLOCK_SIZE_M, BLOCK_SIZE_K], dtype=A.dtype)
    shared_B = tl.alloc_shared([BLOCK_SIZE_K, BLOCK_SIZE_N], dtype=B.dtype)
    
    # 获取当前 Block 的索引
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    # 初始化累加器
    acc = tl.zeros([BLOCK_SIZE_M, BLOCK_SIZE_N], dtype=tl.float32)
    
    # 循环加载分块数据
    for k in range(0, K, BLOCK_SIZE_K):
        # 加载 A 的分块到共享内存
        a_tile = tl.load(A[pid_m * BLOCK_SIZE_M : (pid_m + 1) * BLOCK_SIZE_M, 
                           k : k + BLOCK_SIZE_K])
        shared_A[:] = a_tile
        
        # 加载 B 的分块到共享内存
        b_tile = tl.load(B[k : k + BLOCK_SIZE_K, 
                           pid_n * BLOCK_SIZE_N : (pid_n + 1) * BLOCK_SIZE_N])
        shared_B[:] = b_tile
        
        # 等待所有线程同步，确保数据加载完成
        tl.sync_threads()
        
        # 在共享内存中执行矩阵乘累加
        acc += tl.dot(shared_A, shared_B)
        
        tl.sync_threads()
    
    # 将结果写回全局内存
    tl.store(C[pid_m * BLOCK_SIZE_M : (pid_m + 1) * BLOCK_SIZE_M, 
               pid_n * BLOCK_SIZE_N : (pid_n + 1) * BLOCK_SIZE_N], acc)

这段代码看似简单，但有几个关键点直接决定了在 AMD 硬件上的表现：

共享内存显式管理：通过 alloc_shared，我们强制将热点数据放入 LDS。AMD GPU 的 LDS 带宽远高于全局显存，合理复用能大幅降低内存墙效应。
同步机制：tl.sync_threads() 确保了在数据被使用前，所有线程已完成加载。这在避免竞态条件的同时，也隐式地协调了 Wavefront 的执行节奏。
Block 尺寸匹配：这是最关键的一步。代码中选择的 128x128 Block 尺寸并非随意设定，而是为了匹配 MI300X 的硬件特性。AMD 的 Wavefront 包含 64 个线程，将 Block 设置为 Wavefront 尺寸的整数倍，可以确保每个 Wavefront 内的线程负载均衡，避免出现部分线程空闲等待的情况。

调优实战：如何匹配 Wavefront 与 LDS

在实际调试过程中，很多人容易忽视架构参数的细微差别。比如，NVIDIA H100 的 Warp Size 是 32，而 AMD MI250/MI300 系列的 Wavefront Size 固定为 64。如果你在 TileLang 中沿用了针对 NVIDIA 优化的 64x64 分块策略，在 AMD 卡上虽然能运行，但每个 Block 只占用了一个 Wavefront，可能导致计算单元利用率不足。

通过将 Block Size 调整为 128x128 或 256x128，我们可以让每个 Block 启动多个 Wavefront 并行工作。此外，还需要注意 LDS 的容量限制。不同架构的 LDS 大小有限（例如每 CU 64KB 或 128KB），过大的分块会导致寄存器溢出或 occupancy 下降。TileLang 的优势在于它允许我们快速迭代这些参数：只需修改几行配置，重新编译内核，即可观察性能变化，无需深入繁琐的 HIP C++ 代码编写。

在 SGLang 等框架中集成这种自定义算子也非常顺畅。我们可以将编译好的 kernel 注册为后端算子，替换掉默认的注意力机制实现。特别是在处理长序列推理时，这种针对特定硬件定制的分块策略，能有效减少 KV Cache 读取带来的延迟。