diff --git a/csrc/ascend/batch_invariant_logp_ascend.asc b/csrc/ascend/batch_invariant_logp_ascend.asc index dead4cbe..49139240 100644 --- a/csrc/ascend/batch_invariant_logp_ascend.asc +++ b/csrc/ascend/batch_invariant_logp_ascend.asc @@ -307,10 +307,3 @@ std::vector batch_invariant_logp_ascend_forward(torch::Tensor log } return {logp, lse}; } - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("batch_invariant_logp_ascend", - &batch_invariant_logp_ascend_forward, - "Batch-invariant selected-token log-probability (Ascend C forward)"); -} diff --git a/csrc/ascend/distributed/deterministic_collective_ascend.asc b/csrc/ascend/distributed/deterministic_collective_ascend.asc new file mode 100644 index 00000000..a5377732 --- /dev/null +++ b/csrc/ascend/distributed/deterministic_collective_ascend.asc @@ -0,0 +1,444 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Deterministic (TP-invariant) collectives, Ascend NPU reduction kernel. +// +// out = fixed_tree(rank0, rank1, ...) (ordered adds in the input dtype) +// +// Mirrors the CUDA kernel in csrc/cuda/distributed/deterministic_collective.cu: +// TP sizes 1, 2, 4, and 8 reduce with nested prefixes of the same balanced +// tree -- every node evaluates the lower logical subtree before the higher +// one, and each add rounds once in the input dtype (round-to-nearest-even), +// so the result is bitwise identical on every rank and across TP +// configurations when inputs follow the TBIK-compatible subtree contract. +// +// Unlike the CUDA version there is no cross-device IPC staging on Ascend: the +// Python wrapper gathers every rank's staged input with an HCCL all_gather +// (pure data movement, bitwise exact) into a [world_size, N] buffer, and this +// kernel performs the fixed-tree reduction locally. The reduction order is +// still fully deterministic -- it never depends on the HCCL algorithm. +// +// fp16/bf16 tree levels round per-add like the CUDA ordered_add: the fp32 +// partials (exact sums) are cast back with CAST_RINT, the vector unit's +// round-to-nearest-even conversion (fp16/bf16 vector Adds either do not +// exist or are not exposed by the Add API on this CANN). +// +// Supported dtypes: float32, float16, bfloat16 (same gate as the CUDA op). +// +// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by +// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu. +// Python bindings live in csrc/ascend/ops_npu.asc (single PYBIND11_MODULE). + +#include +#include +#include + +#include "kernel_operator.h" + +#include + +#include "torch_npu/csrc/core/npu/NPUStream.h" + +namespace { +constexpr int kMaxWorldSize = 8; +constexpr uint32_t TILE_FP32 = 1024; +constexpr uint32_t TILE_FP16 = 2048; +constexpr int64_t MAX_BLOCKS = 512; + +// In-place balanced-tree sum over `count` elements of the per-rank UB tiles. +// The rank tiles live back-to-back in one UB buffer; the tree folds even/odd +// pairs in place and finishes by adding the final pair into the out tile. +// This is exactly the CUDA fixed_tree_reduce order: +// ws=2: t0 + t1 +// ws=4: (t0 + t1) + (t2 + t3) +// ws=8: ((t0 + t1) + (t2 + t3)) + ((t4 + t5) + (t6 + t7)) +template +class KernelDeterministicCollectiveReduce { +public: + __aicore__ inline KernelDeterministicCollectiveReduce(AscendC::TPipe* pipe) + : pipe_(pipe) {} + + __aicore__ inline void Init(GM_ADDR gathered, + GM_ADDR out, + int64_t sliceElements, + int64_t sliceOffset, + int64_t rankStride, + int64_t worldSize) + { + sliceElements_ = sliceElements; + sliceOffset_ = sliceOffset; + rankStride_ = rankStride; + worldSize_ = worldSize; + gatheredGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(gathered)); + outGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(out)); + + constexpr uint32_t TILE = std::is_same_v ? TILE_FP32 : TILE_FP16; + pipe_->InitBuffer(rankBuf_, kMaxWorldSize * TILE * sizeof(T)); + pipe_->InitBuffer(rankBufF_, kMaxWorldSize * TILE * sizeof(float)); + pipe_->InitBuffer(outBufF_, TILE * sizeof(float)); + pipe_->InitBuffer(outBuf_, TILE * sizeof(T)); + + // Intra-core pipeline events. NOTE: AscendC::SyncAll() is a cross-core + // barrier and deadlocks when more blocks are launched than there are + // physical cores (resident blocks wait for unscheduled ones), so all + // synchronization here uses per-pipe SetFlag/WaitFlag instead. + eventMTE2S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE2_S); + eventVMTE2_ = pipe_->FetchEventID(AscendC::HardEvent::V_MTE2); + eventVMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::V_MTE3); + eventMTE3S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE3_S); + } + + __aicore__ inline void Process() + { + constexpr uint32_t TILE = std::is_same_v ? TILE_FP32 : TILE_FP16; + const int64_t numTiles = (sliceElements_ + TILE - 1) / TILE; + for (int64_t tile = AscendC::GetBlockIdx(); tile < numTiles; + tile += AscendC::GetBlockNum()) { + ProcessTile(tile); + } + } + +private: + // The rank-r tile views inside the shared rank buffers. GetWithOffset is + // the intrinsic-checker-friendly aliasing view (operator[] views are + // rejected by the vector-binary checks on this CANN version). + __aicore__ inline AscendC::LocalTensor RankBuf(int64_t rank) + { + constexpr uint32_t TILE = std::is_same_v ? TILE_FP32 : TILE_FP16; + return rankBuf_.GetWithOffset(TILE, static_cast(rank * TILE * sizeof(T))); + } + __aicore__ inline AscendC::LocalTensor RankBufF(int64_t rank) + { + constexpr uint32_t TILE = std::is_same_v ? TILE_FP32 : TILE_FP16; + return rankBufF_.GetWithOffset( + TILE, static_cast(rank * TILE * sizeof(float))); + } + + __aicore__ inline void ProcessTile(int64_t tile) + { + constexpr uint32_t TILE = std::is_same_v ? TILE_FP32 : TILE_FP16; + const int64_t start = tile * TILE; + const int64_t remaining = sliceElements_ - start; + const uint32_t count = + static_cast(remaining < TILE ? remaining : TILE); + + // gathered holds [world_size, N] with the slice of interest at + // sliceOffset_ + rank * sliceElements_. Every rank's tile must be in + // UB before any Add; the V_MTE2 wait keeps the next tile's loads from + // clobbering buffers the vector pipe is still reading. + for (int64_t rank = 0; rank < worldSize_; ++rank) { + AscendC::SetFlag(eventVMTE2_); + AscendC::WaitFlag(eventVMTE2_); + const int64_t offset = sliceOffset_ + rank * rankStride_ + start; + AscendC::DataCopyExtParams cp{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPadExtParams pp{false, 0, 0, 0}; + AscendC::DataCopyPad(RankBuf(rank), gatheredGm_[offset], cp, pp); + } + AscendC::SetFlag(eventMTE2S_); + AscendC::WaitFlag(eventMTE2S_); + + if constexpr (std::is_same_v) { + AscendC::LocalTensor out = outBufF_.Get(); + if (worldSize_ == 2) { + AscendC::Add(out, RankBuf(0), RankBuf(1), count); + } else if (worldSize_ == 4) { + AscendC::Add(RankBuf(0), RankBuf(0), RankBuf(1), count); + AscendC::Add(RankBuf(2), RankBuf(2), RankBuf(3), count); + AscendC::Add(out, RankBuf(0), RankBuf(2), count); + } else if (worldSize_ == 8) { + AscendC::Add(RankBuf(0), RankBuf(0), RankBuf(1), count); + AscendC::Add(RankBuf(2), RankBuf(2), RankBuf(3), count); + AscendC::Add(RankBuf(4), RankBuf(4), RankBuf(5), count); + AscendC::Add(RankBuf(6), RankBuf(6), RankBuf(7), count); + AscendC::Add(RankBuf(0), RankBuf(0), RankBuf(2), count); + AscendC::Add(RankBuf(4), RankBuf(4), RankBuf(6), count); + AscendC::Add(out, RankBuf(0), RankBuf(4), count); + } + AscendC::SetFlag(eventVMTE3_); + AscendC::WaitFlag(eventVMTE3_); + AscendC::DataCopyExtParams cpOut{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + if (worldSize_ == 1) { + AscendC::DataCopyPad(outGm_[start], RankBuf(0), cpOut); + } else { + AscendC::DataCopyPad(outGm_[start], out, cpOut); + } + } else { + // fp16/bf16: the CUDA ordered_add rounds after EVERY tree add in + // the input dtype, so each level folds the fp32-exact partials + // (CAST_RINT == IEEE round-to-nearest-even on this CANN) before + // the next level adds them. The RankBuf tiles double as the + // per-level partial storage. + for (int64_t rank = 0; rank < worldSize_; ++rank) { + AscendC::Cast(RankBufF(rank), RankBuf(rank), AscendC::RoundMode::CAST_NONE, count); + } + // Level 1: pair folds. + if (worldSize_ >= 2) { + AscendC::Add(RankBufF(0), RankBufF(0), RankBufF(1), count); + AscendC::Cast(RankBuf(0), RankBufF(0), AscendC::RoundMode::CAST_RINT, count); + } + if (worldSize_ >= 4) { + AscendC::Add(RankBufF(2), RankBufF(2), RankBufF(3), count); + AscendC::Cast(RankBuf(2), RankBufF(2), AscendC::RoundMode::CAST_RINT, count); + } + if (worldSize_ == 8) { + AscendC::Add(RankBufF(4), RankBufF(4), RankBufF(5), count); + AscendC::Cast(RankBuf(4), RankBufF(4), AscendC::RoundMode::CAST_RINT, count); + AscendC::Add(RankBufF(6), RankBufF(6), RankBufF(7), count); + AscendC::Cast(RankBuf(6), RankBufF(6), AscendC::RoundMode::CAST_RINT, count); + } + // Level 2: (01 + 23) and (45 + 67). + if (worldSize_ >= 4) { + AscendC::Cast(RankBufF(0), RankBuf(0), AscendC::RoundMode::CAST_NONE, count); + AscendC::Cast(RankBufF(2), RankBuf(2), AscendC::RoundMode::CAST_NONE, count); + AscendC::Add(RankBufF(0), RankBufF(0), RankBufF(2), count); + AscendC::Cast(RankBuf(0), RankBufF(0), AscendC::RoundMode::CAST_RINT, count); + } + if (worldSize_ == 8) { + AscendC::Cast(RankBufF(4), RankBuf(4), AscendC::RoundMode::CAST_NONE, count); + AscendC::Cast(RankBufF(6), RankBuf(6), AscendC::RoundMode::CAST_NONE, count); + AscendC::Add(RankBufF(4), RankBufF(4), RankBufF(6), count); + AscendC::Cast(RankBuf(4), RankBufF(4), AscendC::RoundMode::CAST_RINT, count); + } + // Level 3: (03 + 47). + AscendC::SetFlag(eventVMTE3_); + AscendC::WaitFlag(eventVMTE3_); + AscendC::DataCopyExtParams cpOut{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + if (worldSize_ == 8) { + AscendC::Cast(RankBufF(0), RankBuf(0), AscendC::RoundMode::CAST_NONE, count); + AscendC::Cast(RankBufF(4), RankBuf(4), AscendC::RoundMode::CAST_NONE, count); + AscendC::Add(RankBufF(0), RankBufF(0), RankBufF(4), count); + AscendC::Cast(outBuf_.Get(), RankBufF(0), AscendC::RoundMode::CAST_RINT, count); + AscendC::SetFlag(eventVMTE3_); + AscendC::WaitFlag(eventVMTE3_); + AscendC::DataCopyPad(outGm_[start], outBuf_.Get(), cpOut); + } else { + // ws 1/2/4: the final partial already lives in RankBuf(0) (or + // is the staged input itself for ws == 1); copy it out + // directly (UB->UB DataCopy has no bf16 form on this CANN). + AscendC::DataCopyPad(outGm_[start], RankBuf(0), cpOut); + } + } + + // Drain MTE3 before the next tile } + + // Drain MTE3 before the next tile stages new values into the shared + // buffers; the scalar pipe issues all later MTE2 copies in order, so + // this wait alone orders them after the copy-out. + AscendC::SetFlag(eventMTE3S_); + AscendC::WaitFlag(eventMTE3S_); + } + + AscendC::TPipe* pipe_; + AscendC::GlobalTensor gatheredGm_; + AscendC::GlobalTensor outGm_; + AscendC::TBuf rankBuf_; + AscendC::TBuf rankBufF_; + AscendC::TBuf outBufF_; + AscendC::TBuf outBuf_; + AscendC::TEventID eventMTE2S_; + AscendC::TEventID eventVMTE2_; + AscendC::TEventID eventVMTE3_; + AscendC::TEventID eventMTE3S_; + int64_t sliceElements_; + int64_t sliceOffset_; + int64_t rankStride_; + int64_t worldSize_; +}; + +} // namespace + +extern "C" __global__ __vector__ void deterministic_collective_reduce_kernel_fp32( + GM_ADDR gathered, GM_ADDR out, int64_t sliceElements, int64_t sliceOffset, + int64_t rankStride, int64_t worldSize) +{ + AscendC::TPipe pipe; + KernelDeterministicCollectiveReduce op(&pipe); + op.Init(gathered, out, sliceElements, sliceOffset, rankStride, worldSize); + op.Process(); +} + +extern "C" __global__ __vector__ void deterministic_collective_reduce_kernel_fp16( + GM_ADDR gathered, GM_ADDR out, int64_t sliceElements, int64_t sliceOffset, + int64_t rankStride, int64_t worldSize) +{ + AscendC::TPipe pipe; + KernelDeterministicCollectiveReduce op(&pipe); + op.Init(gathered, out, sliceElements, sliceOffset, rankStride, worldSize); + op.Process(); +} + +extern "C" __global__ __vector__ void deterministic_collective_reduce_kernel_bf16( + GM_ADDR gathered, GM_ADDR out, int64_t sliceElements, int64_t sliceOffset, + int64_t rankStride, int64_t worldSize) +{ + AscendC::TPipe pipe; + KernelDeterministicCollectiveReduce op(&pipe); + op.Init(gathered, out, sliceElements, sliceOffset, rankStride, worldSize); + op.Process(); +} + +namespace { + +// --------------------------------------------------------------------------- +// Host-side state: staging buffer + reduction dispatch. +// --------------------------------------------------------------------------- + +// --------------------------------------------------------------------------- +// Host-side state: staging buffer + reduction dispatch. +// --------------------------------------------------------------------------- + +class DeterministicCollectiveState { +public: + DeterministicCollectiveState(torch::Tensor& staging, int64_t worldSize, int64_t rank) + : rank_(rank), + world_size_(worldSize), + capacity_bytes_(staging.numel() * staging.element_size()) + { + TORCH_CHECK(staging.is_privateuseone(), "collective staging buffer must be NPU"); + TORCH_CHECK(staging.is_contiguous(), "collective staging buffer must be contiguous"); + TORCH_CHECK( + staging.scalar_type() == torch::kUInt8, + "collective staging buffer must have dtype torch.uint8"); + TORCH_CHECK(capacity_bytes_ > 0, "collective staging capacity must be positive"); + TORCH_CHECK( + world_size_ == 1 || world_size_ == 2 || world_size_ == 4 || world_size_ == 8, + "deterministic collectives require world size 1, 2, 4, or 8; got ", + world_size_); + TORCH_CHECK( + rank_ >= 0 && rank_ < world_size_, + "deterministic collective rank must be in [0, ", + world_size_, + ")"); + staging_ = staging; + } + + void stage(torch::Tensor& input) { + TORCH_CHECK(input.is_privateuseone(), "input must be an NPU tensor"); + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK( + input.get_device() == staging_.get_device(), + "input must be on the staging device"); + const int64_t input_bytes = input.numel() * input.element_size(); + TORCH_CHECK( + input_bytes <= capacity_bytes_, + "input requires ", + input_bytes, + " bytes but staging capacity is ", + capacity_bytes_); + if (input_bytes > 0) { + // Copy on the current NPU stream; the wrapper synchronizes the + // device before the cross-rank gather reads the staging buffer. + staging_.narrow(0, 0, input_bytes) + .view(input.scalar_type()) + .copy_(input); + } + staged_bytes_ = input_bytes; + staged_numel_ = input.numel(); + staged_scalar_type_ = input.scalar_type(); + has_staged_input_ = true; + } + + void reduce(torch::Tensor& gathered, torch::Tensor& output, int64_t sliceOffset) const { + TORCH_CHECK(has_staged_input_, "stage() must be called before reduce()"); + TORCH_CHECK( + gathered.is_privateuseone() && output.is_privateuseone(), + "gathered and output must be NPU tensors"); + TORCH_CHECK( + gathered.is_contiguous() && output.is_contiguous(), + "gathered and output must be contiguous"); + TORCH_CHECK( + output.scalar_type() == staged_scalar_type_, + "reduce output dtype must match the staged input dtype"); + TORCH_CHECK( + gathered.scalar_type() == staged_scalar_type_, + "gathered dtype must match the staged input dtype"); + const int64_t slice_elements = output.numel(); + TORCH_CHECK( + gathered.numel() == staged_numel_ * world_size_, + "gathered must contain one staged input per rank: expected ", + staged_numel_ * world_size_, + " elements, got ", + gathered.numel()); + TORCH_CHECK( + sliceOffset >= 0 && sliceOffset + slice_elements <= staged_numel_, + "reduce slice offset is out of the gathered buffer"); + + if (slice_elements == 0) { + return; + } + auto aclStream = c10_npu::getCurrentNPUStream().stream(true); + const uint32_t tile = + staged_scalar_type_ == at::ScalarType::Float ? TILE_FP32 : TILE_FP16; + const int64_t numTiles = (slice_elements + tile - 1) / tile; + const uint32_t blockNum = static_cast(std::min(numTiles, MAX_BLOCKS)); + + uint8_t* gatheredPtr = reinterpret_cast(gathered.mutable_data_ptr()); + uint8_t* outPtr = reinterpret_cast(output.mutable_data_ptr()); + + switch (staged_scalar_type_) { + case at::ScalarType::Float: + deterministic_collective_reduce_kernel_fp32<<>>( + gatheredPtr, outPtr, slice_elements, sliceOffset, staged_numel_, world_size_); + break; + case at::ScalarType::Half: + deterministic_collective_reduce_kernel_fp16<<>>( + gatheredPtr, outPtr, slice_elements, sliceOffset, staged_numel_, world_size_); + break; + case at::ScalarType::BFloat16: + deterministic_collective_reduce_kernel_bf16<<>>( + gatheredPtr, outPtr, slice_elements, sliceOffset, staged_numel_, world_size_); + break; + default: + TORCH_CHECK( + false, + "deterministic reduce supports float32, float16, and bfloat16; got ", + staged_scalar_type_); + } + } + +private: + int64_t rank_; + int64_t world_size_; + int64_t capacity_bytes_; + int64_t staged_bytes_{0}; + int64_t staged_numel_{0}; + at::ScalarType staged_scalar_type_{at::ScalarType::Undefined}; + bool has_staged_input_{false}; + torch::Tensor staging_; +}; + +DeterministicCollectiveState* state_from_handle(int64_t handle) { + TORCH_CHECK(handle != 0, "deterministic collective handle is closed"); + return reinterpret_cast(handle); +} + +} // namespace + +// Host API (registered in csrc/ascend/ops_npu.asc). +int64_t deterministic_collective_create(torch::Tensor staging, int64_t worldSize, int64_t rank) +{ + auto state = std::make_unique(staging, worldSize, rank); + return reinterpret_cast(state.release()); +} + +void deterministic_collective_destroy(int64_t handle) +{ + delete state_from_handle(handle); +} + +void deterministic_collective_stage(int64_t handle, torch::Tensor input) +{ + state_from_handle(handle)->stage(input); +} + +void deterministic_collective_reduce( + int64_t handle, + torch::Tensor gathered, + torch::Tensor output, + int64_t sliceOffset) +{ + state_from_handle(handle)->reduce(gathered, output, sliceOffset); +} diff --git a/csrc/ascend/ops_npu.asc b/csrc/ascend/ops_npu.asc new file mode 100644 index 00000000..2cfadf04 --- /dev/null +++ b/csrc/ascend/ops_npu.asc @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +// Single Python module entry point for the unified Ascend C extension. +// Keep PYBIND11_MODULE in one translation unit as new .asc kernels are added. + +#include + +#include + +std::vector batch_invariant_logp_ascend_forward( + torch::Tensor logits, torch::Tensor target, int64_t ignore_index); + +int64_t deterministic_collective_create( + torch::Tensor staging, int64_t world_size, int64_t rank); +void deterministic_collective_destroy(int64_t handle); +void deterministic_collective_stage(int64_t handle, torch::Tensor input); +void deterministic_collective_reduce( + int64_t handle, torch::Tensor gathered, torch::Tensor output, int64_t slice_offset); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("batch_invariant_logp_ascend", + &batch_invariant_logp_ascend_forward, + "Batch-invariant selected-token log-probability (Ascend C forward)"); + m.def("deterministic_collective_create", + &deterministic_collective_create, + "Deterministic TP-invariant collective state (Ascend)"); + m.def("deterministic_collective_destroy", + &deterministic_collective_destroy, + "Release a deterministic collective state (Ascend)"); + m.def("deterministic_collective_stage", + &deterministic_collective_stage, + "Stage a tensor into the collective staging buffer (Ascend)"); + m.def("deterministic_collective_reduce", + &deterministic_collective_reduce, + "Fixed-tree ordered reduction over gathered rank tensors (Ascend)"); +} diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index bff5e2e7..4b387a80 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,3 +8,16 @@ def batch_invariant_logp_ascend( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... +def deterministic_collective_create( + staging: torch.Tensor, + world_size: int, + rank: int, +) -> int: ... +def deterministic_collective_destroy(handle: int) -> None: ... +def deterministic_collective_stage(handle: int, input: torch.Tensor) -> None: ... +def deterministic_collective_reduce( + handle: int, + gathered: torch.Tensor, + output: torch.Tensor, + slice_offset: int, +) -> None: ... diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py index 54b00fa5..d9674ac3 100644 --- a/rl_engine/distributed/collectives.py +++ b/rl_engine/distributed/collectives.py @@ -16,8 +16,18 @@ _REDUCTION_DTYPES = (torch.float32, torch.float16, torch.bfloat16) +def _npu_available() -> bool: + """torch.npu only exists after torch_npu is imported; probe defensively.""" + try: + import torch_npu # noqa: F401 + + return hasattr(torch, "npu") and torch.npu.is_available() + except Exception: + return False + + class DeterministicCollective: - """Correctness-first TP-invariant CUDA collectives for one eight-GPU node. + """Correctness-first TP-invariant collectives for one eight-device node. TP sizes 1, 2, 4, and 8 use nested prefixes of the same balanced tree. A reduction is cross-TP bitwise invariant when every rank input is the @@ -25,10 +35,15 @@ class DeterministicCollective: reduction, as produced by a TBIK-compatible row-parallel kernel. Every node evaluates the lower logical subtree before the higher one. - One instance owns a symmetric CUDA IPC staging buffer. All ranks must call - its methods in the same order with matching shapes and dtypes. Calls are - host-synchronizing by design; the first version prioritizes determinism and - lifetime safety over overlap or throughput. + CUDA: one instance owns a symmetric CUDA IPC staging buffer; the reduction + kernel reads every peer's staged data directly. Ascend NPU: no device IPC + exists, so each reduction first gathers every rank's staged input with an + HCCL all_gather (bitwise-exact data movement) and then applies the same + fixed-tree kernel locally -- the reduction order never depends on the HCCL + algorithm. All ranks must call the methods in the same order with matching + shapes and dtypes. Calls are host-synchronizing by design; the first + version prioritizes determinism and lifetime safety over overlap or + throughput. """ def __init__( @@ -40,8 +55,6 @@ def __init__( ) -> None: if not dist.is_available() or not dist.is_initialized(): raise RuntimeError("torch.distributed must be initialized before collectives") - if not torch.cuda.is_available(): - raise RuntimeError("deterministic collectives require CUDA") if max_size_bytes <= 0: raise ValueError("max_size_bytes must be positive") @@ -53,23 +66,64 @@ def __init__( "deterministic collectives require world_size in " f"{_SUPPORTED_WORLD_SIZES}, got {self.world_size}" ) + self.max_size_bytes = int(max_size_bytes) - if device is None: + is_cuda = torch.cuda.is_available() + is_npu = _npu_available() + if is_cuda: + self._backend = "cuda" normalized_device = torch.device("cuda", torch.cuda.current_device()) - elif isinstance(device, int): - normalized_device = torch.device("cuda", device) + if device is not None: + normalized_device = ( + torch.device("cuda", device) + if isinstance(device, int) + else torch.device(device) + ) + if normalized_device.type != "cuda": + raise ValueError( + f"deterministic collectives require a CUDA device, got {device!r}" + ) + if normalized_device.index is None: + normalized_device = torch.device("cuda", torch.cuda.current_device()) + if normalized_device.index != torch.cuda.current_device(): + raise ValueError( + "the collective device must be the current CUDA device; call " + f"torch.cuda.set_device({normalized_device.index}) first" + ) + self._load_cuda_extension() + self.device = normalized_device + self._create_cuda_state() + elif is_npu: + self._backend = "npu" + normalized_device = torch.device("npu", torch.npu.current_device()) + if device is not None: + normalized_device = ( + torch.device("npu", device) if isinstance(device, int) else torch.device(device) + ) + if normalized_device.type != "npu": + raise ValueError( + f"deterministic collectives require an NPU device, got {device!r}" + ) + if normalized_device.index is None: + normalized_device = torch.device("npu", torch.npu.current_device()) + if normalized_device.index != torch.npu.current_device(): + raise ValueError( + "the collective device must be the current NPU device; call " + f"torch.npu.set_device({normalized_device.index}) first" + ) + self._load_npu_extension() + self.device = normalized_device + self._create_npu_state() else: - normalized_device = torch.device(device) - if normalized_device.type != "cuda": - raise ValueError(f"deterministic collectives require a CUDA device, got {device!r}") - if normalized_device.index is None: - normalized_device = torch.device("cuda", torch.cuda.current_device()) - if normalized_device.index != torch.cuda.current_device(): - raise ValueError( - "the collective device must be the current CUDA device; call " - f"torch.cuda.set_device({normalized_device.index}) first" - ) + raise RuntimeError("deterministic collectives require CUDA or Ascend NPU devices") + + self._synchronize_ranks() + # ------------------------------------------------------------------ # + # Backend setup + # ------------------------------------------------------------------ # + + def _load_cuda_extension(self) -> None: try: from rl_engine import _C except ImportError as exc: @@ -92,10 +146,31 @@ def __init__( "the RL-Kernel CUDA extension lacks deterministic collectives: " + ", ".join(missing) ) - - self.device = normalized_device - self.max_size_bytes = int(max_size_bytes) self._extension = _C + + def _load_npu_extension(self) -> None: + try: + from rl_engine import _C_npu + except ImportError as exc: + raise RuntimeError( + "the RL-Kernel Ascend extension is required; rebuild with " + "KERNEL_ALIGN_FORCE_ASCEND=1 and `pip install --no-build-isolation -e .`" + ) from exc + required_symbols = ( + "deterministic_collective_create", + "deterministic_collective_destroy", + "deterministic_collective_stage", + "deterministic_collective_reduce", + ) + missing = [name for name in required_symbols if not hasattr(_C_npu, name)] + if missing: + raise RuntimeError( + "the RL-Kernel Ascend extension lacks deterministic collectives: " + + ", ".join(missing) + ) + self._extension = _C_npu + + def _create_cuda_state(self) -> None: self._lock = threading.Lock() self._handle = 0 self._staging = torch.empty( @@ -111,11 +186,50 @@ def __init__( "capacity": self.max_size_bytes, "hostname": socket.gethostname(), } + gathered_meta = self._exchange_meta(local_meta) + self._validate_meta(gathered_meta) + + handles = [meta["handle"] for meta in gathered_meta] + offsets = [meta["offset"] for meta in gathered_meta] + self._handle = self._extension.deterministic_collective_create( + self._staging, + handles, + offsets, + self.rank, + ) + + def _create_npu_state(self) -> None: + self._lock = threading.Lock() + self._handle = 0 + self._staging = torch.empty( + self.max_size_bytes, + dtype=torch.uint8, + device=self.device, + ) + + # No device IPC on Ascend: only the capacity / hostname invariants are + # exchanged (the data itself moves through HCCL all_gather per call). + local_meta = { + "capacity": self.max_size_bytes, + "hostname": socket.gethostname(), + } + gathered_meta = self._exchange_meta(local_meta) + self._validate_meta(gathered_meta) + + self._handle = self._extension.deterministic_collective_create( + self._staging, + self.world_size, + self.rank, + ) + + def _exchange_meta(self, local_meta: dict[str, Any]) -> list[dict[str, Any]]: gathered_meta: list[dict[str, Any] | None] = [None] * self.world_size dist.all_gather_object(gathered_meta, local_meta, group=self.group) if any(meta is None for meta in gathered_meta): - raise RuntimeError("failed to exchange CUDA IPC metadata") - complete_meta = [meta for meta in gathered_meta if meta is not None] + raise RuntimeError("failed to exchange collective metadata") + return [meta for meta in gathered_meta if meta is not None] + + def _validate_meta(self, complete_meta: list[dict[str, Any]]) -> None: hostnames = {meta["hostname"] for meta in complete_meta} if len(hostnames) != 1: raise ValueError("deterministic collectives require all ranks on one host") @@ -123,15 +237,9 @@ def __init__( if capacities != {self.max_size_bytes}: raise ValueError("all ranks must use the same max_size_bytes") - handles = [meta["handle"] for meta in complete_meta] - offsets = [meta["offset"] for meta in complete_meta] - self._handle = self._extension.deterministic_collective_create( - self._staging, - handles, - offsets, - self.rank, - ) - self._synchronize_ranks() + # ------------------------------------------------------------------ # + # Public collectives + # ------------------------------------------------------------------ # def all_reduce( self, @@ -142,8 +250,9 @@ def all_reduce( """Return the TBIK-compatible fixed-tree sum on every rank. Supported dtypes are float32, float16, and bfloat16. ``out`` may alias - ``input``; the input is staged before the output kernel starts. Cross-TP - invariance requires inputs to follow the class-level subtree contract. + ``input``; the input is staged before the reduction kernel starts. + Cross-TP invariance requires inputs to follow the class-level subtree + contract. """ self._check_open() @@ -156,7 +265,7 @@ def all_reduce( self._validate_matching_signature("all_reduce", input) self._extension.deterministic_collective_stage(self._handle, input) self._synchronize_ranks() - self._extension.deterministic_collective_all_reduce(self._handle, out) + self._run_reduction(input, out, slice_offset=0) self._synchronize_ranks() return out @@ -183,7 +292,13 @@ def all_gather( self._validate_matching_signature("all_gather", input) self._extension.deterministic_collective_stage(self._handle, input) self._synchronize_ranks() - self._extension.deterministic_collective_all_gather(self._handle, out) + if self._backend == "cuda": + self._extension.deterministic_collective_all_gather(self._handle, out) + else: + staged = self._staged_view(input) + shard = input.size(0) + slices = [out[index : index + shard] for index in range(0, out.size(0), shard)] + dist.all_gather(slices, staged, group=self.group) self._synchronize_ranks() return out @@ -217,17 +332,57 @@ def reduce_scatter( self._validate_matching_signature("reduce_scatter", input) self._extension.deterministic_collective_stage(self._handle, input) self._synchronize_ranks() - self._extension.deterministic_collective_reduce_scatter(self._handle, out) + if self._backend == "cuda": + self._extension.deterministic_collective_reduce_scatter(self._handle, out) + else: + self._run_reduction(input, out, slice_offset=self.rank * out.numel()) self._synchronize_ranks() return out + # ------------------------------------------------------------------ # + # Backend helpers + # ------------------------------------------------------------------ # + + def _staged_view(self, input: torch.Tensor) -> torch.Tensor: + """The staged input as a typed view of the staging buffer.""" + return self._staging.narrow(0, 0, input.numel() * input.element_size()).view(input.dtype) + + def _run_reduction( + self, + input: torch.Tensor, + out: torch.Tensor, + *, + slice_offset: int, + ) -> None: + """NPU reduction: HCCL-gather every rank's staged input, then apply the + fixed-tree kernel locally over the [world_size, N] gathered buffer.""" + gathered = torch.empty( + (self.world_size, input.numel()), + dtype=input.dtype, + device=input.device, + ) + dist.all_gather( + list(gathered.unbind(0)), + self._staged_view(input), + group=self.group, + ) + self._extension.deterministic_collective_reduce( + self._handle, + gathered, + out, + slice_offset, + ) + def close(self) -> None: - """Release imported CUDA IPC mappings after the last collective call.""" + """Release the collective state (and CUDA IPC mappings) after the last call.""" handle = getattr(self, "_handle", 0) if not handle: return - torch.cuda.synchronize(self.device) + if self._backend == "cuda": + torch.cuda.synchronize(self.device) + else: + torch.npu.synchronize(self.device) self._handle = 0 self._extension.deterministic_collective_destroy(handle) @@ -254,7 +409,7 @@ def _check_open(self) -> None: raise RuntimeError("deterministic collective is closed") def _validate_reduction_input(self, input: torch.Tensor) -> None: - if not input.is_cuda or input.device != self.device: + if input.device != self.device: raise ValueError(f"input must be on {self.device}, got {input.device}") if not input.is_contiguous(): raise ValueError("input must be contiguous") @@ -270,7 +425,7 @@ def _validate_reduction_input(self, input: torch.Tensor) -> None: ) def _validate_gather_input(self, input: torch.Tensor) -> None: - if not input.is_cuda or input.device != self.device: + if input.device != self.device: raise ValueError(f"input must be on {self.device}, got {input.device}") if not input.is_contiguous(): raise ValueError("input must be contiguous") @@ -317,9 +472,13 @@ def _validate_matching_signature(self, op_name: str, input: torch.Tensor) -> Non ) def _synchronize_ranks(self) -> None: - torch.cuda.synchronize(self.device) - backend = dist.get_backend(self.group) - if backend == dist.Backend.NCCL or str(backend).lower() == "nccl": - dist.barrier(group=self.group, device_ids=[self.device.index]) + if self._backend == "cuda": + torch.cuda.synchronize(self.device) + backend = dist.get_backend(self.group) + if backend == dist.Backend.NCCL or str(backend).lower() == "nccl": + dist.barrier(group=self.group, device_ids=[self.device.index]) + else: + dist.barrier(group=self.group) else: + torch.npu.synchronize(self.device) dist.barrier(group=self.group) diff --git a/setup.py b/setup.py index 79f882d9..e3bb9923 100644 --- a/setup.py +++ b/setup.py @@ -3,10 +3,13 @@ import importlib.util import os +import sysconfig import warnings +from distutils.errors import CompileError +from distutils.spawn import find_executable from pathlib import Path -from setuptools import find_packages, setup +from setuptools import Extension, find_packages, setup def _load_envs_module(): @@ -92,6 +95,88 @@ def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: return filtered_flags +def _ascend_extensions(): + """Ascend C (CANN) kernels, built with bisheng. Gated on KERNEL_ALIGN_FORCE_ASCEND=1. + + Follows the official torch_npu cpp_extension_asc pattern: .asc sources + (kernel + host + pybind) are compiled by the CANN bisheng compiler into a + single rl_engine._C_npu extension module. Requires CANN toolkit (bisheng on + PATH or ASCEND_HOME_PATH set) and torch_npu. + """ + if not envs.env_flag(envs.KERNEL_ALIGN_FORCE_ASCEND): + return [] + try: + import torch # noqa: F401 + import torch_npu # noqa: F401 + except ImportError as e: + raise RuntimeError( + "KERNEL_ALIGN_FORCE_ASCEND=1 requires torch and torch_npu to be installed" + ) from e + + asc_srcs = sorted(str(p) for p in Path("csrc/ascend").glob("**/*.asc")) + if not asc_srcs: + raise RuntimeError("KERNEL_ALIGN_FORCE_ASCEND=1 but no .asc sources under csrc/ascend/") + return [Extension(name="rl_engine._C_npu", sources=asc_srcs, language="asc")] + + +def _bisheng_compile_cmd(ext, ext_fullpath): + """Single-command bisheng build for an Ascend C extension (see op-plugin example).""" + import torch + import torch.utils.cpp_extension as cpp_extension + import torch_npu + + if find_executable("bisheng") is None: + raise RuntimeError( + "bisheng compiler not found on PATH; source the CANN toolkit environment first" + ) + + soc = os.environ.get(envs.KERNEL_ALIGN_ASCEND_ARCH, "dav-2201") # A2/A3; A5: dav-3510 + abi_value = "1" if torch._C._GLIBCXX_USE_CXX11_ABI else "0" + module_name = ext.name.rsplit(".", 1)[-1] + + torch_npu_dir = os.path.dirname(os.path.realpath(torch_npu.__file__)) + ascend_home = os.environ.get("ASCEND_HOME_PATH", "/usr/local/Ascend/ascend-toolkit/latest") + + include_dirs = [ + *cpp_extension.include_paths(), + sysconfig.get_config_var("INCLUDEPY"), + os.path.join(torch_npu_dir, "include"), + os.path.join(torch_npu_dir, "include", "third_party", "acl", "inc"), + os.path.join(ascend_home, "include"), + ] + lib_dirs = [ + sysconfig.get_config_var("LIBDIR"), + os.path.join(os.path.dirname(torch.__file__), "lib"), + os.path.join(torch_npu_dir, "lib"), + os.path.join(ascend_home, "lib64"), + ] + + cmd = [ + "bisheng", + "-x", + "asc", + f"--npu-arch={soc}", + "-shared", + "-fPIC", + "-std=c++17", + "-O2", + f"-D_GLIBCXX_USE_CXX11_ABI={abi_value}", + f"-DTORCH_EXTENSION_NAME={module_name}", + "-lascendcl", + "-ltorch_npu", + "-ltorch", + "-ltorch_cpu", + "-ltorch_python", + "-lc10", + *ext.sources, + "-o", + ext_fullpath, + ] + cmd += [f"-I{d}" for d in include_dirs if d] + cmd += [f"-L{d}" for d in lib_dirs if d] + return cmd + + def get_extensions(): torch, _, CUDAExtension = _load_torch_extension_tools() if torch is None: @@ -231,7 +316,7 @@ def get_extensions(): if enable_sm90 and present_sm90: tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant cuda_sources.extend(present_sm90) - nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") + nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") if "-lcuda" not in extra_link_args: extra_link_args.append("-lcuda") @@ -266,6 +351,11 @@ def get_extensions(): ) ) + # Ascend extension is appended before the native-extension check so an + # NPU-only build (KERNEL_ALIGN_FORCE_ASCEND=1) never trips the CUDA/ROCm + # "no build environment" error above. + extensions.extend(_ascend_extensions()) + if _native_extension_required() and not extensions: raise RuntimeError( "rl_engine._C was requested but no CUDA/ROCm build environment is available. " @@ -280,7 +370,22 @@ def get_cmdclass(): _, BuildExtension, _ = _load_torch_extension_tools() if BuildExtension is None: return {} - return {"build_ext": BuildExtension} + + class AscendBuildExtension(BuildExtension): + """torch BuildExtension + bisheng path for language="asc" extensions.""" + + def build_extension(self, ext): + if getattr(ext, "language", None) != "asc": + super().build_extension(ext) + return + ext_fullpath = self.get_ext_fullpath(ext.name) + os.makedirs(os.path.dirname(ext_fullpath), exist_ok=True) + try: + self.spawn(_bisheng_compile_cmd(ext, ext_fullpath)) + except Exception as e: + raise CompileError(str(e)) from e + + return {"build_ext": AscendBuildExtension} setup( diff --git a/tests/distributed/test_deterministic_all_gather_ascend.py b/tests/distributed/test_deterministic_all_gather_ascend.py new file mode 100644 index 00000000..0428e1d9 --- /dev/null +++ b/tests/distributed/test_deterministic_all_gather_ascend.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Cross-TP deterministic all-gather on Ascend NPUs. + +Same contract as the CUDA test: each rank holds a rank-ordered dimension-0 +shard of a global tensor; the gather must reconstruct the global tensor +bitwise, identically on every rank, and repeatably. +""" + +from __future__ import annotations + +import os +import socket +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from rl_engine.distributed import DeterministicCollective + +_MAX_WORLD_SIZE = 8 +_TP_SIZES = (1, 2, 4, 8) +_EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + + +def _npu_visible_count() -> int: + try: + import torch_npu # noqa: F401 + + return torch.npu.device_count() + except Exception: + return 0 + + +def _ascend_extension_available() -> bool: + if _npu_visible_count() == 0: + return False + try: + from rl_engine import _C_npu + + return hasattr(_C_npu, "deterministic_collective_reduce") + except Exception: + return False + + +pytestmark = [ + pytest.mark.skipif( + _EXTERNAL_WORLD_SIZE != 1, + reason="this cross-TP test owns its worker processes; run pytest directly", + ), + pytest.mark.skipif( + _npu_visible_count() < _MAX_WORLD_SIZE, + reason="requires eight visible Ascend NPUs", + ), + pytest.mark.skipif( + not _ascend_extension_available(), + reason="deterministic_collective_reduce not compiled into _C_npu", + ), +] + + +def _all_gather_tensors( + value: torch.Tensor, + group: dist.ProcessGroup, + world_size: int, +) -> list[torch.Tensor]: + gathered = [torch.empty_like(value) for _ in range(world_size)] + dist.all_gather(gathered, value, group=group) + return gathered + + +def _make_global_input(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + if dtype == torch.int64: + return torch.arange( + _MAX_WORLD_SIZE * 13 * 7, + device=device, + dtype=dtype, + ).reshape(_MAX_WORLD_SIZE * 13, 7) + generator = torch.Generator().manual_seed(20260817) + return torch.randn( + _MAX_WORLD_SIZE * 13, + 7, + dtype=torch.float32, + generator=generator, + ).to(device=device, dtype=dtype) + + +def _worker(rank: int, port: int) -> None: + torch.npu.set_device(rank) + device = torch.device("npu", rank) + dist.init_process_group( + backend="hccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=_MAX_WORLD_SIZE, + timeout=timedelta(minutes=5), + ) + try: + groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} + for tp_size, group in groups.items(): + if rank < tp_size: + with DeterministicCollective( + group=group, + device=device, + max_size_bytes=1024 * 1024, + ) as collective: + group_rank = dist.get_rank(group=group) + for dtype in (torch.float32, torch.bfloat16, torch.int64): + expected = _make_global_input(device, dtype) + input = expected.chunk(tp_size, dim=0)[group_rank].contiguous() + + output = collective.all_gather(input) + assert torch.equal(output, expected) + + peer_outputs = _all_gather_tensors(output, group, tp_size) + assert all(torch.equal(peer_output, output) for peer_output in peer_outputs) + + baseline = output.clone() + for _ in range(3): + repeated = collective.all_gather(input) + assert torch.equal(repeated, baseline) + + provided = torch.empty_like(expected) + returned = collective.all_gather(input, out=provided) + assert returned is provided + assert torch.equal(provided, expected) + dist.barrier() + finally: + dist.destroy_process_group() + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_deterministic_all_gather_cross_tp_ascend() -> None: + mp.spawn( + _worker, + args=(_find_free_port(),), + nprocs=_MAX_WORLD_SIZE, + join=True, + ) diff --git a/tests/distributed/test_deterministic_all_reduce_ascend.py b/tests/distributed/test_deterministic_all_reduce_ascend.py new file mode 100644 index 00000000..c5fff0ab --- /dev/null +++ b/tests/distributed/test_deterministic_all_reduce_ascend.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Cross-TP deterministic all-reduce on Ascend NPUs. + +Same contract as the CUDA test: eight ranks reduce with nested prefixes of the +same balanced tree for TP sizes 1/2/4/8; the result must be the canonical +8-leaf tree sum, bitwise identical across ranks, repeatable, and safe when +``out`` aliases ``input``. +""" + +from __future__ import annotations + +import os +import socket +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from rl_engine.distributed import DeterministicCollective + +_MAX_WORLD_SIZE = 8 +_TP_SIZES = (1, 2, 4, 8) +_EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + + +def _npu_visible_count() -> int: + try: + import torch_npu # noqa: F401 + + return torch.npu.device_count() + except Exception: + return 0 + + +def _ascend_extension_available() -> bool: + if _npu_visible_count() == 0: + return False + try: + from rl_engine import _C_npu + + return hasattr(_C_npu, "deterministic_collective_reduce") + except Exception: + return False + + +pytestmark = [ + pytest.mark.skipif( + _EXTERNAL_WORLD_SIZE != 1, + reason="this cross-TP test owns its worker processes; run pytest directly", + ), + pytest.mark.skipif( + _npu_visible_count() < _MAX_WORLD_SIZE, + reason="requires eight visible Ascend NPUs", + ), + pytest.mark.skipif( + not _ascend_extension_available(), + reason="deterministic_collective_reduce not compiled into _C_npu", + ), +] + + +def _fixed_tree_reference(values: list[torch.Tensor]) -> torch.Tensor: + level = values + while len(level) > 1: + level = [level[index] + level[index + 1] for index in range(0, len(level), 2)] + return level[0] + + +def _all_gather_tensors( + value: torch.Tensor, + group: dist.ProcessGroup, + world_size: int, +) -> list[torch.Tensor]: + gathered = [torch.empty_like(value) for _ in range(world_size)] + dist.all_gather(gathered, value, group=group) + return gathered + + +def _worker(rank: int, port: int) -> None: + torch.npu.set_device(rank) + device = torch.device("npu", rank) + dist.init_process_group( + backend="hccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=_MAX_WORLD_SIZE, + timeout=timedelta(minutes=5), + ) + try: + groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} + for tp_size, group in groups.items(): + if rank < tp_size: + with DeterministicCollective( + group=group, + device=device, + max_size_bytes=1024 * 1024, + ) as collective: + group_rank = dist.get_rank(group=group) + leaves_per_rank = _MAX_WORLD_SIZE // tp_size + start = group_rank * leaves_per_rank + for dtype in (torch.float32, torch.float16, torch.bfloat16): + generator = torch.Generator().manual_seed(20260815) + leaves_tensor = torch.randn( + _MAX_WORLD_SIZE, + 257, + dtype=torch.float32, + generator=generator, + ).to(device=device, dtype=dtype) + leaves = list(leaves_tensor.unbind()) + input = _fixed_tree_reference(leaves[start : start + leaves_per_rank]) + expected = _fixed_tree_reference(leaves) + + output = collective.all_reduce(input) + assert torch.equal(output, expected) + + peer_outputs = _all_gather_tensors(output, group, tp_size) + assert all(torch.equal(peer_output, output) for peer_output in peer_outputs) + + baseline = output.clone() + for _ in range(3): + repeated = collective.all_reduce(input) + assert torch.equal(repeated, baseline) + + inplace = input.clone() + returned = collective.all_reduce(inplace, out=inplace) + assert returned is inplace + assert torch.equal(inplace, expected) + dist.barrier() + finally: + dist.destroy_process_group() + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_deterministic_all_reduce_cross_tp_ascend() -> None: + mp.spawn( + _worker, + args=(_find_free_port(),), + nprocs=_MAX_WORLD_SIZE, + join=True, + ) diff --git a/tests/distributed/test_deterministic_reduce_scatter_ascend.py b/tests/distributed/test_deterministic_reduce_scatter_ascend.py new file mode 100644 index 00000000..cebab4cd --- /dev/null +++ b/tests/distributed/test_deterministic_reduce_scatter_ascend.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Cross-TP deterministic reduce-scatter on Ascend NPUs. + +Same contract as the CUDA test: eight ranks reduce with nested prefixes of the +same balanced tree for TP sizes 1/2/4/8; each rank's output must be the +corresponding shard of the canonical 8-leaf tree sum, bitwise identical, +repeatable, and safe when ``out`` is provided. +""" + +from __future__ import annotations + +import os +import socket +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from rl_engine.distributed import DeterministicCollective + +_MAX_WORLD_SIZE = 8 +_TP_SIZES = (1, 2, 4, 8) +_EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + + +def _npu_visible_count() -> int: + try: + import torch_npu # noqa: F401 + + return torch.npu.device_count() + except Exception: + return 0 + + +def _ascend_extension_available() -> bool: + if _npu_visible_count() == 0: + return False + try: + from rl_engine import _C_npu + + return hasattr(_C_npu, "deterministic_collective_reduce") + except Exception: + return False + + +pytestmark = [ + pytest.mark.skipif( + _EXTERNAL_WORLD_SIZE != 1, + reason="this cross-TP test owns its worker processes; run pytest directly", + ), + pytest.mark.skipif( + _npu_visible_count() < _MAX_WORLD_SIZE, + reason="requires eight visible Ascend NPUs", + ), + pytest.mark.skipif( + not _ascend_extension_available(), + reason="deterministic_collective_reduce not compiled into _C_npu", + ), +] + + +def _fixed_tree_reference(values: list[torch.Tensor]) -> torch.Tensor: + level = values + while len(level) > 1: + level = [level[index] + level[index + 1] for index in range(0, len(level), 2)] + return level[0] + + +def _all_gather_tensors( + value: torch.Tensor, + group: dist.ProcessGroup, + world_size: int, +) -> list[torch.Tensor]: + gathered = [torch.empty_like(value) for _ in range(world_size)] + dist.all_gather(gathered, value, group=group) + return gathered + + +def _worker(rank: int, port: int) -> None: + torch.npu.set_device(rank) + device = torch.device("npu", rank) + dist.init_process_group( + backend="hccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=_MAX_WORLD_SIZE, + timeout=timedelta(minutes=5), + ) + try: + groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} + for tp_size, group in groups.items(): + if rank < tp_size: + with DeterministicCollective( + group=group, + device=device, + max_size_bytes=1024 * 1024, + ) as collective: + group_rank = dist.get_rank(group=group) + leaves_per_rank = _MAX_WORLD_SIZE // tp_size + start = group_rank * leaves_per_rank + for dtype in (torch.float32, torch.float16, torch.bfloat16): + generator = torch.Generator().manual_seed(20260816) + leaves_tensor = torch.randn( + _MAX_WORLD_SIZE, + _MAX_WORLD_SIZE * 17, + 19, + dtype=torch.float32, + generator=generator, + ).to(device=device, dtype=dtype) + leaves = list(leaves_tensor.unbind()) + input = _fixed_tree_reference(leaves[start : start + leaves_per_rank]) + reduced = _fixed_tree_reference(leaves) + expected = reduced.chunk(tp_size, dim=0)[group_rank] + + output = collective.reduce_scatter(input) + assert torch.equal(output, expected) + + peer_outputs = _all_gather_tensors(output, group, tp_size) + assert torch.equal(torch.cat(peer_outputs, dim=0), reduced) + + baseline = output.clone() + for _ in range(3): + repeated = collective.reduce_scatter(input) + assert torch.equal(repeated, baseline) + + provided = torch.empty_like(expected) + returned = collective.reduce_scatter(input, out=provided) + assert returned is provided + assert torch.equal(provided, expected) + dist.barrier() + finally: + dist.destroy_process_group() + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_deterministic_reduce_scatter_cross_tp_ascend() -> None: + mp.spawn( + _worker, + args=(_find_free_port(),), + nprocs=_MAX_WORLD_SIZE, + join=True, + )