[WS1][kernels] Deterministic attention Ascend C kernel

[zhangj1an]

Latest Status [19 Aug 2026]

Ready for review.

Summary

Following the previous deterministic attention implementation (CUDA issue #147 DeterministicAttentionOp + Triton reference TritonBatchInvariantAttentionOp, i.e. rlkernel.attention.deterministic_core.v1), this PR implements an Ascend NPU version:

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

  • Each (b, q_head, row) is processed end-to-end by a single AI Core block, using two streaming reduction passes in a fixed 64-key tile order:
    • pass 1 computes the row max;
    • pass 2 computes sum-exp and P·V.
  • No split-K and no secondary (m, l, u) summary merge. Therefore, the reduction tree for each row depends only on Skv / D / mask, and is independent of batch size or block scheduling/allocation. This makes it batch-invariant, corresponding to the num_splits=1 approach described in the blog.
  • All intermediate computations use FP32. Inputs may be bf16/fp16. Supports D=128, GQA, causal masking, and key_padding_mask. The output dtype matches the input dtype and an FP32 LSE is returned. Fully masked rows produce output 0 and lse = -inf, consistent with the guard in the Triton reference.

• Backward uses the VJP of NativeAttentionOp.forward_fp32 (structurally identical to the gtest gold implementation, with zero gradient error), reusing q/k/v/mask saved during forward. Triton is unavailable on NPU, so this follows the same portable-backward idea as the Triton op.

• Integrated using the .asc extension pattern from PR [draft][WS1][kernels] Batch-invariant logprob Ascend C Kernel #297:

  • Added ASCEND_DETERMINISTIC_ATTENTION to the registry.
  • Added "npu" dispatch for attention / attn / kv_cache_attention.
  • Added an ascend candidate to the gtest spec.
  • Added unit tests, benchmarks, documentation, and a _C_npu.pyi stub.

Build notes (same pattern as PR #297)

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 for all operators are consolidated in csrc/ascend/ops_npu.asc, while individual kernel files contain only the kernel and host forward functions.

The .asc glob in setup.py was also changed to recursive (**/*.asc), allowing kernels to follow the CUDA-style directory structure under csrc/ascend/attention/.

Files

Path	Status
csrc/ascend/attention/deterministic_attention_ascend.asc	Ascend C forward kernel (bf16/fp16, two passes with a fixed tile order, one block per row). New.
csrc/ascend/ops_npu.asc	Aggregated _C_npu pybind registration file (logp + 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/deterministic_attn.py	Autograd wrapper: Ascend C forward + FP32 reference backward. New.
rl_engine/kernels/registry.py	Adds ASCEND_DETERMINISTIC_ATTENTION; adds attention / attn / kv_cache_attention dispatch for the "npu" platform.
rl_engine/kernels/gtest/operator_specs.py	Adds the ascend candidate for attention.
tests/test_attention_ascend.py	Ascend unit tests for correctness / batch invariance / dispatch. New.
tests/test_attention.py	Registry dispatch tests now accept the Ascend deterministic op.
benchmarks/benchmark_attention.py	Attention benchmark comparing native vs Ascend vs CUDA. New.
docs/operators/attention.md	Adds the Ascend backend to the Backends table and Dispatch Behavior documentation.
rl_engine/_C_npu.pyi	Adds the deterministic_attention_ascend stub.
setup.py	Changes the .asc glob to **/*.asc.

Test

gtest

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

python scripts/check_operator.py --op attention --candidate ascend --device npu --dtype bf16 --batch 2 --seq 64
python scripts/check_operator.py --op attention --candidate ascend --device npu --dtype fp16 --batch 2 --seq 64
python scripts/check_operator.py --op attention --candidate ascend --device npu --dtype fp16 --batch 2 --seq 32 --check-grad
python scripts/check_operator.py --op attention --candidate ascend --device npu --dtype bf16 --batch 2 --seq 32 --check-grad

pytest

python -m pytest tests/test_attention_ascend.py -v

Test results

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

Test	Result
gtest attention ascend bf16 (forward)	✅ suite=attention passed=True pass_rate=1.0000
gtest attention ascend fp16 (forward)	✅ suite=attention passed=True pass_rate=1.0000
gtest attention ascend fp16 (forward + grad)	✅ suite=attention passed=True pass_rate=1.0000 (q/k/v gradient errors are all 0)
gtest attention ascend bf16 (forward + grad)	✅ suite=attention passed=True pass_rate=1.0000 (q/k/v gradient errors are all 0)
gtest batch_invariant_logp ascend bf16 (regression)	✅ suite=batch_invariant_logp passed=True pass_rate=1.0000
pytest tests/test_attention_ascend.py	✅ 26 passed

gtest: attention ascend bf16 (raw)

INFO 08-19 20:56:56 [RL-Kernel]: Successfully linked to precompiled _C_npu.deterministic_attention_ascend kernel.
suite=attention passed=True pass_rate=1.0000
candidate=ascend-attention backend=ascend passed=True pass_rate=1.0000
  case=attention-torch.bfloat16-2x32x64x128 output=0 shape=(2, 32, 64, 128) dtype=torch.bfloat16 max_abs=7.79938698e-03 mean_abs=3.43122898e-04 max_rel=1.21453263e-01 tol=(atol=5.000e-02, rtol=2.000e-02) passed=True

gtest: attention ascend fp16 (raw)

INFO 08-19 20:57:05 [RL-Kernel]: Successfully linked to precompiled _C_npu.deterministic_attention_ascend kernel.
suite=attention passed=True pass_rate=1.0000
candidate=ascend-attention backend=ascend passed=True pass_rate=1.0000
  case=attention-torch.float16-2x32x64x128 output=0 shape=(2, 32, 64, 128) dtype=torch.float16 max_abs=9.73939896e-04 mean_abs=4.28804415e-05 max_rel=5.62257826e-01 tol=(atol=1.000e-03, rtol=1.000e-03) passed=True

gtest: attention ascend fp16 --check-grad (raw)

INFO 08-19 20:57:14 [RL-Kernel]: Successfully linked to precompiled _C_npu.deterministic_attention_ascend kernel.
suite=attention passed=True pass_rate=1.0000
candidate=ascend-attention backend=ascend passed=True pass_rate=1.0000
  case=attention-torch.float16-2x32x32x128 output=0 shape=(2, 32, 32, 128) dtype=torch.float16 max_abs=9.73939896e-04 mean_abs=5.40761685e-05 max_rel=4.29534391e-02 tol=(atol=1.000e-03, rtol=1.000e-03) passed=True
  case=attention-torch.float16-2x32x32x128 output=1 gradient:q shape=(2, 32, 32, 128) dtype=torch.float16 max_abs=0.00000000e+00 mean_abs=0.00000000e+00 max_rel=0.00000000e+00 tol=(atol=1.000e-03, rtol=1.000e-03) passed=True
  case=attention-torch.float16-2x32x32x128 output=2 gradient:k shape=(2, 8, 32, 128) dtype=torch.float16 max_abs=0.00000000e+00 mean_abs=0.00000000e+00 max_rel=0.00000000e+00 tol=(atol=1.000e-03, rtol=1.000e-03) passed=True
  case=attention-torch.float16-2x32x32x128 output=3 gradient:v shape=(2, 8, 32, 128) dtype=torch.float16 max_abs=0.00000000e+00 mean_abs=0.00000000e+00 max_rel=0.00000000e+00 tol=(atol=1.000e-03, rtol=1.000e-03) passed=True

gtest: attention ascend bf16 --check-grad (raw)

INFO 08-19 20:57:23 [RL-Kernel]: Successfully linked to precompiled _C_npu.deterministic_attention_ascend kernel.
suite=attention passed=True pass_rate=1.0000
candidate=ascend-attention backend=ascend passed=True pass_rate=1.0000
  case=attention-torch.bfloat16-2x32x32x128 output=0 shape=(2, 32, 32, 128) dtype=torch.bfloat16 max_abs=7.79938698e-03 mean_abs=4.31581051e-04 max_rel=4.02779225e-03 tol=(atol=5.000e-02, rtol=2.000e-02) passed=True
  case=attention-torch.bfloat16-2x32x32x128 output=1 gradient:q shape=(2, 32, 32, 128) dtype=torch.bfloat16 max_abs=0.00000000e+00 mean_abs=0.00000000e+00 max_rel=0.00000000e+00 tol=(atol=5.000e-02, rtol=2.000e-02) passed=True
  case=attention-torch.bfloat16-2x32x32x128 output=2 gradient:k shape=(2, 8, 32, 128) dtype=torch.bfloat16 max_abs=0.00000000e+00 mean_abs=0.00000000e+00 max_rel=0.00000000e+00 tol=(atol=5.000e-02, rtol=2.000e-02) passed=True
  case=attention-torch.bfloat16-2x32x32x128 output=3 gradient:v shape=(2, 8, 32, 128) dtype=torch.bfloat16 max_abs=0.00000000e+00 mean_abs=0.00000000e+00 max_rel=0.00000000e+00 tol=(atol=5.000e-02, rtol=2.000e-02) passed=True

gtest: batch_invariant_logp ascend bf16 regression (raw)

INFO 08-19 20:57:32 [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 ../torch_npu/csrc/aten/common/TensorFactories.cpp:340.)
  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_attention_ascend.py (raw)

============================= test session starts ==============================
platform linux -- Python 3.10.14, pytest-9.1.1, pluggy-1.6.0 -- /home/ma-user/work/z84450661/venv310/bin/python
cachedir: .pytest_cache
rootdir: /home/ma-user/work/z84450661/rl-kernel
configfile: pyproject.toml
plugins: anyio-4.14.1
collecting ... collected 26 items

tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_prefill_causal[dtype0] PASSED [  3%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_prefill_causal[dtype1] PASSED [  7%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_gqa[dtype0] PASSED [ 11%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_gqa[dtype1] PASSED [ 15%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_decode_window[dtype0] PASSED [ 19%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_decode_window[dtype1] PASSED [ 23%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_non_causal[dtype0] PASSED [ 26%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_non_causal[dtype1] PASSED [ 30%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_key_padding_mask[dtype0] PASSED [ 34%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_key_padding_mask[dtype1] PASSED [ 38%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_fully_masked_row_is_zero[dtype0] PASSED [ 42%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_fully_masked_row_is_zero[dtype1] PASSED [ 46%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_explicit_scale[dtype0] PASSED [ 50%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_explicit_scale[dtype1] PASSED [ 53%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_forward_with_lse[dtype0] PASSED [ 57%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_forward_with_lse[dtype1] PASSED [ 61%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_backward_grads[dtype0] PASSED [ 65%]
tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_backward_grads[dtype1] PASSED [ 69%]
tests/test_attention_ascend.py::TestAscendAttentionRejects::test_rejects_fp32 PASSED [ 73%]
tests/test_attention_ascend.py::TestAscendAttentionRejects::test_rejects_bad_head_dim PASSED [ 76%]
tests/test_attention_ascend.py::TestAscendAttentionBatchInvariance::test_batch_size_1_vs_n PASSED [ 80%]
tests/test_attention_ascend.py::TestAscendAttentionBatchInvariance::test_different_positions_in_batch PASSED [ 84%]
tests/test_attention_ascend.py::TestAscendAttentionBatchInvariance::test_block_striding PASSED [ 88%]
tests/test_attention_ascend.py::TestAscendAttentionBatchInvariance::test_repeated_runs_deterministic PASSED [ 92%]
tests/test_attention_ascend.py::TestAscendRegistryDispatch::test_get_op_attention PASSED [ 96%]
tests/test_attention_ascend.py::TestAscendRegistryDispatch::test_get_op_attn_falls_back_to_sdpa PASSED [100%]

=============================== warnings summary ===============================
../venv310/lib/python3.10/site-packages/torch/jit/_script.py:362: 14 warnings
  /home/ma-user/work/z84450661/venv310/lib/python3.10/site-packages/torch/jit/_script.py:362: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(

tests/test_attention_ascend.py::TestAscendAttentionCorrectness::test_key_padding_mask[dtype0]
  /home/ma-user/work/z84450661/rl-kernel/rl_engine/kernels/ops/pytorch/attention/standard_attn.py:177: UserWarning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (Triggered internally at ../torch_npu/csrc/aten/common/TensorFactories.cpp:340.)
    probs = torch.where(all_masked, torch.zeros_like(probs), probs)

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
======================= 26 passed, 15 warnings in 5.33s ========================

benchmark: benchmarks/benchmark_attention.py (raw)

| shape (B x Hq x S)   |   native fwd ms |   ascend fwd ms | vs native   | cuda fwd ms   |
|----------------------|-----------------|-----------------|-------------|---------------|
| 1x32x512             |           0.278 |          48.757 | 0.01x       | -             |
| 1x32x1024            |           0.840 |         194.460 | 0.00x       | -             |

v1 is a correctness-first, pure vector-unit implementation that computes one dot product per key. It does not yet use the CUBE matrix-multiplication path for Q·Kᵀ. The performance gap is known, and CUBE-based optimization is left for future work.

Notes

• This is a single-device operator and does not involve CP/SP or other communication. Its scope is consistent with the CUDA implementation in issue [WS1][kernels] Batch-invariant attention (standard softmax) #147.
• ruff check passes for all modified 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: Pro Plus

Run ID: 05e64d7a-b0cb-4910-8978-a15220bf57cb

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

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[erfgss]

I tested this PR on Ascend A3. Correctness looks good, but I observed a significant performance issue with the current Ascend attention kernel.

Correctness

Both BF16 and FP16 forward checks passed:

BF16, B=2, S=64:
passed=True, pass_rate=1.0000
max_abs=7.799e-03
FP16, B=2, S=64:

passed=True, pass_rate=1.0000

max_abs=9.739e-04

Forward + gradient checks also passed for both BF16 and FP16, with zero error for q/k/v gradients.

The full Ascend attention test suite also passed:

26 passed, 2 warnings in 5.49s

So I did not find a correctness issue in my environment.

Performance

Environment:

Device: Ascend A3

CANN: 9.0.0

I ran:

python benchmarks/benchmark_attention.py --dtype bf16

python benchmarks/benchmark_attention.py --dtype fp16

BF16 results:

Shape (B × Hq × S)	Native fwd	Ascend fwd
1×32×512	0.246 ms	48.752 ms
1×32×1024	0.840 ms	194.418 ms
1×32×2048	3.516 ms	776.835 ms
4×32×2048	13.298 ms	3107.190 ms

The Ascend kernel is around 200–260× slower than the native path in these test cases.

There is also a clear scaling pattern:

S=512   -> ~48.75 ms
S=1024  -> ~194.4 ms
S=2048  -> ~776.9 ms
B=1→4   -> ~776.9→3107 ms

The measured latency is approximately proportional to B × S².