feat(ascend): add deterministic collective Ascend C kernel

[zhangj1an]

Latest Status [29 Aug 2026]

My 8x910C breaks on libccl_kernel.so, so I cannot test this kernel across ranks. I do not have root access to that machine to re-build pytorch npu. It would be great if the community can help me verify this PR works e2e. Thanks in advance!

Summary

Port of the CUDA deterministic collective (PR #312) to Ascend NPU.

• Reduction kernel (_C_npu.deterministic_collective_reduce): same balanced fixed tree for TP 1/2/4/8 (ws=8: ((t0+t1)+(t2+t3))+((t4+t5)+(t6+t7))). Every level rounds once in the input dtype with round-to-nearest-even, matching CUDA ordered_add (add.rn / __hadd), so results are bitwise identical across ranks, TP configs, and CUDA/Ascend.
  • fp16/bf16 have no usable vector Add on this CANN, so each level adds fp32-exact partials and casts back with CAST_RINT (IEEE RN, verified bitwise against torch including tie lanes).
• Communication: no CUDA-style device IPC on Ascend, so the wrapper gathers all ranks' staged inputs via HCCL all_gather (pure data movement, byte-exact) into [world_size, N] and runs the fixed-tree kernel locally. Reduction order never depends on the HCCL algorithm.
• Wrapper: DeterministicCollective gains an NPU backend with the same public surface (all_reduce / all_gather / reduce_scatter, out= aliasing, shape validation, metadata exchange). CUDA IPC path unchanged.
• Build: csrc/ascend/ops_npu.asc consolidates the single PYBIND11_MODULE (one per .asc file, else Bisheng hits duplicate PyInit__C_npu — same pattern as PR [WS1][kernels] Deterministic attention Ascend C kernel #320); setup.py restores the Ascend extension build, which the current test branch had lost.

Files

Path	Status
csrc/ascend/distributed/deterministic_collective_ascend.asc	New. Fixed-tree reduction kernel (fp32/fp16/bf16) + host staging/dispatch.
csrc/ascend/ops_npu.asc	New. Aggregated _C_npu pybind registration.
csrc/ascend/batch_invariant_logp_ascend.asc	Drops PYBIND11_MODULE only; kernel logic unchanged.
rl_engine/distributed/collectives.py	NPU backend (HCCL all_gather + fixed-tree reduce).
rl_engine/_C_npu.pyi	Collective type stubs.
tests/distributed/test_deterministic_{all_reduce,all_gather,reduce_scatter}_ascend.py	New. 8-rank cross-TP bitwise tests, same contract as the CUDA tests.
setup.py	Restores the Ascend extension build.

Test

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

python -m pytest tests/distributed/test_deterministic_all_reduce_ascend.py -v
python -m pytest tests/distributed/test_deterministic_all_gather_ascend.py -v
python -m pytest tests/distributed/test_deterministic_reduce_scatter_ascend.py -v
python -m pytest tests/test_batch_invariant_logp.py -v   # regression

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

Test	Result
Kernel bitwise: all_reduce, fp32/fp16/bf16 × ws 1/2/4/8 × n ∈ {7, 517, 2048, 2049, 4096} vs torch fixed-tree reference	✅ 0 mismatches
Kernel bitwise: reduce_scatter slices, all dtypes × ws 2/4/8 × every rank slice	✅ 0 mismatches
fp16/bf16 tie lanes vs torch (RN ties-even)	✅ 0 / 100000 mismatches (~10% exact ties)
tests/test_batch_invariant_logp.py regression	✅ 44 passed, 42 skipped (non-NPU backends)
8-rank cross-TP tests	⚠️ not executable on this host — see Notes

Notes

• Multi-rank execution is blocked by a host environment fault, not by this PR: a minimal two-rank dist.all_gather on plain HCCL (no RL-Kernel code) aborts in libccl_kernel.so (RunAicpuKfcResInitV2, aicpu exception). The 8-rank tests are the first thing to run once HCCL is healthy.
• Cross-rank/cross-TP bitwise equality holds by construction: identical kernel, identical all_gather bytes, pinned per-add rounding.
• Single-device operator, no CP/SP — scope consistent with the CUDA DeterministicCollective. ruff check passes.

[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: 656e3008-8934-4430-b4cb-57c5a2394e5b

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