Merge blocked
Hi @jins666, welcome to submitting your first PR to CANN!
PR Merge Steps
1. CLA Signing
If the current PR label includes cann-cla/yes, it means you have signed the CLA and can proceed to the next step. If the label includes cann-cla/no, please sign the CLA first. If you have any questions, please refer to the FAQ.
2. CI Check
Please comment /compile to trigger the CI pipeline check. If the CI run is successful, the PR will be tagged with ci-pipeline-passed and you can proceed to the next step. If the CI run fails, the PR will be tagged with ci-pipeline-failed, please check the CI logs to fix the issues in the PR. If you have any questions, please refer to the FAQ.
3. Code Review
After CI passes, please refer to the PR Approval Progress and proactively @ the committers in the table to review the code. After approval, committers will comment /lgtm and /approve. Once the lgtm and approved labels are successfully added, the PR will be merged automatically.
%!(EXTRA string=https://atomgit.com/cann/community/blob/master/en/sig-infrastructure/command.md)
Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
For more, you also can visit HICANN.
PR Approval Progress
⚠️ This PR does not yet meet the following requirements:lgtm (requires ≥ 2 person(s) per module)、approve (requires ≥ 1 person(s) per module)
Module Approval Details
| module | lgtm status | approve status |
|---|---|---|
| repo-cann/ops-tilelang | ❌ (0/2)(You can also ask: jiawei_zhuang, kevin365, ChaoyangJi, 蒋礼锐, 刘丹) | ❌ (0/1)(You can also ask: ChaoyangJi, 马兵, 刘丹, kevin365, 蒋礼锐) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)
CLA Signature Pass
jins666, thanks for your pull request. All authors of the commits have signed the CLA. 👍
| 🚀 CI 流水线已启动 |
|---|
| 📋 执行详情: 点击查看流水线 |
功能
新增
cann_ops_tilelang.hisa_indexer，为 prefill 阶段的每个 query 选取可见 KV token。支持单序列和打包多序列，输出相对于各行可见窗口起点的 INT32 索引。Q/K 为 FP8 E4M3FN，反量化系数、头权重和打分累加为 FP32。粗排将块分数舍入到 BF16，优先保留可见首尾块。精排使用 FP32 分数，同分时取较小的候选槽位，无效槽位填
-1。8192 个候选选 2048 个 token 时使用 radix-4 专用内核。公开入口支持 eager 和单流、固定形状的
torch.compile推理。图模式注册为cann_ops_tilelang::hisa_indexer，沿用相同的签名和选择规则。最终索引由 PyTorch 的 NPU 分配器创建，TileLang 内核直接写入。接口与调用示例验证
npugraph_ex捕获、原地址换值和新地址换输入重放，逐次检查后端事件与参考索引。另在默认后端用完整 8192 单序列形状验证了三次图输出与 eager 的索引集合一致。用例还覆盖 schema/fake、空张量和与上下游 Torch 操作的组合。环境与复现
依赖 CANN、PyTorch、torch_npu 和 Ascend TileLang。测试覆盖默认 AscendC 后端及配套的 PTO 工具链。
OPS_TILELANG_TEST_LEVEL=2 TILELANG_DEFAULT_TARGET=auto python -m pytest --log-cli-level=INFO tests/attention/test_dsa_hisa.py -k hisa_compile OPS_TILELANG_TEST_LEVEL=2 TILELANG_DEFAULT_TARGET=auto python -m pytest tests/attention/test_dsa_hisa.py -m 'not benchmark' # 在配套的 PTO 环境中运行 OPS_TILELANG_TEST_LEVEL=2 TILELANG_DEFAULT_TARGET=pto python -m pytest tests/attention/test_dsa_hisa.py -m 'not benchmark' TILELANG_DEFAULT_TARGET=auto python -m pytest tests/attention/test_dsa_hisa.py -m benchmark --run-benchmark \ --benchmark-regression-threshold=0.10 --benchmark-min-delta-us=1.0本 PR 使用了 AI 辅助开发和验证。