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 7070728a..0350ad55 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -125,6 +125,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" @@ -690,6 +693,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 ee560db0..74c1c184 100644 --- a/setup.py +++ b/setup.py @@ -331,7 +331,13 @@ def _ascend_extensions(): asc_srcs = sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) if not asc_srcs: raise RuntimeError("KERNEL_ALIGN_FORCE_ASCEND=1 but no .asc sources under csrc/ascend/") - return [Extension(name="rl_engine._C_npu", sources=asc_srcs, language="asc")] + sources: list[str] = asc_srcs + # Some kernels ship a C++ pybind host alongside the .asc sources; include + # it when present (compiled per-source by _bisheng_compile_cmd). + host_cpp = Path("csrc/ascend/npu_module.cpp") + if host_cpp.is_file(): + sources = [str(host_cpp), *asc_srcs] + return [Extension(name="rl_engine._C_npu", sources=sources, language="asc")] def _bisheng_compile_cmd(ext, ext_fullpath): @@ -366,6 +372,49 @@ def _bisheng_compile_cmd(ext, ext_fullpath): os.path.join(ascend_home, "lib64"), ] + if any(not str(src).endswith(".asc") for src in ext.sources): + # Mixed extension (C++ pybind host + .asc kernels): the -x asc driver + # cannot compile C++, so compile each source to an object and link them + # in a second step. + import subprocess + import tempfile + + build_temp = tempfile.mkdtemp(prefix="rl_kernel_ascend_") + objects = [] + for src in ext.sources: + src = str(src) + obj = os.path.join(build_temp, Path(src).name + ".o") + src_cmd = [ + "bisheng", + "-std=c++17", + "-O2", + "-fPIC", + "-c", + f"-D_GLIBCXX_USE_CXX11_ABI={abi_value}", + f"-DTORCH_EXTENSION_NAME={module_name}", + ] + if src.endswith(".asc"): + src_cmd += ["-x", "asc", f"--npu-arch={soc}"] + src_cmd += [f"-I{d}" for d in include_dirs if d] + src_cmd += [src, "-o", obj] + subprocess.check_call(src_cmd) + objects.append(obj) + cmd = [ + "bisheng", + "-shared", + *objects, + "-lascendcl", + "-ltorch_npu", + "-ltorch", + "-ltorch_cpu", + "-ltorch_python", + "-lc10", + "-o", + ext_fullpath, + ] + cmd += [f"-L{d}" for d in lib_dirs if d] + return cmd + cmd = [ "bisheng", "-x", 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"