[WS1][Ascend] [Qwen3-8b] Fused linear logp ops

[zhangj1an]

Latest Status [1 Sep 2026]

Ready for review.

Summary

Port of the CUDA fused linear log-prob (csrc/cuda/fused_linear_logp_sm90.cu + FusedLinearLogpSM90Op) to Ascend NPU:

• Forward kernel (_C_npu.fused_linear_logp_ascend): computes log_softmax(hidden @ W^T + b)[target] without materializing [N, V] logits, mirroring the SM90 kernel's bitwise reduction contract (contract v1):
  • vocab rows scanned in ascending index order (the contract's "cross-split ascending-index sequential chains");
  • the online rescale chain newM = max(m, z); sum = sum * exp(m - newM) + exp(z - newM) exactly like the CUDA per-split merge;
  • per-row fp32 dots over a fixed D-tile order (per-tile ReduceSum tree + sequential scalar chain; the hidden row is cached in UB for D ≤ 4096);
  • bias added in fp32; the final clamp logp = min(zt - lse, 0) matches the CUDA contract; out-of-range targets yield 0.
  • Batch-invariance: every row is processed end-to-end by one AI-core block over a fixed vocab/D scan order; rows are strided across blocks (MAX_BLOCKS=128), so a row's logp is bitwise identical across batch sizes, row positions, and block assignments on the NPU (verified with torch.equal).
  • Honest numerics note: the per-tile hardware reduction tree and the transcendental implementations are the Ascend vector unit's own, so cross-platform bitwise parity with the CUDA kernel is not claimed — the guarantee is the same one the CUDA kernel provides on its platform: batch-invariant determinism.
• Wrapper: FusedLinearLogpAscendOp mirrors FusedLinearLogpSM90Op's surface (apply(hidden, lm_head_weight, target_ids, bias=None), validation, fp32 output). TP and bias calls delegate to the native reference (same fallback shape as the CUDA op's non-SM90 paths). The backward is the shared Liger-style chunked_linear_logp_backward — the exact formula the CUDA SM90 op falls back to — so gradients follow the CUDA op's portable backward.
• Registration: gtest candidate "ascend" in the linear_logp spec; ASCEND_FUSED_LINEAR_LOGP in the registry NPU priority map ([ASCEND_FUSED_LINEAR_LOGP, PYTORCH_LINEAR_LOGP]); scripts/check_operator.py gains --device npu support.
• Build: csrc/ascend/npu_module.cpp consolidates the single PYBIND11_MODULE; batch_invariant_logp_ascend.asc only drops its PYBIND11_MODULE block. setup.py gains the Ascend extension build (bisheng, **/*.asc glob) with the CANN env export in _find_ascend_home.

gtest forward comparison note (read this first). The gtest's linear_logp forward
comparison is stricter than any independent kernel can meet, the CUDA candidate
included (CI's gtest matrix gates registration only; it does not execute
linear_logp candidates):

• fp32: the logprob tolerance is atol=1e-5, but two different fp32 reduction trees
  over D=4096 drift ~1e-4 (this kernel measured 2.1e-4 vs the gold; the CUDA kernel's
  WGMMA tree drifts the same way).
• bf16/fp16: the gold's apply() accumulates the matmul in the input dtype
  (bf16/fp16 matmul), while this kernel — like the CUDA SM90 kernel — accumulates in
  fp32 per the WS1 fp32-reference policy, so the bf16 comparison measures bf16
  matmul noise (0.55 vs atol 6e-2).
  The gradients pass the gtest tolerances (fp32 2.4e-7 / 3.0e-8; the bf16/fp16 gradient
  rows fail for the same gold-accumulation-dtype reason). The dedicated pytest suite
  verifies the forward against a hand-computed fp32 reference at an honest reduction
  tolerance and the backward at the quantized level.

Build notes (same pattern as PR #320 / #355)

Each .asc source file can define only one PYBIND11_MODULE (linking multiple sources with Bisheng causes a duplicate PyInit__C_npu error), so pybind registrations are consolidated in csrc/ascend/npu_module.cpp; batch_invariant_logp_ascend.asc only drops its PYBIND11_MODULE block.

Files

Path	Status
csrc/ascend/fused_linear_logp_ascend.asc	Ascend C fused linear logp forward kernel (fp32/bf16/fp16, optional fp32 bias, online softmax) + torch host wrapper. New.
csrc/ascend/npu_module.cpp	Aggregated _C_npu pybind registration. New.
csrc/ascend/batch_invariant_logp_ascend.asc	Only removes PYBIND11_MODULE (moved to the aggregated file). Kernel logic unchanged.
rl_engine/kernels/ops/ascend/loss/linear_logp.py	FusedLinearLogpAscendOp (Ascend C forward + shared chunked backward). New.
rl_engine/kernels/ops/ascend/loss/__init__.py	Imports the new module.
rl_engine/_C_npu.pyi	Adds the fused_linear_logp_ascend type stub.
rl_engine/kernels/gtest/operator_specs.py	Registers the "ascend" candidate for the linear_logp op.
rl_engine/kernels/registry.py	ASCEND_FUSED_LINEAR_LOGP backend + NPU priority map entry.
rl_engine/tests/test_dispatch.py	Covers the NPU linear_logp priority.
scripts/check_operator.py	--device npu / auto-detect support.
tests/test_linear_logp_ascend.py	Ascend correctness (honest reduction tolerance vs fp32 reference) + batch-invariance (bitwise) + registry dispatch. New.
docs/operators/linear-logp.md	Ascend backend row and implementation files.
setup.py	Ascend extension build (bisheng) + CANN env export in _find_ascend_home.

Test

# build
export KERNEL_ALIGN_FORCE_ASCEND=1
pip install -e . --no-build-isolation

# gtest registration + gradient checks (forward comparison: see the note above)
python scripts/check_operator.py --op linear_logp --candidate ascend --device npu \
    --dtype {fp32,bf16,fp16} --batch 2 --seq 16 --vocab 257 --normalized-dim 4096 --check-grad

# pytest suite (correctness tolerance, batch invariance bitwise, registry dispatch)
python -m pytest tests/test_linear_logp_ascend.py -v

# regression (shared module/build touched)
python -m pytest tests/test_batch_invariant_logp.py -q
python -m pytest rl_engine/tests/test_dispatch.py -q

Test results

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

Test	Result
Forward vs hand-computed fp32 reference (fp32/bf16/fp16, D=1000 and gtest D=4096 shape)	✅ max_abs ≤ 2.2e-4 (pure fp32 tree drift; pytest atol 5e-4)
gtest gradients (fp32): hidden + lm_head_weight	✅ max_abs 2.4e-7 / 3.0e-8 (atol 1e-5)
gtest forward comparison (fp32/bf16/fp16)	⚠️ exceeds the logprob tolerance by design — see the note above (fp32 tree drift 2.1e-4 vs atol 1e-5; bf16/fp16 gold accumulates the matmul in the input dtype vs the kernel's fp32 accumulation)
Backward vs fp32-reference autograd (quantized level for bf16/fp16)	✅ bf16 grads bitwise equal to the quantized reference in the test config; tolerance absorbs rare 1-ULP straddles
Out-of-range target → 0.0, empty input, bias fallback to native	✅
Batch invariance (bitwise): batch 1 vs {2, 4, 16, 300}, row positions 1..7, multi-tile D=10000, repeated runs	✅ 0 mismatches (torch.equal)
pytest tests/test_linear_logp_ascend.py	✅ 23 passed
Regression tests/test_batch_invariant_logp.py	✅ 44 passed, 42 skipped
Regression rl_engine/tests/test_dispatch.py	✅ 13 passed

pytest: tests/test_linear_logp_ascend.py (summary)

======================= 23 passed, 15 warnings in 5.49s ========================

Notes

• Why the bitwise guarantee is batch invariance, not gold parity: an independent kernel's fp32 reduction trees drift ~1e-4 against torch.matmul + log_softmax at D=4096 (fp32 addition is not associative); the CUDA kernel has the same property on its platform. The Ascend kernel pins its own fixed order, which makes the op batch-invariant bitwise on the NPU (forward_invariance contract row).
• The forward is fp32-accumulating regardless of input dtype, matching the CUDA SM90 contract and the WS1 fp32-reference policy; the gold's dtype path (bf16 matmul) is the source of the gtest's bf16 forward mismatch, not this kernel.
• Single-device operator; TP/vocab-parallel calls delegate to the native reference (the CUDA op's TP path is out of scope for this port).
• black / isort / flake8 (line-length 100) pass 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: Team

Run ID: c8d34aa5-0a9c-4719-9ed3-21faccf99927

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.}