TileLang 入门指南十几行 Python 写出接近 cuBLAS 的 GPU 内核【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelangTileLang 是一个用 Python 风格语法编写 GPU 内核的领域特定语言DSL底层构建在 TVM 编译器之上。你只需描述数据如何分块、如何在共享内存与寄存器之间搬运线程划分、同步与指令级优化都由编译器完成。本文讲清它的工作方式、一个最小可运行示例以及与 cuBLAS 的实测差距帮你判断是否值得引入。定位给 GPU 内核开发一套 Python 语法TileLang 是面向 GPU/NPU/CPU 算子的 DSL。可以类比成只报菜名的点单方式你声明输入、输出和分块策略后厨流程不用自己写。严格地说你写的是一段受限 Python 子集经 TVM 编译为 CUDA、HIP 或 Metal 设备代码。它替你处理线程布局、访存向量化、bank conflict 规避等细节同时保留显式控制入口比如共享内存的分配方式和流水线级数这一点比 Triton 更接近硬件语义。后端覆盖NVIDIA CUDA主后端、AMD ROCm、Apple Metal、华为 Ascend 950CPU 与 WebGPU 为实验性。安装只需pip install tilelang。核心原理把大矩阵拆成三级搬运先看 TileLang 的分块计算模型矩阵乘法的瓶颈在数据搬运。TileLang 的做法是把 C 切成小块按三级缓存组织数据流全局显存像楼下的仓库大但慢共享内存像小组共用的一块小白板快但只有几十 KB寄存器像手边的草稿纸最快但放不下多少。落到代码上就是T.copy把 A、B 的 tile 从显存搬进共享内存T.gemm调用张量核心做乘累加结果存在 FP32 寄存器 fragment 里K 循环结束后一次性写回。T.Pipelined(num_stages3)开启软件流水线类似后厨边炒菜边备下一盘菜第 k 轮的计算与第 k1 轮的数据加载重叠用来隐藏显存延迟。最小可运行示例一个 FP16 GEMMtilelang.jit装饰后函数在首次调用时按形状特化并编译。核心结构如下tilelang.jit def matmul(A, B, block_M128, block_N128, block_K32): M, N, K T.const(M, N, K) 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) T.copy(C_local, C[by * block_M, bx * block_N]) return C传入 torch 的 CUDA 张量即可执行。kernel.get_profiler().do_bench()测延迟kernel.get_kernel_source()导出生成的 CUDA 源码供审查。epilogue 想融合 ReLU 之类逐元素操作加一个T.Parallel循环即可。性能实测与 cuBLAS、FlashAttention-3 对比H100 上的算子基准。纵轴为延迟 / TileLang 延迟1 表示与 TileLang 持平FP16 GEMM与 cuBLAS、PyTorch 基本持平多数形状落在 0.9–1.15 区间。Conv2D比 PyTorch 快约 1.5 倍部分配置比 Triton 快近 2 倍。FP4 GEMM比 PyTorchbitsandbytes快约 1.5 倍。FlashAttentionFA-3 在部分形状仍领先约 1 倍但 TileLang 稳定快于 Triton 与 PyTorch。MHA 与 Mamba-2 的延迟对比基线同为 TileLangTriton、PyTorch 实现在部分配置下延迟达到 TileLang 的 2 倍左右FlashAttention-3 仍是注意力算子的标杆。跨平台 FP16 GEMM纵轴为相对 cuBLAS/rocBLAS 的加速比NVIDIA 各平台基本追平 cuBLASAMD MI300X 上最高约 1.7 倍于 rocBLAS。量化推理混合精度 GEMV 最高约 8 倍LLM 推理常见的权重量化 GEMVA100 实测纵轴为相对 cuBLAS-FP16 的加速比BitBLAS-TileLang 的 INT2 量化NT8 排布最高约 8.4 倍INT4、NF4 等格式约 3–5 倍。这类收益来自量化格式与硬件搬运方式的精细组合手写 CUDA 同样能做到但要为每种格式单独维护内核。适合谁不适合谁适合写 cuBLAS 覆盖不到的融合算子量化 GEMM、注意力变体、Mamba 类状态空间模型同一份代码跨 CUDA / ROCm / Metal / Ascend 部署想要接近手写 CUDA 的性能同时控制维护成本不适合只需要标准 GEMM——直接调 cuBLAS 即可期望零 GPU 知识入门——分块、共享内存、流水线这些概念仍要自己建立追求单点极致——FlashAttention-3 这类深度手调的库在特定形状上仍更强版本提示当前仍是 0.1.x历次发布有不兼容的 API 变更CPU 与 WebGPU 后端为实验性生产环境建议只用 CUDA 路径。继续深入文档docs/算子示例GEMM、FlashAttention、MLA、量化examples/可当参考用例的测试套件testing/自动搜索块大小与流水线级数看 examples/gemm 的 autotune 示例想逐 pass 审查编译器改动docs/tools/ 里有 Pass Diff 与 IR Lower Trace 的说明。TileLang 没有取代手写 CUDA它把分块 流水这套被验证过的性能模型做成了 Python 里的一等公民。如果你要的是接近库性能、又能长期维护的内核建议先从 examples/gemm 跑起。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考