用途
该技能面向 xLLM 仓库中的 TileLang Ascend 内核开发工作。当任务涉及以下内容时使用:
- xllm/compiler/tilelang/targets/ascend/kernels/*.py
- xllm/core/kernels/npu/tilelang/*_wrapper.cpp
- xllm/core/kernels/npu/tilelang/CMakeLists.txt
- 生成的 TileLang 产物,如 manifest.json、registry.inc 或特化 .cpp
运行环境与入口
- 构建和测试命令需在 NPU 容器内运行。
- TileLang 命令需从 xLLM 仓库根目录运行,而不是从已安装包的环境运行。
- 端到端 TileLang 编译流程优先使用:python xllm/compiler/tilelang_launcher.py ...
- 从 xLLM 仓库根目录设置:export TL_ROOT=$PWD/third_party/tilelang-ascend,并验证 test -f "$TL_ROOT/tilelang/init.py"。
- 任何原始脚本在 import tilelang 之前,需先执行 export TL_ROOT=$PWD/third_party/tilelang-ascend && source third_party/tilelang-ascend/set_env.sh,再执行脚本。
- 不要直接用 python rope.py 运行内核文件;这些文件依赖相对导入,应使用模块方式执行。
直接调试内核脚本时,以模块方式运行并传入所需 CLI 参数:
cd xllm
python -m compiler.tilelang.targets.ascend.kernels.rope
--output ../.tmp/rope.cpp
预期输出:[INFO] RoPE output matches torch reference
同样的模块式规则适用于 xllm/compiler/tilelang/targets/ascend/kernels/ 下的其他内核文件。
主要参考与模式偏好
主要参考:third_party/tilelang-ascend/docs/TileLang-Ascend Programming Guide.md
当任务依赖 T.tile.xxx 语义(如 compare、select、cast 或其他向量内建)时,使用:third_party/tilelang-ascend/.agents/skills/tilelang-custom-skill/tilelang-api-best-practices/references/api-tile-ops.md
xLLM Ascend 内核默认使用 Expert 模式:
- 优先使用 T.tile.xxx、显式 UB/shared 分配和显式 T.copy
- 优先使用显式 T.serial 控制行/块遍历
- 除非内核是明确的 tile-local 逐元素表达式,且改动不会降低对 UB 使用、临时缓冲区或精确运行时语义的控制,否则不要引入 Developer 模式的 T.Parallel
- 翻译 Triton 内核时,先保留 Triton 运行时语义,再选择与之匹配的最小 Expert 模式 lowering
Triton 到 TileLang-Ascend 的常见语义映射
翻译 Triton 内核时可参考以下快速语义映射:
- x + y、x - y、x * y、x / y → T.tile.add/sub/mul/div。在 Expert 风格向量代码中优先使用 tile 操作,而不是手写标量循环。
- tl.exp(x)、tl.log(x)、tl.abs(x) → T.tile.exp、T.tile.ln、T.tile.abs。TileLang 使用 ln,而不是 log。
- x.to(tl.float32) 或 tl.cast(...) → T.tile.cast(dst, src, "CAST_NONE", count)。仅在需要舍入语义时选择非默认 cast 模式。
- x <= y、x < y、x >= y、x == y → T.tile.compare(mask, x, y, "LE"/"LT"/"GE"/"EQ")。T.tile.compare 产生位掩码,而不是浮点张量。
- tl.where(cond, a, b) → T.tile.select(dst, selMask, a, b, selMode)。API 层面匹配。如果 cond 是比较表达式(如 x <= y),先用 T.tile.compare(...) 物化 selMask;对张量-张量或张量-标量选择使用匹配的 VSEL_* 模式。
- tl.full(shape, value, dtype) → 分配缓冲区 + T.tile.fill(dst, value)。将分配与初始化分离。
- tl.arange(0, N) → T.tile.createvecindex(dst, 0) 或显式循环索引。仅当内核确实需要向量索引张量时才优先使用 createvecindex。
保持语义的 lowering 规则:
- 保留 Triton 控制流、掩码和参数语义。除非在支持的输入域内运行时可见行为不变,否则不要替换为数值相近的公式。
- 保持内核 ABI 与 lowering 对齐。每个运行时参数必须要么参与 TileLang 实现,要么从接口中移除。
- 为分支、掩码和边界行为添加针对性测试。如果某些路径仅在特定值下命中,不要只依赖随机输入。
- 根据源操作数选择正确的 VSEL_* 模式。VSEL_CMPMASK_SPR 是 T.tile.compare 产生的掩码的自然匹配;VSEL_TENSOR_SCALAR_MODE 和 VSEL_TENSOR_TENSOR_MODE 用于显式张量/标量或张量/张量选择模式。
新增内核
按以下顺序进行:
- 实现 build_<kernel>_kernel(...)
- 实现 generate_source(...)
- 声明 DISPATCH_SCHEMA 和 SPECIALIZATIONS
- 运行一次 TileLang 编译并检查 registry.inc
- 添加或更新 <kernel>_wrapper.cpp
- 在 xllm/core/kernels/npu/tilelang/CMakeLists.txt 中注册内核:
tilelang_register_runtime_kernel(NAME <kernel> WRAPPER_SRCS <srcs...>)
wrapper 工作要点:
- 先在 Python 侧完成内核精度对齐(build_<kernel>_kernel(...) / generate_source(...)),而不是在 C++ wrapper 中。
- 手写张量检查、布局变换和 build_runtime_specialization(...)。
- 使用生成的 make_<kernel>specialization(...) 和 find<kernel>_kernel_entry(...)。
- 不要手写内核特定的特化结构体或内核函数 typedef。
- 绝不要在 wrapper 中使用 permute、contiguous、transpose、reshape(非平凡时)、clone 或任何触发设备内存拷贝的操作。wrapper 只能传递指向现有张量的指针。如果内核需要不同布局,在内核内部处理,或要求调用方提供正确布局。
新增特化
当内核逻辑和 wrapper ABI 保持不变时使用此路径:
- 更新现有内核的 SPECIALIZATIONS
- 确认每个运行时 dispatch 字段仍与 DISPATCH_SCHEMA 匹配
- 重新运行 TileLang 编译
- 检查 registry.inc 包含新条目
- 检查 wrapper 的 build_runtime_specialization(...) 仍构造匹配的值
调试生成的 Ascend-C
当任务是检查代码生成或比较两个内核实现时,使用隔离的输出根目录:
python xllm/compiler/tilelang_launcher.py compile-kernels
--target ascend
--device a3
--output-root .tmp/tilelang_debug
--kernels <kernel>
--force
然后检查:
- .tmp/tilelang_debug/targets/ascend/<kernel>/<variant_key>/<kernel>_<variant_key>_kernel.cpp
- .tmp/tilelang_debug/targets/ascend/<kernel>/registry.inc
- .tmp/tilelang_debug/targets/ascend/<kernel>/manifest.json
当需要理解生成的符号、字段顺序或 ABI 时,在修改 wrapper 代码之前使用此路径。
验证
优先使用最窄的命令:
python -m py_compile xllm/compiler/tilelang/targets/ascend/kernels/<kernel>.py
cd xllm && python -m compiler.tilelang.targets.ascend.kernels.<kernel> --output ../.tmp/<kernel>.cpp
python xllm/compiler/tilelang_launcher.py prepare-ascend
python setup.py test --test-name <wrapper_test_target> --device npu
参考
- 机制细节阅读 docs/en/dev_guide/tilelang_ascend_kernel_dev.md
- 以 rope 作为具体模板:
- xllm/compiler/tilelang/targets/ascend/kernels/rope.py
- xllm/core/kernels/npu/tilelang/rope_wrapper.cpp
- xllm/core/kernels/npu/tilelang/CMakeLists.txt
限制
资料未提供该技能的收费、价格、平台支持范围、性能效果或具体使用步骤之外的额外信息。