[WS1][Ascend] [Qwen3-8b] LM head ops

[zhangj1an]

Latest Status [1 Sep 2026]

Ready for review.

Summary

Port of the CUDA deterministic LM-head (csrc/cuda/embedding_lm_head_sm90.cu + SM90LMHeadOp) to Ascend NPU:

Bitwise consistency status (read this first). This op is not bitwise-consistent
with the PyTorch reference, and cannot be — the reference
(NativeLMHeadOp.forward_fp32) computes per-row torch.mv (a GEMV), and the internal
reduction order of torch_npu's torch.mv is a private implementation detail (like
cuBLAS's); fp32 addition is not associative, so any independent kernel with a
different summation order drifts by ULPs. Measured drift vs the reference at
2x16x4096x257: fp32 max_abs ≈ 2.4e-4 (contract atol/rtol = 1e-4); bf16/fp16 drift is
ULP-scale after the output cast (max_rel 3.9e-3 / 9.2e-4). The gtest comparison is
therefore tolerance-based per the reduction contract. What is bitwise (verified
with torch.equal): batch invariance on the NPU — a row's logits are bitwise
identical across batch sizes, row positions, and block assignments.

• Forward kernel (_C_npu.lm_head_ascend): mirrors the SM90 CUDA kernel's structure — out[n, v] = dot(hidden[n, :], weight[v, :]) (+ bias[v]) with one output element per block iteration, the full hidden-dimension reduction inside that block over a fixed tile order (products → per-tile sum → sequential scalar accumulation, all fp32), bias added in fp32, final cast to the output dtype with round-to-nearest. No Split-K, no algorithm selection.
  • Batch-invariance: elements are strided across blocks (MAX_BLOCKS=128); the instruction sequence for an element depends only on H, never on N or block assignment, so a row's logits are bitwise identical across batch sizes, row positions, and block assignments on the NPU (verified with torch.equal).
  • Honest numerics note: the per-tile sums use the Ascend vector unit's fixed hardware reduction tree instead of CUDA's warp-shuffle tree, and the PyTorch reference (torch.mv) has its own unspecified internal order, so bitwise parity with either is not claimed for this reduction op — the gtest comparison is tolerance-based per the reduction contract (fp32 drift measured ~2.4e-4 vs atol/rtol 1e-4/1e-4 at the 2x16x4096x257 shape... see Test results; bf16/fp16 drift is ULP-scale after the output cast).
• Wrapper: AscendLMHeadOp mirrors SM90LMHeadOp (forward / forward_fp32, bias support, dtype gate fp32/bf16/fp16, hidden[-1] == weight[1] check). The backward is the standard linear VJP computed in fp32 on the NPU (grad_hidden = grad @ W, grad_weight = grad^T @ H, bias = fixed-order row sum) then cast to the input dtypes — the CUDA op routes through _C.det_gemm, which has no NPU equivalent; gradients compare against the native backward at the gradient contract tolerance.
• Registration: gtest candidate "ascend" in operator_specs.py; ASCEND_LM_HEAD in the registry with the NPU priority [ASCEND_LM_HEAD, PYTORCH_NATIVE_LM_HEAD]; scripts/check_operator.py gains --device npu support.
• Build: csrc/ascend/npu_module.cpp consolidates the single PYBIND11_MODULE (batch_invariant_logp + lm_head); 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.

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/lm_head_ascend.asc	Ascend C LM-head forward kernel (fp32/bf16/fp16, optional fp32 bias) + torch host wrapper. New.
csrc/ascend/npu_module.cpp	Aggregated _C_npu pybind registration (batch_invariant_logp + lm_head). 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/linear/lm_head.py	AscendLMHeadOp (Ascend C forward + fp32-formula VJP backward). New.
rl_engine/kernels/ops/ascend/linear/__init__.py	Package init. New.
rl_engine/kernels/ops/ascend/__init__.py	Imports the new linear subpackage.
rl_engine/_C_npu.pyi	Adds the lm_head_ascend type stub.
rl_engine/kernels/gtest/operator_specs.py	Registers the "ascend" candidate for the lm_head op.
rl_engine/kernels/registry.py	ASCEND_LM_HEAD backend + NPU priority map entry.
rl_engine/tests/test_dispatch.py	Covers the NPU lm_head priority.
scripts/check_operator.py	--device npu / auto-detect support.
tests/test_lm_head_ascend.py	Ascend correctness (contract tolerance) + batch-invariance (bitwise) + registry dispatch. New.
docs/operators/lm_head.md	Ascend backend row, NPU dispatch behavior, tests 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 lm_head --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_lm_head_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 lm_head ascend candidate, fp32 × 2x16x4096x257, output + both gradients	✅ output max_abs=2.4e-4, grads ≤3.1e-5 (tol atol/rtol 1e-4)
gtest lm_head ascend candidate, bf16 × 2x16x4096x257, output + both gradients	✅ output max_abs=9.1e-1 (ULP-scale on
gtest lm_head ascend candidate, fp16 × 2x16x4096x257, output + both gradients	✅ output max_abs=1.1e-1 (ULP-scale), grads ≤1.6e-2 (tol atol 1e-3)
Bias path (fp32/bf16/fp16), lead shapes, empty input	✅
Batch invariance (bitwise): batch 1 vs {2, 4, 16, 300}, row positions 1..7, multi-tile H=10000, repeated runs	✅ 0 mismatches (torch.equal)
pytest tests/test_lm_head_ascend.py	✅ 23 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=lm_head passed=True pass_rate=1.0000
  case=lm_head-torch.float32-2x16x4096x257 output=0 shape=(2, 16, 257) dtype=torch.float32 max_abs=2.44140625e-04 mean_abs=2.66864972e-05 max_rel=6.10409654e-04 tol=(atol=1.000e-04, rtol=1.000e-04) passed=True
  case=lm_head-torch.float32-2x16x4096x257 output=1 gradient:hidden shape=(2, 16, 4096) dtype=torch.float32 max_abs=3.05175781e-05 mean_abs=1.56305839e-06 max_rel=4.10287666e+00 tol=(atol=1.000e-04, rtol=1.000e-04) passed=True
  case=lm_head-torch.float32-2x16x4096x257 output=2 gradient:weight shape=(257, 4096) dtype=torch.float32 max_abs=7.62939453e-06 mean_abs=4.67133901e-07 max_rel=3.36826779e-02 tol=(atol=1.000e-04, rtol=1.000e-04) passed=True

pytest: tests/test_lm_head_ascend.py (summary)

======================= 23 passed, 14 warnings in 5.55s ========================

Notes

• Why tolerance vs the gold, bitwise only for invariance: an LM-head projection is a reduction over H; fp32 addition is not associative, and both the CUDA warp-shuffle tree and torch_npu's torch.mv internal order are unspecified platform internals, so no independent kernel can match them bit for bit. The Ascend kernel pins its own fixed order, which makes the op batch-invariant bitwise on the NPU — the property the WS1 contract's forward_invariance row requires (bitwise) — while forward_accuracy is tolerance-based. Note the bf16 gtest numbers look large in absolute terms but are ULP-scale on logits of magnitude ~40 (max_rel 3.9e-3).
• The backward is the standard linear VJP in fp32 (torch.matmul on NPU), cast back to the input dtypes; gradients compare against the native backward at the gradient contract tolerance. NativeLMHeadOp.forward's bf16 dtype path is itself unavailable on this CANN (torch.mv rejects bf16), which is why the reference comparisons go through forward_fp32.
• Single-device operator; no TP/vocab-parallel path (consistent with the SM90 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: b4149f81-07b7-4ebf-b46c-f9947fbdbf0a

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