feat(ascend): add prefix-shared attention Ascend C kernel

[zhangj1an]

Latest Status [25 Aug 2026]

Will fix CI error then mark as ready for review

Summary

Port of the CUDA prefix-shared fused attention (csrc/cuda/attention/prefix_shared_attention.cu, the GRPO decode workload) to an Ascend NPU version:

• Forward is implemented as an Ascend C (CANN) kernel, _C_npu.prefix_shared_attention_ascend:

  • Layout q [bs, G, len_q, 128], k/v [bs, len_kv, 128] — the G generated responses share one prompt-prefix KV sequence, stored once per batch instead of once per group (same surface as the CUDA PrefixSharedAttentionOp: bf16 only, D=128, non-causal, no key-padding mask, forward-only).
  • Each (bs, g, 64-row query block) is processed end-to-end by one AI-core block, streaming the keys in a fixed 64-key tile order with fp32 online-softmax accumulation (per-row max / sum-exp rescaling per tile — the same flash-style single pass as the CUDA kernel).
  • No split-K and no cross-block summary merge, so the per-row reduction order depends only on len_kv — outputs are batch-invariant: bitwise identical regardless of batch size, batch position, or how many blocks were launched (items are strided across blocks).
  • UB budget ~146 KB < 192 KB. Correctness-first pure vector-unit implementation (one dot product per key); the CUBE Q·Kᵀ path is left for future work, same as the deterministic attention kernel in PR [WS1][kernels] Deterministic attention Ascend C kernel #320.

• Integration:

  • gtest spec prefix_shared_attention (op class attention, reusing its tolerance contract) with pytorch / cuda / ascend candidates and a GtestPrefixSharedAttentionOp gold that expands the shared K/V over G and reuses NativeAttentionOp.forward_fp32 (non-causal, default scale).
  • scripts/check_operator.py learns --device npu (defensive torch_npu probe) so the ascend candidates run in the gtest harness on NPU hosts.
  • setup.py .asc glob changed to recursive **/*.asc, allowing kernels to follow the CUDA-style directory structure under csrc/ascend/attention/.

Build notes (same pattern as PR #320)

Each .asc source file can define only one PYBIND11_MODULE, because linking multiple sources with Bisheng causes a duplicate PyInit__C_npu error. Therefore pybind registrations are consolidated in csrc/ascend/ops_npu.asc (logp + prefix-shared attention), while individual kernel files contain only the kernel and host forward functions.

Files

Path	Status
csrc/ascend/attention/prefix_shared_attention_ascend.asc	Ascend C forward kernel (bf16, fp32 online softmax, fixed 64-key tile order, one block per (bs, g, query block)). New.
csrc/ascend/ops_npu.asc	Aggregated _C_npu pybind registration file (logp + prefix-shared attention). New.
csrc/ascend/batch_invariant_logp_ascend.asc	Only removes PYBIND11_MODULE, which is moved to the aggregated registration file. Kernel logic is unchanged.
rl_engine/kernels/ops/ascend/attention/prefix_shared_attn.py	PrefixSharedAttentionAscendOp wrapper, same surface as the CUDA PrefixSharedAttentionOp (op(q, k, v) -> out) with bf16/D=128 validation. New.
rl_engine/kernels/ops/ascend/attention/__init__.py	Package exports. New.
rl_engine/kernels/gtest/operator_specs.py	Adds the prefix_shared_attention spec (pytorch/cuda/ascend candidates) and the GtestPrefixSharedAttentionOp gold.
rl_engine/kernels/gtest/operator_inputs.py	Adds the prefix-shared input maker and shape name.
scripts/check_operator.py	--device npu support (defensive torch_npu probe).
tests/test_prefix_shared_attention_ascend.py	Ascend unit tests for correctness / batch invariance / validation. New.
tests/test_ws1_gtest_gpu.py	Adds prefix_shared_attention to the registered-ops coverage set.
setup.py	Changes the .asc glob to **/*.asc.

Test

