80 行 Python 写出 H100 级 MLA 解码算子:TileLang 是怎么做到的

一、TileLang 到底解决了什么问题

现代 AI 模型（尤其 MHA / GQA / MLA / Linear Attention 这类变体）越来越依赖融合算子（fused kernels）来榨干硬件。但现状是两难：

• 写 Triton：快，但编译器把共享内存复用、流水线调度、tile 复用这些性能旋钮藏起来了——MLA 的 Triton 实现要 130 行，却只有手写 CUDA（约 500 行）性能的 14.2%。

• 写 CUDA / CUTLASS：性能拉满，但模板机器繁琐、可移植性差、维护成本高。

TileLang 的定位很巧——它处在 Triton 和 CUTLASS 的正中间：你写 Python，但要显式声明"这块 buffer 放共享内存、这条流水线分几级、warp 怎么切"，剩下的线程映射交给编译器的 Layout Inference 自动推。

核心理念一句话：把 tile（张量的超矩形切片）做成一等公民。

• tile 够细，能表达内存层级、warp 划分、流水线；

• tile 又够粗，能形成稳定的可移植 API。

这就是它能兼顾"可编程性"和"性能"的关键颗粒度。

二、为什么 80 行 Python 能做到 FlashMLA 98% 性能

1. 显式的 tile 级原语

TileLang 提供三类显式原语：

• 内存放置：T.alloc_shared/T.alloc_fragment让你自己决定什么进共享内存、什么进寄存器

• 数据搬运：T.copy显式表达 global ↔ shared ↔ fragment 的流向

• 并行调度：T.Pipelined(num_stages=)直接控流水线级数，不像 Triton 那样黑盒

2. Tile Recommendation + Tile Inference 两段式自动补全

光有原语还不够。TileLang 引入统一的融合 tile 级数据流图（Fused Tile-level Dataflow Graph, FTG），配合两步：

1. Tile Recommendation：基于硬件感知给默认 tile 形状、向量化长度、shared memory 大小

2. Tile Inference：通过约束传播自动补全开发者没写的 layout 标注

开发者只需写关键决策，其余编译器推导。这就是为什么论文里多数融合注意力算子能压到70 行以内 Python，代码量较手写最高减 85.5%。

3. 底座仍是 TVM

TileLang 语法层是 Pythonic DSL，编译器基础设施架在 Apache TVM / TVM FFI 之上，所以能复用 TVM 成熟的 lowering 链路，再叠加自己的 tile 抽象。这也是它能在短时间内铺多后端的原因。

4. 官方基准（ICLR'26 论文 + 微软研究院文章双重背书）

⚠️ 注意口径：是"达到 FlashMLA 约 98% 性能"，不是 100% 持平；FlashMLA 本身是 CUTLASS 模板不是纯汇编，自媒体常说的"干翻手写 CUDA"是营销话术，论文措辞是"near hand-written CUDA"。

三、后端覆盖：哪些真测过，哪些是预览

官方 README 的 "Tested Devices" 明确列出的只有：

• NVIDIA：H100（Auto TMA/WGMMA）、A100、V100、RTX 4090、RTX 3090、RTX A6000

• AMD：MI250（Auto MatrixCore）、MI300X（Async Copy）

而以下后端是支持/预览/分支形态：

• Apple Metal：2025-10-07 加入

• 华为昇腾 AscendC / AscendNPU：2025-09-29 加入（含ascendc_pto与npuir两分支，预览态）

• WebGPU：2025-02-15 加入 codegen

• CuTeDSL：2025-12 加入，可编译到 NVIDIA CUTLASS CuTe DSL，覆盖 Blackwell SM100

• NVRTC 后端：显著缩短 cute 模板编译时间

• Z3 定理证明器：2025-12 集成进 TVM Arith Analyzer，做 SMT 符号化自动正确性验证

