TileLang GEMM 教程：从基础矩阵乘法到细粒度 MMA 与自动调优

TileLang GEMM 教程从基础矩阵乘法到细粒度 MMA 与自动调优【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelangTileLang 是一个专为 GPU/CPU/加速器 Kernel 开发设计的领域专用语言DSL通过内置的共享内存分配、调度与分块抽象把 GEMM矩阵乘法这类典型 Kernel 的编写成本大幅降低。本文以 examples/gemm/README.md 为主线结合仓库中 examples/gemm 目录下的完整示例与tilelang源码系统讲解如何用 TileLang 写出一个可运行的 GEMM、如何编译与性能剖析、如何通过 Swizzle/并行拷贝/自动流水线/光栅化优化访存以及如何深入 warp 级ldmatrix/mma/stmatrix细粒度 MMA 编程。读完本文你将能够从零写出并验证一个 GEMM Kernel并能继续向 Persistent Kernel 与自动调优方向进阶。为什么选择 TileLang 编写 GEMM现代 GPU如 NVIDIA 各代架构上获得高性能 GEMM 通常需要手工处理分块tiling、共享内存管理、warp 级 Tensor Core 指令调度、流水线重叠等多层细节。TileLang 的价值在于把这些高频操作抽象为高层的语言原语让开发者专注于算法与性能策略本身。正如 README 所述它提供了对内存分配、调度和分块tiling的高层抽象这些正是现代硬件架构上取得极致性能的关键。同时TileLang 并不封闭——当高层原语如T.gemm无法满足需求时你仍然可以像写 CUDA 一样精确控制每一次 MMA 指令、数据布局与同步点。这种高抽象起步、低抽象兜底的设计使它能同时服务常规开发者与需要榨干每一字节带宽的专家。环境准备与安装根据 README运行 GEMM 示例需要以下环境依赖用途Python 3.8运行 TileLang 脚本NVIDIA GPU 较新的 CUDA Toolkit编译与运行 CUDA KernelPyTorch可选便捷的正确性验证torch.matmul对比tilelang核心 DSL 与编译器bitblas可选进阶示例中的 Swizzle 布局工具安装命令pip install tilelang bitblas若从源码安装或使用其他环境请根据实际情况调整上述命令。安装完成后可以用仓库自带的最小示例快速验证环境python examples/gemm/example_gemm.py该脚本examples/gemm/example_gemm.py会编译一个 1024×1024 的 FP16 GEMM与a b对比后打印 CUDA 源码与耗时正常运行并打印All check passed.即表示环境就绪。基础 GEMM 示例一次完整的 TileLang Kernel 编写README 给出了一个非常简洁的基础示例将两个 1024×1024 的矩阵相乘每个线程块使用 128 个线程。核心代码如下与 examples/gemm/example_gemm.py 中的实现一致import tilelang from tilelang import Profiler import tilelang.language as T def matmul(M, N, K, block_M, block_N, block_K, dtypeT.float16, accum_dtypeT.float): T.prim_func def main( A: T.Tensor((M, K), dtype), B: T.Tensor((K, N), dtype), C: T.Tensor((M, N), dtype), ): # 定义足以覆盖 M×N 的线程块网格 with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads128) as (bx, by): # 为 A、B 的当前分块分配共享内存 A_shared T.alloc_shared((block_M, block_K), dtype) B_shared T.alloc_shared((block_K, block_N), dtype) # 分配寄存器fragment片段用于部分累加 C_local T.alloc_fragment((block_M, block_N), accum_dtype) # 将局部累加缓冲清零 T.clear(C_local) # 沿 K 维以 block_K 为步长循环使用 3 级流水线 for k in T.Pipelined(T.ceildiv(K, block_K), num_stages3): # 从全局内存拷贝到共享内存 T.copy(A[by * block_M, k * block_K], A_shared) T.copy(B[k * block_K, bx * block_N], B_shared) # 对分块执行矩阵乘累加 T.gemm(A_shared, B_shared, C_local) # 将累加结果从局部内存C_local写回全局内存C T.copy(C_local, C[by * block_M, bx * block_N]) return main代码逐段解析定义 Kernel 启动配置with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads128) as (bx, by):这创建了一个二维线程块网格x 维度有ceildiv(N, block_N)个块y 维度有ceildiv(M, block_M)个块每个块 128 个线程。bx/by即当前块在网格中的坐标用于索引各自负责的输出分块。共享内存分配A_shared T.alloc_shared((block_M, block_K), dtype) B_shared T.alloc_shared((block_K, block_N), dtype)A、B 的分块被加载进共享内存以加速访问。从源码看alloc_shared默认的 scope 为shared.dyn动态共享内存见 tilelang/language/allocate.py。局部片段Fragment累加C_local T.alloc_fragment((block_M, block_N), accum_dtype)部分结果保存在寄存器local.fragment作用域见 tilelang/language/allocate.py中从而减少对全局内存的写入。这也是绝大多数高性能 GEMM 的关键策略把频繁累加放在寄存器级。流水线加载与 GEMMfor k in T.Pipelined(T.ceildiv(K, block_K), num_stages3): T.copy(...) T.gemm(...)以分块方式沿 K 维循环并开启最多 3 级流水线让数据搬运与计算重叠减少空转周期。T.Pipelined定义在 tilelang/language/loop.pyT.gemm定义在 tilelang/language/gemm_op.py。结果写回T.copy(C_local, C[by * block_M, bx * block_N])把寄存器/共享内存中的最终分块写回全局内存。编译、运行与性能剖析README 展示了不使用 JIT 装饰器的调用方式func matmul(1024, 1024, 1024, 128, 128, 32) print(func) # 打印 TileLang Kernel 的 IR 表示 artifact tilelang.lower(func) profiler Profiler(artifact.rt_mod, artifact.params, result_idx[2]) import torch a torch.randn(1024, 1024).cuda().half() b torch.randn(1024, 1024).cuda().half() c profiler(a, b) ref_c a b # 验证结果 torch.testing.assert_close(c, ref_c, rtol1e-2, atol1e-2) # 获取 CUDA Kernel 源码 print(artifact.kernel_source)其中tilelang.lower(func)负责把 TileLang 的prim_func下降为可执行运行时模块rt_modProfiler则包装了运行时调用与基准测试能力。仓库中的 examples/gemm/example_gemm.py 采用了另一种更贴近实际使用的tilelang.jit风格tilelang.jit def matmul(A, B, block_M, block_N, block_K, dtypeT.float16, accum_dtypeT.float32): M, N, K T.const(M, N, K) ... return C kernel matmul.compile(M1024, N1024, K1024, block_M128, block_N128, block_K32) c kernel(a, b) print(kernel.get_kernel_source()) # 生成的 CUDA 源码 profiler kernel.get_profiler() latency profiler.do_bench(backendcupti) # 用 CUPTI 计时单位 ms两种方式本质等价前者显式走lowerProfiler后者由 JIT 编译封装。do_bench支持backendcupti依赖 NVIDIA CUPTI或默认的事件计时方式。正确性验证正确性验证可以完全交给 PyTorch生成随机矩阵 → 运行 TileLang Kernel → 与参考实现torch.matmul或运算符对比import torch profiler Profiler(rt_mod, params, result_idx[2]) A torch.randn(1024, 1024).cuda().half() B torch.randn(1024, 1024).cuda().half() C_tilelang profiler(A, B) C_ref A B torch.testing.assert_close(C_tilelang, C_ref, rtol1e-2, atol1e-2) print(Results match!)此外Profiler还提供了assert_allclose(ref_program, ...)这类集成接口例如 examples/gemm/example_gemm_intrinsics.py 中直接用它校验 FP16 细粒度 MMA Kernel 与A B.T的一致性。进阶 GEMM 特性自定义内存布局 / SwizzlingSwizzling数据交错通过重新组织共享内存或全局内存中的排列缓解 bank conflict、改善缓存利用率并更好地匹配 GPU 的 warp 执行模式。TileLang 提供make_swizzle_layout之类的辅助函数来标注缓冲应如何布局。在源码层面make_mma_swizzle_layout的实现tilelang/cuda/intrinsics/layout/mma_layout.py会检查共享缓冲的最后维度是否满足对齐条件shape[-1] * dtype_bits % 512 0只有满足时才真正做 Swizzle 变换否则退回恒等布局——这保证了任意形状下该函数都可安全使用。T.annotate_layout({...})用于把布局标注挂到对应缓冲上其实现位于 tilelang/language/annotations.py它把Layout或可调用对象注册进layout_map供后续布局推断与代码生成阶段消费。并行拷贝与自动流水线并行拷贝Parallel Copy将一块分片的拷贝任务均匀分派给块内所有线程加速全局→共享内存的搬运。自动流水线Auto-Pipelining使用多级 staging 让拷贝与计算重叠减少空闲周期。基础示例中的T.Pipelined(..., num_stages3)即触发该机制。面向 L2 缓存局部性的光栅化Rasterization在 Kernel 级开启swizzle光栅化可以提升数据复用、减少 L2 缓存抖动矩阵规模越大收益越明显。TileLang 通过T.use_swizzle(panel_size, order, enable)提供该能力其实现tilelang/language/annotations.py支持三种顺序row→rasterization2DRow默认column→rasterization2DColumnmlx→rasterization2DMLX当enableFalse时该注解不产生任何效果。panel_size控制光栅化面板尺寸仓库示例中常用值为10。带注解的增强 GEMM 示例README 给出了一个更进阶的片段把内存布局、Swizzle 光栅化与并行拷贝组合起来import tilelang.language as T # make_mma_swizzle_layout 是一个 Python 定义的布局函数 # 用于让数据对齐 MMA矩阵乘累加操作。 from tilelang.cuda.intrinsics import make_mma_swizzle_layout as make_swizzle_layout def matmul(M, N, K, block_M, block_N, block_K, dtypeT.float16, accum_dtypeT.float): T.prim_func def main( A: T.Tensor((M, K), dtype), B: T.Tensor((K, N), dtype), C: T.Tensor((M, N), dtype), ): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads128) as (bx, by): A_shared T.alloc_shared((block_M, block_K), dtype) B_shared T.alloc_shared((block_K, block_N), dtype) C_local T.alloc_fragment((block_M, block_N), accum_dtype) # 标注内存布局 T.annotate_layout({ A_shared: make_swizzle_layout(A_shared), B_shared: make_swizzle_layout(B_shared), }) # 启用基于 swizzle 的光栅化以改善 L2 局部性 T.use_swizzle(panel_size10, enableTrue) T.clear(C_local) for idx in T.Pipelined(T.ceildiv(K, block_K), num_stages3): # 拷贝 A 的分块 T.copy(A[by * block_M, idx * block_K], A_shared) # 并行拷贝 B 的分块 for ko, j in T.Parallel(block_K, block_N): B_shared[ko, j] B[idx * block_K ko, bx * block_N j] T.gemm(A_shared, B_shared, C_local) T.copy(C_local, C[by * block_M, bx * block_N]) return main与基础示例的关键差异T.annotate_layout(...)标注共享内存中的数据组织方式Swizzling。T.use_swizzle(...)启用基于 Swizzle 的光栅化。T.Parallel(...)并行拷贝循环把全局→共享的拷贝分派到所有线程并可能向量化 load/store 指令。其中make_mma_swizzle_layout从tilelang.cuda.intrinsics导出见 tilelang/cuda/intrinsics/layout/init.py是专门为对齐 MMA 操作设计的布局函数。细粒度 MMA 计算像写 CUDA 一样控制 warp对需要完全控制 warp 级矩阵乘操作的高级用户TileLang 允许以接近裸 CUDA 的方式指定细粒度 MMA 计算。虽然T.gemm(...)等高层抽象已能满足多数场景但特定负载例如反量化 GEMM 需要在共享→寄存器阶段做细粒度布局变换可能受益于显式控制每条 MMA 指令、数据布局与同步点。示例工作流simplify_prim_func def tl_matmul(M, N, K, in_dtype, out_dtype, accum_dtype): assert in_dtype in [T.float16, T.int8], Currently only float16 and int8 are supported assert out_dtype in [T.float16, T.float32, T.int32], Currently only float16, float32 and int32 are supported micro_size_x micro_size_y micro_size_k 16 if out_dtype T.int32: micro_size_k 32 # 调试用配置 block_row_warps 2 block_col_warps 2 warp_row_tiles 32 warp_col_tiles 32 chunk 32 shared_scope shared.dyn stage 2 block_M block_row_warps * warp_row_tiles block_N block_col_warps * warp_col_tiles block_K chunk A_shape (M, K) B_shape (N, K) A_shared_shape (block_M, block_K) B_shared_shape (block_N, block_K) C_shared_shape ( block_M // micro_size_x, block_N // micro_size_y, micro_size_x, micro_size_y, ) warp_size 32 threads warp_size * (block_row_warps * block_col_warps) local_size_a (micro_size_x * micro_size_k) // warp_size local_size_b (micro_size_y * micro_size_k) // warp_size local_size_c (micro_size_x * micro_size_y) // warp_size warp_rows warp_row_tiles // micro_size_x warp_cols warp_col_tiles // micro_size_y # MMA 包装器自动生成 MMA 指令代码 mma_emitter TensorCoreIntrinEmitter( a_dtypein_dtype, b_dtypein_dtype, accum_dtypeaccum_dtype, a_transposedFalse, b_transposedTrue, block_row_warpsblock_row_warps, block_col_warpsblock_col_warps, warp_row_tileswarp_row_tiles, warp_col_tileswarp_col_tiles, chunkchunk, ) T.prim_func def main(A: T.Tensor(A_shape, in_dtype), B: T.Tensor(B_shape, in_dtype), C: T.Tensor((M, N), out_dtype)): with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threadsthreads) as (bx, by): A_shared T.alloc_shared(A_shared_shape, in_dtype, scopeshared_scope) B_shared T.alloc_shared(B_shared_shape, in_dtype, scopeshared_scope) C_shared T.alloc_shared(C_shared_shape, out_dtype, scopeshared_scope) A_local T.alloc_local((warp_rows * local_size_a), in_dtype) B_local T.alloc_local((warp_cols * local_size_b), in_dtype) C_local T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype) T.annotate_layout({ A_shared: make_swizzle_layout(A_shared), B_shared: make_swizzle_layout(B_shared), }) # 改善 L2 缓存 T.use_swizzle(panel_size10) T.clear(C_local) for ko in T.Pipelined((K // block_K), num_stagesstage): for i, k in T.Parallel(block_M, block_K): A_shared[i, k] A[by * block_M i, ko * block_K k] for j, k in T.Parallel(block_N, block_K): B_shared[j, k] B[bx * block_N j, ko * block_K k] for ki in T.serial(0, (block_K // micro_size_k)): mma_emitter.ldmatrix_a(A_local, A_shared, ki) mma_emitter.ldmatrix_b(B_local, B_shared, ki) mma_emitter.mma(A_local, B_local, C_local) mma_emitter.stmatrix(C_local, C_shared) for i, j in T.Parallel(block_M, block_N): C[by * block_M i, bx * block_N j] C_shared[ i // micro_size_x, j // micro_size_y, i % micro_size_x, j % micro_size_y, ]这个流程与手写 CUDA 的对应关系如下设置分块尺寸与线程绑定如同 CUDA先确定每个块的 warp/线程数量以及矩阵如何细分。T.Kernel(...)保证正确的线程数每个线程被绑定到具体角色warp ID 或 lane ID。分配 warp 局部片段不用单一共享缓冲做部分和而是分配寄存器片段存放 A、B 的子块A_local T.alloc_local((warp_rows * local_size_a), in_dtype) B_local T.alloc_local((warp_cols * local_size_b), in_dtype) C_local T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype)这些local分配表示每个线程的私有存储区域共同构成 warp 的寄存器分块alloc_local默认作用域为local见 tilelang/language/allocate.py。通过ldmatrix加载数据mma_emitter.ldmatrix_a()/.ldmatrix_b()是 warp 同步内联指令对应 PTXldmatrix的高层封装负责地址计算与数据装载你也可以自行编写加载逻辑for ki in T.serial(0, (block_K // micro_size_k)): mma_emitter.ldmatrix_a(A_local, A_shared, ki) mma_emitter.ldmatrix_b(B_local, B_shared, ki)执行 MMA 指令加载子块后warp 执行mma本质即[ C_{\text{local}} ;; A_{\text{local}} ;\times; B_{\text{local}} ]mma_emitter.mma(A_local, B_local, C_local)底层会翻译为 Tensor Core 指令如 PTX 中的wmma.mma.sync一个 warp 并行处理多个数据元素。通过stmatrix存储结果把 warp 级片段写回共享内存或全局内存mma_emitter.stmatrix(C_local, C_shared)它编排 warp 同步存储确保每个线程把正确的片段元素放到正确位置。小结将 warp 同步内联指令ldmatrix、mma、stmatrix与手动线程绑定、内存分配结合可以在 TileLang 层面复现裸 CUDA 的控制力与性能。这种方式最适合熟悉 GPU warp 级编程的专家用户因为它要求深入理解硬件并发、内存层次与调度但在性能关键路径上收益可能非常可观——每一字节带宽、每一周期延迟都需要精心编排。从仓库示例继续深入Persistent 与自动调优README 聚焦于基础与细粒度 MMA而 examples/gemm 目录还提供了可直接运行的进阶示例可作为本文内容的自然延伸example_gemm_persistent.py演示 Persistent GEMM。它通过driver.get_num_sms()获取 SM 数量将网格压缩为与 SM 数量相当的常驻线程块并使用T.Persistent([...], sm_num, block_id)循环遍历全部输出分块脚本同时提供了非 Persistent 版本并计算加速比。其中还用到了T.use_swizzle(10)与C_shared中转寄存器片段 → 共享 → 全局的两段式写回。example_gemm_autotune.py / example_gemm_advanced_autotune.py演示基于AutoTuner的配置搜索。候选配置涵盖block_M/block_N/block_K、num_stages、thread_num、enable_rasteration等维度--with_roller模式借助MatmulTemplate的 roller 按设备CUDA/CDNA推荐 Tensor Core 友好的分块。进阶版还支持--enable_grouped_compile分组编译加速调优、--use_pipeline编译/基准流水线与--benchmark_multi_gpu多 GPU 并行 benchmark。同时脚本提供了按设备算力SM 80/90预设的启发式配置A100sm_80推荐block_M128, block_N256, block_K32, num_stages2, threads128H100sm_90推荐block_M128, block_N256, block_K64, num_stages3, threads256。example_gemm_intrinsics.py细粒度 MMA 示例的完整可运行版本默认编译 4096×4096×4096 的 FP16 内核并做基准测试与正确性校验其中make_swizzle_layout会根据shape[-1] * bits 512决定是否应用 Swizzle 变换对应 tilelang/cuda/intrinsics/layout/utils.py 中的get_swizzle_layout。测试与回归验证仓库为这些示例配套了自动化测试与性能回归脚本examples/gemm/test_example_gemm.py调用example_gemm.main()与example_gemm_intrinsics.main()进行端到端正确性验证其中 intrinsics 用例因依赖tl.ptx_ldmatrix而标记为requires_cuda。examples/gemm/regression_example_gemm.py通过tilelang.testing.process_func驱动各示例的run_regression_perf用于性能回归跟踪。这意味着你可以在本仓库中直接通过pytest examples/gemm/test_example_gemm.py复现全部验证流程。总结从本文可以看到TileLang 为 GEMM 开发提供了完整的梯度基础路径T.KernelT.alloc_shared/T.alloc_fragmentT.PipelinedT.gemm几十行代码即可获得一个正确的分块流水线 Kernel优化路径T.annotate_layoutmake_mma_swizzle_layout处理共享内存布局T.Parallel并行拷贝T.use_swizzle开启 L2 光栅化专家路径TensorCoreIntrinEmitterldmatrix/mma/stmatrix精确控制 warp 级指令甚至可继续参考仓库中的 Persistent 与自动调优示例逼近手写 CUDA 的控制粒度。无论你是第一次接触 GPU Kernel 开发还是需要在性能关键路径上精细调优都可以在 examples/gemm/README.md 及其配套示例中找到合适的起点。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考