论文精读：TileLang 的可组合平铺编程模型，凭什么在 NVIDIA/AMD 上追平甚至反超手写库？

论文精读TileLang 的可组合平铺编程模型凭什么在 NVIDIA/AMD 上追平甚至反超手写库【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang2025 年 4 月北京大学与微软研究院联合团队在 arXiv 发布《TileLang: A Composable Tiled Programming Model for AI Systems》arXiv:2504.17577。不到半年这篇论文的配套开源项目就完成了从论文代码到社区基础设施的跃迁DeepSeek 官方将其选为昇腾平台算子开发的基座组件海光 DCU 团队发布基于 TileLang 的国产卡编程语言适配国内厂商围绕 TVM/Triton/TileLang 展开的多场技术沙龙把跨硬件统一编译生态推上风口浪尖。一篇 AI 编译器论文能在如此短的时间内搅动芯片、模型、框架三个层面的生态说明它触碰到的不是某个细节技巧而是一个根本性的范式问题——高性能 AI Kernel 到底应该怎么被写出来这篇论文给出的答案是可组合平铺编程模型Composable Tiled Programming Model。本文将从论文核心主张出发结合仓库源码逐层拆解四个关键问题计算与调度为什么必须解耦线程绑定与内存布局如何在编译器里被推理出来软件流水线如何做到手写级调度以及那张最关键的 NVIDIA/AMD 性能成绩单是否经得起推敲。一、核心主张让 Tile 成为一等公民把调度空间从数据流中剥离论文的起点是一个朴素观察几乎所有现代 AI Kernel 都遵循同样的数据流模式——把数据分块Tile从 DRAM 搬进 SRAM在片上执行一串计算再把结果搬回 DRAM。模式清晰但写出高性能实现却极难因为性能由四个相互纠缠的调度维度决定线程绑定Thread Binding块/束/线程之间如何分配数据与计算直接决定负载均衡、同步开销与数据局部性内存布局Memory Layout数据在物理内存中的排布方式决定访存是否合并、是否触发 bank conflict指令张量化Intrinsic TensorizationTensor Core、Matrix Core、异步拷贝等专用指令的使用方式流水线Pipeline计算与数据搬运的重叠策略包括cp.async、TMA、warp specialization 等进阶手段。传统路径只有两个极端要么像手写 CUDA/CUTLASS 那样把所有调度细节显式交给程序员换来极限性能却背负巨大的开发成本要么像 Triton 那样隐藏线程、布局、地址空间用自动策略换取易用性却在低精度量化、自定义布局等场景上束手束脚。论文明确点名了 Triton 的软肋它虽然提供tl.dot和num_stages这类便利接口却不允许用户注册自定义 tile 算子、无法覆盖数据布局、也不能定义完全自定义的流水线。TileLang 的选择是中间路线用户只负责用可组合的 tile 算子T.gemm、T.copy、T.reduce、T.atomic描述数据流编译器接管其余调度空间。当默认优化不够时用户再通过T.Parallel、T.Pipelined、T.annotate_layout、T.use_swizzle等调度原语与注解在前端进行精准干预。数据流与调度解耦是全文设计的轴心。二、从 Python 到 CUDA数据流驱动调度由编译器补齐论文用一个约 20 行的 GEMM 展示了这套模型的全貌仓库的 Quick Start 与此几乎逐行对应见 README.md 与 tilelang/languagetilelang.jit def matmul_relu(A, B, block_M: int 128, block_N: int 128, block_K: int 32): M, N, K T.const(M, N, K) A: T.Tensor((M, K), T.float16) B: T.Tensor((K, N), T.float16) C T.empty((M, N), T.float16) 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), T.float16) B_shared T.alloc_shared((block_K, block_N), T.float16) C_local T.alloc_fragment((block_M, block_N), T.float32) T.clear(C_local) 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) for i, j in T.Parallel(block_M, block_N): C_local[i, j] T.max(C_local[i, j], 0) T.copy(C_local, C[by * block_M, bx * block_N]) return C这份代码里没有任何threadIdx、__syncthreads()、cp.async.wait之类的字样但性能关键的要素一个不少T.alloc_shared把 A、B 的 tile 显式放到片上共享内存T.alloc_fragment在块级声明一个寄存器累加器T.Pipelined(..., num_stages3)声明三段式软件流水线T.gemm把整块矩阵乘映射到目标后端。正如论文所述fragment 的分配表面上是整个块共享同一块寄存器但真正的按线程切分是由编译期的Layout Inference在生成T.Fragment布局对象时完成的。T.Kernel的语义是tile 程序的网格在 CUDA 上一个程序即一个 thread blockbx/by对应blockIdx.x/y在 CPU 上网格退化为外层循环tilelang/language/kernel.py。这种统一 block-and-thread 范式使得同一份前端代码可以横跨多后端——这正是论文第五节结论中延伸多后端支持的设计基础也是今天 TileLang 能同时覆盖 CUDA、ROCm、Metal、Ascend 乃至 LLVM CPU 后端的底层原因。三、计算与调度解耦的三个支点Layout、Fragment 与自动流水线解耦不是一句口号它体现在编译器的三个具体机制上。第一Layout 作为一等的可组合抽象。论文将 Layout 定义为从高维索引到线性地址的映射函数形式化为一个基于 IterVar 的代数表达式。这种抽象支持非双射变换比如 padding可以任意叠加组合出 swizzle 等策略用于消除共享内存 bank conflict。仓库中对应的是 src/layout 下的布局基础设施以及在 src/transform/layout_inference 中维护LayoutMap、按算子优先级自顶向下传播布局推断的完整实现。第二Fragment 把线程视角注入数据布局。Fragment Layout 的输出恒为两维——线程在寄存器堆中的位置 线程本地寄存器索引——从而精确刻画 block 级寄存器文件如何被切分到每个线程。论文以 MMA 的m16k16片元为基座通过repeat、repeat_on_thread、replicate等原语组合出 warp 级、block 级的完整布局。这套机制的价值在于GEMM 对布局与线程绑定有最严格的要求而逐元素算子则宽松得多Layout Inference 据此建立优先级分层先推断约束最强的算子再逐级放宽。仓库的 MLA 教程examples/deepseek_mla/README.md展示了一个非常直观的例子当T.gemm(..., policyT.GemmWarpPolicy.FullCol)声明按列切分 warpgroup 时编译器自动推断出每个 warpgroup 的acc_s_0形状应为[blockM, blockN/2]而后续acc_s V又需要完整的acc_s于是编译器继续向前推断让S_shared与acc_s保持[blockM, blockN]——程序员无需手写任何中间形状tilelang/tileop/base.py 中的GemmWarpPolicySquare/FullRow/FullCol就是这类细粒度控制的入口。第三软件流水线由编译器自动推导。论文详细描述了流水线推断机制编译器分析 Copy 与 GEMM 之间的依赖关系自动交错搬运与计算并在检测到异步机会时自动映射到硬件资源。在 Ampere 上自动插入cp.async/cp.async.wait在 Hopper 上自动做 warp specialization——用活跃变量分析划分 producerTMA 搬运与 consumer计算并自动插入mbarrier同步在 AMD CDNA 上则利用s_waitcnt/buffer_load_dword lds等 HIP 封装原语。仓库中 src/transform/inject_pipeline.cc、src/cuda/codegen/ptx.h内含cp.async系列 PTX 生成与 tilelang/language/copy_op.py 中的T.tma_copy/T.async_copy分阶段接口正是这一设计的工程落地。用户面对的就只有T.Pipelined一个num_stages旋钮复杂度被编译器吸收。四、细粒度控制不缺席线程绑定、swizzle 与指令注入强调编译器自动优化不等于剥夺专家的控制权。论文反复强调的transparent scheduling在源码里可以逐条找到对应物。线程级编程接口T.Kernel提供threads、cluster_dimsSM90 的线程块簇等启动注解T.get_thread_binding()暴露线程索引T.ClusterKernel支持 CUDA 簇语义。文档docs/programming_guides/language_basics.md明确说明这些注解由各后端的语言方言按需解释非 SIMT 目标会自动拒绝不适用特性。内存布局注入T.annotate_layout允许对任意 buffer 指定自定义布局官方还提供make_swizzled_layout便捷原语一键生成 swizzle 布局tilelang/language/annotations.py。T.use_swizzle(panel_size, order)则用一行代码完成 threadblock 级 L2 缓存友好的栅格重排rasterization2DRow/rasterization2DColumn等模式。指令级注入对低精度内核这类极端场景论文展示了 C 源码注入T.import_sourceT.call_extern封装dp4a与T.ptx直接发射内联 PTX如mma.m16n8k32两条通路仓库的 tilelang/language/builtin.py 与 src/cuda/op/builtin.cc 提供了wgmma、tcgen05等全套硬件指令内置函数。论文还客观记录了一个工程权衡默认走 CUTLASS/cute 这类 Tile Library 路线可以快速覆盖新硬件指令但模板膨胀占 CUDA 编译时间的约 90%因此 TileLang 也在逐步用自研指令实现替代这与仓库 README 中构建 self-hosting Tile Library的规划相互印证。五、性能验证NVIDIA/AMD 上的数字到底有多硬论文的实测覆盖 H100、A100、RTX 4090NVIDIA与 MI300XAMD基准包括 cuBLAS/rocBLAS、FlashAttention-3、Triton、Torch、FlashMLA、Marlin、BitsandBytes 等。以下是可以直接引用的事实数据GEMMRTX 4090 / A100 / H100 / MI300X相对厂商库分别取得 1.10×、0.97×、1.00×、1.04× 的加速比即追平甚至小幅反超 cuBLAS/rocBLAS相对 Triton 则为 1.08×、1.03×、1.13×、1.25×在 AMD 上优势更明显。配合 Layout Swizzling论文声明在所有测试设备上实现了无 bank conflict 执行。FlashAttentionH100相对 FlashAttention-3、Triton、PyTorch 分别取得 1.36×、1.41×、1.70× 加速比。值得强调的是TileLang 通过T.Pipelined加自动 warp specialization 就能表达与 FA-3 同等复杂度的流水线方案而 FA-3 是手写 TMA wgmma.mma_async的产物。Mamba-2 Linear AttentionH100chunk-scan 与 chunk-state 相对 Triton 平均加速 1.77× 与 2.10×。MLAH100 / MI300X这是论文最具说服力的案例。H100 上相对 Torch 提速 1075.9×最高达到手写 FlashMLA 的 98%仅需约 70 行 PythonMI300X 上相对 Torch 提速 129.2×达到手写库 AITER 的 95%。仓库的 examples/deepseek_mla/README.md 提供了完整的实现拆解与基准图examples/deepseek_mla/figures/bs64_float16.png、examples/deepseek_mla/figures/bs128_float16.png并展示了通过 Layout Inference 自动推导出的 QK/PV 执行计划examples/deepseek_mla/figures/qk_layout.jpg、examples/deepseek_mla/figures/pv_layout.jpg。要知道 MLA 的 head 维度高达 57651264acc_o极易触发寄存器溢出TileLang 用GemmWarpPolicy.FullCol把输出沿 dim 切给两个 warpgroup、再经共享内存交换部分acc_s的处理方式正是前端一行注解 编译器布局推断协同的巅峰体现。Dequantize GEMMA100相对 cuBLAS(W_FP16/A_FP16) 最高 7.65×W_INT2/A_INT8 配置相对 Marlin 平均 1.04×W_INT4/A_FP16相对 BitsandBytes 平均 1.62×W_NF4/A_FP16。论文将优势归因于线程级接口 PTX 快速数值转换指令 来自 Ladder 的硬件感知张量变换——这些恰是 Triton 难以表达的优化。仓库自身的基准也印证了论文数据benchmark 目录记录的 H800 上 8192×8192 FP16 GEMM 可达 758–766 TFLOPS见 benchmark/matmul/README.mdH100 上的算子性能汇总见 images/mha_performance_h100.png 与 images/op_benchmark_consistent_gemm_fp16.png。六、反超手写库的底气从何而来三个可复盘的工程判断数据背后是三个值得反复咀嚼的判断。其一显式数据流 隐藏调度是比全自动调度更诚实的抽象。TVM 时代调度脚本与计算分离但用户仍需手工编排 scheduleTriton 把调度藏进黑箱换来易用却牺牲表达力。TileLang 把调度空间拆成一组可组合、可注解、可推断的原语恰好落在编译器能自动补全大部分、专家能精准覆盖关键处的甜区。MLA 案例证明编译器推断出的布局甚至能比手写库的手工切分更优雅而用户侧只有 70 行代码。其二布局推断Layout Inference是真正的技术护城河。它把线程绑定、内存布局、向量化、swizzle 四件事统一在一个以优先级驱动的推断框架里让算子即接口Lower InferLayout成为可能。这也解释了为什么社区很快围绕它发展出布局可视化工具链examples/plot_layout与 docs/tools/layout_visualization.md 等配套文档——布局是这类系统性能的物理载体看得见摸得着。其三可组合性决定了生态半径。论文结语提到的三个方向——self-hosting Tile Library、tile 级通信原语与分布式、面向多后端的成本模型——如今都已在仓库中有了雏形或明确路线图CuTe DSL 后端与 tilelang/backend 的模块化后端抽象、Ascend 950 的原生代码生成、Metal 4 协作张量、LLVM CPU 后端……论文中的generalized design已经展开成了一张真实的跨硬件版图。回到标题的问题凭什么追平甚至反超手写库答案不在某一条指令或某一个 pass而在于它把手写库工程师的直觉内存层次、线程划分、流水线重叠系统地编码进了一个可组合、可推断、可注解的编程模型。当手写库为每个新算子重新经历数周的调优周期时TileLang 用户只需重写一份数据流、让编译器去推演调度——这正是 DeepSeek、海光、华为昇腾等生态方相继押注它的原因也是 AI 编译器从自动生成代码走向人机协同优化这条主线上最值得持续跟踪的样本。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考