💡 公众号写作建议：可以写"覆盖 NVIDIA / AMD / Metal / 昇腾 / WebGPU / CuTeDSL 六条后端路径"，但要补一句"NVIDIA 和 AMD 经官方基准验证，其余为支持或预览态"——这样既准确又有信息量。

四、一个能体现"TileLang 思维"的极小例子

下面这段是 GEMM + ReLU 的骨架（依据 Atlas Cloud 教程简化），能看出它和 Triton 的本质差异：

python
import tilelang, tilelang.language as Timport torch@tilelang.jitdef matmul(M, N, K, block_M, block_N, block_K, dtype="float16"):    @T.prim_func    def kernel(A: T.Tensor((M, K), dtype),               B: T.Tensor((K, N), dtype),               C: T.Tensor((M, N), dtype)):        # 1. 显式声明共享内存（Triton 把这步藏起来了）        A_shared = T.alloc_shared((block_M, block_K), dtype)        B_shared = T.alloc_shared((block_K, block_N), dtype)        C_frag   = T.alloc_fragment((block_M, block_N), dtype)        # 2. 显式流水线        with T.Pipelined(K // block_K, num_stages=3):            T.copy(A, A_shared)            T.copy(B, B_shared)            T.gemm(A_shared, B_shared, C_frag)   # 调 Tensor Core        T.copy(C_frag, C)    return kernel

注意三个"显式"：显式 alloc_shared、显式 Pipelined、显式 copy。这正是它比 Triton 多出控制力、又比 CUDA 少写几百行模板的根源。

五、谁在用、意味着什么

已经被采用的信号：

• 2025-09 DeepSeek V3.2-Exp 技术报告明确写"使用 TileLang 做快速原型设计，建议社区使用 TileLang 版本做研究实验"，同日华为、寒武纪、海光宣布支持

• BitBLAS、AttentionEngine 已在生产中使用 TileLang

• ICLR'26 近两万投稿仅 1.18% 获 Oral，TileLang 是其中之一

对行业的真实含义（这部分比"干翻 CUDA"更值得写）：

• 对应用层 / 推理框架团队（vLLM、SGLang、TensorRT-LLM）：可以用 TileLang 快速验证新算子原型，再决定是否值得下沉到 CUDA 手写。门槛从"招百万年薪 Kernel 工程师"降到"懂 Python 的 ML 工程师一周上手"。

• 对算子库厂商：护城河被削弱，但生态位反而扩大——多后端适配成本骤降。

• 对国产硬件（昇腾、寒武纪、海光、MUSA）：原本每款芯片都要养一支 CUDA 移植团队，现在可以用同一份 TileLang 源码 + 各厂后端分支，这是 2025-09 多家国产厂商同日宣布支持的根本原因。

六、给公众号读者的 TL;DR

• TileLang 不是新语言、不是 CUDA 替代品，是架在 TVM 上的 Python 嵌入式 DSL，专写 GPU/加速器 kernel

• 归属是北大杨智组主导 + TileAI/微软研究院，不是 MIT CSAIL，也不是上海 AI Lab

• 80 行 Python 在 H100 上跑 MLA 解码，达到 FlashMLA约 98% 性能；较 Triton 平均 3.02×

• 后端六条路径：NVIDIA/AMD 精测，Metal/昇腾/WebGPU/CuTeDSL 支持或预览

• Z3 形式化验证 + CuTeDSL 桥接已于 2025-12 落地，工程化程度确实领先同类 DSL

• 自媒体爱用的"每秒 1 个 pip install""6.6k stars 碾压"之类数据未见权威源，写稿请谨慎

💡 真正的拐点不在"80 行代码"这个噱头，而在它把 GPU kernel 开发的颗粒度重新定义在 tile 这一层——既给了工程师拧旋钮的权利，又把拧错的成本交给编译器兜底。这才是它能同时被 DeepSeek 写进技术报告、被 ICLR 给 Oral、被国产硬件集体拥抱的原因。