TileLang:用 Python 风格语法编写高性能计算内核
一分钟看懂
TileLang 是一种面向 GPU、CPU 和 NPU 计算内核的领域专用语言。它以接近 Python 的语法描述矩阵乘法、注意力、量化与逐元素计算,再由基于 TVM 的编译基础设施映射到底层硬件,目标是在保留底层优化能力的同时减少手写内核的复杂度。
你可以用它做什么
- 为深度学习模型编写矩阵乘法、FlashAttention 和线性注意力内核。
- 实现量化、反量化和融合算子,并针对不同输入形状调优。
- 在 NVIDIA CUDA、AMD ROCm、Apple Metal 和华为 Ascend 等后端运行内核。
- 使用示例、分析器和布局可视化工具学习高性能内核开发。
主要能力
TileLang 提供 JIT 编译、分块矩阵计算、共享内存与寄存器片段、流水线数据搬运、自动调优和多后端目标选择。主仓库覆盖 CUDA、ROCm、Metal、Ascend 950,并提供实验性的 LLVM CPU、CuTe DSL 与 WebGPU 后端;其他硬件适配器可能在独立仓库维护。
适合哪些人
它适合开发机器学习框架、推理服务和高性能算子的工程师,也适合研究 GPU 编程与编译器优化的学习者。只调用现成模型或常规张量算子的应用开发者,通常不需要直接编写这类内核。
怎么开始使用
稳定版本可通过 pip install tilelang 安装,然后运行官方快速入门示例定义、编译和校验矩阵乘法内核。不同硬件需要相应的驱动和运行时,例如 ROCm 环境需先安装兼容的 PyTorch,Ascend 950 源码构建需要 CANN 与 torch_npu。完整安装和目标配置应以官方文档为准。
使用前需要知道
这是持续演进的编译与运行时项目,各后端支持等级并不一致:CUDA 是主要后端,部分 CPU、WebGPU 和其他硬件路径仍属实验性或生态适配。高性能结果依赖具体设备、算子形状、数据类型和调优参数,不能把官方示例中的性能直接外推到其他环境。夜间构建包含新功能,但稳定性可能低于正式版本;升级时也应检查兼容性说明。项目许可证以仓库 LICENSE 为准。