[Bug] PTO 后端在 ascend_mhc 测试集上 CodeGenPTO 不支持 tir.shift_right / 拒绝 T.Kernel(threads=)

[HecreReed]

背景

在 Ascend A5（Ascend950PR_9589）板卡上，对 examples/ascend_mhc 下的 MHC（Mixture-of-Hyper-Channel）算子测试集做 ascend 原生后端 vs PTO 后端的对比验证时，发现 PTO 后端 (TILELANG_TARGET=pto) 在多个测试文件上无法完成编译/运行。

本 issue 按会议要求重新组织：不以 test 文件为最小归类单位（不同 test 可能调用同类 Kernel，存在同类问题），改为以 Kernel 为粒度归类。包含两张表：

• 表1：每个 test 文件 → 调用的 Kernel 数量 + ascend 状态 + PTO 状态
• 表2：PTO 问题分类（按 Kernel 聚类）

本 issue 不预设责任归属，只汇总 ascend pass/fail + PTO 问题归类 + 复现步骤，供判断属于 PTOAS / tilelang-dsl / tile-kernels 哪一层。

环境

• 板卡：Ascend950PR_9589（A5），device 0（126112MB，vector_core_num=64，cube_core_num=32）；device 1 在连续负载下易 segfault，device 0 更稳定
• OS：aarch64 Linux
• conda env：rp21-PT2.7（python 3.10，torch 2.7.1+cpu + torch_npu）
• CANN：/home/wenquan/cann29/cann-9.0.0
• tilelang：/home/wangmiao/tilelang-simtvf-pr9-buildable（dev build）
• ptoas：/home/wangmiao/PTOAS-main-1df3fba/install/bin/ptoas
• LLVM：/home/ptoas-board-monitor-a5/llvm-project-vpto-feature/build-shared
• tile-kernels：/home/wangmiao/tile-kernels
• 测试目录：/home/wangmiao/tilelang-simtvf-pr9-buildable/examples/ascend_mhc

复现步骤

source /home/anaconda3/etc/profile.d/conda.sh
conda activate rp21-PT2.7
source /home/wenquan/cann29/cann-9.0.0/aarch64-linux/bin/setenv.bash

export PYTHONPATH=/home/wangmiao/tilelang-simtvf-pr9-buildable/3rdparty/tvm/3rdparty/tvm-ffi/python:/home/wangmiao/tilelang-simtvf-pr9-buildable/3rdparty/tvm/python:/home/wangmiao:/home/wangmiao/tile-kernels:/home/wangmiao/tilelang-simtvf-pr9-buildable:${PYTHONPATH:-}
export LD_LIBRARY_PATH=/home/wangmiao/tilelang-simtvf-pr9-buildable/build/lib:/home/ptoas-board-monitor-a5/llvm-project-vpto-feature/build-shared/lib:/home/wangmiao/PTOAS-main-1df3fba/install/lib:${LD_LIBRARY_PATH:-}
export PATH=/home/wangmiao/PTOAS-main-1df3fba/build/tools/ptoas:/home/wangmiao/PTOAS-main-1df3fba/install/bin:${PATH}
export PTOAS_BIN=/home/wangmiao/PTOAS-main-1df3fba/install/bin/ptoas
export LLVM_HOME=/home/ptoas-board-monitor-a5/llvm-project-vpto-feature/build-shared
export ASCEND_RT_VISIBLE_DEVICES=0
export TORCH_DEVICE_BACKEND_AUTOLOAD=0
export PYTHONFAULTHANDLER=1
cd /home/wangmiao/tilelang-simtvf-pr9-buildable/examples/ascend_mhc

# ascend 后端（逐文件，带 transfer_to_npu）
export TILELANG_TARGET=ascend
export PYTEST_ADDOPTS="-p torch_npu.contrib.transfer_to_npu"
for f in test_*.py; do
  python -c "import torch,torch_npu;torch.npu.set_device(0);torch.zeros(1,device=0);torch.npu.synchronize();from tile_kernels.config import get_device_num_sms,get_num_sms;get_device_num_sms.cache_clear();import pytest,sys;sys.exit(pytest.main(['-q','-p','no:cacheprovider']+sys.argv[1:]))" "$f"
done

# pto 后端（逐文件，同样带 transfer_to_npu）
export TILELANG_TARGET=pto
for f in test_*.py; do
  python -c "import torch,torch_npu;torch.npu.set_device(0);torch.zeros(1,device=0);torch.npu.synchronize();from tile_kernels.config import get_device_num_sms,get_num_sms;get_device_num_sms.cache_clear();import pytest,sys;sys.exit(pytest.main(['-q','-p','no:cacheprovider']+sys.argv[1:]))" "$f"
done

注1：TILELANG_TARGET=pto 时必须带 -p torch_npu.contrib.transfer_to_npu，否则测试文件里硬编码的 device='cuda' 参考实现会因 torch_npu 未编 CUDA 而抛 AssertionError: Torch not compiled with CUDA enabled，掩盖真正的 PTO 编译错误。
注2：为避免 transfer_to_npu 启动早期 NPU runtime 未就绪时 tile_kernels.config.get_num_sms() 命中 aclrtGetDeviceInfo 507899 返回 None 导致 T.Kernel(None) segfault，需在同一进程内先做 NPU 预热 + get_device_num_sms.cache_clear() 再跑 pytest（故用 python -c 一起调，不能用 python -m pytest）。每个文件建议单独进程跑。
注3：device 1 在连续负载下存在 segfault 抖动（crash 帧在 tvm::tl::KernelLaunch，libtilelang.so）。device 0 更稳定，下表数字均为 device 0 实测；如复现时遇到 segfault，冷却后重试即可拿到真实结果，不影响 PTO 错误归类判定。

表1：每个 test 文件的 Kernel 数量与 ascend / PTO 状态

下表数字为 device 0（ASCEND_RT_VISIBLE_DEVICES=0）实测；device 1 在连续负载下易 segfault，device 0 更稳定。

测试文件	调用的 Kernel 数	调用的 Kernel	ascend	PTO
test_expand.py	2	expand_to_mhc_fwd_tl_ascend / expand_to_mhc_bwd_tl_ascend	✅ 36 passed	❌ 36 failed（A）
test_head_compute_mix.py	2	mhc_head_compute_mix_fwd_ascend / mhc_head_compute_mix_bwd_ascend	✅ 4 passed	❌ 4 failed（A）
test_pre_split_mixes.py	2	mhc_pre_split_mixes_fwd_ascend / mhc_pre_split_mixes_bwd_ascend	✅ 4 passed	❌ 4 failed（A）
test_sinkhorn.py	2	mhc_sinkhorn_fwd_ascend / mhc_sinkhorn_bwd_ascend	✅ 6 passed	❌ 6 failed（A）
test_pre_apply_mix.py	2	mhc_pre_apply_mix_fwd_ascend / mhc_pre_apply_mix_bwd_ascend	❌ 12 failed（ascend Bisheng 编译失败：__bfloat1622float2 类型不匹配，详见注¹）	❌ 12 failed（A）
test_post.py	2	mhc_post_fwd / mhc_post_bwd	❌ 6 failed + 4 error（ascend 拒绝 threads=，详见注²）	❌ 6 failed + 4 error（B）
test_norm_fn.py	6	mhc_gemm_with_sqrsum_fwd_ascend / mhc_gemm_with_sqrsum_bwd / mhc_reduce_partials_and_rmsnorm_fwd / mhc_reduce_partials_and_rmsnorm_bwd / mhc_fn_normw_merge_fwd / mhc_fn_normw_merge_bwd	❌ 48 failed（ascend：mhc_gemm_with_sqrsum_fwd_ascend 本身能过，失败来自 mhc_reduce_partials_and_rmsnorm_* 的 threads=，详见注³）	❌ 48 failed（B + C）
test_pre_big_fuse.py	3	mhc_gemm_with_sqrsum_fwd_ascend / mhc_gemm_with_sqrsum_bwd / mhc_pre_big_fuse_kernel	❌ 12 failed（ascend Bisheng compile failed + threads=）	❌ 12 failed（C）
test_multilayer_recompute.py	1	mhc_multilayer_recompute	❌ 8 error（ascend TIR 构建阶段 TypeError：tir.Buffer 收到 ffi.Function，详见注¹）	❌ 8 error（A）

