TileLang (tile-lang) 

在 AI 大模型时代，高性能算子（如 GEMM、FlashAttention）是决定推理速度的关键。然而，传统的手写 CUDA/Triton 内核不仅开发门槛高，且难以在不同硬件间移植。

TileLang (tile-lang) 是一种简洁的领域特定语言（DSL），通过 Pythonic 语法 和基于 TVM 的强大编译器基础设施，简化高性能 GPU/CPU 内核的开发。它让开发者能够专注于算法逻辑，同时自动完成底层优化，实现媲美手工汇编的性能。

为什么选择 TileLang？

1. 极简语法，极致性能

TileLang 允许你用类似 Python 的代码描述复杂的并行计算、内存布局和流水线策略。

• 示例：仅用 80 行 Python 代码，即可实现支持 MLA Decoding 的高性能内核，在 NVIDIA H100 上达到与 FlashMLA 相当的性能。
• 自动化优化：内置布局注释、并行复制、Swizzling（改善 L2 缓存局部性）等高级特性，无需手动调整汇编指令。

2. 广泛的硬件支持

TileLang 致力于跨平台移植，目前已验证支持：

• NVIDIA GPU：H100 (Auto TMA/WGMMA), A100, V100, RTX 4090/3090 等。
• AMD GPU：MI250 (Auto MatrixCore), MI300X (Async Copy)。
• 华为昇腾：支持 AscendC 和 AscendNPU IR 后端。
• 其他：Apple Metal, WebGPU, CPU。

3. 丰富的算子库

提供构建各种复杂算子的基石，包括：

• 矩阵乘法：标准 GEMM、反量化 GEMM、稀疏 GEMM (2:4 Sparsity)。
• 注意力机制：FlashAttention, Flash Linear Attention, Native Sparse Attention, Flash MLA Decoding。
• 其他：卷积 (IM2Col)、RetNet/Mamba 实现等。

最新动态：生态持续进化

TileLang 社区活跃，近期更新频繁：

• 2026/02/02：推出 TileLang Puzzles，通过 10 个互动谜题趣味学习编程。
• 2025/12/18：新增 CuTeDSL 后端，支持编译到 NVIDIA CUTLASS CuTe DSL。
• 2025/12/17：集成 Z3 定理证明器，增强符号推理和自动正确性验证。
• 2025/10/31：迁移至 apache-tvm-ffi，显著降低 CPU 开销。
• 2025/09/29：支持 华为昇腾芯片 (AscendC/NPUIR)。
• 2025/04/14：为 AMD MI300X 添加高性能 FlashMLA 实现，媲美手工优化汇编。

快速开始：一个带注释的 GEMM 示例

以下代码展示了如何使用 TileLang 编写一个带有 ReLU激活、共享内存优化和流水线并行的矩阵乘法内核。

import tilelang as T
import torch

# @tilelang.jit 装饰器自动处理编译和缓存
# target 可指定为 "cuda", "hip", "cpu" 等，默认根据输入推断
@tilelang.jit
def matmul_relu(
    A, B,
    block_M: int = 64,
    block_N: int = 64,
    block_K: int = 64,
    dtype: T.dtype = T.float16,
    accum_dtype: T.dtype = T.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)

    # 定义 Kernel 网格和线程块
    with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) 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)

        # 【可选】启用 Swizzle 以优化 L2 缓存局部性
        # T.use_swizzle(panel_size=10, enable=True)

        # 清空局部累加器
        T.clear(C_local)

        # 流水线循环，num_stages=3 实现重叠计算与数据传输
        for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
            # 并行拷贝 A 的分块到共享内存
            T.copy(A[by * block_M, ko * block_K], A_shared)
            # 并行拷贝 B 的分块到共享内存
            T.copy(B[ko * block_K, bx * block_N], B_shared)
            
            # 执行_tile_级别的 GEMM (底层调度至 CuTe/HIP)
            T.gemm(A_shared, B_shared, C_local)

        # Element-wise ReLU 激活
        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

# --- 测试与验证 ---
M, N, K = 1024, 1024, 1024
a = torch.randn(M, K, device="cuda", dtype=torch.float16)
b = torch.randn(K, N, device="cuda", dtype=torch.float16)
c_ref = torch.relu(a @ b)

# 直接调用 JIT 编译后的函数
c = matmul_relu(a, b)
torch.testing.assert_close(c, c_ref, rtol=1e-2, atol=1e-2)

# 性能剖析
kernel = matmul_relu.compile(a, b)
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Normal)
latency = profiler.do_bench()
print(f"Latency: {latency} ms")

安装指南

方法一：Pip 安装（推荐）

pip install tilelang

方法二：从源码构建

适合需要最新特性或贡献代码的用户。

git clone https://github.com/tile-ai/tilelang.git
cd tilelang
pip install -e . -v

注意：需预先安装 CMake、LLVM 等依赖。

方法三：Nightly 版本

体验最新功能（可能不稳定）：

pip install tilelang -f https://tile-ai.github.io/whl/nightly

适用场景

• AI 框架开发者：快速原型化新型算子（如 Mamba, RetNet）。
• 高性能计算工程师：在不牺牲性能的前提下，大幅缩短 CUDA/HIP 内核开发周期。
• 研究人员：探索新的矩阵乘法算法、注意力机制优化策略。
• 跨平台部署团队：一套代码，多处运行（NVIDIA/AMD/Ascend）。