TileLang 多后端实战Target 探测、ROCm、Docker 一网打尽【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang当一份深度学习编译器仓库里同时躺着cuda、rocm、metal、cpu、ascend五套后端目录并且每个后端都配齐了语言方言、Pass 流水线、代码生成与原生算子实现时它就不再是一个CUDA 专用工具而是一个真正意义上的多后端编译器。TileLang 恰好处在这个位置上从 2025 年初开源时的单后端形态到如今被 DeepSeek、海光、昇腾等国产芯片团队相继接入围绕同一份 Kernel 源码在不同硬件上编译出接近手工优化性能这件事它给出了一套完整的工程答案。本文不空谈架构而是把多后端落到三个可以直接上手的实操面Target 探测与显式指定、ROCm 环境搭建、以及用 Docker 快速复现 CUDA/ROCm 两套开发环境。所有代码与路径均来自仓库真实源码读者可以在本地直接对照验证。多后端架构从 BackendContext 到注册清单TileLang 的多后端能力建立在 TVM 的目标抽象之上Target决定选用哪套代码生成器CUDA、HIP、Metal、LLVM……再配以设备相关选项如 GPU 架构旗标。在 tilelang/backend/README.md 中项目给出了它的后端设计模型一个目标后端target backend拥有四块实现领地——语言方言language dialect、上下文贡献target 检测/归一化与兼容执行后端声明、Pass 流水线、主机/设备代码生成生成的代码再交给独立的执行后端execution backend完成 Build、Load、Launch。这套设计的关键在于后端注册清单。每个后端的backend.py会导出一个BackendModule清单声明自己拥有哪些 TVM target kind、用哪条 PassPipeline 和哪套设备代码生成器、以及兼容哪些共享执行后端。例如 ROCm 后端在 tilelang/rocm/backend.py 中的声明BACKEND register_backend( BackendModule( namerocm, target_kinds(hip,), pipelines{hip: PassPipeline(hip, pipeline.ROCMPassPipelineBody)}, device_codegens{hip: DeviceCodegen(hip, buildcodegen.build_hip, ...)}, execution_backendsexecution_backend.EXECUTION_BACKENDS, ... ) )值得注意的一个细节tilelang/rocm目录名虽是 rocm但它声明的 target kind 是hip参见 tilelang/backend/README.md 中的所有权规则——不要从目录名推断 target 归属。CUDA 与 CuTeDSL 两个清单共享cudatarget kind靠supports_target谓词区分变体ROCm 的hip、CPU 的c/llvm、Metal 的metal、WebGPU 的webgpu、Ascend 的ascend则各自独占。每次编译时编译器入口会创建一次不可变的BackendContext其职责写死在 tilelang/backend/module.py 的create_backend_context中归一化设备与主机 Target → 按target.kind.name精确选出唯一一个BackendModule→ 解析一个可用的执行后端。同一份BackendContext贯穿缓存、lowering、代码生成与 JIT后续任何阶段都不允许再次推断后端。仓库自带的回归测试 testing/python/backend/test_tilelang_backend_module.py 直接把这份注册清单固化成了断言例如rocm后端拥有(hip,)且兼容[tvm_ffi, cython]可见后端清单即契约是被测试强制约束的。Target 自动探测一次pip install处处自动识别对绝大多数使用者来说最常打交道的是target参数。TileLang 的auto目标会依次探测 CUDA → HIP → Metal这在同一脚本跑在多台机器的场景下极为实用见 docs/get_started/targets.md。探测机制本身是可插拔的tilelang/backend/target.py维护着全局的目标探测器TargetDetector与目标归一化器TargetNormalizer注册表determine_target(auto)会遍历所有已注册探测器返回第一个命中者。具体实现分布在各个后端目录中CUDA 探测器在 tilelang/cuda/target.py先通过nvcc.find_cuda_path()确认 CUDA 工具链存在再用torch.cuda.is_available()确认真的有一块可用的 NVIDIA 设备最后用torch.cuda.get_device_capability()推算出sm_80、sm_90这样的架构令牌。注释中特意说明了一个反例仅装了 CUDA toolkit 的交叉编译宿主机或 NPU 机器会被排除避免误探测遮蔽 Ascend。HIP 探测器在 tilelang/rocm/target.py通过find_rocm_path()定位 ROCm 安装再用torch.cuda.get_device_properties(0).gcnArchName读出gfx942、gfx950这类架构名取不到架构时回退为裸hip。Metal 探测器在 tilelang/metal/target.py限定 macOS arm64还能额外探测 Metal 4 支持macOS 26 Apple M5 及以上并往 target 的keys里追加metal4。Ascend 探测器在 tilelang/ascend/target.py通过torch.npu.is_available()判断昇腾 NPU 可用性。把探测器接到注册表的方式高度统一例如 tilelang/cuda/target.py 末尾的register_target_detector(cuda, _detect_cuda_target, overrideTrue) register_target_normalizer(cuda, normalize_cuda_target, overrideTrue)归一化器normalizer则负责把用户输入的 target 补全成规范形态。一个典型例子ROCm 的normalize_rocm_target会根据mcpu架构自动补上mtripleamdgcn-amd-amdhsa-hcc与thread_warp_sizeCDNA 架构 warp 为 64RDNA 为 32这些细节在 src/rocm/target_utils.cc 的原生侧也有对应的TargetIsCDNA、TargetIsRDNA、TargetIsGfx950等判定函数。显式指定 target 的三种常见写法同样在 docs/get_started/targets.md 中有完整对照target auto # 自动探测 CUDA / HIP / Metal target cuda # 裸字符串 target {kind: cuda, arch: sm_90} # 配置字典可携带架构选项若不想在代码里传参可以通过环境变量TILELANG_DEFAULT_TARGET设定默认 target。它在 tilelang/env.py 中被读取支持裸字符串或 JSON 配置串export TILELANG_DEFAULT_TARGETcuda export TILELANG_DEFAULT_TARGET{kind: cuda, arch: sm_100f, code: [sm_100a, sm_103a]}对于arch/code的语义文档也解释得很清楚arch传给 NVCC 作为单一-archsm_90code只有在你需要显式-gencode多代码实例生成 fatbin时才使用且必须是精确的 SM 令牌列表。遇到no kernel image is available这类运行时错误时首选排查方向就是arch与实际 GPU 计算能力不匹配。ROCm 实战同一个 wheel两条差异化路径AMD 用户的体验与 CUDA 用户差异不大这在 docs/get_started/Installation.md 的 Installing on AMD GPUs (ROCm) 一节写得很明确PyPI 上的 Linux x86_64 wheel 是 CUDAROCm 双后端 fat build同一个pip install tilelang即可无需源码编译。真正的差异只有两点运行时必须有宿主机 ROCm 安装因为 Kernel 使用宿主机的hipccJIT 编译——与 CUDA 有 pip 提供的nvccextra 不同ROCm 没有 pip 工具链。必须先装 ROCm 版 PyTorch否则默认的 CUDA 版 torch 无法识别 AMD 设备pip install torch --index-url https://download.pytorch.org/whl/rocm7.0 pip install tilelang装好后用一行命令验证 HIP target 是否被正确探测python -c import tilelang; from tilelang.backend.target import determine_target; print(tilelang.__version__, determine_target(return_objectTrue))正常输出是一个带mcpugfx942或你的实际架构的hipTarget。如果 ROCm 装在非标准位置或 PATH 上有多个 HIP 工具链用ROCM_PATH锁定export ROCM_PATH/opt/rocmfind_hipcc()的解析顺序在 tilelang/contrib/rocm.py 中写得很细先看$ROCM_PATH/bin/hipcc再按 PATH 顺序找hipcc最后兜底/opt/rocm/bin/hipccPATH 与兜底候选还必须携带公共 HIP 头文件include/hip/hip_runtime.h才会被接受否则跳过——这是为了防止 ROCm 7 某些安装会往 PATH 里塞不完整的预览编译器。ROCM 后端的语言方言也值得单独提一笔。tilelang.rocm.language是公共语言面 ROCm 扩展的组合见 tilelang/rocm/language/init.py其中最典型的扩展是T.gemm的k_pack旋钮——在 CDNA/RDNA 的 MFMA/WMMA 降低中沿 K 打包多个矩阵核运算这是其他后端没有的性能开关见 tilelang/rocm/language/gemm_op.py。仓库中的 AMD FlashAttention 示例 examples/amd/example_amd_flash_attn_fwd.py 就同时用到了这套方言通过IsRDNA()分支区分 gfx11/gfx12 与 CDNA 的调优空间并在 RDNA WMMA 路径下刻意用 shared memory 中转 softmax 分数以规避 D/A 寄存器布局冲突。面向 gfx950 的算子则有更严格的目标绑定例如 examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_cdna4.py 通过tilelang.jit(out_idx[-1], targethip -mcpugfx950)显式锁定 CDNA4 架构因为底层的hip_fp4.h只在__gfx950__下编译。Docker 复现CUDA 与 ROCm 两条镜像链路多后端开发最烦人的是环境漂移。仓库在 docker/README.md 和 docs/get_started/Installation.md 的 Install Using Docker 一节提供了现成方案docker/目录下并列躺着 8 个 CUDA 版本 Dockerfilecu118 至 cu128和 1 个 ROCm Dockerfile。CUDA 链路非常直接git clone --recursive https://github.com/tile-ai/tilelang TileLang cd TileLang/docker docker build -t tilelang_workspace -f Dockerfile.cu124 . # NVIDIA 设备 docker run -it --gpus all --shm-size4G --cap-addSYS_PTRACE \ --security-opt seccompunconfined --security-opt apparmorunconfined \ --name tilelang_test tilelang_workspace bashAMD 链路的启动命令换成设备节点映射把/dev/kfd与/dev/dri透传给容器即可--device/dev/kfd --device/dev/dri无需--gpus。ROCm 镜像本身值得细看。docker/Dockerfile.rocm 基于官方的rocm/pytorch:rocm7.2_ubuntu22.04_py3.10_pytorch_release_2.10.0关键配置集中在几行环境变量上ENV USE_ROCM1 ENV USE_CUDA0 ENV ROCM_HOME/opt/rocm ENV HIP_PLATFORMamd ENV PYTORCH_ROCM_ARCHgfx90a;gfx942;gfx950;gfx1201;gfx1100PYTORCH_ROCM_ARCH一口气覆盖了 CDNA3gfx90a、CDNA4gfx942、CDNA4 新核gfx950、RDNA4gfx1201与 RDNA3gfx1100这份架构枚举本身就是 TileLang 对 AMD 硬件覆盖面的直观体现。镜像构建时以USE_ROCM1 USE_CUDA0 pip install -e .从源码可编辑安装并在构建前用 sed 从pyproject.toml里临时剔除torch依赖避免 pip 把 ROCm torch 解析回 CUDA 版。如果用户已经在别的 ROCm 容器比如 sglang 镜像里只想原地重建 TileLangdocs/get_started/Installation.md 也给出了一套手工流程补系统库、装setuptools80与 Cython、用循环脚本定位/opt/rocm/llvm/bin/llvm-config找不到则装 LLVM 18、最后pip install -e . -v --no-build-isolation --no-deps并手动补齐apache-tvm-ffi与z3-solver。其中特别提醒ROCm 容器里不要装torch-c-dlpack-ext它的 wheel 依赖 CUDA 库若装过并报libtorch_cuda.so错误需要卸载。后端兼容性一览与示例算子库把前面几节的机制串起来看仓库 README 的 Platform and Backend Support 表格就是一张完整的能力地图CUDA 是 Primary 后端SM70–SM120 全覆盖TMA/WGMMA/TMEM 按架构启用ROCm/HIP 是 SupportedLinux CDNA/RDNA含 gfx942/gfx950 路径CI 跑在自建的 MI300X runner 上Ascend 950 是 Supported需USE_ASCENDON源码构建 CANN torch_npuMetal 是 SupportedApple SiliconM5 可用 Metal 4 cooperative tensorsLLVM CPU、CuTeDSL、WebGPU 处于 Experimental。此外还有一批 Ecosystem 适配器存在于独立仓库昇腾 A2/A3 的ascendc/pto/npuir、MetaX MACA、摩尔线程 MUSA、海光 HCU、矽琻 TANG——海光团队近期把 v0.1.6.post2 版本推向国产 DCU 生态的社区动态正是这张表格在仓库之外的延伸。同一个 Kernel 源码跑多个后端的最佳验证样本是 GEMM。快速入门 examples/quickstart.py 顶部注释点明了规则target可为cuda或hip或cpu不指定则在编译期从输入张量推断。核心代码很短tilelang.jit def matmul(A, B, block_M: int, block_N: int, block_K: int): M, N, K T.const(M, N, K) dtype T.float16 accum_dtype T.float32 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 ko in T.Pipelined(T.ceildiv(K, block_K), num_stages3): T.copy(A[by * block_M, ko * block_K], A_shared) T.copy(B[ko * 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同样的T.gemm、T.copy、T.Pipelined高层原语在 CUDA 后端会落到 cute/MMA 路径在 HIP 后端落到 MFMA/WMMA 路径在 CPU 后端则走标量 tile-op 实现——用户无需改 Kernel 主体。这正是文档中描述的分层编程接口的Developer Level借助 Tile Library 预定义模式无需触碰线程级细节。性能方面仓库的 benchmark 图片给出了多后端一致性的证据。images/op_benchmark_consistent_gemm_fp16.png 展示的是 RTX 4090、A100、H100、MI300X 四类 GPU 上的 FP16 GEMM 吞吐对比——NVIDIA 与 AMD 硬件都被覆盖在同一张图里images/op_benchmark_mi300_fp16_gemm_normalized_latency.png 则专门给出了 MI300X 上归一化延迟的细分。小结回过头看TileLang 多后端能力的三根支柱都清晰可见BackendModule注册清单把一个后端拥有什么做成显式契约并被测试锁定register_target_detector/register_target_normalizer让auto探测成为可插拔的机制CUDA/HIP/Metal/Ascend 各自贡献探测器而同一套 Linux wheel 内置 CUDAROCm 双后端配合 Docker 里从 cu118 到 cu128 再到 rocm7.2 的镜像矩阵让环境复现从手工折腾变成一条命令。对想快速上手的人来说从 examples/quickstart.py 的 GEMM 开始把target在auto、cuda、hip之间切换各跑一遍就是理解整套多后端机制成本最低的路径。【免费下载链接】tilelangDomain-specific language designed to streamline the development of high-performance GPU/CPU/Accelerators kernels项目地址: https://gitcode.com/GitHub_Trending/ti/tilelang创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考