¹ test_pre_apply_mix.py 与 test_multilayer_recompute.py 的 ascend fail 均已抓到完整堆栈（绕开 pytest 的 -p transfer_to_npu segfault，直接调 wrapper 复现）：

• test_pre_apply_mix.py：mhc_pre_apply_mix_fwd_ascend 在 Bisheng 编译期失败，根因是 kernel body 中 float32 += float32 * cast(bf16→f32) 的 fused multiply-add，TileLang ascend lowering 向量化时生成了 __bfloat1622float2(*(uint1*)(...))，Bisheng 找不到匹配重载：

  /tmp/tmpXXX.asc:24:99: error: no matching function for call to '__bfloat1622float2'
  /tmp/tmpXXX.asc:31:38: error: assigning to 'uint1' from incompatible type 'bfloat16x2_t'
  RuntimeError: Compilation Failed! [bisheng --npu-arch=dav-3510 ...]
  

  归属：tilelang-dsl 的 ascend lowering（bf16→f32 vector cast 代码生成）。其余能过 kernel 均无此 bf16→f32+fma 混合 pattern。
• test_multilayer_recompute.py：mhc_multilayer_recompute 是 ptr-list indirection kernel（T.Kernel(n, threads=n_thr)，无 ascend 特化版），ascend 端在 TIR 构建阶段 抛 TypeError，非 Bisheng 编译失败：

  TypeError: Mismatched type on argument #1 when calling tir.Buffer(...). Expected DataType but got ffi.Function
  

  归属待确认（可能 tilelang-dsl 前端对 T.ptr/ptr-list tensor 的 TIR signature 构造）。
  ² test_post.py 的 ascend 与 PTO 均为 does not support threads= 前端拒绝（kernel.py:343），ascend / PTO 两侧一致。其余文件中的 benchmark_timer / benchmark_record 相关 error（未安装 pytest-benchmark）与后端能力无关，不计入 codegen 问题归类。
  ³ test_norm_fn.py ascend 归属更正：原表 1 将该行 ascend fail 标为 Bisheng compile failed + threads=，实测后发现归属需细分——带 _ascend 的 mhc_gemm_with_sqrsum_fwd_ascend 本身能在 ascend 编译通过（经 MHCGemmWithSqrsum.apply 实测 OK）；真正报 threads= 错的是不带 _ascend 的 mhc_reduce_partials_and_rmsnorm_fwd/bwd（该 kernel 用了 T.Kernel(num_tokens, threads=n_thr)，被前端拒绝）。
  ⁴ 下表「11 个带 _ascend 后缀的 Kernel 在 ascend 端的实测状态」按 Kernel 粒度复核，全部经原始 test body 真实调用绕开 pytest segfault 复现（device 0，2026-06-24）。两版本对比：pr9 版（issue 复现环境）与主线版 wenbo/asc 6ade682。
  ⁵ mhc_pre_apply_mix_bwd_ascend 的 no available layout found 详解（主线版 fwd 修好后暴露，pr9 版 bwd 单独调也报同错）：
  • 失败位置：tilelang.transform.LayoutInference → LayoutInferencer::Substitute → BufferUseDefCollector::InferInFreeMode → src/transform/layout_inference.cc:1288 的 ICHECK(min_reg_num < INT64_MAX) "no available layout found"。
  • 机制：InferInFreeMode 遍历每个 fragment 作为 infer_root 候选，对每个候选尝试推导完整 layout，期间可能抛 LayoutConflictException（src/op/parallel.cc / src/op/reduce.cc）、LoopLayoutInjectiveException（src/op/parallel.cc:536）、NormalizeIterException（src/layout/layout.cc:630）三种之一；推导成功的才算 reg_num 选最小。所有候选都抛异常 → min_reg_num 保持 INT64_MAX → ICHECK 失败。
  • 触发条件：bwd kernel（pre_apply_mix_kernel.py:189-204）里 og_frag = T.alloc_fragment((h,), T.float32) 被3 个不同维度/访问模式的 T.Parallel 循环复用——①og_frag[i_h]=og_ub[i_h]（1维写）、②mgl[i_mhc]+=og_frag[i_h]*...（2维 reduce 读）、③xg_ub[i_mhc,i_h]=mix[..]*og_frag[i_h]（2维 elementwise 读）。三种访问对 fragment 的线程铺维要求互相矛盾，选任何 layout 都让某个循环触发冲突。
  • 对比 fwd（能过）：只有 1 个 fragment ol，只在 1 个 T.Parallel 循环里访问，不存在跨循环复用冲突。
  • 归属：tilelang-dsl 的 LayoutInference pass（src/transform/layout_inference.cc + src/op/parallel.cc + src/layout/layout.cc），是通用 pass（tilelang.transform. 前缀，非 ascend 特化），对"单 fragment 跨多循环、多维度访问模式复用"推导无解。kernel 写法本身是 Ascend 常见的寄存器复用优化（注释明说 "Fragment trick for og eliminates repeated UB loads"），属 DSL 能力限制而非 kernel 源码错误。两版本（pr9 / 主线）的 LayoutInference 实现一致，故 bwd 都挂。

表1.5：11 个带 _ascend 后缀的 Kernel 在 ascend 端的实测状态（pr9 版 vs 主线版）

