[WS1][Ascend] [Qwen3-8b] Fused logp ops

csrc/ascend/batch_invariant_logp_ascend.asc

@@ -308,9 +308,5 @@ std::vector<torch::Tensor> 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.

csrc/ascend/fused_logp_ascend.asc

@@ -0,0 +1,332 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+//
+// Batch-invariant fused selected-token log-probability, Ascend C (CANN)
+// forward kernel.
+//
+//   logp[n] = logits[n, target[n]] - logsumexp(logits[n, :])
+//
+// Mirrors the deterministic CUDA kernel in csrc/deterministic_logp_kernel.cu
+// (DeterministicLogpCUDAOp):
+//   - input  : logits [N, V] contiguous, bf16 / fp16 / fp32; target [N]
+//              int64, one per row
+//   - output : logp [N] fp32 (the CUDA deterministic op always returns fp32)
+//   - target[n] outside [0, V) -> logp[n] = 0.0 (same as the CUDA kernel)
+//
+// The math follows the same two-pass fixed-order reduction as the CUDA
+// kernel: row max over a fixed tile order, then sum(exp(x - max)) over the
+// same fixed tile order, lse = max + log(sum), logp = selected - lse. The
+// fp32 accumulation and the formula match the CUDA kernel exactly; the
+// hardware reduction trees and the transcendental implementations are the
+// Ascend vector unit's own (fixed per V), so cross-platform bitwise parity
+// with the CUDA kernel is not claimed -- the guarantee here is the same one
+// the CUDA kernel provides on its platform: batch-invariant determinism.
+//
+// Batch-invariance: every row is processed end-to-end by exactly one AI core
+// block with a fixed tile size and a fixed reduction order. The instruction
+// sequence for a row depends only on V, never on N or on the block the row
+// happens to land on, so a row's output is bitwise identical across batch
+// sizes, row positions, and block assignments on the NPU.
+//
+// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by
+// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu.
+
+#include <type_traits>
+
+#include "kernel_operator.h"
+
+#include <torch/extension.h>
+
+#include "torch_npu/csrc/core/npu/NPUStream.h"
+
+namespace {
+
+// Elements per vocab tile. Fixed for all rows and batch sizes; this is what
+// makes the reduction order batch-invariant. UB budget (input tile + fp32
+// tile + reduce scratch) stays well under the 192 KB UB of current SoCs.
+constexpr uint32_t TILE_LENGTH = 4096;
+// Cap on launched blocks. Rows are strided across blocks, so launching fewer
+// blocks than rows is fine and never changes per-row numerics.
+constexpr int64_t MAX_BLOCKS = 128;
+constexpr float NEG_INF = -3.402823466e+38f;  // -FLT_MAX
+
+template <typename T>
+class KernelFusedLogp {
+public:
+    __aicore__ inline KernelFusedLogp(AscendC::TPipe* pipe) : pipe_(pipe) {}
+
+    __aicore__ inline void Init(GM_ADDR logits,
+                                GM_ADDR target,
+                                GM_ADDR logp,
+                                int64_t numRows,
+                                int64_t vocabSize)
+    {
+        numRows_ = numRows;
+        vocabSize_ = vocabSize;
+        logitsGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(logits));
+        targetGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t*>(target));
+        logpGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(logp));
+        pipe_->InitBuffer(inQueue_, 1, TILE_LENGTH * sizeof(T));
+        pipe_->InitBuffer(fp32Buf_, TILE_LENGTH * sizeof(float));
+        pipe_->InitBuffer(reduceBuf_, TILE_LENGTH * sizeof(float));
+        // 64 B: floats [0,8) for reduce/log scratch, floats [8,16) for the
+        // output staging slots (both halves 32-byte aligned for DataCopyPad).
+        pipe_->InitBuffer(scalarBuf_, 64);
+        // 32 B window for reading target[row] via DataCopyPad (GM scalar
+        // GetValue/SetValue are unreliable on hardware; see cannbot
+        // ascendc-precision-debug common-traps).
+        pipe_->InitBuffer(targetBuf_, 32);
+
+        // Intra-core pipeline events. NOTE: AscendC::SyncAll() is a cross-core
+        // barrier and deadlocks when more blocks are launched than there are
+        // physical cores (resident blocks wait for unscheduled ones), so all
+        // synchronization here uses per-pipe SetFlag/WaitFlag instead.
+        eventVS_ = pipe_->FetchEventID(AscendC::HardEvent::V_S);
+        eventSV_ = pipe_->FetchEventID(AscendC::HardEvent::S_V);
+        eventMTE2S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE2_S);
+        eventSMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE3);
+        eventMTE3S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE3_S);
+    }
+
+    __aicore__ inline void Process()
+    {
+        for (int64_t row = AscendC::GetBlockIdx(); row < numRows_;
+             row += AscendC::GetBlockNum()) {
+            ProcessRow(row);
+        }
+    }
+
+private:
+    // Load one vocab tile into UB and return its fp32 view. When T is fp32 the
+    // queue buffer is used in place; otherwise the tile is cast into fp32Buf_.
+    __aicore__ inline AscendC::LocalTensor<float> LoadTileFp32(int64_t row,
+                                                               int64_t start,
+                                                               uint32_t count)
+    {
+        AscendC::LocalTensor<T> xLocal = inQueue_.AllocTensor<T>();
+        AscendC::DataCopyExtParams copyParams{
+            1, static_cast<uint32_t>(count * sizeof(T)), 0, 0, 0};
+        AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
+        AscendC::DataCopyPad(xLocal, logitsGm_[row * vocabSize_ + start], copyParams, padParams);
+        inQueue_.EnQue(xLocal);
+        xLocal = inQueue_.DeQue<T>();
+        inTile_ = xLocal;
+
+        if constexpr (std::is_same_v<T, float>) {
+            return xLocal;
+        } else {
+            AscendC::LocalTensor<float> fLocal = fp32Buf_.Get<float>();
+            AscendC::Cast(fLocal, xLocal, AscendC::RoundMode::CAST_NONE, count);
+            return fLocal;
+        }
+    }
+
+    // Release the queue-owned input buffer of the current tile.
+    __aicore__ inline void FreeTile()
+    {
+        inQueue_.FreeTensor(inTile_);
+    }
+
+    __aicore__ inline void ProcessRow(int64_t row)
+    {
+        const int64_t target = LoadTarget(row);
+        const bool valid = target >= 0 && target < vocabSize_;
+        const int64_t tileCount = (vocabSize_ + TILE_LENGTH - 1) / TILE_LENGTH;
+        float selected = 0.0f;
+
+        // Pass 1: row max (fixed tile order). Also grab logits[target] on the fly.
+        float rowMax = NEG_INF;
+        for (int64_t tile = 0; tile < tileCount; ++tile) {
+            const int64_t start = tile * TILE_LENGTH;
+            const uint32_t count = TileCount(start);
+            AscendC::LocalTensor<float> fLocal = LoadTileFp32(row, start, count);
+
+            AscendC::LocalTensor<float> rTmp = reduceBuf_.Get<float>();
+            AscendC::LocalTensor<float> scalar = scalarBuf_.Get<float>();
+            AscendC::ReduceMax<float>(scalar, fLocal, rTmp, static_cast<int32_t>(count), false);
+            WaitVector();  // vector -> scalar read
+            const float tileMax = scalar.GetValue(0);
+            rowMax = tileMax > rowMax ? tileMax : rowMax;
+
+            if (valid && target >= start && target < start + count) {
+                selected = fLocal.GetValue(static_cast<uint32_t>(target - start));
+            }
+            FreeTile();
+        }
+
+        // Pass 2: sum(exp(x - rowMax)) with the same fixed tile order.
+        float sumExp = 0.0f;
+        for (int64_t tile = 0; tile < tileCount; ++tile) {
+            const int64_t start = tile * TILE_LENGTH;
+            const uint32_t count = TileCount(start);
+            AscendC::LocalTensor<float> fLocal = LoadTileFp32(row, start, count);
+
+            AscendC::Adds(fLocal, fLocal, -rowMax, count);  // x - rowMax
+            AscendC::Exp(fLocal, fLocal, count);
+            AscendC::LocalTensor<float> rTmp = reduceBuf_.Get<float>();
+            AscendC::LocalTensor<float> scalar = scalarBuf_.Get<float>();
+            AscendC::ReduceSum<float, true>(scalar, fLocal, rTmp, static_cast<int32_t>(count));
+            WaitVector();  // vector -> scalar read
+            sumExp += scalar.GetValue(0);
+            FreeTile();
+        }
+
+        // lse = rowMax + log(sumExp). The scalar unit has no log, so run a
+        // 1-element vector Log (count padded to 8; scalarBuf_ is 32 B aligned).
+        AscendC::LocalTensor<float> scalar = scalarBuf_.Get<float>();
+        scalar.SetValue(0, sumExp);
+        AscendC::SetFlag<AscendC::HardEvent::S_V>(eventSV_);  // scalar write -> vector op
+        AscendC::WaitFlag<AscendC::HardEvent::S_V>(eventSV_);
+        AscendC::Log(scalar, scalar, 8);
+        WaitVector();  // vector -> scalar read
+        const float lse = rowMax + scalar.GetValue(0);
+
+        // Stage the output in UB, then DataCopyPad to GM. GlobalTensor.SetValue
+        // is unreliable on hardware (cannbot ascendc-precision-debug
+        // common-traps), so scalar GM stores are not used. Out-of-range
+        // targets produce 0.0, matching the CUDA deterministic kernel.
+        scalar.SetValue(0, valid ? (selected - lse) : 0.0f);
+        AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(eventSMTE3_);  // scalar write -> copy-out
+        AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(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<AscendC::HardEvent::MTE3_S>(eventMTE3S_);
+        AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(eventMTE3S_);
+    }
+
+    // Wait until all outstanding vector-pipe results are readable as scalars.
+    __aicore__ inline void WaitVector()
+    {
+        AscendC::SetFlag<AscendC::HardEvent::V_S>(eventVS_);
+        AscendC::WaitFlag<AscendC::HardEvent::V_S>(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<uint32_t>(remaining < 4 ? remaining : 4);
+        AscendC::LocalTensor<int64_t> tLocal = targetBuf_.Get<int64_t>();
+        AscendC::DataCopyExtParams copyParams{
+            1, static_cast<uint32_t>(winCount * sizeof(int64_t)), 0, 0, 0};
+        AscendC::DataCopyPadExtParams<int64_t> padParams{false, 0, 0, 0};
+        AscendC::DataCopyPad(tLocal, targetGm_[alignedRow], copyParams, padParams);
+        AscendC::SetFlag<AscendC::HardEvent::MTE2_S>(eventMTE2S_);  // copy-in -> scalar read
+        AscendC::WaitFlag<AscendC::HardEvent::MTE2_S>(eventMTE2S_);
+        return static_cast<int64_t>(tLocal.GetValue(static_cast<uint32_t>(row - alignedRow)));
+    }
+
+    __aicore__ inline uint32_t TileCount(int64_t start) const
+    {
+        const int64_t remaining = vocabSize_ - start;
+        return static_cast<uint32_t>(remaining < TILE_LENGTH ? remaining : TILE_LENGTH);
+    }
+
+    AscendC::TPipe* pipe_;
+    AscendC::GlobalTensor<T> logitsGm_;
+    AscendC::GlobalTensor<int64_t> targetGm_;
+    AscendC::GlobalTensor<float> logpGm_;
+    AscendC::TQue<AscendC::TPosition::VECIN, 1> inQueue_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> fp32Buf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> reduceBuf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> scalarBuf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> targetBuf_;
+    AscendC::LocalTensor<T> inTile_;
+    AscendC::TEventID eventVS_;
+    AscendC::TEventID eventSV_;
+    AscendC::TEventID eventMTE2S_;
+    AscendC::TEventID eventSMTE3_;
+    AscendC::TEventID eventMTE3S_;
+    int64_t numRows_;
+    int64_t vocabSize_;
+};
+
+}  // namespace
+
+extern "C" __global__ __vector__ void fused_logp_ascend_kernel_fp32(
+    GM_ADDR logits, GM_ADDR target, GM_ADDR logp,
+    int64_t numRows, int64_t vocabSize)
+{
+    AscendC::TPipe pipe;
+    KernelFusedLogp<float> op(&pipe);
+    op.Init(logits, target, logp, numRows, vocabSize);
+    op.Process();
+}
+
+extern "C" __global__ __vector__ void fused_logp_ascend_kernel_bf16(
+    GM_ADDR logits, GM_ADDR target, GM_ADDR logp,
+    int64_t numRows, int64_t vocabSize)
+{
+    AscendC::TPipe pipe;
+    KernelFusedLogp<bfloat16_t> op(&pipe);
+    op.Init(logits, target, logp, numRows, vocabSize);
+    op.Process();
+}
+
+extern "C" __global__ __vector__ void fused_logp_ascend_kernel_fp16(
+    GM_ADDR logits, GM_ADDR target, GM_ADDR logp,
+    int64_t numRows, int64_t vocabSize)
+{
+    AscendC::TPipe pipe;
+    KernelFusedLogp<half> op(&pipe);
+    op.Init(logits, target, logp, numRows, vocabSize);
+    op.Process();
+}
+
+torch::Tensor fused_logp_ascend_forward(torch::Tensor logits, torch::Tensor target)
+{
+    TORCH_CHECK(logits.is_privateuseone(), "logits must be on an NPU device");
+    TORCH_CHECK(logits.dim() == 2, "logits must be 2-D [N, V]");
+    TORCH_CHECK(logits.is_contiguous(), "logits must be contiguous");
+    TORCH_CHECK(logits.scalar_type() == at::kBFloat16 || logits.scalar_type() == at::kFloat ||
+                    logits.scalar_type() == at::kHalf,
+                "fused_logp_ascend supports fp32, fp16, and bf16 logits");
+    TORCH_CHECK(logits.size(-1) > 0, "vocab size must be positive");
+    TORCH_CHECK(target.is_privateuseone(), "target must be on the same NPU device as logits");
+    TORCH_CHECK(target.dim() == 1, "target must be 1-D [N]");
+    TORCH_CHECK(target.scalar_type() == at::kLong, "target must be int64");
+    TORCH_CHECK(target.numel() == logits.size(0), "target must have one entry per row");
+
+    const int64_t numRows = logits.size(0);
+    const int64_t vocabSize = logits.size(1);
+
+    // fp32 output, matching the CUDA deterministic logp op's contract.
+    torch::Tensor logp = at::empty({numRows}, logits.options().dtype(at::kFloat));
+    if (numRows == 0) {
+        return logp;
+    }
+
+    torch::Tensor targetContig = target.contiguous();
+
+    // stream(true): flush the task queue before launch so the kernel cannot
+    // overtake earlier NPU work; outputs were allocated with at::empty (no
+    // queued initializer) for the same reason.
+    auto aclStream = c10_npu::getCurrentNPUStream().stream(true);
+    const uint32_t blockNum = static_cast<uint32_t>(std::min(numRows, MAX_BLOCKS));
+
+    if (logits.scalar_type() == at::kBFloat16) {
+        fused_logp_ascend_kernel_bf16<<<blockNum, nullptr, aclStream>>>(
+            reinterpret_cast<uint8_t*>(logits.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(targetContig.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(logp.mutable_data_ptr()),
+            numRows, vocabSize);
+    } else if (logits.scalar_type() == at::kHalf) {
+        fused_logp_ascend_kernel_fp16<<<blockNum, nullptr, aclStream>>>(
+            reinterpret_cast<uint8_t*>(logits.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(targetContig.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(logp.mutable_data_ptr()),
+            numRows, vocabSize);
+    } else {
+        fused_logp_ascend_kernel_fp32<<<blockNum, nullptr, aclStream>>>(
+            reinterpret_cast<uint8_t*>(logits.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(targetContig.mutable_data_ptr()),
+            reinterpret_cast<uint8_t*>(logp.mutable_data_ptr()),
+            numRows, vocabSize);
+    }
+    return logp;
+}
+
+// The PYBIND11_MODULE for rl_engine._C_npu lives in npu_module.cpp so that
+// every Ascend op shares one compiled module.

csrc/ascend/npu_module.cpp

@@ -0,0 +1,28 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+//
+// Pybind entry point for the rl_engine._C_npu extension. The Ascend C kernels
+// and their torch host wrappers live in the sibling *.asc files; this TU only
+// declares and binds them so every Ascend op shares one compiled module.
+//
+// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by
+// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu.
+
+#include <torch/extension.h>
+
+std::vector<torch::Tensor> batch_invariant_logp_ascend_forward(torch::Tensor logits,
+                                                               torch::Tensor target,
+                                                               int64_t ignore_index);
+
+torch::Tensor fused_logp_ascend_forward(torch::Tensor logits,
+                                        torch::Tensor target);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("batch_invariant_logp_ascend",
+          &batch_invariant_logp_ascend_forward,
+          "Batch-invariant selected-token log-probability (Ascend C forward)");
+    m.def("fused_logp_ascend",
+          &fused_logp_ascend_forward,
+          "Batch-invariant fused selected-token log-probability (Ascend C forward)");
+}

docs/operators/fused-logp.md

@@ -31,6 +31,7 @@ reference = logp_ref.forward_fp32(logits, token_ids)
 | --- | --- | --- | --- |
 | CUDA SM90 | `FusedLogpSM90Op` | `_C.fused_logp_sm90` | Experimental TMA-oriented path for 2D contiguous bf16 logits on Hopper-class GPUs. It is disabled by default and requires `RL_KERNEL_ENABLE_EXPERIMENTAL_SM90_LOGP=1`; otherwise the wrapper delegates to the CUDA generic fallback. |
 | CUDA generic | `FusedLogpGenericOp` | `_C.fused_logp` | Generic compiled extension fallback. |
+| Ascend NPU | `FusedLogpAscendOp` | `_C_npu.fused_logp_ascend` | Batch-invariant Ascend C forward: two-pass (row max, then sum-exp) fp32 reduction with a fixed tile order, mirroring the CUDA deterministic kernel. Output is fp32, matching `DeterministicLogpCUDAOp`'s contract; out-of-range targets yield 0.0. |
 | PyTorch native | `NativeLogpOp` | None | PyTorch baseline/reference path. |
 
 ## Tensor Contract
@@ -62,9 +63,14 @@ operator accuracy tests continue to validate native/CUDA fused API compatibility
 ## Implementation Files
 
 - `rl_engine/kernels/registry.py`
-- `rl_engine/kernels/ops/pytorch/loss/logp.py`
-- `rl_engine/kernels/ops/cuda/loss/logp.py`
+- `rl_engine/kernels/ops/pytorch/loss/logp.py` — PyTorch native reference
+- `rl_engine/kernels/ops/cuda/loss/logp.py` — CUDA fused LogP (SM90 + generic)
+- `rl_engine/kernels/ops/ascend/loss/logp.py` — Ascend deterministic op
 - `csrc/ops.cpp`
 - `csrc/fused_logp_kernel.cu`
 - `csrc/cuda/fused_logp_sm90.cu`
+- `csrc/deterministic_logp_kernel.cu` — CUDA deterministic kernel (reference reduction)
+- `csrc/ascend/fused_logp_ascend.asc` — Ascend C forward kernel
+- `csrc/ascend/npu_module.cpp` — shared pybind entry for `rl_engine._C_npu`
 - `tests/test_logp.py`
+- `tests/test_logp_ascend.py` — Ascend correctness + batch-invariance tests