tilelang

环境要求

• Python 环境与可用的 pip
• 支持 CUDA 的 GPU 环境（快速开始示例使用 PyTorch CUDA 张量）
• AMD GPU 需主机安装 ROCm，并先安装 ROCm 版 PyTorch
• 已安装 PyTorch（运行快速开始示例所需）

安装与启动步骤

1. 1安装稳定版

   从 PyPI 安装最新的 tilelang 稳定版，这是 README 推荐的首选安装方式。

   pip install tilelang

   复制

2. 2验证安装

   导入 tilelang 并打印版本号，能输出版本即安装成功。

   python -c "import tilelang; print(tilelang.__version__)"

   复制

3. 3可选装 nightly

   nightly wheel 包含下一个稳定版之前的新特性与修复，但稳定性可能不如正式版。

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

   复制

4. 4AMD 环境准备

   AMD GPU 上 Linux wheel 开箱可用，但需先装 ROCm 版 PyTorch，且运行时要求主机已安装 ROCm。

   pip install torch --index-url https://download.pytorch.org/whl/rocm7.0
   pip install tilelang

   复制

5. 5运行快速开始示例

   把示例代码保存后运行，它会定义、编译、执行并校验一个 FP16 GEMM 加 ReLU 的 kernel，成功时打印 GEMM + ReLU passed.

   import torch
   import tilelang
   import tilelang.language as T
   
   
   @tilelang.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), threads=128) 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_stages=3):
               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
   
   
   M = N = K = 1024
   a = torch.randn((M, K), device="cuda", dtype=torch.float16)
   b = torch.randn((K, N), device="cuda", dtype=torch.float16)
   c = matmul_relu(a, b)
   torch.testing.assert_clos

   复制

如何确认成功

执行 python -c "import tilelang; print(tilelang.__version__)" 能打印版本号；运行快速开始示例打印出 GEMM + ReLU passed.

常见问题

注意事项

• nightly 构建的稳定性可能低于官方正式发布版本。
• 源码构建、可编辑安装、Docker、ROCm 细节、pip 提供的 CUDA 工具链及自定义 TVM 均需参考官方安装指南。
• 运行快速开始示例需要可用的 CUDA/ROCm GPU 环境与 PyTorch。