注：以下 11 个是 tile_kernels/mhc/*_kernel.py 中 def *_ascend(...) 真正定义的 kernel（wrapper 层 _ascend 转发不算）。全部经原始 test body 真实调用、绕开 pytest -p transfer_to_npu 的 segfault、直接调 wrapper 复现（device 0，TILELANG_TARGET=ascend，2026-06-24）。

• pr9 版 = /home/wangmiao/tilelang-simtvf-pr9-buildable（issue 复现环境用的版本，见上「环境」段）
• 主线版 = wenbo/asc 分支 6ade682（/home/wangmiao/tilelang-ds-asc-latest，独立 build，USE_PTO=OFF）。两版本共用同一份 tile-kernels 源码，差异只在 tilelang-dsl 侧。
• 主线版 fwd/bwd 是否通过，不是跑完整 pytest 得出（pytest test_pre_apply_mix.py 带 -p transfer_to_npu 会 segfault，且 wrapper 是 fwd→bwd 顺序、fwd 挂就停，直接跑分不清哪步失败）。改用隔离 probe 分别确认：forward 走 torch.no_grad() 只调 MHCPreApplyMix.apply 不触发 backward；backward 直接调底层 mhc_pre_apply_mix_bwd_ascend（tilelang 的 @tilelang.jit 每 kernel 独立编译，fwd/bwd 互不依赖）。详见 docs(dma-copy): simplify to effective bitwidth tables, remove CCE references #10/refactor(abs): simplify kernel with comments and remove dead code #11 行的「测试方法」。

#	Kernel	归属 test 文件	pr9 版 ascend	主线版 ascend	说明
1	expand_to_mhc_fwd_tl_ascend	test_expand.py	✅	✅	—
2	expand_to_mhc_bwd_tl_ascend	test_expand.py	✅	✅	—
3	mhc_head_compute_mix_fwd_ascend	test_head_compute_mix.py	✅	✅	—
4	mhc_head_compute_mix_bwd_ascend	test_head_compute_mix.py	✅	✅	—
5	mhc_pre_split_mixes_fwd_ascend	test_pre_split_mixes.py	✅	✅	—
6	mhc_pre_split_mixes_bwd_ascend	test_pre_split_mixes.py	✅	✅	—
7	mhc_sinkhorn_fwd_ascend	test_sinkhorn.py	✅	✅	—
8	mhc_sinkhorn_bwd_ascend	test_sinkhorn.py	✅	✅	—
9	mhc_gemm_with_sqrsum_fwd_ascend	test_norm_fn.py	✅	✅	走 MHCGemmWithSqrsum.apply 实测 OK；原表 1 误标为 ascend fail，已更正（见注³）
10	mhc_pre_apply_mix_fwd_ascend	test_pre_apply_mix.py	❌ Bisheng 编译失败	✅ 编译+运行通过	测试方法：因跑完整 pytest test_pre_apply_mix.py 会 segfault、且 wrapper 是 fwd→bwd 顺序（fwd 挂就停），故用隔离 probe 单独验证——①直接调底层 mhc_pre_apply_mix_fwd_ascend（带 wrapper 真实参数 h_blk=256）；②mhc_pre_apply_mix(x,mix) 在 torch.no_grad() 下只走 forward。两步均 OK（sum=655360）。结论：fwd 能过。pr9 失败原因：bf16→f32 vector cast 生成 __bfloat1622float2(*(uint1*)(...))，Bisheng 找不到重载（详见注¹）。主线版已修：codegen 改用 PrintVectorizedCast 显式约束 src 类型为 bfloat16x2_t（codegen_ascend.cc 从 src/target/ 重构到 src/ascend/）
11	mhc_pre_apply_mix_bwd_ascend	test_pre_apply_mix.py	❌（pr9 fwd 先挂，bwd 走不到）	❌ no available layout found	测试方法：tilelang 的 @tilelang.jit 是每个 kernel 函数独立编译的（fwd/bwd 互不依赖），故用隔离 probe 单独调 mhc_pre_apply_mix_bwd_ascend(og,x,mix,xg,num_sms,n_thr=512,h_blk=512)，完全不碰 fwd，3 次均确定性失败。另测"fwd 先跑成功→再调 bwd"结果仍 FAIL，证明 bwd 失败与 fwd 无关。结论：bwd 不能过。失败位置：LayoutInference pass 的 InferInFreeMode 对 og_frag 跨 3 个 T.Parallel 循环（1维写/2维reduce读/2维elementwise读）复用推导无解，min_reg_num 保持 INT64_MAX（详见注⁵）。两版本 bwd 均独立失败（pr9 单独调 bwd 也报同样 layout 错）

结论：

• pr9 版：11 个中 9 个通过，2 个失败（mhc_pre_apply_mix_fwd/bwd_ascend）——fwd 卡 Bisheng bf16 cast，bwd 卡 layout 推导（但 fwd 先挂，bwd 实际走不到）。
• 主线版：11 个中 10 个通过，仅 mhc_pre_apply_mix_bwd_ascend 1 个失败。主线版修好了 fwd 的 bf16 cast codegen 问题，但 bwd 的 layout 推导问题两版本都没修，是 tilelang-dsl LayoutInference pass 的能力限制（见注⁵）。
• 归属：fwd 问题（已修）在 tilelang-dsl ascend codegen；bwd 问题（未修）在 tilelang-dsl tilelang.transform.LayoutInference（src/transform/layout_inference.cc + src/op/parallel.cc），非 PTOAS、非 tile-kernels kernel 源码。

表2：PTO 问题分类（按 Kernel 聚类）

类别 A：CodeGenPTO 不支持 tir.shift_right

• 报错：

tvm.error.InternalError: Unsupported call: Op(tir.shift_right)
<unknown>:0: InternalError
  in tvm::codegen::CodeGenPTO::VisitExpr_(tvm::tir::CallNode const*)
  in tvm::codegen::CodeGenPTO::VisitStmt_(tvm::tir::AllocateNode / DeclBufferNode / AttrStmtNode / ForNode const*)

• 受影响 Kernel：
  • expand_to_mhc_fwd_tl_ascend / expand_to_mhc_bwd_tl_ascend
  • mhc_head_compute_mix_fwd_ascend / mhc_head_compute_mix_bwd_ascend
  • mhc_pre_split_mixes_fwd_ascend / mhc_pre_split_mixes_bwd_ascend
  • mhc_sinkhorn_fwd_ascend / mhc_sinkhorn_bwd_ascend
  • mhc_pre_apply_mix_fwd_ascend / mhc_pre_apply_mix_bwd_ascend
  • mhc_gemm_with_sqrsum_fwd_ascend / mhc_gemm_with_sqrsum_bwd（norm_fn 的 GEMM 部分）
  • mhc_multilayer_recompute
• 受影响 test 文件：test_expand.py / test_head_compute_mix.py / test_pre_split_mixes.py / test_sinkhorn.py / test_pre_apply_mix.py / test_norm_fn.py（部分）/ test_multilayer_recompute.py
• 复现：TILELANG_TARGET=pto PYTEST_ADDOPTS="-p torch_npu.contrib.transfer_to_npu" python -c "import torch,torch_npu;torch.npu.set_device(0);torch.zeros(1,device=0);torch.npu.synchronize();from tile_kernels.config import get_device_num_sms,get_num_sms;get_device_num_sms.cache_clear();import pytest,sys;sys.exit(pytest.main(['-q','-p','no:cacheprovider']+sys.argv[1:]))" test_sinkhorn.py
• 关键观察：shift_right 不在这些 kernel 的源码中直接出现，而是 TileLang lowering 阶段生成（索引计算 / 树形归约 lowering 产生 tir.shift_right），CodeGenPTO::VisitExpr_(CallNode) 未注册对应 handler。归属可能在 PTOAS（codegen 补 visitor）或 tilelang-dsl（lowering 避免生成 tir.shift_right）。

类别 B：PTO 后端不接受 T.Kernel(..., threads=)

• 报错：

ValueError: Ascend backend does not support `threads=` in `T.Kernel(...)`. Use `T.SimtVF(threads=...)` to define thread domains.
  /home/wangmiao/tilelang-simtvf-pr9-buildable/tilelang/language/kernel.py:343: ValueError

• 受影响 Kernel：
  • mhc_post_fwd / mhc_post_bwd
  • mhc_reduce_partials_and_rmsnorm_fwd / mhc_reduce_partials_and_rmsnorm_bwd（norm_fn 的 reduce/RMSNorm 部分）
• 受影响 test 文件：test_post.py / test_norm_fn.py（reduce 部分）
• 复现：TILELANG_TARGET=pto PYTEST_ADDOPTS="-p torch_npu.contrib.transfer_to_npu" python -c "import torch,torch_npu;torch.npu.set_device(0);torch.zeros(1,device=0);torch.npu.synchronize();from tile_kernels.config import get_device_num_sms,get_num_sms;get_device_num_sms.cache_clear();import pytest,sys;sys.exit(pytest.main(['-q','-p','no:cacheprovider']+sys.argv[1:]))" test_post.py
• 说明：错误信息已建议改用 T.SimtVF(threads=...)。需确认这是 PTO 后端能力边界（永远不支持 threads=），还是 tilelang-dsl/PTOAS 可放宽。

类别 C：CodeGenPTO 不支持 AllocateNode scope 'local.var'

• 报错：

tvm.error.InternalError: Check failed: (false) is false: AllocateNode scope 'local.var' not yet supported; only UB (shared/shared.dyn) is supported
<unknown>:0: InternalError

• 受影响 Kernel：
  • mhc_pre_big_fuse_kernel
  • mhc_reduce_partials_and_rmsnorm_fwd / mhc_reduce_partials_and_rmsnorm_bwd（norm_fn 的 reduce/RMSNorm 部分，实测 test_norm_fn.py pto 命中 local.var）
• 受影响 test 文件：test_pre_big_fuse.py / test_norm_fn.py（部分）
• 复现：TILELANG_TARGET=pto PYTEST_ADDOPTS="-p torch_npu.contrib.transfer_to_npu" python -c "import torch,torch_npu;torch.npu.set_device(0);torch.zeros(1,device=0);torch.npu.synchronize();from tile_kernels.config import get_device_num_sms,get_num_sms;get_device_num_sms.cache_clear();import pytest,sys;sys.exit(pytest.main(['-q','-p','no:cacheprovider']+sys.argv[1:]))" test_pre_big_fuse.py
• 关键观察：local.var 来自 fragment 级别内存分配的 lowering。错误信息表明 PTO 后端目前只支持 UB（shared/shared.dyn）scope，local.var 未支持。归属可能在 PTOAS（codegen 支持更多 buffer scope）或 tilelang-dsl（lowering 时将 local.var 提升到 UB）。

归属待确认的问题

类别	问题	可能归属	待确认
A	tir.shift_right 未支持	PTOAS（补 visitor）/ tilelang-dsl（lowering 避免）	shift_right 是 lowering 生成还是 kernel 源码直接写？
B	T.Kernel(threads=) 被前端拒绝	tilelang-dsl（语义边界）/ tile-kernels（改用 SimtVF）/ PTOAS（支持 threads=）	PTO 是否计划支持 threads=？
C	local.var scope 未支持	PTOAS（支持更多 scope）/ tilelang-dsl（lowering 提升到 UB）/ tile-kernels（改用 UB 分配）	local.var 是否计划支持？（norm_fn reduce kernel 实测同时命中 B/C）

期望

希望确认类别 A、B、C 属于哪一层：

• 类别 A：若属 PTOAS codegen 能力补齐，请指认需要补的 visitor；若属 tilelang-dsl lowering 应避免的 pattern，请转移。
• 类别 B：threads= 是否 PTO 后端永远不支持、必须走 SimtVF？还是可放宽？
• 类别 C：local.var 是否计划支持？或应改用 UB 分配？（实测 test_norm_fn.py pto 同时命中 threads= 与 local.var，说明 norm_fn 的 reduce kernel 既有 T.Kernel(threads=) 又有 fragment 分配。）

谢谢！

[KurrinQu]

• PTO Codegen支持 tir.shift_right

tir.shift_right(lhs, rhs)

映射到：

arith.shrsi  // signed int
arith.shrui  // unsigned int

语义与 LLVM 一致：signed 用 arithmetic shift right，unsigned 用 logical shift right。

• PTO Codegen支持local.var

TIR 形态：

v = T.alloc_buffer((1,), dtype, scope="local.var")

MLIR 中应按线程私有scalar 处理，先 lowering 成 stack scalar，和fragment类似，使用llvm.alloca。后续可交给 MLIR/LLVM 的 mem2reg 优化成 SSA

• T.Kernel(threads=xxx），ascendc后端同样不支持，需要对齐该场景是否未改写？

[HecreReed]

补充:zhangyc/restore-pr-9 分支 PTO 后端 11 个 kernel 实测结果(commit 55c154f)

在 issue 原始复现环境(pr9 版 tilelang-simtvf-pr9-buildable)基础上,另用 WenboCodes/tilelang-deepseek 的 zhangyc/restore-pr-9 分支(commit 55c154f,USE_PTO=ON,build 在 a5 /home/wangmiao/tilelang-restore-pr9/)重跑 11 个带 _ascend 后缀的 kernel。复现脚本已留 a5 /home/wangmiao/probe_pto_mhc_full.py(env wrapper:run_pto_env_pr9.sh / run_pto_env_restore.sh)。

测试方法(与 pr9 一致)

pytest + tests/mhc/conftest.py 的 import torch_npu.contrib.transfer_to_npu 在 PTO 后端 KernelLaunch 必段错误(exit 139,栈顶 tvm::tl::KernelLaunch),无法直接跑。故 import 各 test 模块、复用其 generate_*_test_data() + _tester() + assert_close 逻辑,device 直接用 'npu'(torch_npu 已注册 npu 设备,不需要 monkeypatch),与 test 脚本逻辑等价。fwd 用 torch.no_grad() 单独跑(只触发 fwd kernel 编译),bwd 调 test 的 _tester(先 fwd 构图再 torch.autograd.backward(),触发 bwd kernel 编译),fwd/bwd 独立两次跑。device 0,rp21-PT2.7。

restore-pr-9 版完整报错(11 个 kernel,逐个)

A′ 类 — expand_fwd_ascend / expand_bwd_ascend(2 个,完整):

PTO micro-instruction host link failed.
Command: /home/wenquan/cann29/cann-9.0.0/bin/bisheng -c -fPIC -xcce ... --cce-aicore-arch=dav-c310-vec ... /tmp/tilelang_pto_lk5v0_ov/launch.cpp -o .../launch.o
stderr:
/tmp/tilelang_pto_lk5v0_ov/launch.cpp:25:74: error: __bf16 is not supported on this target
  expand_to_mhc_fwd_tl_ascend_kernel_kernel<<<64, 256, stream>>>((__gm__ __bf16 *)arg1, (__gm__ __bf16 *)arg0, (int32_t)arg2);
                                                                         ^
/tmp/tilelang_pto_lk5v0_ov/launch.cpp:25:97: error: __bf16 is not supported on this target
  expand_to_mhc_fwd_tl_ascend_kernel_kernel<<<64, 256, stream>>>((__gm__ __bf16 *)arg1, (__gm__ __bf16 *)arg0, (int32_t)arg2);
                                                                                                ^
2 errors generated.

B′ 类 — head_compute_mix / pre_split_mixes / sinkhorn 的 fwd+bwd(6 个):

Unsupported call: Op(tir.call_pure_extern)

C′ 类 — gemm_with_sqrsum_fwd_ascend(1 个):

Check failed: (false) is false: AllocateNode scope 'shared.l1' not yet supported; only UB (shared/shared.dyn) is supported

D′ 类 — pre_apply_mix_fwd_ascend / pre_apply_mix_bwd_ascend(2 个,完整):

PTO micro-instruction lowering failed.
Command: /home/wangmiao/PTOAS-main-1df3fba/install/bin/ptoas --pto-arch=a5 --pto-backend=vpto /tmp/tilelang_pto_c6ck_r3d/kernel.pto -o /tmp/tilelang_pto_c6ck_r3d/kernel.fatobj.o
stderr:
loc("/tmp/tilelang_pto_c6ck_r3d/kernel.pto":126:40): error: use of value '%load_205' expects different type than prior uses: 'vector<4xf32>' vs 'vector<4xbf16>'
Error: Failed to parse MLIR.

注:kernel.pto 报错定位为第 126 行第 40 列(:126:40),涉及第 125、126 行——同一 SSA value %load_205 在第 125 行以 vector<4xf32> 使用,第 126 行以 vector<4xbf16> 使用,MLIR 要求同一 value 类型一致故解析失败。即错误跨 125/126 两行,定位点在 126:40。

错误类别对比表(逐 kernel,pr9 版 vs restore-pr-9 版)

#	Kernel	pr9 版	restore-pr-9 版	是否变化
1	expand_fwd_ascend	A tir.shift_right	A′ __bf16 host link	变
2	expand_bwd_ascend	A tir.shift_right	A′ __bf16 host link	变
3	mhc_head_compute_mix_fwd_ascend	A tir.shift_right	B′ tir.call_pure_extern	变
4	mhc_head_compute_mix_bwd_ascend	A tir.shift_right	B′ tir.call_pure_extern(被 fwd 拖累)	变
5	mhc_pre_split_mixes_fwd_ascend	A tir.shift_right	B′ tir.call_pure_extern	变
6	mhc_pre_split_mixes_bwd_ascend	A tir.shift_right	B′ tir.call_pure_extern(被 fwd 拖累)	变
7	mhc_sinkhorn_fwd_ascend	A tir.shift_right	B′ tir.call_pure_extern	变
8	mhc_sinkhorn_bwd_ascend	A tir.shift_right	B′ tir.call_pure_extern(被 fwd 拖累)	变
9	mhc_gemm_with_sqrsum_fwd_ascend	C local.var	C′ shared.l1	变(同类,scope 换)
10	mhc_pre_apply_mix_fwd_ascend	A tir.shift_right	D′ bf16 MLIR 类型冲突	变
11	mhc_pre_apply_mix_bwd_ascend	A tir.shift_right	D′ bf16 MLIR 类型冲突	变

"被 fwd 拖累":对 head_compute_mix / pre_split_mixes / sinkhorn 这 3 个,bwd 的 _tester 必须先跑通 forward 才能反向,而 fwd 编译就失败,故 bwd 连 bwd kernel 都没走到,挂在同一个 fwd kernel 上。对 expand / pre_apply_mix,fwd 和 bwd 各自独立失败(同类错误)。

类别演化

pr9 类别	restore-pr-9 状态	说明
A tir.shift_right(10 个)	✅ 已修复,不再出现	commit 补了 shift_right visitor,10 个 kernel 能往下走,暴露更深问题
C local.var(gemm 1 个)	✅ 已修复,但 → C′ shared.l1	同类问题(PTO 只支持 UB),scope 从 local.var 换成 shared.l1
B threads=	➖ 本次不涉及	11 个 _ascend kernel 都不含 threads=(B 只影响 post / reduce_partials 这类非 _ascend kernel)

结论

restore-pr-9 修复了 issue 原始的 A(shift_right)和 C(local.var) 两类,但 11 个 kernel 仍全 FAIL,0 通过——错误位置都从 pr9 的"最前一道坎"推进到了更后面,暴露出 4 类新问题:

• A′ __bf16 host link(expand 2 个):PTO 生成的 host 端 launch.cpp 用 __bf16 类型,Bisheng host 编译器不支持。
• B′ tir.call_pure_extern(6 个):shift_right 修好后,lowering 产生的下一个 TIR call tir.call_pure_extern 又没 visitor(与原 A 同属"PTO codegen 缺 visitor",只是 op 换了)。
• C′ shared.l1 scope(gemm 1 个):原 C 的同类,scope 从 local.var 换成 shared.l1。
• D′ bf16 MLIR 类型冲突(pre_apply_mix 2 个):ptoas 能生成 MLIR,但同一 SSA value %load_205 在 125/126 行两处类型不一致(vector<4xf32> vs vector<4xbf16>),MLIR 解析失败。

复现脚本:a5:/home/wangmiao/probe_pto_mhc_full.py + run_pto_env_pr9.sh / run_pto_env_restore.sh。

[HecreReed]

补充:zhangyc/restore-pr-9 新版(commit 70ab1b6)PTO 后端 11 个 kernel 实测,对比上一版 55c154f

在上一条评论(55c154f)的同一环境上,把 WenboCodes/tilelang-deepseek 的 zhangyc/restore-pr-9 分支更新到最新 commit 70ab1b6(分支新增一个 commit Fix PTO bf16 launch and type lowering,改 codegen_pto.cc +219/-75 和 wrapper.py),在 a5 重 build(USE_PTO=ON,build 在 /home/wangmiao/tilelang-restore-pr9/,旧版备份 tilelang-restore-pr9-old/),用同一套 probe(probe_pto_mhc_full.py + run_pto_env_restore.sh)重跑 11 个 _ascend kernel。全程 PTO 后端(TILELANG_TARGET=pto,build 入口 target.build.tilelang_pto_without_compile)。

新版改动要点

70ab1b6 主要修了 bf16 相关:

• scalar_dtype.is_bfloat16() 判断(原来误用 dtype.is_bfloat16())
• pto.castptr 改用 elemPtrType(不再固定 gmPtrType 导致 bf16 不识别)
• vector<...xbf16> 类型生成
• 新增 shared.l1/l0a/l0b/l0c/bt/fp 的 scope→memory tag 映射

四类问题变化总表(55c154f → 70ab1b6)

旧类别(55c154f)	70ab1b6 状态	新错误(若有)	新类别
A′ __bf16 host link	✅ 已修复	—	—
B′ tir.call_pure_extern(expf)	❌ 未修复	不变	B′
C′ shared.l1 scope	✅ scope 已修复 → 暴露新错	Unsupported thread_tag: cthread	C′′
D′ bf16 MLIR 类型冲突	✅ fwd 已修复 → bwd 暴露新错	no available layout found	D′′

编号约定:原四类 A′~D′;修好后往下暴露的新错用双撇号续编号(C′′、D′′),不另起 E/F。

逐 kernel 对照(11 个)

#	Kernel	55c154f	70ab1b6	变化
1	expand_fwd_ascend	❌ A′ __bf16 host link	✅ PASS	A′ 修好
2	expand_bwd_ascend	❌ A′ __bf16 host link	✅ PASS	A′ 修好
3	mhc_head_compute_mix_fwd_ascend	❌ B′ call_pure_extern	❌ B′ call_pure_extern	未变
4	mhc_head_compute_mix_bwd_ascend	❌ B′(被 fwd 拖累)	❌ B′(被 fwd 拖累)	未变
5	mhc_pre_split_mixes_fwd_ascend	❌ B′ call_pure_extern	❌ B′ call_pure_extern	未变
6	mhc_pre_split_mixes_bwd_ascend	❌ B′(被 fwd 拖累)	❌ B′(被 fwd 拖累)	未变
7	mhc_sinkhorn_fwd_ascend	❌ B′ call_pure_extern	❌ B′ call_pure_extern	未变
8	mhc_sinkhorn_bwd_ascend	❌ B′(被 fwd 拖累)	❌ B′(被 fwd 拖累)	未变
9	mhc_gemm_with_sqrsum_fwd_ascend	❌ C′ shared.l1	❌ C′′ cthread	C′ 修好→C′′
10	mhc_pre_apply_mix_fwd_ascend	❌ D′ bf16 MLIR	✅ PASS	D′ 修好
11	mhc_pre_apply_mix_bwd_ascend	❌ D′ bf16 MLIR	❌ D′′ no available layout	D′ 修好→D′′

两类新错误明细

C′′ — Unsupported thread_tag: cthread(mhc_gemm_with_sqrsum_fwd_ascend,1 个)

tvm.error.InternalError: Unsupported thread_tag: cthread
  in target.build.tilelang_pto_without_compile

• cthread 是 ascend NPU 的 subblock id(子块编号),DSL 里即 T.MixedKernel(num_sms) as (pid, sid) 的 sid,硬件上映射 get_subblockid(),是 ascend Cube/Vector 单元特有的并行层次。
• 来源:src/ir.cc:413-420 的 MixedKernelLaunch 把 T.MixedKernel 展开成 3 个 frame,其中 CreateEnvThread("sid", "cthread", ...) 创建 thread_tag=cthread 的环境线程。
• PTO 未处理:src/target/codegen_pto.cc:407-430 的 VisitStmt_(AttrStmtNode) 处理 thread_extent 时只认 blockIdx.x(→pto.get_block_idx)和 threadIdx(→pto.get_tid_x/y/z),cthread 落到 LOG(FATAL) << "Unsupported thread_tag"。
• ascend 后端对比:src/target/codegen_ascend.cc:1702 把 cthread 映射到 get_subblockid()。
• 与 C′ 同源:gemm kernel 专为 ascend Cube 单元写(alloc_l1/alloc_l0c/cthread/ascend_gemm_l1),旧版卡在第一道门(L1 scope),新版修了 L1,卡在第二道门(cthread)。
• 修复方向:PTO codegen 补 cthread handler(类比 get_subblockid(),加 pto.get_sub_block_id builtin,需 PTO 运行时支持 subblock);或 DSL 侧在 PTO target 下把 MixedKernel 退化成普通 T.Kernel。

D′′ — no available layout found(mhc_pre_apply_mix_bwd_ascend,1 个)

Check failed: (min_reg_num < INT64_MAX) is false: no available layout found

• D′ bf16 MLIR 类型冲突修好后,fwd 过了,bwd 往下走到 tilelang LayoutInference pass 的 InferInFreeMode,撞上 og_frag 跨 3 个 T.Parallel 复用导致 layout 推不出。
• 非 PTO 特有:ascend 原生后端 bwd 也有同一个错——tilelang LayoutInference pass 对 og_frag 跨多 T.Parallel 复用的处理缺陷,不在 PTOAS 范围内。

结论

新版 70ab1b6 修好了 A′ 和 D′(fwd) 两类 bf16 问题,通过率 0/11 → 3/11(expand_fwd/expand_bwd/pre_apply_mix_fwd 全过)。

剩余 8 个失败根因:

数量	类别	根因	修复位置
6	B′	PTO codegen 缺 expf handler(只认 rsqrtf)	PTOAS codegen_pto.cc,补一个 handler
1	C′′	PTO codegen 不认 cthread(需支持 ascend Cube subblock)	PTOAS codegen_pto.cc,或 DSL 侧退化 MixedKernel
1	D′′	tilelang LayoutInference bwd 缺陷(ascend 后端也有)	tilelang LayoutInference pass,非 PTOAS 范围

复现脚本同上一条:a5:/home/wangmiao/probe_pto_mhc_full.py + run_pto_env_restore.sh(env wrapper 指向新版 build)。

[HecreReed]

新版本在 a5 PTO 后端复现结果

背景

• 分支:WenboCodes/tilelang-deepseek 的 zhangyc/restore-pr-9
• 新版 HEAD:cfdf8b2 Support cthread in PTO codegen(比 70ab1b6 多 1 个 commit,只改 src/target/codegen_pto.cc +11 行)
• 后端:全程 PTO 后端(TILELANG_TARGET=pto,入口 target.build.tilelang_pto_without_compile)

---

一、11 kernel 对照表(70ab1b6 → cfdf8b2)

#	Kernel	70ab1b6(上一版)	cfdf8b2(新版)	变化
1	expand_fwd	✅ 编译通过	✅ 编译通过	未变
2	expand_bwd	✅ 编译通过	✅ 编译通过	未变
3	head_compute_mix_fwd	❌ B′ call_pure_extern	❌ B′	未变
4	head_compute_mix_bwd	❌ B′(被 fwd 拖累)	❌ B′	未变
5	pre_split_mixes_fwd	❌ B′	❌ B′	未变
6	pre_split_mixes_bwd	❌ B′(被 fwd 拖累)	❌ B′	未变
7	sinkhorn_fwd	❌ B′	❌ B′	未变
8	sinkhorn_bwd	❌ B′(被 fwd 拖累)	❌ B′	未变
9	gemm_with_sqrsum_fwd	❌ C′′ Unsupported thread_tag: cthread	❌ G′ Unsupported call: Op(tl.ascend_copy_gm_to_cbuf)	cthread 修好→暴露新错
10	pre_apply_mix_fwd	✅ 编译通过	✅ 编译通过	未变
11	pre_apply_mix_bwd	❌ D′′ no available layout	❌ D′′ no available layout	未变

说明:expand_fwd / expand_bwd / pre_apply_mix_fwd 这 3 个能编译通过(不抛编译期异常),但全参数化 assert_close 数值校验下算出的结果与 ref 不一致(见第三节数值报错分析)。

---

二、gemm_with_sqrsum_fwd 现在的具体问题(G′ 类)

错误信息

tvm.error.InternalError: Unsupported call: Op(tl.ascend_copy_gm_to_cbuf)
  in target.build.tilelang_pto_without_compile
  → CodeGenPTO::VisitExpr_(CallNode) → LOG(FATAL) "Unsupported call" (codegen_pto.cc:2319)

tl.ascend_copy_gm_to_cbuf

mhc_gemm_with_sqrsum_fwd_ascend kernel 用 T.alloc_l1 在 Cube Buffer(L1/CBuf)开 buffer,然后用 T.copy 把数据从 GM/UB 搬进 L1:

fn_l1 = T.alloc_l1((mhc_mult3_pad16, hidden_split_size), T.float32)   # L1 = CBuf
...
T.copy(fn[0, blk_h * hidden_split_size], fn_l1)            # GM → L1(CBuf)
T.copy(x_ub, x_l1[sid * (token_block // 2), :], nd2nz=True)  # UB → L1(CBuf, nd2nz)

lowering 层按 src/dst 的 scope 选 DMA path:

• src/op/copy.cc:767 GetAscendDMAPath:GM + L1 ⇒ dma_path == 2
• src/op/copy.cc:978 LowerAscendDMACopy(dma_path==2):构造 Call(ascend_copy_gm_to_cbuf, {dst, src, sid, ..., transpose, ...})

这是 ascend Cube 单元的 GM→CBuf 数据搬运指令。Op 本身定义在 src/op/builtin.h:658 / 注册在 src/op/builtin.cc:505。

PTO codegen 没处理的地方

codegen_pto.cc 的 VisitExpr_(CallNode)(:1511)只实现了 Vector 单元(UB)的两个搬运 op:

op	PTO codegen	行号
tl.ascend_copy_gm_to_ubuf(GM→UB)	✅ 已实现	:1877
tl.ascend_copy_ubuf_to_gm(UB→GM)	✅ 已实现	:1957
tl.ascend_copy_gm_to_cbuf(GM→CBuf/L1)	❌ 未实现 → fatal	(缺)
tl.ascend_copy_cbuf_to_gm	❌ 未实现	(缺)
tl.ascend_gemm_l1(Cube 矩阵乘) 这个也涉及到，目前看源码	❌ 未实现	(缺)

PTO codegen 只支持 UB(Unified Buffer,Vector 单元)的数据搬运,完全没有 Cube 单元(CBuf/L1/L0C)的任何指令 → 落到 2319行代码 的 LOG(FATAL) << "Unsupported call"。

ascend 原生后端对比

• src/target/codegen_ascend.cc:947 实现了 ascend_copy_gm_to_cbuf(发成 copy_gm_to_cbuf_multi_nd2nz/dn2nz 硬件 intrinsic,带 nd2nz/dn2nz 数据排布转换)
• src/target/codegen_ascend.cc:1176 实现了配套的 ascend_gemm_l1(Cube 矩阵乘)

gemm_with_sqrsum_fwd已修复问题和即将面临的问题

阶段	卡点	版本
1	shared.l1 scope(AllocateNode)	70ab1b6 已修(C′→C′′)
2	cthread(MixedKernel 的 sid)	cfdf8b2 已修
3	ascend_copy_gm_to_cbuf(GM→CBuf 搬运)	现在问题
4	ascend_gemm_l1(Cube 矩阵乘)	应该即将会面对

要跑通这个 kernel,PTO 后端需要补整套 Cube 数据通路(CBuf 搬运 + Cube 矩阵乘 + L0C 搬运),而不是补一两个 handler

---

三、全参数化数值报错分析

用 ascend_mhc 目录的原始 test 脚本(仅把 device='cuda'→'npu'、去掉 transfer_to_npu monkeypatch,逻辑/参数化/assert_close 一字未改),使用最新版本run

总览

104 failed, 0 passed in 291.52s (0:04:51)

真实通过率 0/104。

范围说明:严格按 11 个 _ascend kernel 相关的 test 用例。test_norm_fn.py 里 test_correctness[True-...](generate_normw)那 6 组走的是 mhc_fn_normw_merge reduce kernel(2-D grid),不在 11 个 kernel 内,已排除;其余 6 组 test_correctness[False] 调 mhc_pre_norm_fn(含 gemm)算在内。

测试范围(11 kernel ↔ test 文件)

#	kernel	test 文件	参数化组数
1-2	expand_fwd/bwd	test_expand.py	2×2×3×3 = 36
3-4	head_compute_mix_fwd/bwd	test_head_compute_mix.py	2×2×1 = 4
5-6	pre_split_mixes_fwd/bwd	test_pre_split_mixes.py	2×2×1 = 4
7-8	sinkhorn_fwd/bwd	test_sinkhorn.py	2×3×1 = 6
9-10	pre_apply_mix_fwd/bwd	test_pre_apply_mix.py	2×2×3 = 12
11	gemm_with_sqrsum_fwd(+联动)	test_norm_fn.py	6+12+8+8+8 = 42(排除 normw 的 6 组)
		合计	104

各文件 × 根因交叉表

test 文件	组数	失败	根因分布
test_expand	36	36	A′:结果错位 + 少量 NaN(assert_close 失败,mismatch 85%~97%)。fwd 能编译能跑但算错,详见下文
test_head_compute_mix	4	4	B′:Unsupported call: Op(tir.call_pure_extern)(expf) 上表的编译问题
test_pre_split_mixes	4	4	B′:Unsupported call: Op(tir.call_pure_extern)(expf) 上表的编译问题
test_sinkhorn	6	6	B′:Unsupported call: Op(tir.call_pure_extern)(expf) 上表的编译问题
test_pre_apply_mix	12	12	D′ 4 组(h=1280):Check failed: no available layout found(LayoutInference 阶段) H′ 8 组(h=2560/7680):Unsupported call: Op(tir.tvm_storage_sync)(PTO codegen 阶段)
test_norm_fn	42	42	G′ 42 组(全部):Unsupported call: Op(tl.ascend_copy_gm_to_cbuf)(gemm 相关:test_correctness[False] / test_split_k_correctness / test_fwd_matmul / test_fwd_norm / test_full_chain_backward 都调 MHCGemmWithSqrsum.apply)
合计	104	104

详细描述

• A′ expand 结果错位(36 组):expand fwd kernel expand_to_mhc_fwd_tl_ascend 能编译能跑,但算出的结果整体错位。最小 shape(n0=1,n1=1024,mm=2,h=1280)实测:

  • assert_close 失败,mismatch 97.2%(2548225/2621440),Greatest absolute difference: nan at index (0,0,0,638)。
  • out 共 196 个 NaN(占总数 0.0075%,不是全 NaN),散落在 h 方向固定位置 [638, 854, 960, 1234, ...]。
  • 其余非 NaN 的值也错:输入 x[0,0,:8] = [-0.3008, 0.5195, -0.1572, ...](ref 正确复制),但 out[0,0,0,:8] = [-8.8e-24, -0.6016, 6.4e-4, -0.2451, -2.5e-21, 1.92, 28185722880, 0.0251] —— 从第一个值就错位,还出现 28185722880、-8.8e-24 这类垃圾/巨大值。
  • assert_close 报的 "max abs diff = nan" 是因为 |NaN - ref| = nan,nan 来自那 196 个 NaN 位置,不是全 NaN。

  根因在 PTO codegen 对该 kernel 的 GM 地址计算 / 数据搬运有误(该 kernel 用 T.Persistent + token_block_size + get_num_sms() 的 persistent 调度,host launch 参数与 kernel 形参的映射 / GM 偏移计算出了问题),具体机制需进一步 dump MLIR 定位,目前只确认"能编译能跑但数据搬运错位"。所有 36 组 case 都在 fwd assert_close(test_expand.py:31)抛 AssertionError: Tensor-likes are not close!。

• B′ expf(14 组 kernel / 28 次 InternalError):head_compute_mix / pre_split_mixes / sinkhorn 三个 kernel 用了 T.sigmoid/T.exp,lowering 成 tir.call_pure_extern("expf", ...),PTO codegen 的 CallNode handler 只认 rsqrtf,缺 expf handler → Unsupported call: Op(tir.call_pure_extern)。fwd 编译失败,bwd 被 fwd 拖累挂在同一个 fwd kernel。

• G′ gemm Cube 指令(42 组):test_norm_fn 里所有 gemm 相关子测试(test_correctness[False]、test_split_k_correctness、test_fwd_matmul、test_fwd_norm、test_full_chain_backward)都调 MHCGemmWithSqrsum.apply,卡在 tl.ascend_copy_gm_to_cbuf(详见第二节)。

• D′ no available layout(4 组,h=1280):pre_apply_mix 在 h=1280 时,bwd kernel 走到 tilelang LayoutInference 的 InferInFreeMode,推不出 layout → Check failed: (min_reg_num < INT64_MAX) is false: no available layout found。这一阶段在 PTO codegen 之前,所以 h=1280 根本走不到后面的 storage_sync 就挂了。逐 h 隔离子进程复跑确认:h=1280 一律报 D′,n0/n1 维度不影响(只影响 token 数,不改 h 方向的 layout 推导)。这是 tilelang LayoutInference pass 对 og_frag 跨多个 T.Parallel 复用的处理缺陷,非 PTO 特有,ascend 原生后端 bwd 也有同样错误。

• H′ tvm_storage_sync(8 组,h=2560/7680):pre_apply_mix 在 h≥2560 时,LayoutInference 能推出 layout(越过 D′),继续往下走到 PTO codegen 阶段,kernel 用的双缓冲流水线产生 tir.tvm_storage_sync builtin,PTO codegen 不认 → Unsupported call: Op(tir.tvm_storage_sync)。根因是 bwd kernel 的双缓冲:

  # pre_apply_mix_kernel.py:180
  T.annotate_buffer_versions({og_ub: 2, x_ub: 2, xg_ub: 1})   # 多版本 buffer
  for tok in T.Pipelined(..., num_stages=2):                    # 双缓冲流水线

  T.Pipelined(num_stages=2) + 多版本 buffer 的 lowering 必然产生 tir.tvm_storage_sync(流水线 stage 之间的 barrier 同步)。所以 h=1280 卡在前面的 LayoutInference(D′),h≥2560 越过 layout 后卡在后面的 PTO codegen(H′)——同一个 bwd kernel,不同 h 在不同编译阶段挂。caller 层的 n_thr=512/h_blk=512 是硬编码(见 MHCPreApplyMix.backward),h 变化不改这些参数,改变的是 h 方向 fragment 的尺寸从而影响 LayoutInference 能否成功。

[HecreReed]

补一条当前在 A5 上的实测结论，分成性能和精度两部分。

1. expand 两个 _ascend kernel 的 msprof 性能结果

对比条件完全一致：

• 同一份代码：/home/wangmiao/tilelang-restore-pr9-cfdf8b2
• 同一个 ptoas
• 同一个 tile-kernels
• 同一个 msprof
• fwd 和 bwd 都是 full36

比值定义为 pto / ascend：

• > 1：pto 慢
• < 1：pto 快

两轮汇总结论表：

项目	第 1 轮 avg	第 2 轮 avg	多轮均值	多轮中位	PTO 更快
fwd	1.1651	1.1398	1.1524	1.0037	18 / 36
bwd	0.9710	0.9697	0.9704	0.9985	18 / 36
fwd+bwd	1.0680	1.0548	1.0614	0.9985	36 / 72

按 num_tokens 聚合的多轮均值：

num_tokens	多轮均值 pto / ascend	PTO 更快
1024	0.8893	15 / 18
2048	0.9828	12 / 18
4096	1.1650	6 / 18
8192	1.2086	3 / 18

稳定性：

• fwd 平均 CV = 0.0216
• bwd 平均 CV = 0.0189
• all 平均 CV = 0.0202

这说明两轮之间波动很小，结论已经比较稳。

最终判断：

• bwd 基本打平，略偏 pto 更快
• fwd 是主要差距来源，整体上 pto 稍慢
• 全部 72 个 case 合并后，pto 相对 ascend 的多轮均值约为 1.0614x
• 更直接说：同仓、同码、同口径下，PTO 整体平均慢约 6.1%

2. expand_fwd_ascend / expand_bwd_ascend 精度 mismatch 的修复点

这两个 kernel 之前在 PTO backend 下有精度 mismatch，根因是 PTO 路径被错误地绕开了 autoschedule 这段 pass，因此没有把 autoschedule 生成的同步逻辑带下来。

旧代码在 phase.py 里是：

if allow_autoschedule() and not is_pto_target(target):

修复就是去掉这里的：

and not is_pto_target(target)

也就是改成：

if allow_autoschedule():

这样 PTO target 也会走：

• AnnotateMultiBufferEligible
• IfConditionExtract
• AutoSchedule
• NormalizeMixedKernelSid
• Simplify

修完后：

• PTO dump 里能看到对应的 set_flag / wait_flag 同步
• expand_fwd_ascend / expand_bwd_ascend 这两个 kernel 的精度 mismatch 问题可以修掉

所以这里不是 ptoas 本身没自动同步，而是上游 phase.py 把 PTO target 排除在 autoschedule 之外，导致双缓冲流水线下同步没有正确插入。

[HecreReed]

补充一次 2026-07-15 基于最新版 TileLang asc + 最新 PTOAS main 的 A5 重测结果。这次结果与 issue 正文使用的旧版 pr9 环境不同，建议作为当前主线状态参考。

1. 测试版本与目标确认

• A5：Ascend950PR_9589
• TileLang：WenboCodes/tilelang-deepseek，asc 分支
  • commit：2918b4aab64c1cbb21daa7d2c037571ec22ce121
  • TVM submodule：28a0d34420d2fa9bc71fc891445a3f1396fca759
• PTOAS：hw-native-sys/PTOAS，main
  • commit：a1371bc
  • ptoas 0.51，LLVM 21.1.8
• tile-kernels：现有 ascend-dev 工作树
  • commit：f7ff125472ee1a1f6a5bd96d2e9e507526b0982f
• 测试集：/home/wangmiao/ascend_mhc_npu 对应的 11 个 SIMT kernel

最新版 TileLang 的默认目标变量已改为 TILELANG_DEFAULT_TARGET。为避免旧的 TILELANG_TARGET=pto 实际仍进入 AscendC，本次没有依赖环境变量，而是显式创建 PTO Target：

pto_target = determine_target("pto", return_object=True)
with pto_target:
    ...

运行时确认 Target 为：

kind = c
keys = ["pto", "ascend"]

2. 结论

显式 PTO backend 下 0/11 编译通过。11 个 case 均在 TileLang DSL/TIR → PTODSL 适配阶段失败，没有任何一个执行到 ptoas subprocess。

实际链路是：

TileLang DSL
  -> TileLang TIR / passes
  -> tilelang src/ascend/codegen/codegen_pto.cc 生成 PTODSL Python
  -> TileLang PTO host wrapper
  -> PTODSL kernel.compile() 生成 .pto
  -> ptoas --pto-arch=a5 --pto-backend=vpto

本轮错误均发生在前四步，因此不能归类为“PTOAS 编译 11 个 kernel 失败”，更准确的说法是“TileLang 尚未为这 11 个 kernel 生成可交给 PTODSL/PTOAS 的合法输入”。

3. 11 个 kernel 的具体报错

#	kernel / harness case	具体报错	失败阶段
1	expand_to_mhc	PTO kernel argument count mismatch: device=3, host=2	TileLang PTO host wrapper
2	mhc_head_compute_mix	PTO kernel argument count mismatch: device=5, host=4	TileLang PTO host wrapper
3	MHCGemmWithSqrsum_fwd	TypeError: copy() got an unexpected keyword argument 'nd2nz'	TileLang DSL API
4	MHCReducePartialsAndRmsnorm_fwd	同上：不再接受 nd2nz=True	TileLang DSL API
5	mhc_pre_norm_fn	同上：不再接受 nd2nz=True	TileLang DSL API
6	mhc_post_fwd	ValueError: Ascend backend does not support threads= in T.Kernel(...). Use T.SimtVF(threads=...) to define thread domains.	TileLang DSL/TIR 构建
7	mhc_pre_apply_mix	InternalError: PTO SIMT local allocation currently supports float32 only, got bfloat16	TileLang PTODSL codegen
8	mhc_pre_split_mixes	PTO kernel argument count mismatch: device=7, host=6	TileLang PTO host wrapper
9	sinkhorn_normalize	InternalError: PTO codegen currently maps only sum reductions to pto.simt_allreduce_sum	TileLang PTODSL codegen
10	mhc_pre_big_fuse	TypeError: copy() got an unexpected keyword argument 'nd2nz'	TileLang DSL API
11	mhc_multilayer_recompute	与 #6 相同：T.Kernel(..., threads=...) 被拒绝	TileLang DSL/TIR 构建

按根因聚合：

• 3 个：PTODSL device signature 与 host wrapper ABI 相差一个参数
• 4 个：旧的 nd2nz=True API 与最新版 TileLang 不兼容
• 2 个：旧的 T.Kernel(..., threads=...) API 与最新版 TileLang 不兼容
• 1 个：TileLang PTO codegen 将 SIMT local allocation 硬限制为 float32
• 1 个：TileLang PTO codegen 只实现了 sum reduction 映射

4. 代码侧定位与可行修复建议

A. host/device 参数不一致

最新版 wrapper 会解析生成的 PTODSL 函数签名，然后与 host PrimFunc 参数逐项比较：

device_arg_names = self._parse_ptodsl_kernel_args(kernel_name)
if len(device_arg_names) != len(host_args):
    raise RuntimeError(
        f"PTO kernel {kernel_name} argument count mismatch: "
        f"device={len(device_arg_names)}, host={len(host_args)}"
    )

当前 wrapper 的 PTODSL signature 来自 device module，但 host 参数取自 scheduled module 的 primary PrimFunc。三个动态/特化 kernel 都稳定表现为 device 比 host 多一个参数，说明 host/device split 后有一个动态标量没有被一致保留或常量折叠。

修复建议：

1. 在错误信息里同时 dump device 参数名、host 参数名和原始 PrimFunc 参数，先确定多出的参数是运行时 scalar 还是 specialization-only scalar。
2. 如果是运行时参数，应在 host ABI 中保留并传给 launch stub。
3. 如果是编译期特化参数，应在 device PrimFunc/生成 PTODSL 前完成替换和 DCE，不应留在 PTODSL 函数签名里。
4. wrapper 不应只比较参数数量，建议基于 canonical parameter mapping 生成 prototype/call args，并增加动态 shape/scalar 的单测。

B. nd2nz=True 已删除

最新版 T.copy() 已没有 nd2nz 关键字。当前 asc 分支的设计是：普通 UB→L1 copy 在 LayoutInference 将目标 L1 标记为 NZ 后，由 InsertNd2Nz pass 自动重写为 ascend_nd2nz_scatter + ascend_nd2nz_post_copy。

修复建议：

• 将 tile-kernels 从显式 nd2nz=True 迁移到新版 layout-driven contract。
• 不能只机械删除参数；需要确认目标 L1 buffer 的 NZ layout 能被 LayoutInference 正确推导，并检查 InsertNd2Nz 后的 TIR。
• 给原来的四个 GEMM/norm kernel 增加 pass dump 和数值回归，避免“能编译但 L1 实际仍是 ND”的静默错误。
• 如果需要兼容旧 kernel，可在 T.copy 前端临时接收 nd2nz，并转换成 annotation/layout hint，但不建议长期保留两套语义。

C. T.Kernel(..., threads=...) 迁移

修复建议：

with T.Kernel(blocks) as bx:
    with T.SimtVF(threads=n_thr):
        # 原 SIMT body
        ...

迁移时要保持 fragment allocation、T.Parallel 和 reduction 都位于 SimtVF lexical scope 内，并分别验证 AscendC/PTO 两条路径。

D. SIMT local bfloat16 被 TileLang codegen 硬拦截

当前 codegen_pto.cc 的 EmitPtoBufferAllocation 中存在：

ICHECK(IsFloat32(buffer->dtype))
    << "PTO SIMT local allocation currently supports float32 only, got "
    << buffer->dtype;
stream << vid << " = pto.alloc_buffer((" << size << ",), pto.f32)\n";

修复建议：

• 先用独立 PTODSL 最小用例确认 pto.alloc_buffer(..., pto.bf16) 的支持情况。
• 如果 PTODSL 已支持，TileLang 应使用 PtoTypeName(buffer->dtype) 发出真实 dtype，而不是固定 pto.f32。
• 如果 PTODSL 暂不支持，可在 TileLang pass 中显式将 local fragment 提升为 float32，并在 load/store 边界插入 bf16↔f32 cast；不要在 codegen 阶段无上下文地直接改 dtype。

E. sinkhorn reduction 映射不完整

当前 EmitPTOAllReduceExpr_ 只接受 tl::AscendAllReducetl::SumOp,...，其他 reduction 直接 ICHECK 失败。

修复建议：

• 先 dump sinkhorn 的实际 reduction op，区分 max 与 sum（稳定 softmax/sinkhorn 通常两者都会出现）。
• 如果 PTODSL 有对应 allreduce max/min primitive，按 reducer type 分派到对应 PTODSL API。
• 如果 PTODSL 暂无对应 primitive，需要用 SIMT scratch + barrier/tree reduction 实现，或在进入 PTO codegen 前把该 reduction lower 成 PTODSL 已支持的基础操作。
• 增加 sum/max 各自的非 2 次幂线程数、thread_offset 和 scratch 回归。

5. PTOAS 独立验证

同一套最新版 PTOAS 环境下，直接编译 PTOAS 自带的 A5 V-PTO 用例成功：

ptoas --pto-arch=a5 --pto-backend=vpto \
  /home/wangmiao/PTOAS-main-latest/test/lit/pto/subview_tmatmul_acc_slices_a5.pto \
  -o /tmp/ptoas_latest_smoke_vpto_cann.fatobj.o

输出为有效的 AArch64 ELF relocatable object。因此当前证据支持：

• PTOAS CLI / V-PTO 基础链路正常；
• 当前 0/11 的首要修复位置在 WenboCodes/tilelang-deepseek asc 分支及 tile-kernels API 迁移；
• 修复上述前置问题、真正生成 kernel.pto 后，才能继续判断 PTODSL/PTOAS 是否还有下一阶段错误。

完整重测记录：

• /home/wangmiao/codex_runs/tilelang_asc_ptoas_latest_20260715_1216/pto_explicit_compile_first_run.json
• /home/wangmiao/codex_runs/tilelang_asc_ptoas_latest_20260715_1216/pto_explicit_compile_first_run.log