TileLang GPU 内核开发教程用几行 Python 写出对标 cuBLAS 的 GEMM【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelangTileLang 是一个构建在 TVM 编译器之上的 GPU 内核领域特定语言DSL。你用类 Python 语法描述 GEMM、FlashAttention 等算子的分块逻辑编译器自动完成内存布局、线程调度和流水线生成最终产出可直接运行的 CUDA 内核性能接近 cuBLAS 等库函数。从 Python 到 CUDATileLang 的编译链路长什么样TileLang 的核心抽象是瓦片tile一块有形状的数据单元由 warp 或线程块持有和搬运。你负责描述哪些瓦片在什么时候从哪级内存搬到哪级内存编译器负责把它翻译成硬件指令。整个语言按抽象程度分为三级你可以只停在高层也可以在同一个内核里混用不同层级初学者层硬件无感知不关心共享内存、线程细节直接描述数据搬运和计算逻辑适合先跑通算法正确性开发者层瓦片库显式使用T.alloc_shared、T.gemm等瓦片库原语控制数据在片上内存的驻留位置这是日常开发的主力层级专家层线程原语直接操作线程级原语、同步与数据布局用于压榨最后一点性能。以矩阵乘法为例分块策略贯穿三级内存全局矩阵先切成block_M × block_K的子块拷入共享内存乘积累加放在寄存器里的 fragment 中最后再把结果块写回全局内存。T.Pipelined让搬下一块和算这一块在时间上重叠这是掩盖访存延迟的关键。两步装好环境并跑出第一个 GEMM 内核安装只需要一条命令pip install tilelang python -c import tilelang; print(tilelang.__version__)官方预编译 wheel 覆盖 Linux x86-64/AArch64、Windows x86-64 和 macOS arm64。如果想从源码构建需要自定义 TVM 或特定后端开关再考虑克隆仓库git clone https://gitcode.com/GitHub_Trending/ti/tilelang cd tilelang pip install -e .编写第一个内核。examples/gemm/example_gemm.py里的 FP16 GEMM 完整实现不到 30 行核心逻辑如下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) A: T.Tensor((M, K), dtype) B: T.Tensor((K, N), dtype) C T.empty((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.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逐行解读这段代码tilelang.jit装饰器让函数在首次调用时按具体形状做特化编译T.Kernel声明启动网格每个线程块负责一个block_M × block_N的输出子块T.alloc_shared/T.alloc_fragment把瓦片分别钉在共享内存和寄存器两级T.Pipelined(num_stages3)自动把 K 维循环展开成三级流水T.gemm则映射到目标后端对应的 Tensor Core 指令NVIDIA 上会落到 WGMMA/MMA 路径。编译完成后kernel.get_kernel_source()可以打印生成的 CUDA 源码供审查kernel.get_profiler().do_bench(backendcupti)能直接测出内核延迟不用自己搭计时脚手架。GEMM 之外注意力、量化与多后端的真实覆盖TileLang 不是 GEMM 玩具算子覆盖面很宽。仓库examples/目录里的实现大多可以当作生产级参考FlashAttention 前向/反向、变长序列、GQA见examples/flash_attention/SM100Blackwell上另有专用版本DeepSeek MLA 解码examples/deepseek_mla/提供了 H100 和 AMD MI300X 两版实现是复杂算子用 TileLang 落地的典型样本量化与反量化 GEMMexamples/dequantize_gemm/覆盖 W4A8、FP4、MXFP4 等低比特路径examples/gemm_fp8/覆盖 FP8 矩阵乘2:4 结构化稀疏T.gemm_sp直接利用稀疏 Tensor Core自动调优from tilelang.autotuner import AutoTuner可以把块大小、流水级数、线程数交给搜索。example_gemm_autotune.py中还能通过MatmulTemplate的 roller 拿到针对当前芯片的 Top-K 推荐配置省掉手工枚举。跨平台方面同一份内核代码可通过 Target 对象切换后端cuda主后端SM70–SM120、hipAMD CDNA/RDNA、metalApple 芯片、ascend华为 Ascend 950CPU 走 LLVM 后端实验性。对大多数用户auto目标会自动探测当前设备无需显式指定。基准测试TileLang 与 cuBLAS 差多少结论先说标准 FP16 GEMM 场景下TileLang 与 cuBLAS 基本打平个别尺寸略优低比特量化场景优势更明显。以下为项目维护的基准结果上图横跨 RTX 4090、A100、H100、MI300X 四款 GPU以 cuBLAS/rocBLAS 为 1.0 基准TileLang 在多数矩阵尺寸上落在基准线附近或之上。混合精度推理是 TileLang 拉开差距的区间。在 A100 上的权重量化 GEMV 测试中BitBLAS-TileLang 的 WINT2AFP16 实现最高取得约 7 倍加速注意力场景的结论要客观H100 上 FlashAttention-3 仍是第一TileLang 处于紧随其后的领先梯队images/mha_performance_h100.png有完整曲线。换句话说TileLang 的价值在于一套代码同时覆盖 GEMM、Attention、量化推理和稀疏计算且每类都在各自的第一梯队而不是单项全部称王。写在最后TileLang 把 GPU 内核开发中最脏的活——内存布局、同步、流水——收进了编译器留给你的只是分块逻辑本身。上手成本是一次pip install加一个 GEMM 例子回报是能自己迭代 cuBLAS 之外的定制算子。建议从 examples/gemm/ 读到 examples/flash_attention/再按需查 docs/ 下的编程指南软件流水线、类型系统、自动调优都有专篇。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考