Build first (same as PR #320):

source /usr/local/Ascend/ascend-toolkit/set_env.sh
export KERNEL_ALIGN_FORCE_ASCEND=1
pip install -e . --no-build-isolation --no-deps

gtest

python scripts/check_operator.py --op prefix_shared_attention --candidate ascend --device npu --dtype bf16 --batch 2 --seq 64
python scripts/check_operator.py --op prefix_shared_attention --candidate pytorch --device npu --dtype bf16 --batch 2 --seq 64
python scripts/check_operator.py --op batch_invariant_logp --candidate ascend --device npu --dtype bf16

pytest

python -m pytest tests/test_prefix_shared_attention_ascend.py -v
python -m pytest tests/test_batch_invariant_logp.py -v

Test results

Environment: 8× Ascend 910, CANN 8.5.1 (Bisheng), torch 2.7.1 + torch_npu 2.7.1.

Test	Result
gtest prefix_shared_attention ascend bf16 (forward)	✅ suite=prefix_shared_attention passed=True pass_rate=1.0000
gtest prefix_shared_attention pytorch bf16 (gold wiring)	✅ suite=prefix_shared_attention passed=True pass_rate=1.0000
gtest batch_invariant_logp ascend bf16 (regression)	✅ suite=batch_invariant_logp passed=True pass_rate=1.0000
pytest tests/test_prefix_shared_attention_ascend.py	✅ 14 passed
pytest tests/test_batch_invariant_logp.py (regression)	✅ 44 passed, 42 skipped (non-NPU backends)

gtest: prefix_shared_attention ascend bf16 (raw)

INFO 08-25 14:19:47 [RL-Kernel]: Successfully linked to precompiled _C_npu.prefix_shared_attention_ascend kernel.
suite=prefix_shared_attention passed=True pass_rate=1.0000
candidate=ascend-prefix_shared_attention backend=ascend passed=True pass_rate=1.0000
  case=prefix_shared_attention-torch.bfloat16-2x32x64x128 output=0 shape=(2, 32, 64, 128) dtype=torch.bfloat16 max_abs=3.89850140e-03 mean_abs=2.17124965e-04 max_rel=9.33131874e-02 tol=(atol=5.000e-02, rtol=2.000e-02) passed=True

gtest: prefix_shared_attention pytorch bf16 (raw)

suite=prefix_shared_attention passed=True pass_rate=1.0000
candidate=pytorch-prefix_shared_attention backend=pytorch passed=True pass_rate=1.0000
  case=prefix_shared_attention-torch.bfloat16-2x32x64x128 output=0 shape=(2, 32, 64, 128) dtype=torch.bfloat16 max_abs=1.66380405e-02 mean_abs=6.89244305e-04 max_rel=2.99679004e+03 tol=(atol=5.000e-02, rtol=2.000e-02) passed=True

gtest: batch_invariant_logp ascend bf16 regression (raw)

INFO 08-25 14:20:07 [RL-Kernel]: Successfully linked to precompiled _C_npu.batch_invariant_logp_ascend kernel.
/home/ma-user/work/z84450661/rl-kernel/rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py:105: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at build/CMakeFiles/torch_npu.dir/compiler_depend.ts:335.)
  selected_logp = selected_logp.where(valid_mask, torch.zeros_like(selected_logp))
suite=batch_invariant_logp passed=True pass_rate=1.0000
candidate=ascend-batch_invariant_logp backend=ascend passed=True pass_rate=1.0000
  case=batch_invariant_logp-torch.bfloat16-2x16x257 output=0 shape=(2, 16) dtype=torch.float32 max_abs=4.76837158e-07 mean_abs=1.49011612e-08 max_rel=7.54759455e-08 tol=(atol=6.000e-02, rtol=0.000e+00) passed=True

pytest: tests/test_prefix_shared_attention_ascend.py (raw)

============================= test session starts ==============================
platform linux -- Python 3.10.20, pytest-9.1.1, pluggy-1.6.0
rootdir: /home/ma-user/work/z84450661/rl-kernel
configfile: pyproject.toml
collecting ... collected 14 items

tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionCorrectness::test_basic PASSED [  7%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionCorrectness::test_exact_tiles PASSED [ 14%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionCorrectness::test_decode_window PASSED [ 21%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionCorrectness::test_long_prefix_multi_tile PASSED [ 28%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionCorrectness::test_shared_kv_across_groups PASSED [ 35%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionBatchInvariance::test_batch_size_1_vs_n PASSED [ 42%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionBatchInvariance::test_different_positions_in_batch PASSED [ 50%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionBatchInvariance::test_block_striding PASSED [ 57%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionBatchInvariance::test_repeated_runs_deterministic PASSED [ 64%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionValidation::test_rejects_fp32 PASSED [ 71%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionValidation::test_rejects_fp16 PASSED [ 78%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionValidation::test_rejects_bad_head_dim PASSED [ 85%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionValidation::test_rejects_4d_kv PASSED [ 92%]
tests/test_prefix_shared_attention_ascend.py::TestAscendPrefixSharedAttentionValidation::test_rejects_kv_length_mismatch PASSED [100%]

============================== 14 passed in 5.81s ==============================

pytest: tests/test_batch_invariant_logp.py regression (summary)

================== 44 passed, 42 skipped, 1 warning in 6.37s ===================

(all TestAscend* classes passed; skips are CUDA/Triton-only tests on this NPU host)

Notes

• Single-device operator; no CP/SP or other communication. Scope is consistent with the CUDA PrefixSharedAttentionOp.
• The op is not added to the kernel registry dispatch: the CUDA PrefixSharedAttentionOp is consumed directly (e.g. benchmarks/benchmark_attention.py) and is not registered either; the Ascend op follows the same surface.
• ruff check passes for all modified Python files.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: 2f5f6500-9c78-4326-95df-9befb8fc8c76

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}

[Flink-ddd]

LGTM