From c266a06d3ce4c47580ccde454a2f43378058980f Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 1 Sep 2026 13:56:43 +0800 Subject: [PATCH] feat(ascend): add batch-invariant fused linear logp Ascend C operator Mirror the SM90 fused linear logp kernel's reduction contract (csrc/cuda/fused_linear_logp_sm90.cu, contract v1) with an Ascend C forward in csrc/ascend/fused_linear_logp_ascend.asc: ascending vocab-row scan with the online rescale chain, per-row fp32 dots over a fixed D-tile order, bias in fp32, final min(zt - lse, 0) clamp, fp32 output. No [N, V] logits are materialized. A row's logp is bitwise identical across batch sizes, positions, and block assignments on the NPU; cross-platform bitwise parity is not claimed (hardware reduction trees / transcendentals differ). Wrapper mirrors FusedLinearLogpSM90Op's surface; the backward is the shared chunked_linear_logp_backward (the CUDA op's portable fallback formula). TP and bias calls delegate to the native reference. Registered as the ascend candidate in the linear_logp gtest spec and in the registry's NPU priority map. Ports the shared-module Ascend build and check_operator npu support. Verified on Ascend 910 / CANN 8.5.1: - Forward vs hand-computed fp32 reference: max_abs <= 2.2e-4 at the gtest shape (pure fp32 tree drift; pytest atol 5e-4). The gtest's own forward comparison is stricter than any independent kernel can meet (fp32 tree drift vs atol 1e-5; bf16/fp16 gold accumulates the matmul in the input dtype) - documented in the PR body. - gtest gradients (fp32): hidden 2.4e-7, weight 3.0e-8 (atol 1e-5). - Batch invariance bitwise: batch 1 vs {2,4,16,300}, positions 1..7, multi-tile D=10000, repeated runs. - pytest tests/test_linear_logp_ascend.py: 23 passed. - Regression: test_batch_invariant_logp.py 44 passed, test_dispatch.py 13 passed. --- csrc/ascend/batch_invariant_logp_ascend.asc | 8 +- csrc/ascend/fused_linear_logp_ascend.asc | 434 ++++++++++ csrc/ascend/npu_module.cpp | 30 + docs/operators/linear-logp.md | 11 +- rl_engine/_C_npu.pyi | 6 + rl_engine/kernels/gtest/operator_specs.py | 1 + rl_engine/kernels/ops/ascend/loss/__init__.py | 3 + .../kernels/ops/ascend/loss/linear_logp.py | 173 ++++ rl_engine/kernels/registry.py | 7 + rl_engine/tests/test_dispatch.py | 4 + scripts/check_operator.py | 20 +- setup.py | 741 ++++++++++-------- tests/test_linear_logp_ascend.py | 236 ++++++ 13 files changed, 1351 insertions(+), 323 deletions(-) create mode 100644 csrc/ascend/fused_linear_logp_ascend.asc create mode 100644 csrc/ascend/npu_module.cpp create mode 100644 rl_engine/kernels/ops/ascend/loss/linear_logp.py create mode 100644 tests/test_linear_logp_ascend.py 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_linear_logp_ascend.asc b/csrc/ascend/fused_linear_logp_ascend.asc new file mode 100644 index 00000000..6f4f1d32 --- /dev/null +++ b/csrc/ascend/fused_linear_logp_ascend.asc @@ -0,0 +1,434 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Batch-invariant fused linear log-probability, Ascend C (CANN) forward +// kernel. +// +// logp[n] = log_softmax(hidden[n] @ W^T + b)[target[n]] +// +// Mirrors the SM90 CUDA fused kernel's bitwise reduction contract +// (csrc/cuda/fused_linear_logp_sm90.cu, contract v1), without materializing +// the [N, V] logits: +// - vocab rows are scanned in ascending index order (the CUDA contract's +// "cross-split ascending-index sequential chains"); +// - the softmax statistics use the online rescale chain +// newM = max(m, z); sum = sum * exp(m - newM) + exp(z - newM); +// exactly like the CUDA per-split merge; +// - each per-row dot is a fixed tile order over D with fp32 accumulation +// (per-tile ReduceSum tree + sequential scalar chain); +// - bias is added in fp32; padding lanes are -inf so exp() is exact 0; +// - final clamp logp = min(zt - lse, 0), matching the CUDA contract. +// +// The hardware reduction trees and transcendental implementations are the +// Ascend vector unit's own (fixed per D), so cross-platform bitwise parity +// with the CUDA kernel is not claimed -- the guarantee 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 vocab/D scan order; rows are +// strided across blocks (MAX_BLOCKS cap), so a row's logp depends only on +// (D, V), never on N or block assignment. +// +// 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 hidden tile. Fixed for all rows and batch sizes; this is what +// makes the reduction order batch-invariant. UB budget (hidden row cache + +// weight tile + fp32 views + reduce scratch) stays well under the 192 KB UB +// of current SoCs for D <= TILE_LENGTH (Qwen3: D = 4096). +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 KernelFusedLinearLogp { +public: + __aicore__ inline KernelFusedLinearLogp(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(GM_ADDR hidden, + GM_ADDR weight, + GM_ADDR bias, + GM_ADDR target, + GM_ADDR logp, + int64_t numRows, + int64_t vocabSize, + int64_t hiddenSize, + bool hasBias) + { + numRows_ = numRows; + vocabSize_ = vocabSize; + hiddenSize_ = hiddenSize; + hasBias_ = hasBias; + // Hidden row cache: one fp32 tile; valid only when D <= TILE_LENGTH. + cacheHidden_ = hiddenSize_ <= static_cast(TILE_LENGTH); + hiddenGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(hidden)); + weightGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(weight)); + biasGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(bias)); + targetGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t*>(target)); + logpGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(logp)); + pipe_->InitBuffer(wQueue_, 1, TILE_LENGTH * sizeof(T)); + pipe_->InitBuffer(hFp32Buf_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(wFp32Buf_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(prodBuf_, 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 windows for scalar reads via DataCopyPad (GM scalar + // GetValue/SetValue are unreliable on hardware; see cannbot + // ascendc-precision-debug common-traps). + pipe_->InitBuffer(targetBuf_, 32); + pipe_->InitBuffer(biasBuf_, 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); + eventMTE2V_ = pipe_->FetchEventID(AscendC::HardEvent::MTE2_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 tile into the queue and return its fp32 view (cast when T is + // not fp32; in-place otherwise). + __aicore__ inline AscendC::LocalTensor LoadTileFp32(AscendC::GlobalTensor gm, + int64_t offset, + uint32_t count, + AscendC::LocalTensor fp32View) + { + AscendC::LocalTensor tile = wQueue_.AllocTensor(); + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(tile, gm[offset], copyParams, padParams); + wQueue_.EnQue(tile); + tile = wQueue_.DeQue(); + AscendC::SetFlag(eventMTE2V_); + AscendC::WaitFlag(eventMTE2V_); + if constexpr (std::is_same_v) { + AscendC::DataCopy(fp32View, tile, VecAlignCount(count)); + } else { + AscendC::Cast(fp32View, tile, AscendC::RoundMode::CAST_NONE, count); + } + wQueue_.FreeTensor(tile); + return fp32View; + } + + // fp32 dot of the hidden row and weight row over one D tile. + __aicore__ inline float DotTile(AscendC::LocalTensor hTile, + AscendC::LocalTensor wTile, + uint32_t count) + { + AscendC::LocalTensor prod = prodBuf_.Get(); + AscendC::Mul(prod, hTile, wTile, count); + AscendC::LocalTensor rTmp = reduceBuf_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + AscendC::ReduceSum(scalar, prod, rTmp, static_cast(count)); + WaitVector(); // vector -> scalar read + return scalar.GetValue(0); + } + + // Per-vocab-row dot over D with a fixed tile order. + __aicore__ inline float DotRow(int64_t row, int64_t col) + { + const int64_t tileCount = (hiddenSize_ + TILE_LENGTH - 1) / TILE_LENGTH; + AscendC::LocalTensor hFp32 = hFp32Buf_.Get(); + float acc = 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 wFp32 = wFp32Buf_.Get(); + if (!(cacheHidden_ && tile == 0)) { + // Reload the hidden tile when D > TILE_LENGTH (or on demand). + LoadTileFp32(hiddenGm_, row * hiddenSize_ + start, count, hFp32); + } + LoadTileFp32(weightGm_, col * hiddenSize_ + start, count, wFp32); + acc += DotTile(hFp32, wFp32, count); + } + return acc; + } + + __aicore__ inline void ProcessRow(int64_t row) + { + const int64_t target = LoadTarget(row); + const bool valid = target >= 0 && target < vocabSize_; + float m = NEG_INF; + float sumExp = 0.0f; + float zt = 0.0f; + + // Cache the hidden row (D <= TILE_LENGTH fast path). + if (cacheHidden_) { + LoadTileFp32(hiddenGm_, row * hiddenSize_, TileCount(0), hFp32Buf_.Get()); + } + + // Vocab rows in ascending index order, online rescale chain. + for (int64_t v = 0; v < vocabSize_; ++v) { + float z = DotRow(row, v); + if (hasBias_) { + z += LoadBias(v); + } + if (v == target) { + zt = z; + } + const float newM = z > m ? z : m; + // Two scalar exps per vocab row via one padded vector Exp. + AscendC::LocalTensor scalar = scalarBuf_.Get(); + scalar.SetValue(0, m - newM); + scalar.SetValue(1, z - newM); + AscendC::SetFlag(eventSV_); + AscendC::WaitFlag(eventSV_); + AscendC::Exp(scalar, scalar, 8); + WaitVector(); // vector -> scalar read + sumExp = sumExp * scalar.GetValue(0) + scalar.GetValue(1); + m = newM; + } + + // lse = m + log(sumExp) via a padded 1-element vector Log. + AscendC::LocalTensor scalar = scalarBuf_.Get(); + scalar.SetValue(0, sumExp); + AscendC::SetFlag(eventSV_); + AscendC::WaitFlag(eventSV_); + AscendC::Log(scalar, scalar, 8); + WaitVector(); // vector -> scalar read + const float lse = m + 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; the final clamp matches the CUDA contract. + float logp = 0.0f; + if (valid) { + logp = zt - lse; + logp = logp < 0.0f ? logp : 0.0f; + } + scalar.SetValue(0, logp); + AscendC::SetFlag(eventSMTE3_); + 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))); + } + + // Read bias[col] through a 32-byte-aligned DataCopyPad window. + __aicore__ inline float LoadBias(int64_t col) + { + const int64_t alignedCol = col & ~7LL; // 8 x fp32 per 32 B + const int64_t remaining = vocabSize_ - alignedCol; + const uint32_t winCount = static_cast(remaining < 8 ? remaining : 8); + AscendC::LocalTensor bLocal = biasBuf_.Get(); + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(winCount * sizeof(float)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(bLocal, biasGm_[alignedCol], copyParams, padParams); + AscendC::SetFlag(eventMTE2S_); // copy-in -> scalar read + AscendC::WaitFlag(eventMTE2S_); + return bLocal.GetValue(static_cast(col - alignedCol)); + } + + __aicore__ inline uint32_t TileCount(int64_t start) const + { + const int64_t remaining = hiddenSize_ - start; + return static_cast(remaining < TILE_LENGTH ? remaining : TILE_LENGTH); + } + + // Round an element count up to a 32 B boundary (vector-pipe minimum). + __aicore__ inline uint32_t VecAlignCount(uint32_t count) const + { + constexpr uint32_t elemsPer32B = 32 / sizeof(T); + return (count + elemsPer32B - 1) / elemsPer32B * elemsPer32B; + } + + AscendC::TPipe* pipe_; + AscendC::GlobalTensor hiddenGm_; + AscendC::GlobalTensor weightGm_; + AscendC::GlobalTensor biasGm_; + AscendC::GlobalTensor targetGm_; + AscendC::GlobalTensor logpGm_; + AscendC::TQue wQueue_; + AscendC::TBuf hFp32Buf_; + AscendC::TBuf wFp32Buf_; + AscendC::TBuf prodBuf_; + AscendC::TBuf reduceBuf_; + AscendC::TBuf scalarBuf_; + AscendC::TBuf targetBuf_; + AscendC::TBuf biasBuf_; + AscendC::TEventID eventVS_; + AscendC::TEventID eventSV_; + AscendC::TEventID eventMTE2V_; + AscendC::TEventID eventMTE2S_; + AscendC::TEventID eventSMTE3_; + AscendC::TEventID eventMTE3S_; + int64_t numRows_; + int64_t vocabSize_; + int64_t hiddenSize_; + bool hasBias_; + bool cacheHidden_; +}; + +} // namespace + +extern "C" __global__ __vector__ void fused_linear_logp_ascend_kernel_fp32( + GM_ADDR hidden, GM_ADDR weight, GM_ADDR bias, GM_ADDR target, GM_ADDR logp, + int64_t numRows, int64_t vocabSize, int64_t hiddenSize, bool hasBias) +{ + AscendC::TPipe pipe; + KernelFusedLinearLogp op(&pipe); + op.Init(hidden, weight, bias, target, logp, numRows, vocabSize, hiddenSize, hasBias); + op.Process(); +} + +extern "C" __global__ __vector__ void fused_linear_logp_ascend_kernel_bf16( + GM_ADDR hidden, GM_ADDR weight, GM_ADDR bias, GM_ADDR target, GM_ADDR logp, + int64_t numRows, int64_t vocabSize, int64_t hiddenSize, bool hasBias) +{ + AscendC::TPipe pipe; + KernelFusedLinearLogp op(&pipe); + op.Init(hidden, weight, bias, target, logp, numRows, vocabSize, hiddenSize, hasBias); + op.Process(); +} + +extern "C" __global__ __vector__ void fused_linear_logp_ascend_kernel_fp16( + GM_ADDR hidden, GM_ADDR weight, GM_ADDR bias, GM_ADDR target, GM_ADDR logp, + int64_t numRows, int64_t vocabSize, int64_t hiddenSize, bool hasBias) +{ + AscendC::TPipe pipe; + KernelFusedLinearLogp op(&pipe); + op.Init(hidden, weight, bias, target, logp, numRows, vocabSize, hiddenSize, hasBias); + op.Process(); +} + +torch::Tensor fused_linear_logp_ascend_forward(torch::Tensor hidden, + torch::Tensor weight, + torch::optional bias, + torch::Tensor target) +{ + TORCH_CHECK(hidden.is_privateuseone(), "hidden must be on an NPU device"); + TORCH_CHECK(weight.is_privateuseone(), "lm_head_weight must be on an NPU device"); + TORCH_CHECK(hidden.device() == weight.device(), + "hidden and lm_head_weight must be on the same NPU device"); + TORCH_CHECK(hidden.dim() == 2, "hidden must be 2-D [N, D]"); + TORCH_CHECK(weight.dim() == 2, "lm_head_weight must be 2-D [V, D]"); + TORCH_CHECK(hidden.size(-1) == weight.size(1), "hidden/weight hidden-dim mismatch"); + TORCH_CHECK(hidden.scalar_type() == at::kFloat || hidden.scalar_type() == at::kHalf || + hidden.scalar_type() == at::kBFloat16, + "fused_linear_logp_ascend supports fp32, fp16, and bf16 hidden states"); + TORCH_CHECK(weight.scalar_type() == hidden.scalar_type(), + "fused_linear_logp_ascend requires weight to match the hidden dtype"); + TORCH_CHECK(hidden.size(-1) > 0, "hidden dimension must be positive"); + TORCH_CHECK(target.is_privateuseone(), "target must be on the same NPU device as hidden"); + 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() == hidden.size(0), "target must have one entry per row"); + + const int64_t numRows = hidden.size(0); + const int64_t hiddenSize = hidden.size(1); + const int64_t vocabSize = weight.size(0); + + torch::Tensor biasF; + uint8_t* biasPtr = nullptr; + bool hasBias = false; + if (bias.has_value()) { + TORCH_CHECK(bias->is_privateuseone(), "bias must be on an NPU device"); + TORCH_CHECK(bias->device() == hidden.device(), "bias must be on the same NPU device"); + TORCH_CHECK(bias->dim() == 1 && bias->numel() == vocabSize, + "bias must be 1-D [V]"); + biasF = bias->reshape({vocabSize}).to(at::kFloat).contiguous(); + biasPtr = reinterpret_cast(biasF.mutable_data_ptr()); + hasBias = true; + } + + // fp32 output, matching the gold reference's contract. + torch::Tensor logp = at::empty({numRows}, hidden.options().dtype(at::kFloat)); + if (numRows == 0 || vocabSize == 0) { + return logp; + } + + auto hiddenContig = hidden.contiguous(); + auto weightContig = weight.contiguous(); + auto 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 (hidden.scalar_type() == at::kBFloat16) { + fused_linear_logp_ascend_kernel_bf16<<>>( + reinterpret_cast(hiddenContig.mutable_data_ptr()), + reinterpret_cast(weightContig.mutable_data_ptr()), + biasPtr, + reinterpret_cast(targetContig.mutable_data_ptr()), + reinterpret_cast(logp.mutable_data_ptr()), + numRows, vocabSize, hiddenSize, hasBias); + } else if (hidden.scalar_type() == at::kHalf) { + fused_linear_logp_ascend_kernel_fp16<<>>( + reinterpret_cast(hiddenContig.mutable_data_ptr()), + reinterpret_cast(weightContig.mutable_data_ptr()), + biasPtr, + reinterpret_cast(targetContig.mutable_data_ptr()), + reinterpret_cast(logp.mutable_data_ptr()), + numRows, vocabSize, hiddenSize, hasBias); + } else { + fused_linear_logp_ascend_kernel_fp32<<>>( + reinterpret_cast(hiddenContig.mutable_data_ptr()), + reinterpret_cast(weightContig.mutable_data_ptr()), + biasPtr, + reinterpret_cast(targetContig.mutable_data_ptr()), + reinterpret_cast(logp.mutable_data_ptr()), + numRows, vocabSize, hiddenSize, hasBias); + } + 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..ee6325f7 --- /dev/null +++ b/csrc/ascend/npu_module.cpp @@ -0,0 +1,30 @@ +// 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_linear_logp_ascend_forward(torch::Tensor hidden, + torch::Tensor weight, + torch::optional bias, + 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_linear_logp_ascend", + &fused_linear_logp_ascend_forward, + "Batch-invariant fused linear log-probability (Ascend C forward)"); +} diff --git a/docs/operators/linear-logp.md b/docs/operators/linear-logp.md index 4b5231ef..0dd65bde 100644 --- a/docs/operators/linear-logp.md +++ b/docs/operators/linear-logp.md @@ -39,6 +39,7 @@ logp.sum().backward() # gradients flow into hidden, lm_head_weight, bias | --- | --- | --- | | CUDA SM90 (Hopper) | `FusedLinearLogpSM90Op` | TMA-streamed, Double Buffering, tensor-core forward (`mma.sync.m16n8k16`), online softmax in smem; chunked backward. Compiles for `sm_90a`; validated fp32-accurate on H100. Falls back to Triton/native for fp32/fp16 inputs or hidden dims not divisible by 32. | | CUDA / ROCm (Triton) | `TritonLinearLogpOp` | Triton online-softmax forward; Liger-style chunked backward (cuBLAS matmuls, deterministic). Phase 1. | +| Ascend NPU | `FusedLinearLogpAscendOp` | Batch-invariant Ascend C forward mirroring the SM90 reduction contract: ascending vocab-row scan with the online rescale chain, per-row fp32 dots over a fixed D-tile order, `min(zt - lse, 0)` clamp; the shared chunked backward. Output fp32. | | PyTorch native | `NativeLinearLogpOp` | Naive `F.linear` + `log_softmax` + `gather` reference; CPU / Triton-less fallback. | The SM90 backend (`csrc/cuda/fused_linear_logp_sm90.cu`) streams hidden/weight @@ -168,10 +169,14 @@ For 4-GPU tensor-parallel validation, use ## Implementation Files - `rl_engine/kernels/ops/triton/loss/linear_logp.py` -- `rl_engine/kernels/ops/pytorch/loss/linear_logp.py` -- `rl_engine/kernels/ops/cuda/loss/linear_logp.py` (SM90 wrapper + chunked backward) -- `csrc/cuda/fused_linear_logp_sm90.cu`, `csrc/ops.cpp`, `setup.py` (SM90 kernel + build) +- `rl_engine/kernels/ops/pytorch/loss/linear_logp.py` — native reference, chunked backward, TP helpers +- `rl_engine/kernels/ops/cuda/loss/linear_logp.py` — CUDA fused implementation (SM90 wrapper + chunked backward) +- `rl_engine/kernels/ops/ascend/loss/linear_logp.py` — Ascend deterministic op +- `csrc/cuda/fused_linear_logp_sm90.cu`, `csrc/ops.cpp`, `setup.py` — SM90 kernel + build +- `csrc/ascend/fused_linear_logp_ascend.asc` — Ascend C forward kernel +- `csrc/ascend/npu_module.cpp` — shared pybind entry for `rl_engine._C_npu` - `rl_engine/kernels/registry.py` - `tests/test_linear_logp.py` +- `tests/test_linear_logp_ascend.py` — Ascend correctness + batch-invariance tests - `benchmarks/benchmark_linear_logp.py` - `docs/design/fused-linear-logp.md` diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index bff5e2e7..e3fde7d7 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,3 +8,9 @@ def batch_invariant_logp_ascend( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... +def fused_linear_logp_ascend( + hidden: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + 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..13e67280 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -115,6 +115,7 @@ def _load_object(path: str) -> Any: "pytorch": "rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp", "triton": "rl_engine.kernels.ops.triton.loss.linear_logp.TritonLinearLogpOp", "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.linear_logp.FusedLinearLogpSM90Op", + "ascend": "rl_engine.kernels.ops.ascend.loss.linear_logp.FusedLinearLogpAscendOp", }, grad_input_names=("hidden", "lm_head_weight"), ), diff --git a/rl_engine/kernels/ops/ascend/loss/__init__.py b/rl_engine/kernels/ops/ascend/loss/__init__.py index 86cf4c9d..a7302f46 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 linear_logp # noqa: F401 diff --git a/rl_engine/kernels/ops/ascend/loss/linear_logp.py b/rl_engine/kernels/ops/ascend/loss/linear_logp.py new file mode 100644 index 00000000..af44bc90 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/loss/linear_logp.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Any, Optional + +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 + +_SUPPORTED_DTYPES = {torch.float32, torch.float16, torch.bfloat16} + + +class _FusedLinearLogpAscendFunction(torch.autograd.Function): + """Autograd bridge for the Ascend fused linear log-prob forward. + + The backward is the shared Liger-style chunked formula from + ``rl_engine.kernels.ops.pytorch.loss.linear_logp.chunked_linear_logp_backward`` + (the same formula the CUDA SM90 op falls back to), so gradients follow + the CUDA op's portable backward exactly. + """ + + @staticmethod + def forward(ctx, hidden, lm_head_weight, target_ids): + hidden_2d = hidden.reshape(-1, hidden.size(-1)).contiguous() + weight = lm_head_weight.contiguous() + target_1d = ( + target_ids.reshape(-1).to(device=hidden_2d.device, dtype=torch.long).contiguous() + ) + output = _C_npu.fused_linear_logp_ascend(hidden_2d, weight, None, target_1d) + ctx.save_for_backward(hidden_2d, weight, target_1d) + ctx.lead_shape = hidden.shape[:-1] + ctx.hidden_dtype = hidden.dtype + ctx.weight_dtype = lm_head_weight.dtype + return output.reshape(hidden.shape[:-1]) + + @staticmethod + def backward(ctx, grad_logp): + from rl_engine.kernels.ops.pytorch.loss.linear_logp import chunked_linear_logp_backward + + hidden_2d, weight, target_1d = ctx.saved_tensors + grad_hidden, grad_weight, _ = chunked_linear_logp_backward( + grad_logp, + hidden_2d, + weight, + target_1d, + hidden_2d, # bias placeholder; has_bias=False + has_bias=False, + lead_shape=ctx.lead_shape, + hidden_dtype=ctx.hidden_dtype, + weight_dtype=ctx.weight_dtype, + bias_dtype=None, + ) + return grad_hidden, grad_weight, None + + +class FusedLinearLogpAscendOp: + """Batch-invariant fused linear log-prob for Ascend NPU. + + Computes ``log_softmax(hidden @ W^T + b)[target]`` without materializing + the ``[N, V]`` logits. The Ascend C forward mirrors the SM90 kernel's + reduction contract: ascending vocab-row scan with the online rescale + chain, per-row fp32 dots over a fixed D-tile order, final + ``min(zt - lse, 0)`` clamp; the output is fp32. + """ + + is_fused_logp = True + is_batch_invariant = True + + def __init__(self) -> None: + if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "fused_linear_logp_ascend"): + raise RuntimeError( + "fused_linear_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_linear_logp_ascend + logger.info("Successfully linked to precompiled _C_npu.fused_linear_logp_ascend kernel.") + + def __call__( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + tp_group: Any = None, + vocab_start_index: int = 0, + global_vocab_size: Optional[int] = None, + ) -> torch.Tensor: + return self.apply( + hidden, + lm_head_weight, + target_ids, + bias, + tp_group=tp_group, + vocab_start_index=vocab_start_index, + global_vocab_size=global_vocab_size, + ) + + def apply( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + tp_group: Any = None, + vocab_start_index: int = 0, + global_vocab_size: Optional[int] = None, + ) -> torch.Tensor: + from rl_engine.kernels.ops.pytorch.loss.linear_logp import ( + NativeLinearLogpOp, + should_use_tensor_parallel_linear_logp, + ) + + if lm_head_weight.size(-1) != hidden.size(-1): + raise ValueError( + f"hidden dim {hidden.size(-1)} must match lm_head_weight dim " + f"{lm_head_weight.size(-1)}" + ) + if lm_head_weight.device != hidden.device: + raise ValueError( + f"lm_head_weight device {lm_head_weight.device} must match hidden " + f"device {hidden.device}" + ) + if hidden.shape[:-1] != target_ids.shape: + raise ValueError( + f"hidden leading shape {tuple(hidden.shape[:-1])} must match " + f"target_ids shape {tuple(target_ids.shape)}" + ) + # Tensor-parallel and bias paths are not covered by the Ascend forward; + # delegate to the native reference (same fallback as the CUDA op). + if ( + should_use_tensor_parallel_linear_logp( + tp_group, + int(vocab_start_index), + global_vocab_size, + lm_head_weight.size(0), + ) + or bias is not None + ): + return NativeLinearLogpOp().apply( + hidden, + lm_head_weight, + target_ids, + bias, + tp_group=tp_group, + vocab_start_index=vocab_start_index, + global_vocab_size=global_vocab_size, + ) + if not self._ascend_supported(hidden, lm_head_weight): + return NativeLinearLogpOp().apply(hidden, lm_head_weight, target_ids) + return _FusedLinearLogpAscendFunction.apply(hidden, lm_head_weight, target_ids) + + @staticmethod + def _ascend_supported(hidden: torch.Tensor, lm_head_weight: torch.Tensor) -> bool: + return ( + hidden.device.type == "npu" + and lm_head_weight.device.type == "npu" + and hidden.is_contiguous() + and lm_head_weight.is_contiguous() + and hidden.dtype in _SUPPORTED_DTYPES + and lm_head_weight.dtype == hidden.dtype + ) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 12ea9b21..36152c16 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -107,6 +107,9 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ASCEND_BATCH_INVARIANT_LOGP = ( "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAscendOp" ) + ASCEND_FUSED_LINEAR_LOGP = ( + "rl_engine.kernels.ops.ascend.loss.linear_logp.FusedLinearLogpAscendOp" + ) # 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 +625,10 @@ def __init__(self): OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ] + self._priority_map["npu"]["linear_logp"] = [ + OpBackend.ASCEND_FUSED_LINEAR_LOGP, + OpBackend.PYTORCH_LINEAR_LOGP, + ] 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..f3a0361c 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"]["linear_logp"] == [ + OpBackend.ASCEND_FUSED_LINEAR_LOGP, + OpBackend.PYTORCH_LINEAR_LOGP, + ] 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_linear_logp_ascend.py b/tests/test_linear_logp_ascend.py new file mode 100644 index 00000000..7ece62b4 --- /dev/null +++ b/tests/test_linear_logp_ascend.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for the Ascend NPU batch-invariant fused linear log-prob. + +Validates the same two orthogonal properties as the CUDA deterministic op: + +1. **Correctness** - output matches a hand-computed fp32 reference + (``hidden.float() @ weight.float().T`` + ``log_softmax`` + gather + + clamp) within an honest reduction tolerance (~2e-4 at D=4096, pure + fp32-tree drift). The gtest's own forward comparison is stricter than + any independent kernel can meet (see Notes in the PR description): the + fp32 logprob tolerance is 1e-5 while two different fp32 reduction trees + over D=4096 drift ~1e-4, and the gold's dtype path accumulates the + matmul in bf16/fp16 while this kernel (like the CUDA SM90 kernel) + accumulates in fp32. +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 over a fixed vocab scan). +""" + +import pytest +import torch + +_VOCAB = 129 +_HIDDEN = 1000 + +# Honest forward tolerance vs the fp32 reference: pure fp32 reduction-tree +# drift (measured 2.1e-4 at D=4096, V=257). +_FWD_ATOL = 5.0e-4 +_FWD_RTOL = 1.0e-5 +# Gradient tolerance: the chunked backward casts to the input dtype, so +# low-precision grads compare at their own quantization level. +_GRAD_ATOL = {torch.float32: 5.0e-4, torch.bfloat16: 2.0e-2, torch.float16: 1.0e-2} + + +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.linear_logp import _NPU_EXT_AVAILABLE, _C_npu + except Exception: + return False + return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "fused_linear_logp_ascend") + + +requires_ascend = pytest.mark.skipif( + not _ascend_kernel_available(), + reason="fused_linear_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.linear_logp import FusedLinearLogpAscendOp + + return FusedLinearLogpAscendOp() + + +def _make_inputs(shape, vocab=_VOCAB, hidden=_HIDDEN, dtype=torch.float32, seed=0): + generator = torch.Generator(device="cpu").manual_seed(seed) + hidden_t = torch.randn(*shape, hidden, dtype=dtype, generator=generator).to("npu") + weight = torch.randn(vocab, hidden, dtype=dtype, generator=generator).to("npu") + target_ids = torch.randint(0, vocab, shape, generator=generator).long().to("npu") + return hidden_t, weight, target_ids + + +def _ref_fp32(hidden, weight, target_ids, bias=None): + """Hand-written fp32 reference matching the WS1 fp32-reference policy.""" + logits = hidden.float().reshape(-1, hidden.size(-1)) @ weight.float().t() + if bias is not None: + logits = logits + bias.float() + flat = target_ids.reshape(-1) + logp = torch.log_softmax(logits, dim=-1) + selected = logp.gather(1, flat.unsqueeze(1)).squeeze(1) + return selected.clamp(max=0).reshape(hidden.shape[:-1]) + + +# --------------------------------------------------------------------------- +# Correctness +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@requires_ascend +class TestAscendFusedLinearLogpCorrectness: + def test_forward_matches_fp32_reference(self, dtype): + op = _get_op() + hidden, weight, target_ids = _make_inputs((3, 5), dtype=dtype) + out = op(hidden, weight, target_ids) + ref = _ref_fp32(hidden, weight, target_ids) + assert out.dtype == torch.float32 + assert out.shape == (3, 5) + assert torch.allclose(out, ref, atol=_FWD_ATOL, rtol=_FWD_RTOL) + + def test_forward_large_shape(self, dtype): + # gtest shape: D=4096 (single cached hidden tile), V=257. + op = _get_op() + hidden, weight, target_ids = _make_inputs((2, 16), vocab=257, hidden=4096, dtype=dtype) + out = op(hidden, weight, target_ids) + ref = _ref_fp32(hidden, weight, target_ids) + assert torch.allclose(out, ref, atol=_FWD_ATOL, rtol=_FWD_RTOL) + + def test_out_of_range_target_is_zero(self, dtype): + op = _get_op() + hidden, weight, target_ids = _make_inputs((2, 4), vocab=32, dtype=dtype) + target_ids = target_ids.reshape(-1) + target_ids[1] = 32 + 5 # out of [0, V) + out = op(hidden, weight, target_ids.reshape(2, 4)) + assert out.reshape(-1)[1].item() == 0.0 + + def test_bias_falls_back_to_native(self, dtype): + from rl_engine.kernels.ops.pytorch.loss.linear_logp import NativeLinearLogpOp + + op = _get_op() + hidden, weight, target_ids = _make_inputs((2, 3), dtype=dtype) + bias = torch.randn(_VOCAB, device="npu", dtype=dtype) + out = op(hidden, weight, target_ids, bias) + ref = NativeLinearLogpOp().apply(hidden, weight, target_ids, bias) + assert torch.allclose(out.float(), ref.float(), atol=1e-5, rtol=1e-5) + + def test_backward_matches_fp32_reference(self, dtype): + op = _get_op() + hidden, weight, target_ids = _make_inputs((3, 5), dtype=dtype) + grad_out = torch.randn(3, 5, device="npu", dtype=dtype) + + h_a = hidden.clone().requires_grad_() + w_a = weight.clone().requires_grad_() + op(h_a, w_a, target_ids).backward(grad_out) + + h_f = hidden.float().clone().requires_grad_() + w_f = weight.float().clone().requires_grad_() + _ref_fp32(h_f, w_f, target_ids).backward(grad_out.float()) + + # Compare at the quantized level for low-precision inputs: both + # backends compute the VJP in fp32 and cast to the input dtype, so + # the fp32 tree drift collapses into (usually identical) quantized + # bits; the tolerance only absorbs the rare 1-ULP straddle. + assert torch.allclose( + h_a.grad.float(), h_f.grad.to(dtype).float(), atol=_GRAD_ATOL[dtype], rtol=1.0e-4 + ) + assert torch.allclose( + w_a.grad.float(), w_f.grad.to(dtype).float(), atol=_GRAD_ATOL[dtype], 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() + hidden, weight, target_ids = _make_inputs((4, 8), dtype=dtype) + hidden[1].copy_(hidden[0]) + target_ids[1] = target_ids[0] + + h_g = hidden.clone().requires_grad_() + grad_out = torch.randn(4, 8, device="npu", dtype=dtype) + grad_out[1] = grad_out[0] + op(h_g, weight, target_ids).backward(grad_out) + grad = h_g.grad + assert torch.equal(grad[0], grad[1]) + for row in range(2, 4): + assert not torch.equal(grad[0], grad[row]) + + +# --------------------------------------------------------------------------- +# Batch invariance +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendFusedLinearLogpBatchInvariance: + def test_batch_size_1_vs_n(self): + # One fixed row embedded in batches of growing size: its logp must be + # bitwise identical regardless of batch size. + dtype = torch.bfloat16 + op = _get_op() + alone_hidden, weight, alone_ids = _make_inputs((1,), dtype=dtype, seed=7) + alone = op(alone_hidden, weight, alone_ids)[0] + for batch in (2, 4, 16, 300): # 300 rows -> > MAX_BLOCKS strided blocks + hidden, _, target_ids = _make_inputs((batch,), dtype=dtype, seed=7) + hidden[0].copy_(alone_hidden[0]) + target_ids[0] = alone_ids[0] + in_batch = op(hidden, weight, target_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() + hidden, weight, target_ids = _make_inputs((8,), dtype=dtype, seed=11) + base, base_id = hidden[0].clone(), int(target_ids[0]) + for pos in range(1, 8): + hidden[pos].copy_(base) + target_ids[pos] = base_id + out = op(hidden, weight, target_ids) + for pos in range(1, 8): + assert torch.equal(out[pos], out[0]), f"drift at position={pos}" + + def test_multi_tile_rows(self): + # hidden > TILE_LENGTH (4096): the per-row dots span multiple tiles. + dtype = torch.float32 + op = _get_op() + hidden, weight, target_ids = _make_inputs((4,), hidden=10000, dtype=dtype, seed=5) + out = op(hidden, weight, target_ids) + assert torch.equal(out, op(hidden, weight, target_ids)) + + def test_repeated_runs_deterministic(self): + dtype = torch.bfloat16 + hidden, weight, target_ids = _make_inputs((3, 5), dtype=dtype, seed=5) + op = _get_op() + first = op(hidden, weight, target_ids) + for _ in range(3): + again = op(hidden, weight, target_ids) + assert torch.equal(first, again) + + +# --------------------------------------------------------------------------- +# Registry dispatch +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendRegistryDispatch: + def test_get_op_linear_logp(self): + from rl_engine.kernels.registry import kernel_registry + + op = kernel_registry.get_op("linear_logp", device="npu") + assert type(op).__name__ == "FusedLinearLogpAscendOp"