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

[zhangj1an]

Latest Status [1 Sep 2026]

Ready for review.

Summary

Port of the CUDA deterministic fused logp (csrc/deterministic_logp_kernel.cu + DeterministicLogpCUDAOp) to Ascend NPU:

• Forward kernel (_C_npu.fused_logp_ascend): mirrors the CUDA deterministic kernel's math — logp[n] = logits[n, target[n]] - logsumexp(logits[n, :]) — with the same two-pass fixed-order reduction: row max over a fixed tile order, then sum(exp(x - max)) over the same fixed tile order, lse = max + log(sum), logp = selected - lse, all in fp32. Out-of-range targets produce 0.0, matching the CUDA kernel. Output is fp32, matching DeterministicLogpCUDAOp's contract.
  • Batch-invariance: every row is processed end-to-end by exactly one AI-core block with a fixed tile size (4096 elements) and a fixed reduction order; rows are strided across blocks (MAX_BLOCKS=128). The instruction sequence for a row depends only on V, never on N or block assignment, so a row's logp is bitwise identical across batch sizes, row positions, and block assignments on the NPU (verified).
  • Honest numerics note: the fp32 accumulation and formula match the CUDA kernel exactly, but the hardware reduction trees and transcendental implementations are the Ascend vector unit's own (fixed per V) — cross-platform bitwise parity with the CUDA kernel is not claimed (CUDA expf is a software polynomial; the Ascend Exp is a hardware instruction). The guarantee is the same one the CUDA kernel provides on its platform: batch-invariant determinism. Measured drift vs the fp32 gold reference is ~1e-7.
• Wrapper: FusedLogpAscendOp mirrors the CUDA _FusedLogpAutograd bridge — lead-shape support ([..., V] -> [...]), row-local fp32 VJP backward (dlogits = grad * (one_hot(target) - softmax)), native fallback for non-NPU/non-contiguous inputs. Registered as the "ascend" candidate in the logp gtest spec and as ASCEND_FUSED_LOGP in the registry's NPU priority map ([ASCEND_FUSED_LOGP, PYTORCH_NATIVE]).
• Build: csrc/ascend/npu_module.cpp consolidates the single PYBIND11_MODULE (logp + fused logp); 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 (durable fresh-shell fix). scripts/check_operator.py gains --device npu support.

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_logp_ascend.asc	Ascend C two-pass fused logp forward kernel (fp32/bf16/fp16 inputs, fp32 output) + torch host wrapper. New.
csrc/ascend/npu_module.cpp	Aggregated _C_npu pybind registration (batch_invariant_logp + fused logp). 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/logp.py	FusedLogpAscendOp (Ascend C forward + row-local fp32 VJP backward). New.
rl_engine/kernels/ops/ascend/loss/__init__.py	Imports the new module.
rl_engine/_C_npu.pyi	Adds the fused_logp_ascend type stub.
rl_engine/kernels/gtest/operator_specs.py	Registers the "ascend" candidate for the logp op.
rl_engine/kernels/registry.py	ASCEND_FUSED_LOGP backend + NPU priority map entry.
rl_engine/tests/test_dispatch.py	Covers the NPU logp priority.
scripts/check_operator.py	--device npu / auto-detect support.
tests/test_logp_ascend.py	Ascend correctness (contract tolerance) + batch-invariance (bitwise) + registry dispatch. New.
docs/operators/fused-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 (single-op check, ascend candidate vs the PyTorch gold)
python scripts/check_operator.py --op logp --candidate ascend --device npu \
    --dtype {fp32,bf16,fp16} --batch 2 --seq 16 --vocab 257 --check-grad

# pytest suite (correctness tolerance, batch invariance bitwise, registry dispatch)
python -m pytest tests/test_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
gtest logp ascend candidate, fp32 × 2x16x257, output + gradient	✅ output max_abs=9.5e-7, grad max_abs=1.2e-7 (tol atol=1e-5)
gtest logp ascend candidate, bf16 × 2x16x257, output + gradient	✅ output max_abs=4.8e-7, grad max_abs=7.6e-6 (tol atol=6e-2 / 5e-2)
gtest logp ascend candidate, fp16 × 2x16x257, output + gradient	✅ output max_abs=4.8e-7, grad max_abs=1.9e-6 (tol atol=5e-3)
Batch invariance (bitwise): batch 1 vs {2, 4, 16, 300}, row positions 1..7, multi-tile V=10000, repeated runs	✅ 0 mismatches (torch.equal)
Out-of-range target → 0.0, empty input, fp32 output dtype	✅
pytest tests/test_logp_ascend.py	✅ 22 passed
Regression tests/test_batch_invariant_logp.py	✅ 44 passed, 42 skipped
Regression rl_engine/tests/test_dispatch.py	✅ 13 passed

gtest raw output (fp32, representative)

suite=logp passed=True pass_rate=1.0000
  case=logp-torch.float32-2x16x257 output=0 shape=(2, 16) dtype=torch.float32 max_abs=9.53674316e-07 mean_abs=2.23517418e-07 max_rel=1.29912607e-07 tol=(atol=1.000e-05, rtol=0.000e+00) passed=True
  case=logp-torch.float32-2x16x257 output=1 gradient:logits shape=(2, 16, 257) dtype=torch.float32 max_abs=1.19209290e-07 mean_abs=4.15880497e-10 max_rel=6.74039995e-07 tol=(atol=1.000e-05, rtol=0.000e+00) passed=True

pytest: tests/test_logp_ascend.py (summary)

======================= 22 passed, 14 warnings in 5.43s ========================

Notes

• Why tolerance vs the gold, bitwise only for invariance: the CUDA deterministic kernel's reduction tree (strided-thread partials + warp shuffle tree) is CUDA-specific, and the transcendental implementations differ per platform (expf polynomial vs the Ascend vector Exp instruction). The Ascend kernel mirrors the formula and the two-pass fixed-order structure; on the NPU its own fixed tree makes the op batch-invariant bitwise, which is the property the WS1 contract's forward_invariance row requires (bitwise) — the forward_accuracy row is tolerance-based (fp32 drift ~1e-7, well inside atol=1e-5).
• The backward is the same row-local VJP formula as the CUDA wrapper's _FusedLogpAutograd, computed in fp32 and cast back to the input dtype; no cross-token reduction, so gradients are batch-layout independent.
• Single-device operator; no TP/vocab-parallel path (consistent with the CUDA deterministic op's scope).
• 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: a38462e3-a953-428f-b605-bbfdd5383578

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