diff --git a/csrc/ascend/batch_invariant_logp_ascend.asc b/csrc/ascend/batch_invariant_logp_ascend.asc index dead4cbe..3b7e46b9 100644 --- a/csrc/ascend/batch_invariant_logp_ascend.asc +++ b/csrc/ascend/batch_invariant_logp_ascend.asc @@ -308,9 +308,5 @@ 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)"); -} +// The PYBIND11_MODULE for rl_engine._C_npu lives in npu_module.cpp so that +// every Ascend op shares one compiled module. diff --git a/csrc/ascend/fused_logp_ascend.asc b/csrc/ascend/fused_logp_ascend.asc new file mode 100644 index 00000000..0b4c1aec --- /dev/null +++ b/csrc/ascend/fused_logp_ascend.asc @@ -0,0 +1,332 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Batch-invariant fused selected-token log-probability, Ascend C (CANN) +// forward kernel. +// +// logp[n] = logits[n, target[n]] - logsumexp(logits[n, :]) +// +// Mirrors the deterministic CUDA kernel in csrc/deterministic_logp_kernel.cu +// (DeterministicLogpCUDAOp): +// - input : logits [N, V] contiguous, bf16 / fp16 / fp32; target [N] +// int64, one per row +// - output : logp [N] fp32 (the CUDA deterministic op always returns fp32) +// - target[n] outside [0, V) -> logp[n] = 0.0 (same as the CUDA kernel) +// +// The math follows the same two-pass fixed-order reduction as the CUDA +// kernel: 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. The +// fp32 accumulation and the formula match the CUDA kernel exactly; the +// hardware reduction trees and the transcendental implementations are the +// Ascend vector unit's own (fixed per V), so cross-platform bitwise parity +// with the CUDA kernel is not claimed -- the guarantee here is the same one +// the CUDA kernel provides on its platform: batch-invariant determinism. +// +// Batch-invariance: every row is processed end-to-end by exactly one AI core +// block with a fixed tile size and a fixed reduction order. The instruction +// sequence for a row depends only on V, never on N or on the block the row +// happens to land on, so a row's output is bitwise identical across batch +// sizes, row positions, and block assignments on the NPU. +// +// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by +// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu. + +#include + +#include "kernel_operator.h" + +#include + +#include "torch_npu/csrc/core/npu/NPUStream.h" + +namespace { + +// Elements per vocab tile. Fixed for all rows and batch sizes; this is what +// makes the reduction order batch-invariant. UB budget (input tile + fp32 +// tile + reduce scratch) stays well under the 192 KB UB of current SoCs. +constexpr uint32_t TILE_LENGTH = 4096; +// Cap on launched blocks. Rows are strided across blocks, so launching fewer +// blocks than rows is fine and never changes per-row numerics. +constexpr int64_t MAX_BLOCKS = 128; +constexpr float NEG_INF = -3.402823466e+38f; // -FLT_MAX + +template +class KernelFusedLogp { +public: + __aicore__ inline KernelFusedLogp(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(GM_ADDR logits, + GM_ADDR target, + GM_ADDR logp, + int64_t numRows, + int64_t vocabSize) + { + numRows_ = numRows; + vocabSize_ = vocabSize; + logitsGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(logits)); + targetGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t*>(target)); + logpGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(logp)); + pipe_->InitBuffer(inQueue_, 1, TILE_LENGTH * sizeof(T)); + pipe_->InitBuffer(fp32Buf_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(reduceBuf_, TILE_LENGTH * sizeof(float)); + // 64 B: floats [0,8) for reduce/log scratch, floats [8,16) for the + // output staging slots (both halves 32-byte aligned for DataCopyPad). + pipe_->InitBuffer(scalarBuf_, 64); + // 32 B window for reading target[row] via DataCopyPad (GM scalar + // GetValue/SetValue are unreliable on hardware; see cannbot + // ascendc-precision-debug common-traps). + pipe_->InitBuffer(targetBuf_, 32); + + // 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. + eventVS_ = pipe_->FetchEventID(AscendC::HardEvent::V_S); + eventSV_ = pipe_->FetchEventID(AscendC::HardEvent::S_V); + eventMTE2S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE2_S); + eventSMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE3); + eventMTE3S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE3_S); + } + + __aicore__ inline void Process() + { + for (int64_t row = AscendC::GetBlockIdx(); row < numRows_; + row += AscendC::GetBlockNum()) { + ProcessRow(row); + } + } + +private: + // Load one vocab tile into UB and return its fp32 view. When T is fp32 the + // queue buffer is used in place; otherwise the tile is cast into fp32Buf_. + __aicore__ inline AscendC::LocalTensor LoadTileFp32(int64_t row, + int64_t start, + uint32_t count) + { + AscendC::LocalTensor xLocal = inQueue_.AllocTensor(); + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(xLocal, logitsGm_[row * vocabSize_ + start], copyParams, padParams); + inQueue_.EnQue(xLocal); + xLocal = inQueue_.DeQue(); + inTile_ = xLocal; + + if constexpr (std::is_same_v) { + return xLocal; + } else { + AscendC::LocalTensor fLocal = fp32Buf_.Get(); + AscendC::Cast(fLocal, xLocal, AscendC::RoundMode::CAST_NONE, count); + return fLocal; + } + } + + // Release the queue-owned input buffer of the current tile. + __aicore__ inline void FreeTile() + { + inQueue_.FreeTensor(inTile_); + } + + __aicore__ inline void ProcessRow(int64_t row) + { + const int64_t target = LoadTarget(row); + const bool valid = target >= 0 && target < vocabSize_; + const int64_t tileCount = (vocabSize_ + TILE_LENGTH - 1) / TILE_LENGTH; + float selected = 0.0f; + + // Pass 1: row max (fixed tile order). Also grab logits[target] on the fly. + float rowMax = NEG_INF; + for (int64_t tile = 0; tile < tileCount; ++tile) { + const int64_t start = tile * TILE_LENGTH; + const uint32_t count = TileCount(start); + AscendC::LocalTensor fLocal = LoadTileFp32(row, start, count); + + AscendC::LocalTensor rTmp = reduceBuf_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + AscendC::ReduceMax(scalar, fLocal, rTmp, static_cast(count), false); + WaitVector(); // vector -> scalar read + const float tileMax = scalar.GetValue(0); + rowMax = tileMax > rowMax ? tileMax : rowMax; + + if (valid && target >= start && target < start + count) { + selected = fLocal.GetValue(static_cast(target - start)); + } + FreeTile(); + } + + // Pass 2: sum(exp(x - rowMax)) with the same fixed tile order. + float sumExp = 0.0f; + for (int64_t tile = 0; tile < tileCount; ++tile) { + const int64_t start = tile * TILE_LENGTH; + const uint32_t count = TileCount(start); + AscendC::LocalTensor fLocal = LoadTileFp32(row, start, count); + + AscendC::Adds(fLocal, fLocal, -rowMax, count); // x - rowMax + AscendC::Exp(fLocal, fLocal, count); + AscendC::LocalTensor rTmp = reduceBuf_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + AscendC::ReduceSum(scalar, fLocal, rTmp, static_cast(count)); + WaitVector(); // vector -> scalar read + sumExp += scalar.GetValue(0); + FreeTile(); + } + + // lse = rowMax + log(sumExp). The scalar unit has no log, so run a + // 1-element vector Log (count padded to 8; scalarBuf_ is 32 B aligned). + AscendC::LocalTensor scalar = scalarBuf_.Get(); + scalar.SetValue(0, sumExp); + AscendC::SetFlag(eventSV_); // scalar write -> vector op + AscendC::WaitFlag(eventSV_); + AscendC::Log(scalar, scalar, 8); + WaitVector(); // vector -> scalar read + const float lse = rowMax + scalar.GetValue(0); + + // Stage the output in UB, then DataCopyPad to GM. GlobalTensor.SetValue + // is unreliable on hardware (cannbot ascendc-precision-debug + // common-traps), so scalar GM stores are not used. Out-of-range + // targets produce 0.0, matching the CUDA deterministic kernel. + scalar.SetValue(0, valid ? (selected - lse) : 0.0f); + AscendC::SetFlag(eventSMTE3_); // scalar write -> copy-out + AscendC::WaitFlag(eventSMTE3_); + AscendC::DataCopyExtParams outParams{1, sizeof(float), 0, 0, 0}; + AscendC::DataCopyPad(logpGm_[row], scalar[0], outParams); + // Drain MTE3 before the next row stages new values into scalarBuf_. + AscendC::SetFlag(eventMTE3S_); + AscendC::WaitFlag(eventMTE3S_); + } + + // Wait until all outstanding vector-pipe results are readable as scalars. + __aicore__ inline void WaitVector() + { + AscendC::SetFlag(eventVS_); + AscendC::WaitFlag(eventVS_); + } + + // Read target[row] through a 32-byte-aligned DataCopyPad window. + __aicore__ inline int64_t LoadTarget(int64_t row) + { + const int64_t alignedRow = row & ~3LL; // 4 x int64 per 32 B + const int64_t remaining = numRows_ - alignedRow; + const uint32_t winCount = static_cast(remaining < 4 ? remaining : 4); + AscendC::LocalTensor tLocal = targetBuf_.Get(); + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(winCount * sizeof(int64_t)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(tLocal, targetGm_[alignedRow], copyParams, padParams); + AscendC::SetFlag(eventMTE2S_); // copy-in -> scalar read + AscendC::WaitFlag(eventMTE2S_); + return static_cast(tLocal.GetValue(static_cast(row - alignedRow))); + } + + __aicore__ inline uint32_t TileCount(int64_t start) const + { + const int64_t remaining = vocabSize_ - start; + return static_cast(remaining < TILE_LENGTH ? remaining : TILE_LENGTH); + } + + AscendC::TPipe* pipe_; + AscendC::GlobalTensor logitsGm_; + AscendC::GlobalTensor targetGm_; + AscendC::GlobalTensor logpGm_; + AscendC::TQue inQueue_; + AscendC::TBuf fp32Buf_; + AscendC::TBuf reduceBuf_; + AscendC::TBuf scalarBuf_; + AscendC::TBuf targetBuf_; + AscendC::LocalTensor inTile_; + AscendC::TEventID eventVS_; + AscendC::TEventID eventSV_; + AscendC::TEventID eventMTE2S_; + AscendC::TEventID eventSMTE3_; + AscendC::TEventID eventMTE3S_; + int64_t numRows_; + int64_t vocabSize_; +}; + +} // namespace + +extern "C" __global__ __vector__ void fused_logp_ascend_kernel_fp32( + GM_ADDR logits, GM_ADDR target, GM_ADDR logp, + int64_t numRows, int64_t vocabSize) +{ + AscendC::TPipe pipe; + KernelFusedLogp op(&pipe); + op.Init(logits, target, logp, numRows, vocabSize); + op.Process(); +} + +extern "C" __global__ __vector__ void fused_logp_ascend_kernel_bf16( + GM_ADDR logits, GM_ADDR target, GM_ADDR logp, + int64_t numRows, int64_t vocabSize) +{ + AscendC::TPipe pipe; + KernelFusedLogp op(&pipe); + op.Init(logits, target, logp, numRows, vocabSize); + op.Process(); +} + +extern "C" __global__ __vector__ void fused_logp_ascend_kernel_fp16( + GM_ADDR logits, GM_ADDR target, GM_ADDR logp, + int64_t numRows, int64_t vocabSize) +{ + AscendC::TPipe pipe; + KernelFusedLogp op(&pipe); + op.Init(logits, target, logp, numRows, vocabSize); + op.Process(); +} + +torch::Tensor fused_logp_ascend_forward(torch::Tensor logits, torch::Tensor target) +{ + TORCH_CHECK(logits.is_privateuseone(), "logits must be on an NPU device"); + TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [N, V]"); + TORCH_CHECK(logits.is_contiguous(), "logits must be contiguous"); + TORCH_CHECK(logits.scalar_type() == at::kBFloat16 || logits.scalar_type() == at::kFloat || + logits.scalar_type() == at::kHalf, + "fused_logp_ascend supports fp32, fp16, and bf16 logits"); + TORCH_CHECK(logits.size(-1) > 0, "vocab size must be positive"); + TORCH_CHECK(target.is_privateuseone(), "target must be on the same NPU device as logits"); + TORCH_CHECK(target.dim() == 1, "target must be 1-D [N]"); + TORCH_CHECK(target.scalar_type() == at::kLong, "target must be int64"); + TORCH_CHECK(target.numel() == logits.size(0), "target must have one entry per row"); + + const int64_t numRows = logits.size(0); + const int64_t vocabSize = logits.size(1); + + // fp32 output, matching the CUDA deterministic logp op's contract. + torch::Tensor logp = at::empty({numRows}, logits.options().dtype(at::kFloat)); + if (numRows == 0) { + return logp; + } + + torch::Tensor targetContig = target.contiguous(); + + // stream(true): flush the task queue before launch so the kernel cannot + // overtake earlier NPU work; outputs were allocated with at::empty (no + // queued initializer) for the same reason. + auto aclStream = c10_npu::getCurrentNPUStream().stream(true); + const uint32_t blockNum = static_cast(std::min(numRows, MAX_BLOCKS)); + + if (logits.scalar_type() == at::kBFloat16) { + fused_logp_ascend_kernel_bf16<<>>( + reinterpret_cast(logits.mutable_data_ptr()), + reinterpret_cast(targetContig.mutable_data_ptr()), + reinterpret_cast(logp.mutable_data_ptr()), + numRows, vocabSize); + } else if (logits.scalar_type() == at::kHalf) { + fused_logp_ascend_kernel_fp16<<>>( + reinterpret_cast(logits.mutable_data_ptr()), + reinterpret_cast(targetContig.mutable_data_ptr()), + reinterpret_cast(logp.mutable_data_ptr()), + numRows, vocabSize); + } else { + fused_logp_ascend_kernel_fp32<<>>( + reinterpret_cast(logits.mutable_data_ptr()), + reinterpret_cast(targetContig.mutable_data_ptr()), + reinterpret_cast(logp.mutable_data_ptr()), + numRows, vocabSize); + } + return logp; +} + +// The PYBIND11_MODULE for rl_engine._C_npu lives in npu_module.cpp so that +// every Ascend op shares one compiled module. diff --git a/csrc/ascend/npu_module.cpp b/csrc/ascend/npu_module.cpp new file mode 100644 index 00000000..36c377d3 --- /dev/null +++ b/csrc/ascend/npu_module.cpp @@ -0,0 +1,28 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Pybind entry point for the rl_engine._C_npu extension. The Ascend C kernels +// and their torch host wrappers live in the sibling *.asc files; this TU only +// declares and binds them so every Ascend op shares one compiled module. +// +// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by +// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu. + +#include + +std::vector batch_invariant_logp_ascend_forward(torch::Tensor logits, + torch::Tensor target, + int64_t ignore_index); + +torch::Tensor fused_logp_ascend_forward(torch::Tensor logits, + torch::Tensor target); + +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("fused_logp_ascend", + &fused_logp_ascend_forward, + "Batch-invariant fused selected-token log-probability (Ascend C forward)"); +} diff --git a/docs/operators/fused-logp.md b/docs/operators/fused-logp.md index 7fc8008a..912cecb7 100644 --- a/docs/operators/fused-logp.md +++ b/docs/operators/fused-logp.md @@ -31,6 +31,7 @@ reference = logp_ref.forward_fp32(logits, token_ids) | --- | --- | --- | --- | | CUDA SM90 | `FusedLogpSM90Op` | `_C.fused_logp_sm90` | Experimental TMA-oriented path for 2D contiguous bf16 logits on Hopper-class GPUs. It is disabled by default and requires `RL_KERNEL_ENABLE_EXPERIMENTAL_SM90_LOGP=1`; otherwise the wrapper delegates to the CUDA generic fallback. | | CUDA generic | `FusedLogpGenericOp` | `_C.fused_logp` | Generic compiled extension fallback. | +| Ascend NPU | `FusedLogpAscendOp` | `_C_npu.fused_logp_ascend` | Batch-invariant Ascend C forward: two-pass (row max, then sum-exp) fp32 reduction with a fixed tile order, mirroring the CUDA deterministic kernel. Output is fp32, matching `DeterministicLogpCUDAOp`'s contract; out-of-range targets yield 0.0. | | PyTorch native | `NativeLogpOp` | None | PyTorch baseline/reference path. | ## Tensor Contract @@ -62,9 +63,14 @@ operator accuracy tests continue to validate native/CUDA fused API compatibility ## Implementation Files - `rl_engine/kernels/registry.py` -- `rl_engine/kernels/ops/pytorch/loss/logp.py` -- `rl_engine/kernels/ops/cuda/loss/logp.py` +- `rl_engine/kernels/ops/pytorch/loss/logp.py` — PyTorch native reference +- `rl_engine/kernels/ops/cuda/loss/logp.py` — CUDA fused LogP (SM90 + generic) +- `rl_engine/kernels/ops/ascend/loss/logp.py` — Ascend deterministic op - `csrc/ops.cpp` - `csrc/fused_logp_kernel.cu` - `csrc/cuda/fused_logp_sm90.cu` +- `csrc/deterministic_logp_kernel.cu` — CUDA deterministic kernel (reference reduction) +- `csrc/ascend/fused_logp_ascend.asc` — Ascend C forward kernel +- `csrc/ascend/npu_module.cpp` — shared pybind entry for `rl_engine._C_npu` - `tests/test_logp.py` +- `tests/test_logp_ascend.py` — Ascend correctness + batch-invariance tests diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index bff5e2e7..cc7a64a1 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,3 +8,7 @@ def batch_invariant_logp_ascend( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... +def fused_logp_ascend( + logits: torch.Tensor, + target: torch.Tensor, +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index ca4a462e..8f7c2568 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -103,6 +103,7 @@ def _load_object(path: str) -> Any: "cuda": "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp", "cuda-generic": "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp", "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpSM90Op", + "ascend": "rl_engine.kernels.ops.ascend.loss.logp.FusedLogpAscendOp", }, grad_input_names=("logits",), ), diff --git a/rl_engine/kernels/ops/ascend/loss/__init__.py b/rl_engine/kernels/ops/ascend/loss/__init__.py index 86cf4c9d..100f6068 100644 --- a/rl_engine/kernels/ops/ascend/loss/__init__.py +++ b/rl_engine/kernels/ops/ascend/loss/__init__.py @@ -1,2 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors + +from . import batch_invariant_logp # noqa: F401 +from . import logp # noqa: F401 diff --git a/rl_engine/kernels/ops/ascend/loss/logp.py b/rl_engine/kernels/ops/ascend/loss/logp.py new file mode 100644 index 00000000..17084086 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/loss/logp.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Any + +import torch + +from rl_engine.utils.logger import logger + +_C_npu: Any = None +try: + from rl_engine import _C_npu + + _NPU_EXT_AVAILABLE = True +except ImportError: # pragma: no cover - Ascend extension not built + _NPU_EXT_AVAILABLE = False + + +class _FusedLogpAscendAutograd(torch.autograd.Function): + """Autograd bridge for the Ascend fused selected-logprob forward. + + Mirrors the CUDA ``_FusedLogpAutograd``: the VJP is row-local + (``dlogits = grad * (one_hot(target) - softmax)``), computed in FP32 and + cast only the final input VJP back to the input dtype. There is no + cross-token reduction, so Batch/Chunk layout cannot change the result. + """ + + @staticmethod + def forward(ctx, logits: torch.Tensor, token_ids: torch.Tensor): + logits_2d = logits.reshape(-1, logits.size(-1)).contiguous() + labels = token_ids.reshape(-1).to(device=logits.device, dtype=torch.long).contiguous() + output = _C_npu.fused_logp_ascend(logits_2d, labels) + ctx.save_for_backward(logits_2d, labels) + ctx.input_shape = tuple(logits.shape) + ctx.input_dtype = logits.dtype + return output.reshape(logits.shape[:-1]) + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + logits, labels = ctx.saved_tensors + probs = torch.softmax(logits.float(), dim=-1) + rows = torch.arange(logits.size(0), device=logits.device) + probs[rows, labels] -= 1.0 + grad = -grad_output.reshape(-1, 1).float() * probs + return grad.to(ctx.input_dtype).reshape(ctx.input_shape), None + + +class FusedLogpAscendOp: + """Batch-invariant fused LogP for Ascend NPU. + + The Ascend C forward mirrors the deterministic CUDA kernel's two-pass + (row max, then sum-exp) fp32 reduction with a fixed tile order; the + output is fp32, matching ``DeterministicLogpCUDAOp``'s contract. + """ + + is_fused_logp = True + is_batch_invariant = True + + def __init__(self): + if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "fused_logp_ascend"): + raise RuntimeError( + "fused_logp_ascend is not compiled into the extension. Rebuild with " + "KERNEL_ALIGN_FORCE_ASCEND=1 on an Ascend NPU host: 'pip install -e .'" + ) + self.op = _C_npu.fused_logp_ascend + logger.info("Successfully linked to precompiled _C_npu.fused_logp_ascend kernel.") + + def __call__(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: + return self.apply(logits, token_ids) + + def _ascend_supported(self, logits: torch.Tensor) -> bool: + """NPU tensors only; bf16/fp16/fp32 (mirrors the CUDA kernel's gate).""" + return ( + logits.device.type == "npu" + and logits.is_contiguous() + and logits.dtype in (torch.bfloat16, torch.float16, torch.float32) + ) + + def apply(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: + if not self._ascend_supported(logits): + from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp + + return NativeLogpOp()(logits, token_ids) + return _FusedLogpAscendAutograd.apply(logits, token_ids) + + def apply_fp32(self, logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: + if not self._ascend_supported(logits): + from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp + + return NativeLogpOp().forward_fp32(logits, token_ids) + return _FusedLogpAscendAutograd.apply(logits, token_ids) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 12ea9b21..a5586d9c 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -107,6 +107,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ASCEND_BATCH_INVARIANT_LOGP = ( "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAscendOp" ) + ASCEND_FUSED_LOGP = "rl_engine.kernels.ops.ascend.loss.logp.FusedLogpAscendOp" # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3) PYTORCH_VOCAB_PARALLEL_LOGP = ( "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" @@ -622,6 +623,10 @@ def __init__(self): OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ] + self._priority_map["npu"]["logp"] = [ + OpBackend.ASCEND_FUSED_LOGP, + OpBackend.PYTORCH_NATIVE, + ] logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") self._adjust_priority_for_hardware() self._adjust_priority_from_env() diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index 388c4aec..ede6d95a 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -166,6 +166,10 @@ def fake_load_backend(backend): OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ] + assert registry._priority_map["npu"]["logp"] == [ + OpBackend.ASCEND_FUSED_LOGP, + OpBackend.PYTORCH_NATIVE, + ] def test_npu_available_handles_runtime_failure(monkeypatch): diff --git a/scripts/check_operator.py b/scripts/check_operator.py index 9dbca48d..ccf18a28 100644 --- a/scripts/check_operator.py +++ b/scripts/check_operator.py @@ -35,9 +35,24 @@ def _parse_dtype(value: str) -> torch.dtype: raise ValueError(f"unsupported dtype: {value}") +def _npu_available() -> bool: + """torch.npu only exists after torch_npu is imported; probe defensively.""" + try: + import torch_npu # noqa: F401 + except ImportError: + return False + return hasattr(torch, "npu") and torch.npu.is_available() + + def _select_device(value: str) -> torch.device: if value == "auto": - return torch.device("cuda" if torch.cuda.is_available() else "cpu") + if torch.cuda.is_available(): + return torch.device("cuda") + if _npu_available(): + return torch.device("npu") + return torch.device("cpu") + if value == "npu" and not _npu_available(): + raise RuntimeError("--device npu was requested, but no Ascend NPU is available") device = torch.device(value) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("--device cuda was requested, but CUDA is not available") @@ -73,7 +88,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--candidate", default="pytorch", - help="Candidate backend to validate, for example pytorch, cuda, cuda-sm90, triton.", + help="Candidate backend to validate, for example pytorch, cuda, cuda-sm90, triton, " + "ascend.", ) parser.add_argument("--dtype", choices=("fp32", "bf16", "fp16"), default="fp32") parser.add_argument("--device", default="auto") diff --git a/setup.py b/setup.py index 79f882d9..9c7cbd02 100644 --- a/setup.py +++ b/setup.py @@ -1,313 +1,430 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -import importlib.util -import os -import warnings -from pathlib import Path - -from setuptools import find_packages, setup - - -def _load_envs_module(): - envs_path = Path(__file__).with_name("envs.py") - spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"failed to load environment helpers from {envs_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -envs = _load_envs_module() - - -def _load_torch_extension_tools(): - try: - import torch - except ModuleNotFoundError as exc: - if exc.name != "torch": - raise - return None, None, None - - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - - # CUDAExtension is also the supported extension entry point for ROCm - # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when - # torch.version.hip is set. - return torch, BuildExtension, CUDAExtension - - -def _native_extension_required() -> bool: - """Whether the caller explicitly requested a native extension build.""" - return ( - envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) - or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) - or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) - or envs.env_flag("FORCE_CUDA") - ) - - -def _cuda_define_from_env(name: str, macro: str) -> list[str]: - value = os.environ.get(name) - if value is None: - return [] - parsed = int(value) - if parsed <= 0: - raise ValueError(f"{name} must be positive, got {value!r}") - return [f"-D{macro}={parsed}"] - - -_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( - "-Xfatbin", - "-compress-all", - "-gencode", - "--generate-code", - "--expt-", - "-lineinfo", - "-allow-unsupported-compiler", - "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", -) -_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { - "-Xfatbin", - "-gencode", - "--generate-code", -} - - -def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: - """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" - filtered_flags = [] - skip_next = False - for flag in flags: - if skip_next: - skip_next = False - continue - if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: - skip_next = True - continue - if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): - continue - filtered_flags.append(flag) - return filtered_flags - - -def get_extensions(): - torch, _, CUDAExtension = _load_torch_extension_tools() - if torch is None: - message = ( - "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " - "CUDA/ROCm PyTorch build first, then run " - "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." - ) - if _native_extension_required(): - raise RuntimeError(message) - warnings.warn( - f"{message} Continuing with the pure-Python fallback because no native extension " - "was explicitly requested.", - RuntimeWarning, - stacklevel=2, - ) - return [] - - extensions = [] - torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") - torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] - if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": - torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") - is_rocm = getattr(torch.version, "hip", None) is not None - - # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, - # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also - # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add - # --offload-arch. Do not require a visible GPU when a ROCm target was - # explicitly selected. - no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() - if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: - raise RuntimeError( - "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " - "Set one or more ';'-separated targets, for example " - "PYTORCH_ROCM_ARCH='gfx942;gfx950'." - ) - - if is_rocm or torch.cuda.is_available(): - cuda_sources = [ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cu", - "csrc/deterministic_logp_kernel.cu", - "csrc/cuda/gemm/det_gemm_kernel.cu", - "csrc/cuda/rmsnorm.cu", - "csrc/cuda/activation.cu", - "csrc/cuda/attention/deterministic_attention.cu", - "csrc/cuda/distributed/deterministic_collective.cu", - ] - if not is_rocm: - # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). - # The ROCm dispatcher falls back to PyTorch SDPA for this operator. - cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") - - nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): - nvcc_flags.append("--use_fast_math") - if not is_rocm: - cc_major, cc_minor = torch.cuda.get_device_capability() - enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" - if not enable_sm90: - # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. - nvcc_flags.append( - f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" - ) - nvcc_flags.append("--expt-relaxed-constexpr") - nvcc_flags.append("--expt-extended-lambda") - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - ) - ) - if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): - nvcc_flags.append("-lineinfo") - if ( - not is_rocm - and os.name == "nt" - and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) - ): - nvcc_flags.append("-allow-unsupported-compiler") - nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") - - cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] - extra_link_args = list(torch_rpath) - if os.name != "nt": - # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). - extra_link_args.append("-lcuda") - - if not is_rocm: - sm90_srcs = [ - "csrc/cuda/fused_logp_sm90.cu", - "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob - "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp - "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - # Single-card batch-invariant embedding/lm-head. - "csrc/cuda/embedding_lm_head_sm90.cu", - ] - enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) - present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] - 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) +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import importlib.util +import os +import platform +import subprocess +import sysconfig +import warnings +from pathlib import Path + +from setuptools import Extension, find_packages, setup + + +def _load_envs_module(): + envs_path = Path(__file__).with_name("envs.py") + spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"failed to load environment helpers from {envs_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +envs = _load_envs_module() + + +def _load_torch_extension_tools(): + try: + import torch + except ModuleNotFoundError as exc: + if exc.name != "torch": + raise + return None, None, None + + from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + # CUDAExtension is also the supported extension entry point for ROCm + # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when + # torch.version.hip is set. + return torch, BuildExtension, CUDAExtension + + +def _native_extension_required() -> bool: + """Whether the caller explicitly requested a native extension build.""" + return ( + envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) + or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) + or envs.env_flag("FORCE_CUDA") + ) + + +def _cuda_define_from_env(name: str, macro: str) -> list[str]: + value = os.environ.get(name) + if value is None: + return [] + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive, got {value!r}") + return [f"-D{macro}={parsed}"] + + +_ASCEND_EXTENSION_NAME = "rl_engine._C_npu" +_ASCEND_CPU_DIRS = {"aarch64": "aarch64-linux", "x86_64": "x86_64-linux"} + + +def _find_ascend_home() -> str: + """Locate the CANN toolkit root (must contain bin/bisheng).""" + candidates = [ + os.environ.get("ASCEND_HOME_PATH"), + os.environ.get("ASCEND_TOOLKIT_HOME"), + ] + candidates += [str(p) for p in sorted(Path.home().glob("Ascend/cann-*"), reverse=True)] + candidates.append("/usr/local/Ascend/ascend-toolkit/latest") + for cand in candidates: + if cand and (Path(cand) / "bin" / "bisheng").is_file(): + # The bisheng driver and its Ascend C plugin resolve toolkit data + # (impl include dirs, stub JSON generation) through these env + # vars. Without ASCEND_HOME_PATH the plugin crashes (segfault) + # while compiling any kernel TU, so export them for the compiler + # subprocesses once we know where the toolkit lives. + os.environ["ASCEND_HOME_PATH"] = cand + os.environ.setdefault("ASCEND_TOOLKIT_HOME", cand) + return cand + raise RuntimeError( + "KERNEL_ALIGN_FORCE_ASCEND=1 was requested but no CANN toolkit with bin/bisheng " + "was found. Set ASCEND_HOME_PATH to the toolkit root." + ) + + +def _ascend_extension_spec() -> Extension: + sources = ["csrc/ascend/npu_module.cpp"] + sources += sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) + ext = Extension(name=_ASCEND_EXTENSION_NAME, sources=sources) + ext._rl_kernel_ascend = True # intercepted by the custom build_ext below + return ext + + +def _compile_ascend_extension(build_ext, ext) -> None: + """Compile the Ascend C extension with bisheng (torch's BuildExtension + does not know the .asc language, so we drive the compiler directly).""" + torch, _, _ = _load_torch_extension_tools() + try: + import torch_npu + except ModuleNotFoundError as exc: + raise RuntimeError( + "KERNEL_ALIGN_FORCE_ASCEND=1 requires torch_npu. Install a matching " + "torch_npu build first." + ) from exc + + ascend_home = _find_ascend_home() + cpu_dir = _ASCEND_CPU_DIRS.get(platform.machine()) + if cpu_dir is None: + raise RuntimeError(f"unsupported Ascend host architecture: {platform.machine()}") + arch = os.environ.get(envs.KERNEL_ALIGN_ASCEND_ARCH, "dav-c220") + + bisheng = os.path.join(ascend_home, "bin", "bisheng") + torch_dir = os.path.dirname(torch.__file__) + tnpu_dir = os.path.dirname(torch_npu.__file__) + + includes = [ + f"-I{os.path.join(ascend_home, cpu_dir, 'asc', 'include')}", + f"-I{os.path.join(torch_dir, 'include')}", + f"-I{os.path.join(torch_dir, 'include', 'torch', 'csrc', 'api', 'include')}", + f"-I{os.path.join(tnpu_dir, 'include')}", + f"-I{sysconfig.get_paths()['include']}", + ] + defines = [f"-DTORCH_EXTENSION_NAME={_ASCEND_EXTENSION_NAME.rsplit('.', 1)[-1]}"] + + build_temp = os.path.join(build_ext.build_temp, "ascend") + os.makedirs(build_temp, exist_ok=True) + + objects = [] + for src in ext.sources: + obj = os.path.join(build_temp, Path(src).name + ".o") + cmd = [bisheng, "-std=c++17", "-O2", "-fPIC", "-c"] + if src.endswith(".asc"): + cmd += ["-x", "asc", f"--cce-aicore-arch={arch}"] + cmd += includes + defines + [src, "-o", obj] + subprocess.check_call(cmd) + objects.append(obj) + + out_path = build_ext.get_ext_fullpath(ext.name) + os.makedirs(os.path.dirname(out_path), exist_ok=True) + link = [bisheng, "-shared", *objects] + for lib_dir, libs in ( + (os.path.join(torch_dir, "lib"), ["torch", "torch_cpu", "torch_python", "c10"]), + (os.path.join(tnpu_dir, "lib"), ["torch_npu"]), + (os.path.join(ascend_home, "runtime", "lib64"), ["ascendcl"]), + (os.path.join(ascend_home, cpu_dir, "lib64"), ["runtime"]), + ): + link.append(f"-L{lib_dir}") + link += [f"-l{name}" for name in libs] + link += [ + f"-Wl,-rpath,{os.path.join(torch_dir, 'lib')}", + f"-Wl,-rpath,{os.path.join(tnpu_dir, 'lib')}", + "-o", + out_path, + ] + subprocess.check_call(link) + + +def _make_build_extension(BuildExtension): + class AscendAwareBuildExtension(BuildExtension): + def build_extension(self, ext): + if getattr(ext, "_rl_kernel_ascend", False): + _compile_ascend_extension(self, ext) + return + super().build_extension(ext) + + return AscendAwareBuildExtension + + +_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( + "-Xfatbin", + "-compress-all", + "-gencode", + "--generate-code", + "--expt-", + "-lineinfo", + "-allow-unsupported-compiler", + "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", +) +_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { + "-Xfatbin", + "-gencode", + "--generate-code", +} + + +def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: + """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" + filtered_flags = [] + skip_next = False + for flag in flags: + if skip_next: + skip_next = False + continue + if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: + skip_next = True + continue + if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): + continue + filtered_flags.append(flag) + return filtered_flags + + +def get_extensions(): + torch, _, CUDAExtension = _load_torch_extension_tools() + if torch is None: + message = ( + "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " + "CUDA/ROCm PyTorch build first, then run " + "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." + ) + if _native_extension_required(): + raise RuntimeError(message) + warnings.warn( + f"{message} Continuing with the pure-Python fallback because no native extension " + "was explicitly requested.", + RuntimeWarning, + stacklevel=2, + ) + return [] + + extensions = [] + torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") + torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] + if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": + torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") + is_rocm = getattr(torch.version, "hip", None) is not None + + # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, + # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also + # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add + # --offload-arch. Do not require a visible GPU when a ROCm target was + # explicitly selected. + no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() + if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: + raise RuntimeError( + "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " + "Set one or more ';'-separated targets, for example " + "PYTORCH_ROCM_ARCH='gfx942;gfx950'." + ) + + if is_rocm or torch.cuda.is_available(): + cuda_sources = [ + "csrc/ops.cpp", + "csrc/fused_logp_kernel.cu", + "csrc/deterministic_logp_kernel.cu", + "csrc/cuda/gemm/det_gemm_kernel.cu", + "csrc/cuda/rmsnorm.cu", + "csrc/cuda/activation.cu", + "csrc/cuda/attention/deterministic_attention.cu", + "csrc/cuda/distributed/deterministic_collective.cu", + ] + if not is_rocm: + # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). + # The ROCm dispatcher falls back to PyTorch SDPA for this operator. + cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") + + nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] + if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + nvcc_flags.append("--use_fast_math") + if not is_rocm: + cc_major, cc_minor = torch.cuda.get_device_capability() + enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" + if not enable_sm90: + # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. + nvcc_flags.append( + f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" + ) + nvcc_flags.append("--expt-relaxed-constexpr") + nvcc_flags.append("--expt-extended-lambda") + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + ) + ) + if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): + nvcc_flags.append("-lineinfo") + if ( + not is_rocm + and os.name == "nt" + and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) + ): + nvcc_flags.append("-allow-unsupported-compiler") + nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") + + cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] + extra_link_args = list(torch_rpath) + if os.name != "nt": + # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). + extra_link_args.append("-lcuda") + + if not is_rocm: + sm90_srcs = [ + "csrc/cuda/fused_logp_sm90.cu", + "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp + "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build + # Single-card batch-invariant embedding/lm-head. + "csrc/cuda/embedding_lm_head_sm90.cu", + ] + enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) + present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] + 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}") - cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - - # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp - # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in - # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. - enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" - if enable_det_gemm_sm90: - tma_arch = f"{cc_major}{cc_minor}a" - arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" - if arch_flag not in nvcc_flags: - nvcc_flags.append(arch_flag) - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") - cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") - - if is_rocm: - nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) - - extensions.append( - CUDAExtension( - name="rl_engine._C", - sources=cuda_sources, - include_dirs=[], - extra_compile_args={ - "cxx": cxx_flags, - "nvcc": nvcc_flags, - }, - extra_link_args=extra_link_args, - ) - ) - - if _native_extension_required() and not extensions: - raise RuntimeError( - "rl_engine._C was requested but no CUDA/ROCm build environment is available. " - "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " - "PYTORCH_ROCM_ARCH to the target architecture." - ) - - return extensions - - -def get_cmdclass(): - _, BuildExtension, _ = _load_torch_extension_tools() - if BuildExtension is None: - return {} - return {"build_ext": BuildExtension} - - -setup( - name="rl-engine", - version="0.1.0", - packages=find_packages(include=["rl_engine", "rl_engine.*"]), - install_requires=[ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", - ], - ext_modules=get_extensions(), - cmdclass=get_cmdclass(), - extras_require={ - "cuda": ["flashinfer"], - "rocm": ["aiter"], - "vllm": ["vllm>=0.6.0"], - "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], - }, - entry_points={ - "console_scripts": [ - "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", - ], - }, - python_requires=">=3.10", - include_package_data=True, - zip_safe=False, -) + cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + + # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp + # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in + # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. + enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" + if enable_det_gemm_sm90: + tma_arch = f"{cc_major}{cc_minor}a" + arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" + if arch_flag not in nvcc_flags: + nvcc_flags.append(arch_flag) + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") + cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + if is_rocm: + nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) + + extensions.append( + CUDAExtension( + name="rl_engine._C", + sources=cuda_sources, + include_dirs=[], + extra_compile_args={ + "cxx": cxx_flags, + "nvcc": nvcc_flags, + }, + extra_link_args=extra_link_args, + ) + ) + + if envs.env_flag(envs.KERNEL_ALIGN_FORCE_ASCEND): + extensions.append(_ascend_extension_spec()) + + if _native_extension_required() and not extensions: + raise RuntimeError( + "rl_engine._C was requested but no CUDA/ROCm build environment is available. " + "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " + "PYTORCH_ROCM_ARCH to the target architecture." + ) + + return extensions + + +def get_cmdclass(): + _, BuildExtension, _ = _load_torch_extension_tools() + if BuildExtension is None: + return {} + return {"build_ext": _make_build_extension(BuildExtension)} + + +setup( + name="rl-engine", + version="0.1.0", + packages=find_packages(include=["rl_engine", "rl_engine.*"]), + install_requires=[ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", + ], + ext_modules=get_extensions(), + cmdclass=get_cmdclass(), + extras_require={ + "cuda": ["flashinfer"], + "rocm": ["aiter"], + "vllm": ["vllm>=0.6.0"], + "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], + }, + entry_points={ + "console_scripts": [ + "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", + ], + }, + python_requires=">=3.10", + include_package_data=True, + zip_safe=False, +) diff --git a/tests/test_logp_ascend.py b/tests/test_logp_ascend.py new file mode 100644 index 00000000..9197c433 --- /dev/null +++ b/tests/test_logp_ascend.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for the Ascend NPU batch-invariant fused selected-token logp. + +Validates the same two orthogonal properties as the CUDA deterministic op: + +1. **Correctness** - output matches the ``NativeLogpOp.forward_fp32`` ground + truth within the logprob contract tolerance. The Ascend C kernel mirrors + the CUDA deterministic kernel's two-pass (row max, then sum-exp) fp32 + reduction with a fixed tile order; the hardware reduction trees differ + from CUDA's, so the comparison is tolerance-based (fp32 drift ~1e-7). +2. **Batch-invariance** - a row's logp is bitwise identical regardless of + batch size, batch position, or how many AI-core blocks were launched + (each row is reduced end-to-end by one block; no split-K merge exists). +""" + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp + +# Accuracy tolerances from the gtest contract, "logprob" op class. +_ATOL = { + torch.float32: 1.0e-5, + torch.bfloat16: 6.0e-2, + torch.float16: 5.0e-3, +} + + +def _npu_available() -> bool: + try: + import torch_npu # noqa: F401 + except ImportError: + return False + return hasattr(torch, "npu") and torch.npu.is_available() + + +def _ascend_kernel_available() -> bool: + if not _npu_available(): + return False + try: + from rl_engine.kernels.ops.ascend.loss.logp import _NPU_EXT_AVAILABLE, _C_npu + except Exception: + return False + return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "fused_logp_ascend") + + +requires_ascend = pytest.mark.skipif( + not _ascend_kernel_available(), + reason="fused_logp_ascend kernel not compiled " + "(needs KERNEL_ALIGN_FORCE_ASCEND=1 on an Ascend NPU host).", +) + + +def _get_op(): + from rl_engine.kernels.ops.ascend.loss.logp import FusedLogpAscendOp + + return FusedLogpAscendOp() + + +def _make_inputs(shape, vocab, dtype, seed=0): + generator = torch.Generator(device="cpu").manual_seed(seed) + logits = torch.randn(*shape, vocab, dtype=dtype, generator=generator).to("npu") + token_ids = torch.randint(0, vocab, shape, generator=generator).long().to("npu") + return logits, token_ids + + +# --------------------------------------------------------------------------- +# Correctness +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@requires_ascend +class TestAscendFusedLogpCorrectness: + def test_forward_matches_pytorch_reference(self, dtype): + op = _get_op() + logits, token_ids = _make_inputs((3, 5), 257, dtype) + out = op(logits, token_ids) + ref = NativeLogpOp().forward_fp32(logits, token_ids) + assert out.dtype == torch.float32 # matches DeterministicLogpCUDAOp's contract + assert out.shape == (3, 5) + assert torch.allclose(out.float(), ref, atol=_ATOL[dtype], rtol=0.0) + + def test_apply_fp32_matches_reference(self, dtype): + op = _get_op() + logits, token_ids = _make_inputs((3, 5), 257, dtype) + out = op.apply_fp32(logits, token_ids) + ref = NativeLogpOp().forward_fp32(logits, token_ids) + assert torch.allclose(out.float(), ref, atol=_ATOL[dtype], rtol=0.0) + + def test_out_of_range_target_is_zero(self, dtype): + op = _get_op() + logits, token_ids = _make_inputs((2, 4), 32, dtype) + token_ids = token_ids.reshape(-1) + token_ids[1] = 32 + 5 # out of [0, V) + out = op(logits, token_ids.reshape(2, 4)) + assert out.reshape(-1)[1].item() == 0.0 + + def test_backward_matches_native_reference(self, dtype): + op = _get_op() + logits, token_ids = _make_inputs((3, 5), 257, dtype) + + logits_a = logits.clone().requires_grad_() + op(logits_a, token_ids).backward(torch.ones(3, 5, device="npu", dtype=dtype)) + + logits_n = logits.clone().requires_grad_() + NativeLogpOp()(logits_n, token_ids).backward(torch.ones(3, 5, device="npu", dtype=dtype)) + + assert torch.allclose( + logits_a.grad.float(), logits_n.grad.float(), atol=1.0e-4, rtol=1.0e-4 + ) + + def test_backward_has_no_cross_row_leak(self, dtype): + """Row-local VJP: the same row's grad is bitwise identical wherever the + row sits in the batch.""" + op = _get_op() + logits, token_ids = _make_inputs((4, 8), 257, dtype) + logits[1].copy_(logits[0]) + token_ids[1] = token_ids[0] + + logits_g = logits.clone().requires_grad_() + grad_out = torch.randn(4, 8, device="npu", dtype=dtype) + grad_out[1] = grad_out[0] # identical (logits, target, dy) triples + op(logits_g, token_ids).backward(grad_out) + grad = logits_g.grad + # Rows 0 and 1 saw identical inputs, so their VJPs must be bitwise + # identical; rows 2/3 stay independent. + assert torch.equal(grad[0], grad[1]) + for row in range(2, 4): + assert not torch.equal(grad[0], grad[row]) + +# --------------------------------------------------------------------------- +# Fallback +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendFusedLogpFallback: + def test_rejects_non_npu_falls_back_to_native(self): + op = _get_op() + logits, token_ids = _make_inputs((2, 3), 17, torch.float32) + out = op(logits.cpu(), token_ids.cpu()) + ref = NativeLogpOp()(logits.cpu(), token_ids.cpu()) + assert torch.equal(out, ref) + + +# --------------------------------------------------------------------------- +# Batch invariance +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendFusedLogpBatchInvariance: + def _run_row(self, batch, vocab, dtype, pos, seed=7): + """One fixed row embedded at position `pos` of a random batch.""" + op = _get_op() + logits, token_ids = _make_inputs((batch,), vocab, dtype, seed=seed) + out = op(logits, token_ids) + return out[pos].clone() + + def test_batch_size_1_vs_n(self): + # One fixed (row, target) pair embedded in batches of growing size: + # its logp must be bitwise identical regardless of batch size. + dtype = torch.bfloat16 + op = _get_op() + alone_logits, alone_ids = _make_inputs((1,), 257, dtype, seed=7) + alone = op(alone_logits, alone_ids)[0] + for batch in (2, 4, 16, 300): # 300 > MAX_BLOCKS -> strided blocks + logits, token_ids = _make_inputs((batch,), 257, dtype, seed=7) + logits[0].copy_(alone_logits[0]) + token_ids[0] = alone_ids[0] + in_batch = op(logits, token_ids)[0] + assert torch.equal(alone, in_batch), f"drift at batch_size={batch}" + + def test_different_positions_in_batch(self): + # The same row content copied to every position reduces bitwise + # identically regardless of where it lands. + dtype = torch.float16 + op = _get_op() + logits, token_ids = _make_inputs((8,), 257, dtype, seed=11) + base, base_id = logits[0].clone(), int(token_ids[0]) + for pos in range(1, 8): + logits[pos].copy_(base) + token_ids[pos] = base_id + out = op(logits, token_ids) + for pos in range(1, 8): + assert torch.equal(out[pos], out[0]), f"drift at position={pos}" + + def test_block_striding(self): + # 300 rows > MAX_BLOCKS (128): rows are strided across blocks, so + # numerics must not depend on block assignment. + dtype = torch.bfloat16 + op = _get_op() + logits, token_ids = _make_inputs((300,), 257, dtype, seed=13) + out = op(logits, token_ids) + again = op(logits, token_ids) + assert torch.equal(out, again) + + def test_multi_tile_rows(self): + # vocab > TILE_LENGTH (4096): rows span multiple fixed-order tiles. + dtype = torch.float32 + op = _get_op() + logits, token_ids = _make_inputs((4,), 10000, dtype, seed=5) + out = op(logits, token_ids) + assert torch.equal(out, op(logits, token_ids)) + + def test_repeated_runs_deterministic(self): + dtype = torch.bfloat16 + logits, token_ids = _make_inputs((3, 5), 257, dtype, seed=5) + op = _get_op() + first = op(logits, token_ids) + for _ in range(3): + again = op(logits, token_ids) + assert torch.equal(first, again) + + +# --------------------------------------------------------------------------- +# Registry dispatch +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendRegistryDispatch: + def test_get_op_logp(self): + from rl_engine.kernels.registry import kernel_registry + + op = kernel_registry.get_op("logp", device="npu") + assert type(op).__name__ == "FusedLogpAscendOp"