From 73f4deb2e122b6ba5f8349098de85b220711e309 Mon Sep 17 00:00:00 2001 From: chenyang <2082464740@qq.com> Date: Wed, 12 Aug 2026 17:59:08 +0800 Subject: [PATCH 1/5] feat: ascend_ Signed-off-by: chenyang <2082464740@qq.com> (cherry picked from commit 093339d4e7533e789d5bf41f09f4d52695011957) --- benchmarks/benchmark_batch_invariant_logp.py | 163 ++++++--- csrc/ascend/batch_invariant_logp_ascend.asc | 310 ++++++++++++++++++ docs/operators/batch-invariant-logp.md | 32 +- envs.py | 2 + rl_engine/_C_npu.pyi | 10 + rl_engine/kernels/gtest/operator_specs.py | 2 + rl_engine/kernels/ops/ascend/__init__.py | 4 + rl_engine/kernels/ops/ascend/loss/__init__.py | 2 + .../ops/ascend/loss/batch_invariant_logp.py | 137 ++++++++ rl_engine/kernels/registry.py | 13 + rl_engine/platforms/device.py | 27 +- setup.py | 105 +++++- tests/test_batch_invariant_logp.py | 221 +++++++++++++ 13 files changed, 962 insertions(+), 66 deletions(-) create mode 100644 csrc/ascend/batch_invariant_logp_ascend.asc create mode 100644 rl_engine/_C_npu.pyi create mode 100644 rl_engine/kernels/ops/ascend/__init__.py create mode 100644 rl_engine/kernels/ops/ascend/loss/__init__.py create mode 100644 rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py diff --git a/benchmarks/benchmark_batch_invariant_logp.py b/benchmarks/benchmark_batch_invariant_logp.py index 4b39540b..f38af22d 100644 --- a/benchmarks/benchmark_batch_invariant_logp.py +++ b/benchmarks/benchmark_batch_invariant_logp.py @@ -1,21 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Benchmark batch-invariant logp: Native vs Triton vs CUDA (SM90 TMA). +"""Benchmark batch-invariant logp: Native vs Triton vs CUDA (SM90 TMA) vs Ascend. -All three backends compute ``logits[t, target[t]] - logsumexp(logits[t, :])`` -from a materialized ``[N, V]`` logits tensor with a locked, per-row reduction -order (batch-invariant). The comparison here is latency and peak VRAM across a -vocab sweep: +All backends compute ``logits[t, target[t]] - logsumexp(logits[t, :])`` from a +materialized ``[N, V]`` logits tensor with a locked, per-row reduction order +(batch-invariant). The comparison here is latency and peak device memory across +a vocab sweep: - Native materializes ``log_softmax`` over the full ``[N, V]`` tensor. - Triton streams the vocab through an online softmax (grid = one program/row). - CUDA is the Hopper TMA online-softmax kernel (one CTA/row); only present when the extension is built with ``KERNEL_ALIGN_FORCE_SM90=1`` on an SM90 device. +- Ascend is the CANN two-pass streaming kernel (one AI core block/row); only + present when the extension is built with ``KERNEL_ALIGN_FORCE_ASCEND=1``. By default only the forward pass is timed; pass ``--backward`` to also emit the -forward+backward table (grad w.r.t. logits). Timing uses ``torch.cuda`` events, -so the benchmark runs on CUDA/ROCm devices. +forward+backward table (grad w.r.t. logits). Timing and memory accounting +dispatch through the active accelerator (``torch.cuda`` or ``torch.npu``), so +the benchmark runs on CUDA/ROCm/NPU devices. Usage: python benchmarks/benchmark_batch_invariant_logp.py @@ -29,11 +32,30 @@ from tabulate import tabulate from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp -from rl_engine.kernels.ops.triton.loss.batch_invariant_logp import TritonBatchInvariantLogpOp from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger +def _accel(): + """The active accelerator module (torch.npu on Ascend, torch.cuda otherwise).""" + if device_ctx.device_type == "npu": + return torch.npu + return torch.cuda + + +def _maybe_triton_op(): + """The Triton op, or None when unavailable (no Triton / NPU-only host).""" + if device_ctx.device_type == "npu": + return None + try: + from rl_engine.kernels.ops.triton.loss.batch_invariant_logp import ( + TritonBatchInvariantLogpOp, + ) + except ImportError: + return None + return TritonBatchInvariantLogpOp() + + def _maybe_sm90_op(): """The Hopper TMA op, or None when unavailable (non-Hopper / not built).""" from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE @@ -50,6 +72,30 @@ def _maybe_sm90_op(): return BatchInvariantLogpSM90Op() +def _maybe_ascend_op(): + """The Ascend C op, or None when unavailable (no NPU / not built).""" + if device_ctx.device_type != "npu": + return None + try: + from rl_engine.kernels.ops.ascend.loss.batch_invariant_logp import ( + BatchInvariantLogpAscendOp, + ) + except (ImportError, RuntimeError): + return None + return BatchInvariantLogpAscendOp() + + +def _maybe_accelerated_op(): + """(label, op) for the fastest hardware kernel on this host, else (None, None).""" + sm90_op = _maybe_sm90_op() + if sm90_op is not None: + return "cuda", sm90_op + ascend_op = _maybe_ascend_op() + if ascend_op is not None: + return "ascend", ascend_op + return None, None + + # (num_tokens, vocab); vocab kept a multiple of 8 so the bf16 TMA path runs. DEFAULT_CONFIGS = [ (4096, 32768), @@ -66,30 +112,32 @@ def _make_inputs(num_tokens, vocab, device, dtype): def _time_ms(fn, warmup, iters): + acc = _accel() for _ in range(warmup): fn() - torch.cuda.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) + acc.synchronize() + start = acc.Event(enable_timing=True) + end = acc.Event(enable_timing=True) start.record() for _ in range(iters): fn() end.record() - torch.cuda.synchronize() + acc.synchronize() return start.elapsed_time(end) / iters def _peak_vram_gb(fn, warmup, iters): + acc = _accel() for _ in range(warmup): fn() - torch.cuda.synchronize() - torch.cuda.empty_cache() - baseline = torch.cuda.memory_allocated() - torch.cuda.reset_peak_memory_stats() + acc.synchronize() + acc.empty_cache() + baseline = acc.memory_allocated() + acc.reset_peak_memory_stats() for _ in range(iters): fn() - torch.cuda.synchronize() - return (torch.cuda.max_memory_allocated() - baseline) / (1024**3) + acc.synchronize() + return (acc.max_memory_allocated() - baseline) / (1024**3) def _forward_closure(op, logits, target): @@ -111,65 +159,74 @@ def run(): def _bench_table(configs, backends, closure_factory, device, dtype, warmup, iters): - """Build a github-markdown table timing each backend, plus CUDA speedups. + """Build a github-markdown table timing each backend, plus accel speedups. - ``backends`` is ``(native_op, triton_op, sm90_op_or_None)``; ``closure_factory`` - maps ``(op, logits, target) -> callable`` (forward-only or forward+backward). + ``backends`` is ``(native_op, triton_op_or_None, (accel_label, accel_op))``; + ``closure_factory`` maps ``(op, logits, target) -> callable`` (forward-only + or forward+backward). """ - native, triton_op, sm90_op = backends + native, triton_op, (accel_label, accel_op) = backends label = "fwd" if closure_factory is _forward_closure else "fwd+bwd" rows = [] for num_tokens, vocab in configs: logits, target = _make_inputs(num_tokens, vocab, device, dtype) n_c = closure_factory(native, logits, target) - t_c = closure_factory(triton_op, logits, target) n_ms = _time_ms(n_c, warmup, iters) - t_ms = _time_ms(t_c, warmup, iters) n_mb = _peak_vram_gb(n_c, warmup, iters) * 1024 - t_mb = _peak_vram_gb(t_c, warmup, iters) * 1024 - - row = [ - f"{num_tokens}x{vocab}", - f"{n_ms:.3f}", - f"{t_ms:.3f}", - f"{n_ms/t_ms:.2f}x", - f"{n_mb:.0f}", - f"{t_mb:.0f}", - ] - if sm90_op is not None: - s_c = closure_factory(sm90_op, logits, target) - s_ms = _time_ms(s_c, warmup, iters) - s_mb = _peak_vram_gb(s_c, warmup, iters) * 1024 - row += [f"{s_ms:.3f}", f"{n_ms/s_ms:.2f}x", f"{t_ms/s_ms:.2f}x", f"{s_mb:.0f}"] + + row = [f"{num_tokens}x{vocab}", f"{n_ms:.3f}"] + t_ms = None + if triton_op is not None: + t_c = closure_factory(triton_op, logits, target) + t_ms = _time_ms(t_c, warmup, iters) + t_mb = _peak_vram_gb(t_c, warmup, iters) * 1024 + row += [f"{t_ms:.3f}", f"{n_ms/t_ms:.2f}x", f"{n_mb:.0f}", f"{t_mb:.0f}"] + else: + row += [f"{n_mb:.0f}"] + if accel_op is not None: + a_c = closure_factory(accel_op, logits, target) + a_ms = _time_ms(a_c, warmup, iters) + a_mb = _peak_vram_gb(a_c, warmup, iters) * 1024 + row += [f"{a_ms:.3f}", f"{n_ms/a_ms:.2f}x"] + if t_ms is not None: + row += [f"{t_ms/a_ms:.2f}x"] + row += [f"{a_mb:.0f}"] rows.append(row) - headers = [ - "shape (N x V)", - f"native {label} ms", - f"triton {label} ms", - f"{label} speedup", - f"native {label} MB", - f"triton {label} MB", - ] - if sm90_op is not None: - headers += [f"cuda {label} ms", "cuda vs native", "cuda vs triton", f"cuda {label} MB"] + headers = ["shape (N x V)", f"native {label} ms"] + if triton_op is not None: + headers += [ + f"triton {label} ms", + f"{label} speedup", + f"native {label} MB", + f"triton {label} MB", + ] + else: + headers += [f"native {label} MB"] + if accel_op is not None: + headers += [f"{accel_label} {label} ms", f"{accel_label} vs native"] + if triton_op is not None: + headers += [f"{accel_label} vs triton"] + headers += [f"{accel_label} {label} MB"] print(tabulate(rows, headers=headers, tablefmt="github")) def run_benchmark(args): - if device_ctx.device_type not in ["cuda", "hip"]: + if device_ctx.device_type not in ["cuda", "hip", "npu"]: raise RuntimeError( - "batch_invariant_logp benchmark requires a CUDA/ROCm GPU (uses torch.cuda timing)." + "batch_invariant_logp benchmark requires a CUDA/ROCm/NPU device " + "(uses accelerator-event timing)." ) device = device_ctx.device dtype = torch.bfloat16 - backends = (NativeBatchInvariantLogpOp(), TritonBatchInvariantLogpOp(), _maybe_sm90_op()) + accel_label, accel_op = _maybe_accelerated_op() + backends = (NativeBatchInvariantLogpOp(), _maybe_triton_op(), (accel_label, accel_op)) logger.info( f"batch_invariant_logp benchmark on {device} (dtype={dtype}); " - f"SM90 TMA backend {'enabled' if backends[2] is not None else 'unavailable'}" + f"accelerated backend: {accel_label or 'unavailable'}" ) print("Forward") diff --git a/csrc/ascend/batch_invariant_logp_ascend.asc b/csrc/ascend/batch_invariant_logp_ascend.asc new file mode 100644 index 00000000..09226c85 --- /dev/null +++ b/csrc/ascend/batch_invariant_logp_ascend.asc @@ -0,0 +1,310 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Batch-invariant selected-token log-probability, Ascend C (CANN) forward kernel. +// +// logp[n] = logits[n, target[n]] - logsumexp(logits[n, :]) +// +// Mirrors the SM90 CUDA kernel in csrc/cuda/batch_invariant_logp_kernel_sm90.cu: +// - input : logits [N, V] contiguous, bf16 or fp32; target [N] (cast to int32) +// - output : logp [N] fp32, lse [N] fp32 +// - target[n] == ignore_index -> logp[n] = 0, but lse is always computed. +// +// 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 (two-pass +// max -> sum-exp). The instruction sequence for a row depends only on V, +// never on N or on the block the row happens to land on. +// +// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by +// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu. + +#include + +#include "kernel_operator.h" + +#include + +#include "torch_npu/csrc/core/npu/NPUStream.h" + +namespace { + +// Elements per vocab tile. Fixed for all rows and batch sizes; this is what +// makes the reduction order batch-invariant. UB budget (input tile + fp32 tile +// + reduce scratch) stays well under the 192 KB UB of current SoCs. +constexpr uint32_t TILE_LENGTH = 4096; +// Cap on launched blocks. Rows are strided across blocks, so launching fewer +// blocks than rows is fine and never changes per-row numerics. +constexpr int64_t MAX_BLOCKS = 128; +constexpr float NEG_INF = -3.402823466e+38f; // -FLT_MAX + +template +class KernelBatchInvariantLogp { +public: + __aicore__ inline KernelBatchInvariantLogp(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(GM_ADDR logits, + GM_ADDR target, + GM_ADDR logp, + GM_ADDR lse, + int64_t numRows, + int64_t vocabSize, + int64_t ignoreIndex) + { + numRows_ = numRows; + vocabSize_ = vocabSize; + ignoreIndex_ = ignoreIndex; + logitsGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(logits)); + targetGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(target)); + logpGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(logp)); + lseGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(lse)); + pipe_->InitBuffer(inQueue_, 1, TILE_LENGTH * sizeof(T)); + pipe_->InitBuffer(fp32Buf_, TILE_LENGTH * sizeof(float)); + pipe_->InitBuffer(reduceBuf_, TILE_LENGTH * sizeof(float)); + // 64 B: floats [0,8) for reduce/log scratch, floats [8,16) for the + // output staging slots (both halves 32-byte aligned for DataCopyPad). + pipe_->InitBuffer(scalarBuf_, 64); + // 32 B window for reading target[row] via DataCopyPad (GM scalar + // GetValue/SetValue are unreliable on hardware; see cannbot + // ascendc-precision-debug common-traps). + pipe_->InitBuffer(targetBuf_, 32); + + // Intra-core pipeline events. NOTE: AscendC::SyncAll() is a cross-core + // barrier and deadlocks when more blocks are launched than there are + // physical cores (resident blocks wait for unscheduled ones), so all + // synchronization here uses per-pipe SetFlag/WaitFlag instead. + eventVS_ = pipe_->FetchEventID(AscendC::HardEvent::V_S); + eventSV_ = pipe_->FetchEventID(AscendC::HardEvent::S_V); + eventMTE2S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE2_S); + eventSMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE3); + eventMTE3S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE3_S); + } + + __aicore__ inline void Process() + { + for (int64_t row = AscendC::GetBlockIdx(); row < numRows_; row += AscendC::GetBlockNum()) { + ProcessRow(row); + } + } + +private: + // Load one vocab tile into UB and return its fp32 view. When T is fp32 the + // queue buffer is used in place; otherwise the tile is cast into fp32Buf_. + __aicore__ inline AscendC::LocalTensor LoadTileFp32(int64_t row, + int64_t start, + uint32_t count) + { + AscendC::LocalTensor xLocal = inQueue_.AllocTensor(); + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(count * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(xLocal, logitsGm_[row * vocabSize_ + start], copyParams, padParams); + inQueue_.EnQue(xLocal); + xLocal = inQueue_.DeQue(); + inTile_ = xLocal; + + if constexpr (std::is_same_v) { + return xLocal; + } else { + AscendC::LocalTensor fLocal = fp32Buf_.Get(); + AscendC::Cast(fLocal, xLocal, AscendC::RoundMode::CAST_NONE, count); + return fLocal; + } + } + + // Release the queue-owned input buffer of the current tile. + __aicore__ inline void FreeTile() + { + inQueue_.FreeTensor(inTile_); + } + + __aicore__ inline void ProcessRow(int64_t row) + { + const int64_t target = LoadTarget(row); + const bool valid = target != ignoreIndex_; + const int64_t tileCount = (vocabSize_ + TILE_LENGTH - 1) / TILE_LENGTH; + float selected = 0.0f; + + // Pass 1: row max (fixed tile order). Also grab logits[target] on the fly. + float rowMax = NEG_INF; + for (int64_t tile = 0; tile < tileCount; ++tile) { + const int64_t start = tile * TILE_LENGTH; + const uint32_t count = TileCount(start); + AscendC::LocalTensor fLocal = LoadTileFp32(row, start, count); + + AscendC::LocalTensor rTmp = reduceBuf_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + AscendC::ReduceMax(scalar, fLocal, rTmp, static_cast(count), false); + WaitVector(); // vector -> scalar read + const float tileMax = scalar.GetValue(0); + rowMax = tileMax > rowMax ? tileMax : rowMax; + + if (valid && target >= start && target < start + count) { + selected = fLocal.GetValue(static_cast(target - start)); + } + FreeTile(); + } + + // Pass 2: sum(exp(x - rowMax)) with the same fixed tile order. + float sumExp = 0.0f; + for (int64_t tile = 0; tile < tileCount; ++tile) { + const int64_t start = tile * TILE_LENGTH; + const uint32_t count = TileCount(start); + AscendC::LocalTensor fLocal = LoadTileFp32(row, start, count); + + AscendC::Adds(fLocal, fLocal, -rowMax, count); // x - rowMax + AscendC::Exp(fLocal, fLocal, count); + AscendC::LocalTensor rTmp = reduceBuf_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + AscendC::ReduceSum(scalar, fLocal, rTmp, static_cast(count)); + WaitVector(); // vector -> scalar read + sumExp += scalar.GetValue(0); + FreeTile(); + } + + // lse = rowMax + log(sumExp). The scalar unit has no log, so run a + // 1-element vector Log (count padded to 8; scalarBuf_ is 32 B aligned). + AscendC::LocalTensor scalar = scalarBuf_.Get(); + scalar.SetValue(0, sumExp); + AscendC::SetFlag(eventSV_); // scalar write -> vector op + AscendC::WaitFlag(eventSV_); + AscendC::Log(scalar, scalar, 8); + WaitVector(); // vector -> scalar read + const float lse = rowMax + scalar.GetValue(0); + + // Stage outputs 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. + scalar.SetValue(0, valid ? (selected - lse) : 0.0f); + scalar.SetValue(8, lse); + AscendC::SetFlag(eventSMTE3_); // scalar write -> copy-out + AscendC::WaitFlag(eventSMTE3_); + AscendC::DataCopyExtParams outParams{1, sizeof(float), 0, 0, 0}; + AscendC::DataCopyPad(logpGm_[row], scalar[0], outParams); + AscendC::DataCopyPad(lseGm_[row], scalar[8], 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 & ~7LL; // 8 x int32 per 32 B + const int64_t remaining = numRows_ - alignedRow; + const uint32_t winCount = static_cast(remaining < 8 ? remaining : 8); + AscendC::LocalTensor tLocal = targetBuf_.Get(); + AscendC::DataCopyExtParams copyParams{ + 1, static_cast(winCount * sizeof(int32_t)), 0, 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(tLocal, targetGm_[alignedRow], copyParams, padParams); + AscendC::SetFlag(eventMTE2S_); // copy-in -> scalar read + AscendC::WaitFlag(eventMTE2S_); + return static_cast(tLocal.GetValue(static_cast(row - alignedRow))); + } + + __aicore__ inline uint32_t TileCount(int64_t start) const + { + const int64_t remaining = vocabSize_ - start; + return static_cast(remaining < TILE_LENGTH ? remaining : TILE_LENGTH); + } + + AscendC::TPipe* pipe_; + AscendC::GlobalTensor logitsGm_; + AscendC::GlobalTensor targetGm_; + AscendC::GlobalTensor logpGm_; + AscendC::GlobalTensor lseGm_; + AscendC::TQue inQueue_; + AscendC::TBuf fp32Buf_; + AscendC::TBuf reduceBuf_; + AscendC::TBuf scalarBuf_; + AscendC::TBuf targetBuf_; + AscendC::LocalTensor inTile_; + AscendC::TEventID eventVS_; + AscendC::TEventID eventSV_; + AscendC::TEventID eventMTE2S_; + AscendC::TEventID eventSMTE3_; + AscendC::TEventID eventMTE3S_; + int64_t numRows_; + int64_t vocabSize_; + int64_t ignoreIndex_; +}; + +} // namespace + +extern "C" __global__ __vector__ void batch_invariant_logp_ascend_kernel_fp32( + GM_ADDR logits, GM_ADDR target, GM_ADDR logp, GM_ADDR lse, + int64_t numRows, int64_t vocabSize, int64_t ignoreIndex) +{ + AscendC::TPipe pipe; + KernelBatchInvariantLogp op(&pipe); + op.Init(logits, target, logp, lse, numRows, vocabSize, ignoreIndex); + op.Process(); +} + +extern "C" __global__ __vector__ void batch_invariant_logp_ascend_kernel_bf16( + GM_ADDR logits, GM_ADDR target, GM_ADDR logp, GM_ADDR lse, + int64_t numRows, int64_t vocabSize, int64_t ignoreIndex) +{ + AscendC::TPipe pipe; + KernelBatchInvariantLogp op(&pipe); + op.Init(logits, target, logp, lse, numRows, vocabSize, ignoreIndex); + op.Process(); +} + +std::vector batch_invariant_logp_ascend_forward(torch::Tensor logits, + torch::Tensor target, + int64_t ignore_index) +{ + 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 must be bf16 or fp32"); + 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.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); + + torch::Tensor targetI32 = target.to(torch::kInt32).contiguous(); + torch::Tensor logp = at::empty({numRows}, logits.options().dtype(at::kFloat)); + torch::Tensor lse = at::empty({numRows}, logits.options().dtype(at::kFloat)); + + // stream(true): flush the task queue before launch so the kernel cannot + // overtake earlier NPU work; outputs were allocated with at::empty (no + // queued initializer) for the same reason. + auto aclStream = c10_npu::getCurrentNPUStream().stream(true); + const uint32_t blockNum = static_cast(std::min(numRows, MAX_BLOCKS)); + + if (logits.scalar_type() == at::kBFloat16) { + batch_invariant_logp_ascend_kernel_bf16<<>>( + reinterpret_cast(logits.mutable_data_ptr()), + reinterpret_cast(targetI32.mutable_data_ptr()), + reinterpret_cast(logp.mutable_data_ptr()), + reinterpret_cast(lse.mutable_data_ptr()), + numRows, vocabSize, ignore_index); + } else { + batch_invariant_logp_ascend_kernel_fp32<<>>( + reinterpret_cast(logits.mutable_data_ptr()), + reinterpret_cast(targetI32.mutable_data_ptr()), + reinterpret_cast(logp.mutable_data_ptr()), + reinterpret_cast(lse.mutable_data_ptr()), + numRows, vocabSize, ignore_index); + } + return {logp, lse}; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("batch_invariant_logp_ascend", + &batch_invariant_logp_ascend_forward, + "Batch-invariant selected-token log-probability (Ascend C forward)"); +} diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index fbc0e9f1..ab4b0942 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -38,6 +38,7 @@ logp.sum().backward() # gradients flow into logits only | Backend | Wrapper | Status | | --- | --- | --- | | CUDA (SM90 TMA) | `BatchInvariantLogpSM90Op` | Hopper TMA online-softmax forward. | +| Ascend (CANN) | `BatchInvariantLogpAscendOp` | Ascend C two-pass streaming forward; PyTorch-formula backward. | | CUDA / ROCm (Triton) | `TritonBatchInvariantLogpOp` | Triton online-softmax forward and tile-wise backward. Requires a GPU tensor. | | PyTorch native | `NativeBatchInvariantLogpOp` | FP32 reference path; CPU fallback and Triton-less fallback. | @@ -46,6 +47,7 @@ Current dispatch: ```text CUDA (Hopper, SM90 kernel compiled): CUDA (SM90 TMA) -> Triton -> PyTorch CUDA / ROCm (otherwise): Triton -> PyTorch +Ascend NPU: Ascend -> PyTorch CPU: PyTorch ``` @@ -54,6 +56,14 @@ CUDA priority list when the extension exposes `_C.batch_invariant_logp_sm90` (built with `KERNEL_ALIGN_FORCE_SM90=1`) on an SM90 device. On any other build or device, dispatch is unchanged (Triton -> PyTorch). +The Ascend backend lives on the `npu` platform key and is available when the +extension exposes `_C_npu.batch_invariant_logp_ascend` (built with +`KERNEL_ALIGN_FORCE_ASCEND=1` on a CANN + torch_npu host; `npu-arch` defaults +to `dav-2201`, override with `KERNEL_ALIGN_ASCEND_ARCH`). When the extension is +not compiled, instantiation fails and dispatch falls through to PyTorch native. +bf16/fp32 NPU tensors run the Ascend C kernel; anything else (e.g. fp16) +silently falls back to the native op. + ## Benchmarks `benchmarks/benchmark_batch_invariant_logp.py` compares Native, Triton, and the @@ -159,6 +169,10 @@ The operator is designed so each row is computed independently: - Triton backward uses `grid=(num_tokens, vocab_tiles)` and writes one row tile per program. It reuses the forward-saved per-row `lse`, so no backward reduction crosses row boundaries. +- The Ascend forward strides rows across blocks, so one AI core block owns + exactly one row; the vocab is scanned left-to-right in fixed + `TILE_LENGTH=4096` tiles with a two-pass (max, then sum-exp) fixed-order + reduction. - No atomic writes are used. These constraints ensure the result for a row depends only on that row's logits @@ -206,21 +220,25 @@ out.sum().backward() python -m pytest tests/test_batch_invariant_logp.py -q -rs ``` -All backends (Native, Triton) are tested in a single file. Coverage includes: -correctness, leading-shape preservation, batch-invariance (bitwise), validation, -ignore-index behavior, backward correctness, CUDA smoke cases, registry -dispatch, and Triton-specific fp32/fp16/bf16 correctness, large vocab, backward -gradient batch-invariance, and ignored-row zero gradients. +All backends (Native, Triton, SM90, Ascend) are tested in a single file. +Coverage includes: correctness, leading-shape preservation, batch-invariance +(bitwise), validation, ignore-index behavior, backward correctness, CUDA smoke +cases, registry dispatch, and Triton-specific fp32/fp16/bf16 correctness, large +vocab, backward gradient batch-invariance, and ignored-row zero gradients. -Triton tests skip when Triton or CUDA is unavailable. On Windows, run via -WSL/Linux with CUDA. +Triton tests skip when Triton or CUDA is unavailable. SM90 tests skip without a +Hopper build; Ascend tests skip without an NPU + `_C_npu` build. On Windows, run +via WSL/Linux with CUDA. ## Implementation Files - `rl_engine/kernels/ops/pytorch/loss/batch_invariant_logp.py` - `rl_engine/kernels/ops/triton/loss/batch_invariant_logp.py` - `rl_engine/kernels/ops/cuda/loss/batch_invariant_logp.py` +- `rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py` - `csrc/cuda/batch_invariant_logp_kernel_sm90.cu` +- `csrc/ascend/batch_invariant_logp_ascend.asc` - `rl_engine/kernels/registry.py` +- `rl_engine/platforms/device.py` - `tests/test_batch_invariant_logp.py` - `benchmarks/benchmark_batch_invariant_logp.py` diff --git a/envs.py b/envs.py index 34aded7c..833f00ac 100644 --- a/envs.py +++ b/envs.py @@ -26,3 +26,5 @@ def env_flag(name: str, default: bool = False) -> bool: KERNEL_ALIGN_NCU_LINEINFO = "KERNEL_ALIGN_NCU_LINEINFO" KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC = "KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC" KERNEL_ALIGN_FORCE_SM90 = "KERNEL_ALIGN_FORCE_SM90" +KERNEL_ALIGN_FORCE_ASCEND = "KERNEL_ALIGN_FORCE_ASCEND" +KERNEL_ALIGN_ASCEND_ARCH = "KERNEL_ALIGN_ASCEND_ARCH" diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi new file mode 100644 index 00000000..bff5e2e7 --- /dev/null +++ b/rl_engine/_C_npu.pyi @@ -0,0 +1,10 @@ +# rl_engine/_C_npu.pyi +# Type stub for the compiled Ascend C (CANN) extension module. +# Built only when KERNEL_ALIGN_FORCE_ASCEND=1 on a machine with CANN + torch_npu. +import torch + +def batch_invariant_logp_ascend( + logits: torch.Tensor, + target: torch.Tensor, + ignore_index: int, +) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index df17c911..bde87edb 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -187,6 +187,8 @@ def _load_object(path: str) -> Any: "TritonBatchInvariantLogpOp", "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp." "BatchInvariantLogpSM90Op", + "ascend": "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp." + "BatchInvariantLogpAscendOp", }, grad_input_names=("logits",), ), diff --git a/rl_engine/kernels/ops/ascend/__init__.py b/rl_engine/kernels/ops/ascend/__init__.py new file mode 100644 index 00000000..ab85458d --- /dev/null +++ b/rl_engine/kernels/ops/ascend/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from . import loss # noqa: F401 diff --git a/rl_engine/kernels/ops/ascend/loss/__init__.py b/rl_engine/kernels/ops/ascend/loss/__init__.py new file mode 100644 index 00000000..86cf4c9d --- /dev/null +++ b/rl_engine/kernels/ops/ascend/loss/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py new file mode 100644 index 00000000..657ee127 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import torch + +from rl_engine.utils.logger import logger + +_C_npu = 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 + + +def _ascend_supported(logits: torch.Tensor) -> bool: + """Whether the Ascend C forward can run these logits directly. + + NPU tensors only, bf16/fp32 only (mirrors the SM90 kernel's dtype gate). + """ + return logits.device.type == "npu" and logits.dtype in ( + torch.bfloat16, + torch.float32, + ) + + +def _fallback_op(): + """Portable op for inputs the Ascend forward cannot take. + + Triton rejects non-CUDA devices, so on NPU the only fallback is native. + """ + from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp + + return NativeBatchInvariantLogpOp() + + +class _BatchInvariantLogpAscendFunction(torch.autograd.Function): + # Autograd wrapper: Ascend C forward + PyTorch-formula backward. + # The SM90 op reuses Triton's tile-wise backward; Triton is unavailable on + # NPU, so the backward uses the same onehot - softmax formula as the SM90 + # wrapper's portable branch, reusing the forward-saved lse. + + @staticmethod + def forward(ctx, logits, target_ids, ignore_index): + lead_shape = logits.shape[:-1] + vocab_size = logits.size(-1) + + logits_2d = logits.reshape(-1, vocab_size).contiguous() + target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.int64).contiguous() + + logp, lse = _C_npu.batch_invariant_logp_ascend(logits_2d, target_1d, int(ignore_index)) + + ctx.save_for_backward(logits_2d, target_1d, lse) + ctx.ignore_index = ignore_index + ctx.lead_shape = lead_shape + ctx.vocab_size = vocab_size + return logp.reshape(lead_shape) + + @staticmethod + def backward(ctx, grad_output): + logits_2d, target_1d, lse = ctx.saved_tensors + ignore_index = ctx.ignore_index + vocab_size = ctx.vocab_size + + grad_flat = grad_output.reshape(-1).contiguous().to(torch.float32) + + valid = target_1d != ignore_index + safe_target = torch.where(valid, target_1d, torch.zeros_like(target_1d)) + probs = torch.exp(logits_2d.float() - lse.unsqueeze(1)) + onehot = torch.zeros_like(probs) + onehot.scatter_(1, safe_target.unsqueeze(1), 1.0) + grad = grad_flat.unsqueeze(1) * (onehot - probs) + grad = torch.where(valid.unsqueeze(1), grad, torch.zeros_like(grad)) + grad_logits = grad.to(logits_2d.dtype) + + grad_logits = grad_logits.reshape(ctx.lead_shape + (vocab_size,)) + return grad_logits, None, None + + +class BatchInvariantLogpAscendOp: + # Ascend C batch-invariant selected-token log-probability (forward kernel). + + def __init__(self) -> None: + if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "batch_invariant_logp_ascend"): + raise RuntimeError( + "batch_invariant_logp_ascend is not compiled into the extension. Rebuild with " + "KERNEL_ALIGN_FORCE_ASCEND=1 on an Ascend NPU host: 'pip install -e .'" + ) + logger.info("Successfully linked to precompiled _C_npu.batch_invariant_logp_ascend kernel.") + + def __call__( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int = -100, + *, + validate: bool = False, + ) -> torch.Tensor: + return self.apply(logits, target_ids, ignore_index=ignore_index, validate=validate) + + def apply( + self, + logits: torch.Tensor, + target_ids: torch.Tensor, + ignore_index: int = -100, + *, + validate: bool = False, + ) -> torch.Tensor: + if logits.dim() < 2: + raise ValueError( + f"logits must be at least 2-D ([*lead, V]), got shape {tuple(logits.shape)}" + ) + if logits.shape[:-1] != target_ids.shape: + raise ValueError( + f"logits leading shape {tuple(logits.shape[:-1])} must match " + f"target_ids shape {tuple(target_ids.shape)}" + ) + + if not _ascend_supported(logits): + return _fallback_op()(logits, target_ids, ignore_index=ignore_index, validate=validate) + + if validate: + vocab_size = logits.size(-1) + target_flat = target_ids.reshape(-1) + valid_targets = target_flat[target_flat != ignore_index] + if valid_targets.numel() > 0 and ( + (valid_targets < 0).any() or (valid_targets >= vocab_size).any() + ): + bad = valid_targets[(valid_targets < 0) | (valid_targets >= vocab_size)] + raise ValueError( + f"target_ids contains values outside [0, {vocab_size}): {bad.tolist()}" + ) + + return _BatchInvariantLogpAscendFunction.apply(logits, target_ids, ignore_index) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index efde5c25..4020cd01 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -74,6 +74,9 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_BATCH_INVARIANT_LOGP_SM90 = ( "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90Op" ) + ASCEND_BATCH_INVARIANT_LOGP = ( + "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAscendOp" + ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" @@ -287,6 +290,14 @@ def __init__(self): "silu": [OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], }, + # Ascend NPU: op types without an entry fall back to PYTORCH_NATIVE + # (see get_op's default), so only Ascend-accelerated ops are listed. + "npu": { + "batch_invariant_logp": [ + OpBackend.ASCEND_BATCH_INVARIANT_LOGP, + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + ], + }, } logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") self._adjust_priority_for_hardware() @@ -407,6 +418,8 @@ def _platform_for_device(self, device: torch.device | str | None) -> str: return "rocm" if device_ctx.device_type == "cuda": return "cuda" + if device_ctx.device_type == "npu": + return "npu" return "cpu" resolved = torch.device(device) diff --git a/rl_engine/platforms/device.py b/rl_engine/platforms/device.py index 5a494d71..98f84811 100644 --- a/rl_engine/platforms/device.py +++ b/rl_engine/platforms/device.py @@ -7,12 +7,22 @@ from rl_engine.utils.logger import logger +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() + + class DeviceContext: """ Hardware-aware context manager for high-performance RL tasks. - Provides transparent support for both AMD (ROCm/HIP) and NVIDIA (CUDA) - architectures to ensure backend-agnostic scaling for RL operators. + Provides transparent support for AMD (ROCm/HIP), NVIDIA (CUDA) and Huawei + Ascend (NPU/CANN) architectures to ensure backend-agnostic scaling for RL + operators. """ def __init__(self): @@ -41,8 +51,17 @@ def __init__(self): f" (Version: {self.backend_version})" ) else: - self.device_type = DeviceType.CPU.value - logger.warning("No GPU detected. RL-Engine is falling back to CPU mode.") + if _npu_available(): + self.device = torch.device(DeviceType.NPU.value) + self.device_type = DeviceType.NPU.value + self.backend_version = getattr(torch.version, "cann", None) or "N/A" + logger.info_once( + f"RL-Engine initialized with Huawei Ascend NPU backend" + f" (CANN Version: {self.backend_version})" + ) + else: + self.device_type = DeviceType.CPU.value + logger.warning("No GPU detected. RL-Engine is falling back to CPU mode.") def get_preferred_dtype(self): """ diff --git a/setup.py b/setup.py index 2c7af89e..57c98070 100644 --- a/setup.py +++ b/setup.py @@ -3,9 +3,12 @@ import importlib.util import os +import sysconfig +from distutils.errors import CompileError +from distutils.spawn import find_executable from pathlib import Path -from setuptools import find_packages, setup +from setuptools import Extension, find_packages, setup def _load_envs_module(): @@ -195,14 +198,112 @@ def get_extensions(): extra_link_args=extra_link_args, ) ) + extensions.extend(_ascend_extensions()) return extensions +def _ascend_extensions(): + """Ascend C (CANN) kernels, built with bisheng. Gated on KERNEL_ALIGN_FORCE_ASCEND=1. + + Follows the official torch_npu cpp_extension_asc pattern: .asc sources + (kernel + host + pybind) are compiled by the CANN bisheng compiler into a + single rl_engine._C_npu extension module. Requires CANN toolkit (bisheng on + PATH or ASCEND_HOME_PATH set) and torch_npu. + """ + if not envs.env_flag(envs.KERNEL_ALIGN_FORCE_ASCEND): + return [] + try: + import torch # noqa: F401 + import torch_npu # noqa: F401 + except ImportError as e: + raise RuntimeError( + "KERNEL_ALIGN_FORCE_ASCEND=1 requires torch and torch_npu to be installed" + ) from e + + asc_srcs = sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) + if not asc_srcs: + raise RuntimeError("KERNEL_ALIGN_FORCE_ASCEND=1 but no .asc sources under csrc/ascend/") + return [Extension(name="rl_engine._C_npu", sources=asc_srcs, language="asc")] + + +def _bisheng_compile_cmd(ext, ext_fullpath): + """Single-command bisheng build for an Ascend C extension (see op-plugin example).""" + import torch + import torch.utils.cpp_extension as cpp_extension + import torch_npu + + if find_executable("bisheng") is None: + raise RuntimeError( + "bisheng compiler not found on PATH; source the CANN toolkit environment first" + ) + + soc = os.environ.get(envs.KERNEL_ALIGN_ASCEND_ARCH, "dav-2201") # A2/A3; A5: dav-3510 + abi_value = "1" if torch._C._GLIBCXX_USE_CXX11_ABI else "0" + module_name = ext.name.rsplit(".", 1)[-1] + + torch_npu_dir = os.path.dirname(os.path.realpath(torch_npu.__file__)) + ascend_home = os.environ.get("ASCEND_HOME_PATH", "/usr/local/Ascend/ascend-toolkit/latest") + + include_dirs = [ + *cpp_extension.include_paths(), + sysconfig.get_config_var("INCLUDEPY"), + os.path.join(torch_npu_dir, "include"), + os.path.join(torch_npu_dir, "include", "third_party", "acl", "inc"), + os.path.join(ascend_home, "include"), + ] + lib_dirs = [ + sysconfig.get_config_var("LIBDIR"), + os.path.join(os.path.dirname(torch.__file__), "lib"), + os.path.join(torch_npu_dir, "lib"), + os.path.join(ascend_home, "lib64"), + ] + + cmd = [ + "bisheng", + "-x", + "asc", + f"--npu-arch={soc}", + "-shared", + "-fPIC", + "-std=c++17", + "-O2", + f"-D_GLIBCXX_USE_CXX11_ABI={abi_value}", + f"-DTORCH_EXTENSION_NAME={module_name}", + "-lascendcl", + "-ltorch_npu", + "-ltorch", + "-ltorch_cpu", + "-ltorch_python", + "-lc10", + *ext.sources, + "-o", + ext_fullpath, + ] + cmd += [f"-I{d}" for d in include_dirs if d] + cmd += [f"-L{d}" for d in lib_dirs if d] + return cmd + + def get_cmdclass(): _, BuildExtension, _, _ = _load_torch_extension_tools() if BuildExtension is None: return {} - return {"build_ext": BuildExtension} + + class AscendBuildExtension(BuildExtension): + """torch BuildExtension + bisheng path for language="asc" extensions.""" + + def build_extension(self, ext): + if getattr(ext, "language", None) != "asc": + super().build_extension(ext) + return + ext_fullpath = self.get_ext_fullpath(ext.name) + os.makedirs(os.path.dirname(ext_fullpath), exist_ok=True) + try: + self.spawn(_bisheng_compile_cmd(ext, ext_fullpath)) + except Exception as e: + raise CompileError(str(e)) from e + + return {"build_ext": AscendBuildExtension} setup( diff --git a/tests/test_batch_invariant_logp.py b/tests/test_batch_invariant_logp.py index 06e5dbb4..93ec3548 100644 --- a/tests/test_batch_invariant_logp.py +++ b/tests/test_batch_invariant_logp.py @@ -60,6 +60,34 @@ def _sm90_kernel_available() -> bool: "(needs KERNEL_ALIGN_FORCE_SM90=1 on an SM90/Hopper device).", ) + +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.batch_invariant_logp import ( + _NPU_EXT_AVAILABLE, + _C_npu, + ) + except ImportError: + return False + return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "batch_invariant_logp_ascend") + + +requires_ascend = pytest.mark.skipif( + not _ascend_kernel_available(), + reason="batch_invariant_logp_ascend kernel not compiled " + "(needs KERNEL_ALIGN_FORCE_ASCEND=1 on an Ascend NPU host).", +) + # 16-byte-aligned vocab so the TMA forward runs directly (not the fallback): # bf16 needs V % 8 == 0, fp32 needs V % 4 == 0. _VC = 1024 @@ -1022,6 +1050,198 @@ def test_fp16_is_rejected(self): op(logits, target) +# --------------------------------------------------------------------------- +# 7c. Ascend C kernel backend +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendCorrectness: + """Compiled Ascend C kernel output must match log_softmax + gather.""" + + def _get_op(self): + from rl_engine.kernels.ops.ascend.loss.batch_invariant_logp import ( + BatchInvariantLogpAscendOp, + ) + + return BatchInvariantLogpAscendOp() + + def test_matches_reference_fp32(self): + op = self._get_op() + logits = torch.randn(8, _VC, device="npu") + target = torch.randint(0, _VC, (8,), device="npu") + out = op(logits, target) + ref = _reference_logp(logits, target) + assert out.dtype == torch.float32 + assert torch.allclose(out, ref, atol=1e-4) + + def test_matches_reference_bf16(self): + op = self._get_op() + logits = torch.randn(8, _VC, device="npu", dtype=torch.bfloat16) + target = torch.randint(0, _VC, (8,), device="npu") + out = op(logits, target) + ref = _reference_logp(logits.float(), target) + assert out.dtype == torch.float32 + assert torch.allclose(out, ref, atol=1e-3) + + def test_large_vocab(self): + op = self._get_op() + logits = torch.randn(4, 128256, device="npu", dtype=torch.bfloat16) + target = torch.randint(0, 128256, (4,), device="npu") + out = op(logits, target) + ref = _reference_logp(logits.float(), target) + assert torch.allclose(out, ref, atol=2e-3) + + def test_single_token(self): + op = self._get_op() + logits = torch.randn(1, _VC, device="npu") + target = torch.randint(0, _VC, (1,), device="npu") + out = op(logits, target) + ref = _reference_logp(logits, target) + assert torch.allclose(out, ref, atol=1e-4) + + def test_3d_logits(self): + op = self._get_op() + logits = torch.randn(2, 3, _VC, device="npu", dtype=torch.bfloat16) + target = torch.randint(0, _VC, (2, 3), device="npu") + out = op(logits, target) + assert out.shape == (2, 3) + ref = _reference_logp(logits.float(), target) + assert torch.allclose(out, ref, atol=1e-3) + + def test_matches_pytorch_op(self): + op = self._get_op() + pytorch_op = NativeBatchInvariantLogpOp() + logits = torch.randn(16, _VC, device="npu") + target = torch.randint(0, _VC, (16,), device="npu") + assert torch.allclose(op(logits, target), pytorch_op(logits, target), atol=1e-4) + + +@requires_ascend +class TestAscendBatchInvariance: + """Ascend kernel must be bitwise batch-invariant (one block per row).""" + + def _get_op(self): + from rl_engine.kernels.ops.ascend.loss.batch_invariant_logp import ( + BatchInvariantLogpAscendOp, + ) + + return BatchInvariantLogpAscendOp() + + def test_batch_size_1_vs_n(self): + op = self._get_op() + row = _make_row(42, vocab=_VC, device="npu") + target = torch.tensor([7], device="npu") + result_alone = op(row, target).item() + + for batch_size in [2, 4, 8, 16, 32, 64, 128]: + batch_logits = torch.randn(batch_size, _VC, device="npu") + batch_target = torch.randint(0, _VC, (batch_size,), device="npu") + batch_logits[0] = row.squeeze(0) + batch_target[0] = target.squeeze(0) + result_in_batch = op(batch_logits, batch_target)[0].item() + assert result_alone == result_in_batch, ( + f"Ascend drift at batch_size={batch_size}: " + f"alone={result_alone}, in_batch={result_in_batch}" + ) + + def test_different_positions(self): + op = self._get_op() + row = _make_row(99, vocab=_VC, device="npu") + target = torch.tensor([13], device="npu") + batch_size = 16 + results = [] + for pos in range(batch_size): + batch_logits = torch.randn(batch_size, _VC, device="npu") + batch_target = torch.randint(0, _VC, (batch_size,), device="npu") + batch_logits[pos] = row.squeeze(0) + batch_target[pos] = target.squeeze(0) + results.append(op(batch_logits, batch_target)[pos].item()) + assert all( + r == results[0] for r in results + ), f"Ascend position drift: unique = {set(results)}" + + def test_repeated_runs(self): + op = self._get_op() + logits = torch.randn(16, _VC, device="npu", dtype=torch.bfloat16) + target = torch.randint(0, _VC, (16,), device="npu") + results = [op(logits, target) for _ in range(50)] + for i, r in enumerate(results[1:], 1): + assert torch.equal(r, results[0]), f"Ascend run {i} differs from run 0" + + +@requires_ascend +class TestAscendBackward: + """Gradient through the Ascend op must match reference.""" + + def _get_op(self): + from rl_engine.kernels.ops.ascend.loss.batch_invariant_logp import ( + BatchInvariantLogpAscendOp, + ) + + return BatchInvariantLogpAscendOp() + + def test_backward_matches_reference(self): + op = self._get_op() + logits = torch.randn(4, _VC, device="npu", requires_grad=True) + target = torch.randint(0, _VC, (4,), device="npu") + op(logits, target).sum().backward() + grad = logits.grad.detach().clone() + + ref_logits = logits.detach().clone().requires_grad_(True) + _reference_logp(ref_logits, target).sum().backward() + assert torch.allclose(grad, ref_logits.grad, atol=1e-4) + + def test_ignored_row_grad_is_zero(self): + op = self._get_op() + logits = torch.randn(4, _VC, device="npu", requires_grad=True) + target = torch.tensor([0, -100, 2, -100], device="npu") + op(logits, target).sum().backward() + assert torch.equal(logits.grad[1], torch.zeros(_VC, device="npu")) + assert torch.equal(logits.grad[3], torch.zeros(_VC, device="npu")) + + +@requires_ascend +class TestAscendIgnoreIndex: + + def _get_op(self): + from rl_engine.kernels.ops.ascend.loss.batch_invariant_logp import ( + BatchInvariantLogpAscendOp, + ) + + return BatchInvariantLogpAscendOp() + + def test_ignore_outputs_zero(self): + op = self._get_op() + logits = torch.randn(4, _VC, device="npu") + target = torch.tensor([0, -100, 2, -100], device="npu") + out = op(logits, target) + assert out[1].item() == 0.0 + assert out[3].item() == 0.0 + ref = _reference_logp(logits[[0, 2]], target[[0, 2]]) + assert torch.allclose(out[[0, 2]], ref, atol=1e-4) + + +@requires_ascend +class TestAscendFallback: + """Inputs the Ascend path can't take must silently fall back and stay correct.""" + + def _get_op(self): + from rl_engine.kernels.ops.ascend.loss.batch_invariant_logp import ( + BatchInvariantLogpAscendOp, + ) + + return BatchInvariantLogpAscendOp() + + def test_fp16_falls_back(self): + op = self._get_op() + logits = torch.randn(8, _VC, device="npu", dtype=torch.float16) + target = torch.randint(0, _VC, (8,), device="npu") + out = op(logits, target) + ref = _reference_logp(logits.float(), target) + assert torch.allclose(out, ref, atol=1e-3) + + # --------------------------------------------------------------------------- # 8. Registry dispatch test # --------------------------------------------------------------------------- @@ -1035,6 +1255,7 @@ def test_registry_dispatches_correctly(): isinstance(op, NativeBatchInvariantLogpOp) or type(op).__name__ == "TritonBatchInvariantLogpOp" or type(op).__name__ == "BatchInvariantLogpSM90Op" + or type(op).__name__ == "BatchInvariantLogpAscendOp" ) logits = torch.randn(4, _V, device="cuda" if torch.cuda.is_available() else "cpu") target = torch.randint(0, _V, (4,), device=logits.device) From 4e06d4a65670ec936bf2119c831de5147f584bfd Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Thu, 13 Aug 2026 21:18:30 +0800 Subject: [PATCH 2/5] fix lint and tested with gtest Signed-off-by: Zhang Jian (cherry picked from commit 7a763dd295ce0541f15c4ef717d3137a48a33319) --- .../ops/ascend/loss/batch_invariant_logp.py | 4 +++- scripts/check_operator.py | 20 +++++++++++++++++-- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py b/rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py index 657ee127..b12a8413 100644 --- a/rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py +++ b/rl_engine/kernels/ops/ascend/loss/batch_invariant_logp.py @@ -3,11 +3,13 @@ from __future__ import annotations +from typing import Any + import torch from rl_engine.utils.logger import logger -_C_npu = None +_C_npu: Any = None try: from rl_engine import _C_npu 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") From 7eef254c141d4be9a727a93a457994e28bdba40b Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Wed, 19 Aug 2026 19:35:11 +0800 Subject: [PATCH 3/5] feat(ws1): ascend deterministic attention (attention op type) - csrc/ascend/attention/deterministic_attention_ascend.asc: Ascend C batch-invariant standard-softmax attention forward (one AI-core block per (b, q_head, row), fixed 64-key tile order, two-pass streaming reduction, no split-K); bf16/fp16, D=128, GQA, causal + key_padding_mask - rl_engine/kernels/ops/ascend/attention/deterministic_attn.py: autograd wrapper (Ascend C forward + reference backward via NativeAttentionOp) - registry: ASCEND_DETERMINISTIC_ATTENTION op backend; npu priority map gains attention/attn/kv_cache_attention entries - gtest spec: attention gains an ascend candidate - tests/test_attention_ascend.py, benchmarks/benchmark_attention.py, docs/operators/attention.md, _C_npu.pyi stub - setup.py: recursive .asc glob so kernels may live in subdirectories (cherry picked from commit 7970761c947bbd30d647d55356701faff6c2d04e) --- benchmarks/benchmark_attention.py | 225 ++++++-- .../deterministic_attention_ascend.asc | 505 ++++++++++++++++++ docs/operators/attention.md | 15 + rl_engine/_C_npu.pyi | 9 + rl_engine/kernels/gtest/operator_specs.py | 4 + .../kernels/ops/ascend/attention/__init__.py | 8 + .../ascend/attention/deterministic_attn.py | 193 +++++++ rl_engine/kernels/registry.py | 14 + setup.py | 2 +- tests/test_attention_ascend.py | 277 ++++++++++ 10 files changed, 1197 insertions(+), 55 deletions(-) create mode 100644 csrc/ascend/attention/deterministic_attention_ascend.asc create mode 100644 rl_engine/kernels/ops/ascend/attention/__init__.py create mode 100644 rl_engine/kernels/ops/ascend/attention/deterministic_attn.py create mode 100644 tests/test_attention_ascend.py diff --git a/benchmarks/benchmark_attention.py b/benchmarks/benchmark_attention.py index c9509a19..0f4c7adc 100644 --- a/benchmarks/benchmark_attention.py +++ b/benchmarks/benchmark_attention.py @@ -1,70 +1,187 @@ -# File: benchmarks/benchmark_attention.py -import pandas as pd -import torch -import triton -from tabulate import tabulate +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors -from rl_engine.kernels.ops.cuda.attention.prefix_shared_attn import PrefixSharedAttentionOp +"""Benchmark deterministic standard-softmax attention across backends. +All backends compute ``softmax(Q K^T * scale + causal mask) @ V`` with a locked, +per-row reduction order (batch-invariant, no split-K). The comparison here is +latency across a sequence sweep: -def run_benchmark(): - bs = 1 - G = 64 - len_q = 512 - dim = 128 +- Native is the pure-PyTorch fp32-accumulating ground-truth reference. +- Ascend is the CANN two-pass streaming kernel (one AI core block/row); only + present when the extension is built with ``KERNEL_ALIGN_FORCE_ASCEND=1``. +- CUDA is the deterministic op (one CTA/row); only present when the extension + is built with ``KERNEL_ALIGN_FORCE_SM90=1`` on an SM90 device. - len_kvs = [1024, 2048, 4096, 8192, 16384] +Timing dispatch through the active accelerator (``torch.cuda`` or +``torch.npu``), so the benchmark runs on CUDA/ROCm/NPU devices. - print("Benchmarking GRPO Prefix-Shared Attention") - print(f"Fixed Shapes: Batch={bs}, Group(Response)={G}, Query_Len={len_q}, Head_Dim={dim}\n") +Usage: + python benchmarks/benchmark_attention.py + python benchmarks/benchmark_attention.py --backward + python benchmarks/benchmark_attention.py --configs "1,8,512;2,8,2048" +""" - prefix_shared_sdpa = PrefixSharedAttentionOp() - results = [] +import argparse - for len_kv in len_kvs: - q = torch.randn(bs, G, len_q, dim, dtype=torch.bfloat16, device="cuda") - k = torch.randn(bs, len_kv, dim, dtype=torch.bfloat16, device="cuda") - v = torch.randn(bs, len_kv, dim, dtype=torch.bfloat16, device="cuda") +import torch +from tabulate import tabulate - k_exp = k.unsqueeze(1).expand(-1, G, -1, -1).reshape(bs * G, 1, len_kv, dim).contiguous() - v_exp = v.unsqueeze(1).expand(-1, G, -1, -1).reshape(bs * G, 1, len_kv, dim).contiguous() - q_res = q.view(bs * G, 1, len_q, dim) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.platforms.device import device_ctx +from rl_engine.utils.logger import logger - for _ in range(5): - _ = torch.nn.functional.scaled_dot_product_attention(q_res, k_exp, v_exp) - _ = prefix_shared_sdpa(q, k, v) +_D = 128 +_DEFAULT_KV_HEADS = 8 - native_ms = triton.testing.do_bench( - lambda: torch.nn.functional.scaled_dot_product_attention(q_res, k_exp, v_exp), - return_mode="median", - ) - custom_ms = triton.testing.do_bench( - lambda: prefix_shared_sdpa(q, k, v), return_mode="median" - ) +def _accel(): + """The active accelerator module (torch.npu on Ascend, torch.cuda otherwise).""" + if device_ctx.device_type == "npu": + return torch.npu + return torch.cuda - speedup = native_ms / custom_ms - reduction = (native_ms - custom_ms) / native_ms * 100 - - flops = 4 * bs * G * len_q * len_kv * dim - native_tflops = (flops / 1e12) / (native_ms / 1000) - custom_tflops = (flops / 1e12) / (custom_ms / 1000) - - results.append( - { - "Prompt Len": len_kv, - "Native (ms)": f"{native_ms:.3f}", - "RL-Kernel (ms)": f"{custom_ms:.3f}", - "Native TFLOPS": f"{native_tflops:.1f}", - "RL-Kernel TFLOPS": f"{custom_tflops:.1f}", - "Speedup": f"{speedup:.2f}x", - "Time Saved": f"{reduction:.1f}%", - } - ) - df = pd.DataFrame(results) - print(tabulate(df, headers="keys", tablefmt="pretty", stralign="center", showindex=False)) +def _maybe_ascend_op(): + """The Ascend C op, or None when unavailable (no NPU / not built).""" + if device_ctx.device_type != "npu": + return None + try: + from rl_engine.kernels.ops.ascend.attention.deterministic_attn import ( + DeterministicAttentionAscendOp, + ) + except (ImportError, RuntimeError): + return None + return DeterministicAttentionAscendOp() + + +def _maybe_cuda_op(): + """The CUDA deterministic op, or None when unavailable (no CUDA / not built).""" + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + if not ( + torch.cuda.is_available() + and _EXT_AVAILABLE + and hasattr(_C, "deterministic_attention_forward") + ): + return None + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + DeterministicAttentionOp, + ) + + return DeterministicAttentionOp() + + +# (batch, num_q_heads, seqlen); Hq = 32, Hkv = 8 (Qwen3-style GQA, g = 4). +DEFAULT_CONFIGS = [ + (1, 32, 512), + (1, 32, 1024), + (1, 32, 2048), + (4, 32, 2048), +] + + +def _make_inputs(batch, hq, seq, device, dtype): + generator = torch.Generator(device="cpu").manual_seed(0) + q = torch.randn(batch, hq, seq, _D, dtype=dtype, generator=generator).to(device) + k = torch.randn(batch, _DEFAULT_KV_HEADS, seq, _D, dtype=dtype, generator=generator).to(device) + v = torch.randn(batch, _DEFAULT_KV_HEADS, seq, _D, dtype=dtype, generator=generator).to(device) + return q, k, v + + +def _time_ms(fn, warmup, iters): + acc = _accel() + for _ in range(warmup): + fn() + acc.synchronize() + start = acc.Event(enable_timing=True) + end = acc.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + acc.synchronize() + return start.elapsed_time(end) / iters + + +def _forward_closure(op, q, k, v): + def run(): + with torch.no_grad(): + op(q, k, v, causal=True) + + return run + + +def _forward_backward_closure(op, q, k, v): + def run(): + qq = q.detach().requires_grad_(True) + kk = k.detach().requires_grad_(True) + vv = v.detach().requires_grad_(True) + op(qq, kk, vv, causal=True).sum().backward() + + return run + + +def _bench_table(configs, native_op, other_ops, closure_factory, device, dtype, warmup, iters): + label = "fwd" if closure_factory is _forward_closure else "fwd+bwd" + rows = [] + for batch, hq, seq in configs: + q, k, v = _make_inputs(batch, hq, seq, device, dtype) + n_c = closure_factory(native_op, q, k, v) + n_ms = _time_ms(n_c, warmup, iters) + row = [f"{batch}x{hq}x{seq}", f"{n_ms:.3f}"] + for _name, op in other_ops: + if op is None: + row += ["-"] + continue + o_ms = _time_ms(closure_factory(op, q, k, v), warmup, iters) + row += [f"{o_ms:.3f}", f"{n_ms/o_ms:.2f}x"] + rows.append(row) + + headers = ["shape (B x Hq x S)", f"native {label} ms"] + for name, _ in other_ops: + headers += [f"{name} {label} ms", f"vs native"] + logger.info("\n" + tabulate(rows, headers=headers, tablefmt="github")) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--backward", action="store_true", help="also emit the forward+backward table" + ) + parser.add_argument( + "--configs", + default=";".join(",".join(map(str, c)) for c in DEFAULT_CONFIGS), + help="semicolon-separated 'batch,hq,seq' triples", + ) + parser.add_argument("--dtype", choices=("bf16", "fp16"), default="bf16") + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=20) + args = parser.parse_args() + + configs = [tuple(int(x) for x in part.split(",")) for part in args.configs.split(";")] + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + device = torch.device(device_ctx.device_type) + + native_op = NativeAttentionOp() + other_ops = [ + ("ascend", _maybe_ascend_op()), + ("cuda", _maybe_cuda_op()), + ] + + _bench_table(configs, native_op, other_ops, _forward_closure, device, dtype, args.warmup, args.iters) + if args.backward: + _bench_table( + configs, + native_op, + other_ops, + _forward_backward_closure, + device, + dtype, + args.warmup, + args.iters, + ) if __name__ == "__main__": - run_benchmark() + main() diff --git a/csrc/ascend/attention/deterministic_attention_ascend.asc b/csrc/ascend/attention/deterministic_attention_ascend.asc new file mode 100644 index 00000000..ebf6a036 --- /dev/null +++ b/csrc/ascend/attention/deterministic_attention_ascend.asc @@ -0,0 +1,505 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Deterministic (batch-invariant) standard-softmax attention, Ascend C (CANN) +// forward kernel. +// +// out = softmax(Q K^T * scale + masks) @ V, lse = rowmax + log(sum(exp(s - rowmax))) +// +// Mirrors the batch-invariant algorithm of the Triton reference +// (rl_engine/kernels/ops/triton/attention/standard_attn.py, i.e. +// rlkernel.attention.deterministic_core.v1) and the CUDA deterministic op +// (issue #147): +// - layout : q [B, Hq, Sq, D], k/v [B, Hkv, Skv, D] contiguous, D = 128 +// - masks : causal (upper triangle at offset Skv - Sq + 1) and optional +// key_padding_mask [B, Skv] bool, True = keep +// - numerics: all fp32 intermediate; bf16/fp16 inputs are upcast, the output +// row is cast back once at the end +// +// Batch-invariance / no split-K: every (b, q_head, row) is processed +// end-to-end by exactly one AI-core block, with a fixed 64-key tile size and a +// fixed two-pass (max, then sum-exp + P.V) reduction order over the key +// dimension. The instruction sequence for a row depends only on Skv, D and the +// masks -- never on the batch size, the block the row lands on, or how many +// blocks were launched (rows are strided across blocks). No second-pass merge +// of per-split (m, l, u) summaries exists, so the reduction tree is fixed. +// +// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by +// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu. + +#include +#include + +#include "kernel_operator.h" + +#include + +#include "torch_npu/csrc/core/npu/NPUStream.h" + +namespace { + +// Fixed head dimension (matches the CUDA deterministic op gate). +constexpr uint32_t HEAD_DIM = 128; +// Keys per tile. Fixed for all rows and batch sizes; this is what makes the +// reduction order batch-invariant (mirrors the Triton reference _BLOCK_N = 64). +constexpr uint32_t TILE_N = 64; +// Cap on launched blocks. Work items are strided across blocks, so launching +// fewer blocks than items is fine and never changes per-row numerics. +constexpr int64_t MAX_BLOCKS = 512; +constexpr float NEG_INF = -3.402823466e+38f; // -FLT_MAX + +template +class KernelDeterministicAttention { +public: + __aicore__ inline KernelDeterministicAttention(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(GM_ADDR q, + GM_ADDR k, + GM_ADDR v, + GM_ADDR mask, + GM_ADDR out, + GM_ADDR lse, + int64_t B, + int64_t Hq, + int64_t Hkv, + int64_t Sq, + int64_t Skv, + float scale, + int32_t causal, + int32_t hasMask) + { + B_ = B; + Hq_ = Hq; + Hkv_ = Hkv; + Sq_ = Sq; + Skv_ = Skv; + scale_ = scale; + causal_ = causal; + hasMask_ = hasMask; + qGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(q)); + kGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(k)); + vGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(v)); + maskGm_.SetGlobalBuffer(reinterpret_cast<__gm__ uint8_t*>(mask)); + outGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(out)); + lseGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(lse)); + + // UB budget stays well under 192 KB: + // k tile bf16 16 KB + k tile fp32 32 KB + v tile fp32 32 KB + // + q/acc/prod/work/scores/scalar/mask ~4 KB. + pipe_->InitBuffer(qBufF_, HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(accBufF_, HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(prodBufF_, HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(workBufF_, HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(kBufT_, TILE_N * HEAD_DIM * sizeof(T)); + pipe_->InitBuffer(kBufF_, TILE_N * HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(vBufF_, TILE_N * HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(scoresBuf_, TILE_N * sizeof(float)); + // 128 B: mask tile (64 B) + up to 31 B misalignment + 32 B rounding. + pipe_->InitBuffer(maskBuf_, 128); + // 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); + + // 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); + eventSMTE2_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE2); + eventVMTE2_ = pipe_->FetchEventID(AscendC::HardEvent::V_MTE2); + eventSMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::S_MTE3); + eventVMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::V_MTE3); + eventMTE3S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE3_S); + } + + __aicore__ inline void Process() + { + const int64_t items = B_ * Hq_ * Sq_; + for (int64_t item = AscendC::GetBlockIdx(); item < items; + item += AscendC::GetBlockNum()) { + const int64_t row = item % Sq_; + const int64_t qh = (item / Sq_) % Hq_; + const int64_t b = item / (Sq_ * Hq_); + ProcessRow(b, qh, row); + } + } + +private: + __aicore__ inline void LoadQRow(int64_t b, int64_t qh, int64_t row) + { + const int64_t offset = ((b * Hq_ + qh) * Sq_ + row) * HEAD_DIM; + AscendC::LocalTensor qT = kBufT_.Get(); + AscendC::DataCopyExtParams cp{1, HEAD_DIM * sizeof(T), 0, 0, 0}; + AscendC::DataCopyPadExtParams pp{false, 0, 0, 0}; + AscendC::DataCopyPad(qT, qGm_[offset], cp, pp); + AscendC::SetFlag(eventMTE2S_); + AscendC::WaitFlag(eventMTE2S_); + AscendC::Cast(qBufF_.Get(), qT, AscendC::RoundMode::CAST_NONE, HEAD_DIM); + } + + __aicore__ inline void LoadKTile(int64_t b, int64_t kvh, int64_t start, uint32_t count) + { + // kBufT_ staging may hold data the vector pipe is still reading (the + // q cast of the first tile or the v cast of the previous tile). + AscendC::SetFlag(eventVMTE2_); + AscendC::WaitFlag(eventVMTE2_); + const int64_t offset = ((b * Hkv_ + kvh) * Skv_ + start) * HEAD_DIM; + AscendC::LocalTensor kT = kBufT_.Get(); + AscendC::DataCopyExtParams cp{1, count * HEAD_DIM * sizeof(T), 0, 0, 0}; + AscendC::DataCopyPadExtParams pp{false, 0, 0, 0}; + AscendC::DataCopyPad(kT, kGm_[offset], cp, pp); + AscendC::SetFlag(eventMTE2S_); + AscendC::WaitFlag(eventMTE2S_); + AscendC::Cast(kBufF_.Get(), kT, AscendC::RoundMode::CAST_NONE, + count * HEAD_DIM); + } + + __aicore__ inline void LoadVTile(int64_t b, int64_t kvh, int64_t start, uint32_t count) + { + // vBufF_ is still being read by the P . V accumulation of the previous + // tile; kBufT_ staging may be in use by the vector pipe as well. + AscendC::SetFlag(eventVMTE2_); + AscendC::WaitFlag(eventVMTE2_); + const int64_t offset = ((b * Hkv_ + kvh) * Skv_ + start) * HEAD_DIM; + AscendC::LocalTensor vT = kBufT_.Get(); + AscendC::DataCopyExtParams cp{1, count * HEAD_DIM * sizeof(T), 0, 0, 0}; + AscendC::DataCopyPadExtParams pp{false, 0, 0, 0}; + AscendC::DataCopyPad(vT, vGm_[offset], cp, pp); + AscendC::SetFlag(eventMTE2S_); + AscendC::WaitFlag(eventMTE2S_); + AscendC::Cast(vBufF_.Get(), vT, AscendC::RoundMode::CAST_NONE, + count * HEAD_DIM); + } + + // Load the mask window covering keys [start, start + count) of batch b. + // The window is read from the last 32-byte-aligned GM address at or below + // the tile start; returns the byte offset of the tile inside the window. + __aicore__ inline uint32_t LoadMaskTile(int64_t b, int64_t start, uint32_t count) + { + // maskBuf_ holds values the scalar pipe read for the previous tile; + // wait for those reads before MTE2 overwrites the window. + AscendC::SetFlag(eventSMTE2_); + AscendC::WaitFlag(eventSMTE2_); + const int64_t base = b * Skv_ + start; + const int64_t aligned = base & ~31LL; + const uint32_t offset = static_cast(base - aligned); + const int64_t remaining = (b + 1) * Skv_ - aligned; + uint32_t alignedCount = (offset + count + 31) & ~31u; + if (alignedCount > remaining) { + alignedCount = static_cast(remaining); + } + AscendC::LocalTensor m = maskBuf_.Get(); + AscendC::DataCopyExtParams cp{1, alignedCount, 0, 0, 0}; + AscendC::DataCopyPadExtParams pp{false, 0, 0, 0}; + AscendC::DataCopyPad(m, maskGm_[aligned], cp, pp); + AscendC::SetFlag(eventMTE2S_); + AscendC::WaitFlag(eventMTE2S_); + return offset; + } + + // scores[j] = scale * (q . k[start + j]) for j in [0, count), with the + // causal and key-padding masks applied; masked lanes become NEG_INF. + __aicore__ inline void ComputeScores(int64_t b, + int64_t row, + int64_t start, + uint32_t count, + uint32_t maskOffset) + { + AscendC::LocalTensor qRow = qBufF_.Get(); + AscendC::LocalTensor kTile = kBufF_.Get(); + AscendC::LocalTensor scores = scoresBuf_.Get(); + AscendC::LocalTensor prod = prodBufF_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + AscendC::LocalTensor maskTile = maskBuf_.Get(); + + const int64_t causalKeep = row + Skv_ - Sq_; + for (uint32_t j = 0; j < count; ++j) { + AscendC::Mul(prod, qRow, kTile[j * HEAD_DIM], HEAD_DIM); + AscendC::ReduceSum(scalar, prod, workBufF_.Get(), HEAD_DIM); + WaitVector(); + float s = scalar.GetValue(0) * scale_; + const int64_t jGlobal = start + j; + if (causal_ && jGlobal > causalKeep) { + s = NEG_INF; + } + if (hasMask_ && maskTile.GetValue(maskOffset + j) == 0) { + s = NEG_INF; + } + scores.SetValue(j, s); + } + } + + __aicore__ inline void ProcessRow(int64_t b, int64_t qh, int64_t row) + { + const int64_t kvh = qh / (Hq_ / Hkv_); // GQA: query head h -> KV head h / g + const int64_t tileCount = (Skv_ + TILE_N - 1) / TILE_N; + AscendC::LocalTensor scores = scoresBuf_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + + LoadQRow(b, qh, row); + + // Pass 1: row max with a fixed tile order. + float rowMax = NEG_INF; + for (int64_t tile = 0; tile < tileCount; ++tile) { + const int64_t start = tile * TILE_N; + const uint32_t count = TileCount(start); + LoadKTile(b, kvh, start, count); + uint32_t maskOffset = 0; + if (hasMask_) { + maskOffset = LoadMaskTile(b, start, count); + } + ComputeScores(b, row, start, count, maskOffset); + for (uint32_t j = count; j < TILE_N; ++j) { + scores.SetValue(j, NEG_INF); + } + AscendC::SetFlag(eventSV_); + AscendC::WaitFlag(eventSV_); + AscendC::ReduceMax(scalar, scores, workBufF_.Get(), TILE_N, false); + WaitVector(); + const float tileMax = scalar.GetValue(0); + rowMax = tileMax > rowMax ? tileMax : rowMax; + } + + // Pass 2: sum(exp(s - rowMax)) and P . V with the same fixed tile + // order (see the batch-invariance note at the top of the file). + float sumExp = 0.0f; + AscendC::LocalTensor acc = accBufF_.Get(); + // Real zeroing: accBuf_ starts as uninitialized UB and 0 * inf == NaN. + AscendC::Duplicate(acc, 0.0f, HEAD_DIM); + for (int64_t tile = 0; tile < tileCount; ++tile) { + const int64_t start = tile * TILE_N; + const uint32_t count = TileCount(start); + LoadKTile(b, kvh, start, count); + uint32_t maskOffset = 0; + if (hasMask_) { + maskOffset = LoadMaskTile(b, start, count); + } + ComputeScores(b, row, start, count, maskOffset); + for (uint32_t j = count; j < TILE_N; ++j) { + scores.SetValue(j, NEG_INF); + } + AscendC::SetFlag(eventSV_); + AscendC::WaitFlag(eventSV_); + AscendC::Adds(scores, scores, -rowMax, TILE_N); + AscendC::Exp(scores, scores, TILE_N); + AscendC::ReduceSum(scalar, scores, workBufF_.Get(), TILE_N); + WaitVector(); // vector -> scalar read; also covers the Exp above + sumExp += scalar.GetValue(0); + + LoadVTile(b, kvh, start, count); + AscendC::LocalTensor vTile = vBufF_.Get(); + AscendC::LocalTensor prod = prodBufF_.Get(); + for (uint32_t j = 0; j < count; ++j) { + const float pj = scores.GetValue(j); + if (pj == 0.0f) { + continue; + } + AscendC::Muls(prod, vTile[j * HEAD_DIM], pj, HEAD_DIM); + AscendC::Add(acc, acc, prod, HEAD_DIM); + } + } + + // out = acc / sumExp. A fully-masked row keeps rowMax == NEG_INF, so + // sumExp is 0 and the output row is defined as 0 with lse = -inf + // (mirrors the Triton reference's denom > 0 guard). + const float invDenom = (sumExp > 0.0f) ? (1.0f / sumExp) : 0.0f; + AscendC::Muls(acc, acc, invDenom, HEAD_DIM); + WriteOutputs(b, qh, row, rowMax, sumExp, acc); + } + + __aicore__ inline void WriteOutputs(int64_t b, + int64_t qh, + int64_t row, + float rowMax, + float sumExp, + AscendC::LocalTensor acc) + { + AscendC::LocalTensor scalar = scalarBuf_.Get(); + float lse = NEG_INF; + if (rowMax > NEG_INF) { + scalar.SetValue(0, sumExp); + AscendC::SetFlag(eventSV_); // scalar write -> vector op + AscendC::WaitFlag(eventSV_); + AscendC::Log(scalar, scalar, 8); + WaitVector(); // vector -> scalar read + lse = rowMax + scalar.GetValue(0); + } + // Stage outputs 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. + scalar.SetValue(0, lse); + AscendC::LocalTensor outT = kBufT_.Get(); + AscendC::Cast(outT, acc, AscendC::RoundMode::CAST_ROUND, HEAD_DIM); + AscendC::SetFlag(eventVMTE3_); // vector write -> copy-out + AscendC::WaitFlag(eventVMTE3_); + AscendC::SetFlag(eventSMTE3_); // scalar write -> copy-out + AscendC::WaitFlag(eventSMTE3_); + AscendC::DataCopyExtParams p4{1, sizeof(float), 0, 0, 0}; + AscendC::DataCopyPad(lseGm_[(b * Hq_ + qh) * Sq_ + row], scalar[0], p4); + AscendC::DataCopyExtParams outCp{1, HEAD_DIM * sizeof(T), 0, 0, 0}; + AscendC::DataCopyPadExtParams outPp{false, 0, 0, 0}; + AscendC::DataCopyPad(outGm_[((b * Hq_ + qh) * Sq_ + row) * HEAD_DIM], outT, outCp, + outPp); + // Drain MTE3 before the next row stages new values into the shared + // buffers; the scalar pipe issues all later MTE2 copies in order, so + // this wait alone orders them after the copy-outs. + 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_); + } + + __aicore__ inline uint32_t TileCount(int64_t start) const + { + const int64_t remaining = Skv_ - start; + return static_cast(remaining < TILE_N ? remaining : TILE_N); + } + + AscendC::TPipe* pipe_; + AscendC::GlobalTensor qGm_; + AscendC::GlobalTensor kGm_; + AscendC::GlobalTensor vGm_; + AscendC::GlobalTensor maskGm_; + AscendC::GlobalTensor outGm_; + AscendC::GlobalTensor lseGm_; + AscendC::TBuf qBufF_; + AscendC::TBuf accBufF_; + AscendC::TBuf prodBufF_; + AscendC::TBuf workBufF_; + AscendC::TBuf kBufT_; + AscendC::TBuf kBufF_; + AscendC::TBuf vBufF_; + AscendC::TBuf scoresBuf_; + AscendC::TBuf maskBuf_; + AscendC::TBuf scalarBuf_; + AscendC::TEventID eventVS_; + AscendC::TEventID eventSV_; + AscendC::TEventID eventMTE2S_; + AscendC::TEventID eventSMTE2_; + AscendC::TEventID eventVMTE2_; + AscendC::TEventID eventSMTE3_; + AscendC::TEventID eventVMTE3_; + AscendC::TEventID eventMTE3S_; + int64_t B_; + int64_t Hq_; + int64_t Hkv_; + int64_t Sq_; + int64_t Skv_; + float scale_; + int32_t causal_; + int32_t hasMask_; +}; + +} // namespace + +extern "C" __global__ __vector__ void deterministic_attention_ascend_kernel_bf16( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR mask, GM_ADDR out, GM_ADDR lse, + int64_t B, int64_t Hq, int64_t Hkv, int64_t Sq, int64_t Skv, + float scale, int32_t causal, int32_t hasMask) +{ + AscendC::TPipe pipe; + KernelDeterministicAttention op(&pipe); + op.Init(q, k, v, mask, out, lse, B, Hq, Hkv, Sq, Skv, scale, causal, hasMask); + op.Process(); +} + +extern "C" __global__ __vector__ void deterministic_attention_ascend_kernel_fp16( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR mask, GM_ADDR out, GM_ADDR lse, + int64_t B, int64_t Hq, int64_t Hkv, int64_t Sq, int64_t Skv, + float scale, int32_t causal, int32_t hasMask) +{ + AscendC::TPipe pipe; + KernelDeterministicAttention op(&pipe); + op.Init(q, k, v, mask, out, lse, B, Hq, Hkv, Sq, Skv, scale, causal, hasMask); + op.Process(); +} + +std::vector deterministic_attention_ascend_forward( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + bool causal, + double scale, + c10::optional key_padding_mask) +{ + TORCH_CHECK(q.is_privateuseone() && k.is_privateuseone() && v.is_privateuseone(), + "q, k, v must be on an NPU device"); + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, + "q, k, v must be 4-D [B, H, S, D]"); + TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(), + "q, k, v must be contiguous"); + TORCH_CHECK(q.scalar_type() == at::kBFloat16 || q.scalar_type() == at::kHalf, + "q must be bf16 or fp16"); + TORCH_CHECK(k.scalar_type() == q.scalar_type() && v.scalar_type() == q.scalar_type(), + "q, k, v must share the same dtype"); + TORCH_CHECK(q.size(3) == HEAD_DIM && k.size(3) == HEAD_DIM && v.size(3) == HEAD_DIM, + "head dim D must be 128"); + TORCH_CHECK(q.size(0) == k.size(0) && q.size(0) == v.size(0), + "batch size mismatch between q/k/v"); + TORCH_CHECK(k.size(1) == v.size(1) && k.size(2) == v.size(2), + "k/v must have the same head and key-length layout"); + TORCH_CHECK(q.size(1) % k.size(1) == 0, "Hq not divisible by Hkv (GQA group)"); + + const int64_t B = q.size(0); + const int64_t Hq = q.size(1); + const int64_t Hkv = k.size(1); + const int64_t Sq = q.size(2); + const int64_t Skv = k.size(2); + + torch::Tensor mask; + bool hasMask = key_padding_mask.has_value() && key_padding_mask->defined(); + if (hasMask) { + mask = key_padding_mask->to(torch::kBool).contiguous(); + TORCH_CHECK(mask.is_privateuseone(), "key_padding_mask must be on an NPU device"); + TORCH_CHECK(mask.dim() == 2 && mask.size(0) == B && mask.size(1) == Skv, + "key_padding_mask must be [B, Skv]"); + } + + torch::Tensor out = at::empty({B, Hq, Sq, HEAD_DIM}, q.options()); + torch::Tensor lse = at::empty({B, Hq, Sq}, q.options().dtype(at::kFloat)); + + // 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(B * Hq * Sq, MAX_BLOCKS)); + uint8_t* maskPtr = hasMask ? reinterpret_cast(mask.mutable_data_ptr()) : nullptr; + + if (q.scalar_type() == at::kBFloat16) { + deterministic_attention_ascend_kernel_bf16<<>>( + reinterpret_cast(q.mutable_data_ptr()), + reinterpret_cast(k.mutable_data_ptr()), + reinterpret_cast(v.mutable_data_ptr()), + reinterpret_cast(maskPtr), + reinterpret_cast(out.mutable_data_ptr()), + reinterpret_cast(lse.mutable_data_ptr()), + B, Hq, Hkv, Sq, Skv, static_cast(scale), causal ? 1 : 0, + hasMask ? 1 : 0); + } else { + deterministic_attention_ascend_kernel_fp16<<>>( + reinterpret_cast(q.mutable_data_ptr()), + reinterpret_cast(k.mutable_data_ptr()), + reinterpret_cast(v.mutable_data_ptr()), + reinterpret_cast(maskPtr), + reinterpret_cast(out.mutable_data_ptr()), + reinterpret_cast(lse.mutable_data_ptr()), + B, Hq, Hkv, Sq, Skv, static_cast(scale), causal ? 1 : 0, + hasMask ? 1 : 0); + } + return {out, lse}; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("deterministic_attention_ascend", + &deterministic_attention_ascend_forward, + "Deterministic batch-invariant standard-softmax attention (Ascend C forward)"); +} diff --git a/docs/operators/attention.md b/docs/operators/attention.md index ebff9a58..4edc44b4 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -49,6 +49,7 @@ The op exposes the WS1 dual-path contract: | --- | --- | --- | --- | | PyTorch fallback | `NativeAttentionOp` | None | fp32 ground-truth reference; CPU and any GPU. | | CUDA deterministic | `DeterministicAttentionOp` | `_C.deterministic_attention_forward/backward` | Batch-invariant CUDA implementation (issue #147). | +| Ascend NPU deterministic | `DeterministicAttentionAscendOp` | `_C_npu.deterministic_attention_ascend` | Batch-invariant Ascend C implementation (issue #147). | ## Tensor Contract @@ -81,6 +82,20 @@ the inputs' device. 1. `CUDA_DETERMINISTIC_ATTENTION` — `DeterministicAttentionOp` (batch-invariant, fixed-order). 2. `PYTORCH_NATIVE_ATTENTION` — `NativeAttentionOp` (fallback). +On `npu` the priority is: + +1. `ASCEND_DETERMINISTIC_ATTENTION` — `DeterministicAttentionAscendOp` (batch-invariant, + fixed-order Ascend C forward; bf16/fp16 inputs, head dim 128). +2. `PYTORCH_NATIVE_ATTENTION` — `NativeAttentionOp` (fallback). + +The Ascend kernel implements the same algorithm as the Triton reference and the CUDA op: +every `(b, q_head, row)` is processed end-to-end by exactly one AI-core block, streaming the +keys in a fixed 64-key tile order with a two-pass (max, then sum-exp + P·V) reduction. There +is **no split-K** and no second-pass merge of per-split `(m, l, u)` summaries, so the +reduction tree for a row depends only on `Skv`, `D` and the masks — never on batch size or +block assignment. The backward recomputes the native reference forward under autograd +(Triton is unavailable on NPU), matching the Triton op's portable backward. + Calling it (`__call__` -> `forward(...)`) computes in the input dtype; `forward_fp32(...)` is the explicit fp32 golden path (NativeAttentionOp only). The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index bff5e2e7..6ced124a 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,3 +8,12 @@ def batch_invariant_logp_ascend( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... + +def deterministic_attention_ascend( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + causal: bool, + scale: float, + key_padding_mask: torch.Tensor | None, +) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index bde87edb..49d39830 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -73,6 +73,10 @@ def _load_object(path: str) -> Any: "rl_engine.kernels.ops.cuda.attention.deterministic_attn." "DeterministicAttentionOp" ), + "ascend": ( + "rl_engine.kernels.ops.ascend.attention.deterministic_attn." + "DeterministicAttentionAscendOp" + ), }, grad_input_names=("q", "k", "v"), ), diff --git a/rl_engine/kernels/ops/ascend/attention/__init__.py b/rl_engine/kernels/ops/ascend/attention/__init__.py new file mode 100644 index 00000000..07b76408 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/attention/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from rl_engine.kernels.ops.ascend.attention.deterministic_attn import ( + DeterministicAttentionAscendOp, +) + +__all__ = ["DeterministicAttentionAscendOp"] diff --git a/rl_engine/kernels/ops/ascend/attention/deterministic_attn.py b/rl_engine/kernels/ops/ascend/attention/deterministic_attn.py new file mode 100644 index 00000000..28895c5b --- /dev/null +++ b/rl_engine/kernels/ops/ascend/attention/deterministic_attn.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Ascend NPU deterministic standard-softmax attention (issue #147). + +Forward: QK -> masked softmax+LSE -> PV (all FP32 intermediate) on an +Ascend C kernel (`_C_npu.deterministic_attention_ascend`). Every +(b, q_head, row) is reduced end-to-end by one AI-core block with a fixed +64-key tile order and no split-K merge, so per-row numerics are +batch-invariant (the same algorithm as the Triton reference and the CUDA +deterministic op). + +Backward: Triton is unavailable on NPU, so the backward mirrors the +Triton op's portable branch -- it recomputes the reference forward with +`NativeAttentionOp` under autograd and VJPs the upstream gradient through +it, reusing the forward-saved q/k/v/mask. +""" + +from __future__ import annotations + +import math +from typing import Any, Optional + +import torch +from torch.autograd import Function +from torch.autograd.function import once_differentiable + +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +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 + +_HEAD_DIM = 128 + + +class _DeterministicAttentionAscendFn(Function): + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + causal: bool, + scale: float, + key_padding_mask: Optional[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + q_c = q.contiguous() + k_c = k.contiguous() + v_c = v.contiguous() + mask_c = key_padding_mask.contiguous() if key_padding_mask is not None else None + + out, lse = _C_npu.deterministic_attention_ascend(q_c, k_c, v_c, causal, float(scale), mask_c) + + ctx.save_for_backward(q_c, k_c, v_c, mask_c) + ctx.causal = causal + ctx.scale = scale + ctx.has_mask = mask_c is not None + ctx.mark_non_differentiable(lse) + return out, lse + + @staticmethod + @once_differentiable + def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): + del grad_lse # lse is non-differentiable; always None upstream + q, k, v, mask = ctx.saved_tensors + with torch.enable_grad(): + q_ref = q.detach().requires_grad_(True) + k_ref = k.detach().requires_grad_(True) + v_ref = v.detach().requires_grad_(True) + out = NativeAttentionOp().forward( + q_ref, + k_ref, + v_ref, + causal=ctx.causal, + scale=ctx.scale, + key_padding_mask=mask if ctx.has_mask else None, + ) + dq, dk, dv = torch.autograd.grad(out, (q_ref, k_ref, v_ref), grad_out) + return dq, dk, dv, None, None, None + + +class DeterministicAttentionAscendOp: + """Batch-invariant standard softmax attention on Ascend NPU. + + Public surface matches ``NativeAttentionOp`` / ``DeterministicAttentionOp`` + so the #108 harness can call ``forward(**inputs)`` with ``key_padding_mask``. + Inputs the Ascend C kernel cannot take fall back to the native reference. + """ + + def __init__(self) -> None: + if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "deterministic_attention_ascend"): + raise RuntimeError( + "deterministic_attention_ascend is not compiled into the extension. " + "Rebuild with KERNEL_ALIGN_FORCE_ASCEND=1 on an Ascend NPU host: " + "'pip install -e .'" + ) + logger.info( + "Successfully linked to precompiled _C_npu.deterministic_attention_ascend kernel." + ) + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.forward(q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Harness / registry main path: return out only. Differentiable.""" + out, _lse = self.forward_with_lse( + q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask + ) + return out + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return (out, lse) with FP32 LSE for debug / handoff hooks.""" + self._validate_inputs(q, k, v, key_padding_mask) + resolved_scale = scale if scale is not None else (1.0 / math.sqrt(q.shape[-1])) + out, lse = _DeterministicAttentionAscendFn.apply( + q, k, v, causal, resolved_scale, key_padding_mask + ) + return out, lse + + @staticmethod + def _validate_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + key_padding_mask: Optional[torch.Tensor], + ) -> None: + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + raise ValueError( + f"q/k/v must be 4-D [B, H, S, D], got q={tuple(q.shape)}, " + f"k={tuple(k.shape)}, v={tuple(v.shape)}" + ) + b, hq, sq, d = q.shape + hkv, skv = k.shape[1], k.shape[2] + if k.shape[0] != b or v.shape[0] != b: + raise ValueError("batch size mismatch between q/k/v") + if v.shape[1] != hkv or v.shape[2] != skv or k.shape[3] != d or v.shape[3] != d: + raise ValueError( + f"k/v shape mismatch: k={tuple(k.shape)}, v={tuple(v.shape)}, " + f"expected k/v [B={b}, Hkv, Skv, D={d}]" + ) + if d != _HEAD_DIM: + raise ValueError(f"head dim D must be {_HEAD_DIM}, got {d}") + if hq % hkv != 0: + raise ValueError(f"Hq={hq} not divisible by Hkv={hkv} (GQA group)") + if q.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"only FP16/BF16 supported, got {q.dtype}") + if k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("q, k, v must share the same dtype") + if not (q.device.type == "npu" and k.device.type == "npu" and v.device.type == "npu"): + raise ValueError("q, k, v must be NPU tensors") + if key_padding_mask is not None: + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + if key_padding_mask.shape != (b, skv): + raise ValueError( + f"key_padding_mask must be [B, Skv]=[{b}, {skv}], " + f"got {tuple(key_padding_mask.shape)}" + ) + if sq < 1 or skv < 1: + raise ValueError(f"Sq and Skv must be positive, got Sq={sq}, Skv={skv}") diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 4020cd01..e611b1be 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -77,6 +77,11 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ASCEND_BATCH_INVARIANT_LOGP = ( "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAscendOp" ) + # Ascend NPU deterministic (batch-invariant, no split-K) standard-softmax + # attention (issue #147); Ascend C forward + reference backward. + ASCEND_DETERMINISTIC_ATTENTION = ( + "rl_engine.kernels.ops.ascend.attention.deterministic_attn.DeterministicAttentionAscendOp" + ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" @@ -297,6 +302,15 @@ def __init__(self): OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ], + "attention": [ + OpBackend.ASCEND_DETERMINISTIC_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], + # Production SDPA-layout paths: no Ascend-accelerated candidate + # yet, so pin the explicit native fallbacks (the registry-wide + # default would resolve to a logp op, which is the wrong kind). + "attn": [OpBackend.PYTORCH_ATTN], + "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], }, } logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") diff --git a/setup.py b/setup.py index 57c98070..1205b7f5 100644 --- a/setup.py +++ b/setup.py @@ -220,7 +220,7 @@ def _ascend_extensions(): "KERNEL_ALIGN_FORCE_ASCEND=1 requires torch and torch_npu to be installed" ) from e - asc_srcs = sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) + 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")] diff --git a/tests/test_attention_ascend.py b/tests/test_attention_ascend.py new file mode 100644 index 00000000..7facd37d --- /dev/null +++ b/tests/test_attention_ascend.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for the Ascend NPU deterministic standard-softmax attention. + +Validates the same two orthogonal properties as the CUDA deterministic op: +1. **Correctness** - output matches the ``NativeAttentionOp.forward_fp32`` + ground truth within the reduction tolerances. +2. **Batch-invariance** - a query row's output is bitwise identical regardless + of batch size, batch position, or how many AI-core blocks were launched + (each row is reduced end-to-end by one block; no split-K merge exists). +""" + +import math + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + +_D = 128 + +# Accuracy tolerance from the gtest contract, "attention" op class. +_ATOL = {torch.bfloat16: 5.0e-2, torch.float16: 1.0e-3} +_RTOL = {torch.bfloat16: 2.0e-2, torch.float16: 1.0e-3} + + +def _npu_available() -> bool: + try: + import torch_npu # noqa: F401 + except ImportError: + return False + return hasattr(torch, "npu") and torch.npu.is_available() + + +def _ascend_kernel_available() -> bool: + if not _npu_available(): + return False + try: + from rl_engine.kernels.ops.ascend.attention.deterministic_attn import ( + _C_npu, + _NPU_EXT_AVAILABLE, + ) + except Exception: + return False + return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "deterministic_attention_ascend") + + +requires_ascend = pytest.mark.skipif( + not _ascend_kernel_available(), + reason="deterministic_attention_ascend kernel not compiled " + "(needs KERNEL_ALIGN_FORCE_ASCEND=1 on an Ascend NPU host).", +) + + +def _get_op(): + from rl_engine.kernels.ops.ascend.attention.deterministic_attn import ( + DeterministicAttentionAscendOp, + ) + + return DeterministicAttentionAscendOp() + + +def _gold(q, k, v, causal=True, scale=None, key_padding_mask=None): + """fp32 ground truth: NativeAttentionOp.forward_fp32.""" + return NativeAttentionOp().forward_fp32( + q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask + ) + + +def _make_qkv(batch, hq, hkv, sq, skv, dtype, seed=0): + generator = torch.Generator(device="cpu").manual_seed(seed) + q = torch.randn(batch, hq, sq, _D, dtype=dtype, generator=generator).to("npu") + k = torch.randn(batch, hkv, skv, _D, dtype=dtype, generator=generator).to("npu") + v = torch.randn(batch, hkv, skv, _D, dtype=dtype, generator=generator).to("npu") + return q, k, v + + +# --------------------------------------------------------------------------- +# Correctness +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@requires_ascend +class TestAscendAttentionCorrectness: + def test_prefill_causal(self, dtype): + op = _get_op() + q, k, v = _make_qkv(2, 8, 2, 128, 128, dtype) + out = op(q, k, v, causal=True) + gold = _gold(q, k, v, causal=True) + assert out.dtype == dtype + assert torch.allclose(out.float(), gold, atol=_ATOL[dtype], rtol=_RTOL[dtype]) + + def test_gqa(self, dtype): + op = _get_op() + q, k, v = _make_qkv(2, 8, 2, 64, 64, dtype) # g = 8/2 = 4 + out = op(q, k, v, causal=True) + gold = _gold(q, k, v, causal=True) + assert torch.allclose(out.float(), gold, atol=_ATOL[dtype], rtol=_RTOL[dtype]) + + def test_decode_window(self, dtype): + op = _get_op() + q, k, v = _make_qkv(2, 8, 2, 4, 96, dtype) # Sq < Skv + out = op(q, k, v, causal=True) + gold = _gold(q, k, v, causal=True) + assert torch.allclose(out.float(), gold, atol=_ATOL[dtype], rtol=_RTOL[dtype]) + + def test_non_causal(self, dtype): + op = _get_op() + q, k, v = _make_qkv(2, 8, 2, 64, 64, dtype) + out = op(q, k, v, causal=False) + gold = _gold(q, k, v, causal=False) + assert torch.allclose(out.float(), gold, atol=_ATOL[dtype], rtol=_RTOL[dtype]) + + def test_key_padding_mask(self, dtype): + op = _get_op() + q, k, v = _make_qkv(2, 8, 2, 64, 100, dtype) # Skv not a multiple of 64 + mask = torch.ones(2, 100, dtype=torch.bool, device="npu") + mask[:, 80:] = False + out = op(q, k, v, causal=True, key_padding_mask=mask) + gold = _gold(q, k, v, causal=True, key_padding_mask=mask) + assert torch.allclose(out.float(), gold, atol=_ATOL[dtype], rtol=_RTOL[dtype]) + + def test_fully_masked_row_is_zero(self, dtype): + op = _get_op() + q, k, v = _make_qkv(2, 4, 4, 2, 32, dtype) + mask = torch.ones(2, 32, dtype=torch.bool, device="npu") + mask[1, :] = False # batch 1 has no valid key at all + out, lse = op.forward_with_lse(q, k, v, causal=True, key_padding_mask=mask) + # Batch 1 has zero valid keys -> defined as 0, lse = -inf. + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert torch.all(lse[1] == float("-inf")) + # Batch 0 still finite. + assert torch.isfinite(out[0]).all() + assert torch.isfinite(lse[0]).all() + + def test_explicit_scale(self, dtype): + op = _get_op() + q, k, v = _make_qkv(1, 4, 4, 32, 32, dtype) + out = op(q, k, v, causal=False, scale=0.05) + gold = _gold(q, k, v, causal=False, scale=0.05) + assert torch.allclose(out.float(), gold, atol=_ATOL[dtype], rtol=_RTOL[dtype]) + + def test_forward_with_lse(self, dtype): + op = _get_op() + q, k, v = _make_qkv(1, 4, 4, 64, 64, dtype) + out, lse = op.forward_with_lse(q, k, v, causal=True) + scale = 1.0 / math.sqrt(_D) + qf, kf = q.float(), k.float() + scores = qf @ kf.transpose(-1, -2) * scale + cm = torch.triu(torch.ones(64, 64, dtype=torch.bool, device="npu"), 1) + scores = scores.masked_fill(cm, float("-inf")) + ref_lse = torch.logsumexp(scores, dim=-1) + assert lse.dtype == torch.float32 + assert lse.shape == (1, 4, 64) + assert torch.allclose(lse, ref_lse, atol=1e-3, rtol=1e-3) + + def test_backward_grads(self, dtype): + op = _get_op() + q, k, v = _make_qkv(1, 4, 4, 32, 48, dtype) + q.requires_grad_(True) + k.requires_grad_(True) + v.requires_grad_(True) + out = op(q, k, v, causal=True) + grad_out = torch.randn_like(out) + out.backward(grad_out) + assert all(g is not None for g in (q.grad, k.grad, v.grad)) + assert all(torch.isfinite(g).all() for g in (q.grad, k.grad, v.grad)) + + # The backward is the VJP of the native reference forward; compare. + with torch.enable_grad(): + q_ref = q.detach().requires_grad_(True) + k_ref = k.detach().requires_grad_(True) + v_ref = v.detach().requires_grad_(True) + ref_out = NativeAttentionOp().forward(q_ref, k_ref, v_ref, causal=True) + dq_ref, dk_ref, dv_ref = torch.autograd.grad( + ref_out, (q_ref, k_ref, v_ref), grad_out + ) + # The backward recomputes the same reference forward, so the VJPs + # match to numerical noise. + assert torch.allclose(q.grad.float(), dq_ref.float(), atol=1e-6, rtol=1e-5) + assert torch.allclose(k.grad.float(), dk_ref.float(), atol=1e-6, rtol=1e-5) + assert torch.allclose(v.grad.float(), dv_ref.float(), atol=1e-6, rtol=1e-5) + + def test_rejects_fp32(self): + op = _get_op() + q, k, v = _make_qkv(1, 4, 4, 32, 32, torch.float32) + with pytest.raises(ValueError, match="only FP16/BF16"): + op(q, k, v) + + def test_rejects_bad_head_dim(self): + op = _get_op() + q = torch.randn(1, 4, 32, 64, device="npu", dtype=torch.float16) + k = torch.randn(1, 4, 32, 64, device="npu", dtype=torch.float16) + v = torch.randn(1, 4, 32, 64, device="npu", dtype=torch.float16) + with pytest.raises(ValueError, match="head dim D must be 128"): + op(q, k, v) + + +# --------------------------------------------------------------------------- +# Batch invariance +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendAttentionBatchInvariance: + def _run_row(self, batch, hq, hkv, sq, skv, dtype, pos, seed=7): + """One fixed query row embedded at position `pos` of a random batch.""" + op = _get_op() + q, k, v = _make_qkv(batch, hq, hkv, sq, skv, dtype, seed=seed) + out = op(q, k, v, causal=True) + return out[0, 0, pos, :].clone() + + def test_batch_size_1_vs_n(self): + dtype = torch.float16 + alone = self._run_row(1, 8, 2, 64, 64, dtype, pos=0, seed=7) + for batch in (2, 4, 8): + in_batch = self._run_row(batch, 8, 2, 64, 64, dtype, pos=0, seed=7) + assert torch.equal(alone, in_batch), f"drift at batch_size={batch}" + + def test_different_positions_in_batch(self): + dtype = torch.bfloat16 + baseline = None + for pos in range(16): + row = self._run_row(2, 8, 2, 64, 64, dtype, pos=pos, seed=11) + if baseline is None: + baseline = row + else: + assert torch.equal(baseline, row), f"drift at position={pos}" + + def test_block_striding(self): + # 2 * 8 * 128 = 2048 work items > MAX_BLOCKS (512): rows are strided + # across blocks, so numerics must not depend on block assignment. + # The small run (1 * 4 * 32 = 128 items, one block per item) and the + # strided run must give the bitwise-identical row for the same content. + dtype = torch.float16 + op = _get_op() + small_q, small_k, small_v = _make_qkv(1, 4, 2, 32, 32, dtype, seed=3) + small = op(small_q, small_k, small_v, causal=True) + big_q, big_k, big_v = _make_qkv(2, 8, 2, 128, 128, dtype, seed=3) + big_q[:, 0, 0, :] = small_q[0, 0, 0, :] + big_k[:, 0, :32, :] = small_k[0, 0, :, :] + big_v[:, 0, :32, :] = small_v[0, 0, :, :] + big = op(big_q, big_k, big_v, causal=True) + # Row (0, head 0, pos 0): causal window is j <= 0 in both runs, so the + # other 96 keys cannot influence the result. + assert torch.equal(big[0, 0, 0, :], small[0, 0, 0, :]) + + def test_repeated_runs_deterministic(self): + dtype = torch.bfloat16 + q, k, v = _make_qkv(2, 8, 2, 128, 128, dtype, seed=5) + op = _get_op() + first = op(q, k, v, causal=True) + for _ in range(3): + again = op(q, k, v, causal=True) + assert torch.equal(first, again) + + +# --------------------------------------------------------------------------- +# Registry dispatch +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendRegistryDispatch: + def test_get_op_attention(self): + from rl_engine.kernels.registry import kernel_registry + + op = kernel_registry.get_op("attention", device="npu") + assert type(op).__name__ == "DeterministicAttentionAscendOp" + + def test_get_op_attn_falls_back_to_sdpa(self): + from rl_engine.kernels.registry import kernel_registry + + op = kernel_registry.get_op("attn", device="npu") + assert type(op).__name__ == "NativeAttentionOp" From a67f613cf0f4738f8a386046f6807ce929394bbe Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Wed, 19 Aug 2026 20:56:20 +0800 Subject: [PATCH 4/5] fix(ws1): fully-masked rows, lse -inf, batch the per-tile score syncs - fully-masked rows: track anyValid in pass 1 and short-circuit to out = 0 / lse = -inf instead of letting exp(NEG_INF - NEG_INF) = exp(0) leak weight into P.V (mirrors the Triton reference guards) - lse of invalid rows is true -inf (matches the reference max_score/log(0) path), not -FLT_MAX - ComputeScores: ReduceSum lands directly in the score lanes and the whole tile syncs once (S_V before, V_S after) instead of per-key WaitVector round trips - tests: independent generators per q/k/v (batch size must not shift the k/v stream), non-causal same-content position sweep, dispatch test now accepts the Ascend deterministic op --- benchmarks/benchmark_attention.py | 8 ++- .../deterministic_attention_ascend.asc | 63 ++++++++++++------- csrc/ascend/batch_invariant_logp_ascend.asc | 7 --- csrc/ascend/ops_npu.asc | 34 ++++++++++ .../ascend/attention/deterministic_attn.py | 20 +++--- tests/test_attention.py | 14 +++-- tests/test_attention_ascend.py | 44 ++++++++----- 7 files changed, 132 insertions(+), 58 deletions(-) create mode 100644 csrc/ascend/ops_npu.asc diff --git a/benchmarks/benchmark_attention.py b/benchmarks/benchmark_attention.py index 0f4c7adc..0aebec22 100644 --- a/benchmarks/benchmark_attention.py +++ b/benchmarks/benchmark_attention.py @@ -135,12 +135,12 @@ def _bench_table(configs, native_op, other_ops, closure_factory, device, dtype, row += ["-"] continue o_ms = _time_ms(closure_factory(op, q, k, v), warmup, iters) - row += [f"{o_ms:.3f}", f"{n_ms/o_ms:.2f}x"] + row += [f"{o_ms:.3f}", f"{n_ms / o_ms:.2f}x"] rows.append(row) headers = ["shape (B x Hq x S)", f"native {label} ms"] for name, _ in other_ops: - headers += [f"{name} {label} ms", f"vs native"] + headers += [f"{name} {label} ms", "vs native"] logger.info("\n" + tabulate(rows, headers=headers, tablefmt="github")) @@ -169,7 +169,9 @@ def main() -> None: ("cuda", _maybe_cuda_op()), ] - _bench_table(configs, native_op, other_ops, _forward_closure, device, dtype, args.warmup, args.iters) + _bench_table( + configs, native_op, other_ops, _forward_closure, device, dtype, args.warmup, args.iters + ) if args.backward: _bench_table( configs, diff --git a/csrc/ascend/attention/deterministic_attention_ascend.asc b/csrc/ascend/attention/deterministic_attention_ascend.asc index ebf6a036..7d26b89b 100644 --- a/csrc/ascend/attention/deterministic_attention_ascend.asc +++ b/csrc/ascend/attention/deterministic_attention_ascend.asc @@ -28,6 +28,7 @@ // KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu. #include +#include #include #include "kernel_operator.h" @@ -46,7 +47,10 @@ constexpr uint32_t TILE_N = 64; // Cap on launched blocks. Work items are strided across blocks, so launching // fewer blocks than items is fine and never changes per-row numerics. constexpr int64_t MAX_BLOCKS = 512; -constexpr float NEG_INF = -3.402823466e+38f; // -FLT_MAX +constexpr float NEG_INF = -3.402823466e+38f; // -FLT_MAX, mask sentinel +// lse of a fully-masked row: true -inf (matches the Triton reference, whose +// max_score == -inf / log(denom=0) == -inf path writes -inf, not -FLT_MAX). +constexpr float LSE_INVALID = -std::numeric_limits::infinity(); template class KernelDeterministicAttention { @@ -131,7 +135,7 @@ private: { const int64_t offset = ((b * Hq_ + qh) * Sq_ + row) * HEAD_DIM; AscendC::LocalTensor qT = kBufT_.Get(); - AscendC::DataCopyExtParams cp{1, HEAD_DIM * sizeof(T), 0, 0, 0}; + AscendC::DataCopyExtParams cp{1, static_cast(HEAD_DIM * sizeof(T)), 0, 0, 0}; AscendC::DataCopyPadExtParams pp{false, 0, 0, 0}; AscendC::DataCopyPad(qT, qGm_[offset], cp, pp); AscendC::SetFlag(eventMTE2S_); @@ -147,7 +151,7 @@ private: AscendC::WaitFlag(eventVMTE2_); const int64_t offset = ((b * Hkv_ + kvh) * Skv_ + start) * HEAD_DIM; AscendC::LocalTensor kT = kBufT_.Get(); - AscendC::DataCopyExtParams cp{1, count * HEAD_DIM * sizeof(T), 0, 0, 0}; + AscendC::DataCopyExtParams cp{1, static_cast(count * HEAD_DIM * sizeof(T)), 0, 0, 0}; AscendC::DataCopyPadExtParams pp{false, 0, 0, 0}; AscendC::DataCopyPad(kT, kGm_[offset], cp, pp); AscendC::SetFlag(eventMTE2S_); @@ -164,7 +168,7 @@ private: AscendC::WaitFlag(eventVMTE2_); const int64_t offset = ((b * Hkv_ + kvh) * Skv_ + start) * HEAD_DIM; AscendC::LocalTensor vT = kBufT_.Get(); - AscendC::DataCopyExtParams cp{1, count * HEAD_DIM * sizeof(T), 0, 0, 0}; + AscendC::DataCopyExtParams cp{1, static_cast(count * HEAD_DIM * sizeof(T)), 0, 0, 0}; AscendC::DataCopyPadExtParams pp{false, 0, 0, 0}; AscendC::DataCopyPad(vT, vGm_[offset], cp, pp); AscendC::SetFlag(eventMTE2S_); @@ -201,6 +205,11 @@ private: // scores[j] = scale * (q . k[start + j]) for j in [0, count), with the // causal and key-padding masks applied; masked lanes become NEG_INF. + // + // Each dot product lands directly in its score lane (ReduceSum's dst + // scalar aliases scores[j]), so the vector pipe runs the whole tile + // without per-key scalar synchronization; a single V_S wait covers all + // lanes before the scalar mask pass. __aicore__ inline void ComputeScores(int64_t b, int64_t row, int64_t start, @@ -211,15 +220,24 @@ private: AscendC::LocalTensor kTile = kBufF_.Get(); AscendC::LocalTensor scores = scoresBuf_.Get(); AscendC::LocalTensor prod = prodBufF_.Get(); - AscendC::LocalTensor scalar = scalarBuf_.Get(); AscendC::LocalTensor maskTile = maskBuf_.Get(); - const int64_t causalKeep = row + Skv_ - Sq_; + // The previous tile's mask pass (and the caller's padding loop) write + // the score lanes on the scalar pipe; drain them before the vector + // ReduceSum targets the same lanes. + AscendC::SetFlag(eventSV_); + AscendC::WaitFlag(eventSV_); for (uint32_t j = 0; j < count; ++j) { AscendC::Mul(prod, qRow, kTile[j * HEAD_DIM], HEAD_DIM); - AscendC::ReduceSum(scalar, prod, workBufF_.Get(), HEAD_DIM); - WaitVector(); - float s = scalar.GetValue(0) * scale_; + AscendC::ReduceSum(scores[j], prod, workBufF_.Get(), HEAD_DIM); + } + WaitVector(); // all dot products visible to the scalar pipe + AscendC::Muls(scores, scores, scale_, count); + WaitVector(); // scaled lanes visible to the scalar pipe + + const int64_t causalKeep = row + Skv_ - Sq_; + for (uint32_t j = 0; j < count; ++j) { + float s = scores.GetValue(j); const int64_t jGlobal = start + j; if (causal_ && jGlobal > causalKeep) { s = NEG_INF; @@ -242,6 +260,7 @@ private: // Pass 1: row max with a fixed tile order. float rowMax = NEG_INF; + bool anyValid = false; for (int64_t tile = 0; tile < tileCount; ++tile) { const int64_t start = tile * TILE_N; const uint32_t count = TileCount(start); @@ -259,6 +278,9 @@ private: AscendC::ReduceMax(scalar, scores, workBufF_.Get(), TILE_N, false); WaitVector(); const float tileMax = scalar.GetValue(0); + if (tileMax > NEG_INF) { + anyValid = true; + } rowMax = tileMax > rowMax ? tileMax : rowMax; } @@ -268,6 +290,14 @@ private: AscendC::LocalTensor acc = accBufF_.Get(); // Real zeroing: accBuf_ starts as uninitialized UB and 0 * inf == NaN. AscendC::Duplicate(acc, 0.0f, HEAD_DIM); + if (!anyValid) { + // Fully-masked row: exp(s - rowMax) would be exp(0) = 1 for the + // masked lanes (NEG_INF - NEG_INF == 0), not exp(-inf) = 0, so the + // row is defined as out = 0, lse = -inf -- mirroring the Triton + // reference's max_score == -inf / denom > 0 guards. + WriteOutputs(b, qh, row, NEG_INF, 0.0f, acc); + return; + } for (int64_t tile = 0; tile < tileCount; ++tile) { const int64_t start = tile * TILE_N; const uint32_t count = TileCount(start); @@ -317,7 +347,7 @@ private: AscendC::LocalTensor acc) { AscendC::LocalTensor scalar = scalarBuf_.Get(); - float lse = NEG_INF; + float lse = LSE_INVALID; if (rowMax > NEG_INF) { scalar.SetValue(0, sumExp); AscendC::SetFlag(eventSV_); // scalar write -> vector op @@ -338,10 +368,8 @@ private: AscendC::WaitFlag(eventSMTE3_); AscendC::DataCopyExtParams p4{1, sizeof(float), 0, 0, 0}; AscendC::DataCopyPad(lseGm_[(b * Hq_ + qh) * Sq_ + row], scalar[0], p4); - AscendC::DataCopyExtParams outCp{1, HEAD_DIM * sizeof(T), 0, 0, 0}; - AscendC::DataCopyPadExtParams outPp{false, 0, 0, 0}; - AscendC::DataCopyPad(outGm_[((b * Hq_ + qh) * Sq_ + row) * HEAD_DIM], outT, outCp, - outPp); + AscendC::DataCopyExtParams outCp{1, static_cast(HEAD_DIM * sizeof(T)), 0, 0, 0}; + AscendC::DataCopyPad(outGm_[((b * Hq_ + qh) * Sq_ + row) * HEAD_DIM], outT, outCp); // Drain MTE3 before the next row stages new values into the shared // buffers; the scalar pipe issues all later MTE2 copies in order, so // this wait alone orders them after the copy-outs. @@ -496,10 +524,3 @@ std::vector deterministic_attention_ascend_forward( } return {out, lse}; } - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("deterministic_attention_ascend", - &deterministic_attention_ascend_forward, - "Deterministic batch-invariant standard-softmax attention (Ascend C forward)"); -} diff --git a/csrc/ascend/batch_invariant_logp_ascend.asc b/csrc/ascend/batch_invariant_logp_ascend.asc index 09226c85..c65020af 100644 --- a/csrc/ascend/batch_invariant_logp_ascend.asc +++ b/csrc/ascend/batch_invariant_logp_ascend.asc @@ -301,10 +301,3 @@ std::vector batch_invariant_logp_ascend_forward(torch::Tensor log } return {logp, lse}; } - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("batch_invariant_logp_ascend", - &batch_invariant_logp_ascend_forward, - "Batch-invariant selected-token log-probability (Ascend C forward)"); -} diff --git a/csrc/ascend/ops_npu.asc b/csrc/ascend/ops_npu.asc new file mode 100644 index 00000000..98632a35 --- /dev/null +++ b/csrc/ascend/ops_npu.asc @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Aggregator for the Ascend C (CANN) extension module rl_engine._C_npu. +// +// bisheng's "-x asc" driver compiles every .asc source it is given and links +// them into a single shared object; a PYBIND11_MODULE in more than one source +// would define duplicate PyInit symbols. The pybind module is therefore +// defined exactly once, here, and the per-operator .asc files below only +// provide kernel + host forward functions. + +#include + +std::vector batch_invariant_logp_ascend_forward(torch::Tensor logits, + torch::Tensor target, + int64_t ignore_index); + +std::vector deterministic_attention_ascend_forward( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + bool causal, + double scale, + c10::optional key_padding_mask); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("batch_invariant_logp_ascend", + &batch_invariant_logp_ascend_forward, + "Batch-invariant selected-token log-probability (Ascend C forward)"); + m.def("deterministic_attention_ascend", + &deterministic_attention_ascend_forward, + "Deterministic batch-invariant standard-softmax attention (Ascend C forward)"); +} diff --git a/rl_engine/kernels/ops/ascend/attention/deterministic_attn.py b/rl_engine/kernels/ops/ascend/attention/deterministic_attn.py index 28895c5b..32fc9533 100644 --- a/rl_engine/kernels/ops/ascend/attention/deterministic_attn.py +++ b/rl_engine/kernels/ops/ascend/attention/deterministic_attn.py @@ -10,10 +10,10 @@ batch-invariant (the same algorithm as the Triton reference and the CUDA deterministic op). -Backward: Triton is unavailable on NPU, so the backward mirrors the -Triton op's portable branch -- it recomputes the reference forward with -`NativeAttentionOp` under autograd and VJPs the upstream gradient through -it, reusing the forward-saved q/k/v/mask. +Backward: Triton is unavailable on NPU, so the backward recomputes the +fp32 reference forward (`NativeAttentionOp.forward_fp32`, the same golden +path the forward kernel accumulates in) under autograd and VJPs the +upstream gradient through it, reusing the forward-saved q/k/v/mask. """ from __future__ import annotations @@ -55,7 +55,9 @@ def forward( v_c = v.contiguous() mask_c = key_padding_mask.contiguous() if key_padding_mask is not None else None - out, lse = _C_npu.deterministic_attention_ascend(q_c, k_c, v_c, causal, float(scale), mask_c) + out, lse = _C_npu.deterministic_attention_ascend( + q_c, k_c, v_c, causal, float(scale), mask_c + ) ctx.save_for_backward(q_c, k_c, v_c, mask_c) ctx.causal = causal @@ -69,11 +71,14 @@ def forward( def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): del grad_lse # lse is non-differentiable; always None upstream q, k, v, mask = ctx.saved_tensors + # VJP of the fp32 reference forward: the Ascend C forward accumulates in + # fp32 (like the CUDA deterministic op), so the backward must match the + # fp32 golden path, not the low-precision dtype path. with torch.enable_grad(): q_ref = q.detach().requires_grad_(True) k_ref = k.detach().requires_grad_(True) v_ref = v.detach().requires_grad_(True) - out = NativeAttentionOp().forward( + out = NativeAttentionOp().forward_fp32( q_ref, k_ref, v_ref, @@ -90,7 +95,8 @@ class DeterministicAttentionAscendOp: Public surface matches ``NativeAttentionOp`` / ``DeterministicAttentionOp`` so the #108 harness can call ``forward(**inputs)`` with ``key_padding_mask``. - Inputs the Ascend C kernel cannot take fall back to the native reference. + Out-of-domain inputs are rejected up front (the registry-level + ``PYTORCH_NATIVE_ATTENTION`` entry covers unavailable-kernel fallback). """ def __init__(self) -> None: diff --git a/tests/test_attention.py b/tests/test_attention.py index 469c6d30..77e0e050 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -434,13 +434,19 @@ def test_gradient_matches_reference(): def test_registry_dispatches_native_attention_op(): - """Resolve attention to the deterministic CUDA op or native fallback.""" + """Resolve attention to the deterministic op of the active platform or native fallback.""" op = kernel_registry.get_op("attention") - # On CUDA with the extension built, the registry prefers DeterministicAttentionOp. - # On CPU or without the CUDA extension, it falls back to NativeAttentionOp. + # On CUDA with the extension built, the registry prefers DeterministicAttentionOp; + # on NPU with the Ascend extension, DeterministicAttentionAscendOp. On CPU or + # without the platform extension, it falls back to NativeAttentionOp. + from rl_engine.kernels.ops.ascend.attention.deterministic_attn import ( + DeterministicAttentionAscendOp, + ) from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionOp - assert isinstance(op, (NativeAttentionOp, DeterministicAttentionOp)) + assert isinstance( + op, (NativeAttentionOp, DeterministicAttentionOp, DeterministicAttentionAscendOp) + ) # --------------------------------------------------------------------------- # diff --git a/tests/test_attention_ascend.py b/tests/test_attention_ascend.py index 7facd37d..1d1f614e 100644 --- a/tests/test_attention_ascend.py +++ b/tests/test_attention_ascend.py @@ -69,10 +69,15 @@ def _gold(q, k, v, causal=True, scale=None, key_padding_mask=None): def _make_qkv(batch, hq, hkv, sq, skv, dtype, seed=0): - generator = torch.Generator(device="cpu").manual_seed(seed) - q = torch.randn(batch, hq, sq, _D, dtype=dtype, generator=generator).to("npu") - k = torch.randn(batch, hkv, skv, _D, dtype=dtype, generator=generator).to("npu") - v = torch.randn(batch, hkv, skv, _D, dtype=dtype, generator=generator).to("npu") + # Independent generator per tensor: batch size must not shift the k/v + # content (a shared generator would make k[0] differ between batch sizes, + # breaking the batch-invariance comparisons below). + gq = torch.Generator(device="cpu").manual_seed(seed) + gk = torch.Generator(device="cpu").manual_seed(seed + 1) + gv = torch.Generator(device="cpu").manual_seed(seed + 2) + q = torch.randn(batch, hq, sq, _D, dtype=dtype, generator=gq).to("npu") + k = torch.randn(batch, hkv, skv, _D, dtype=dtype, generator=gk).to("npu") + v = torch.randn(batch, hkv, skv, _D, dtype=dtype, generator=gv).to("npu") return q, k, v @@ -168,21 +173,24 @@ def test_backward_grads(self, dtype): assert all(g is not None for g in (q.grad, k.grad, v.grad)) assert all(torch.isfinite(g).all() for g in (q.grad, k.grad, v.grad)) - # The backward is the VJP of the native reference forward; compare. + # The backward is the VJP of the fp32 reference forward; compare. with torch.enable_grad(): q_ref = q.detach().requires_grad_(True) k_ref = k.detach().requires_grad_(True) v_ref = v.detach().requires_grad_(True) - ref_out = NativeAttentionOp().forward(q_ref, k_ref, v_ref, causal=True) - dq_ref, dk_ref, dv_ref = torch.autograd.grad( - ref_out, (q_ref, k_ref, v_ref), grad_out - ) + ref_out = NativeAttentionOp().forward_fp32(q_ref, k_ref, v_ref, causal=True) + dq_ref, dk_ref, dv_ref = torch.autograd.grad(ref_out, (q_ref, k_ref, v_ref), grad_out) # The backward recomputes the same reference forward, so the VJPs # match to numerical noise. assert torch.allclose(q.grad.float(), dq_ref.float(), atol=1e-6, rtol=1e-5) assert torch.allclose(k.grad.float(), dk_ref.float(), atol=1e-6, rtol=1e-5) assert torch.allclose(v.grad.float(), dv_ref.float(), atol=1e-6, rtol=1e-5) + +@requires_ascend +class TestAscendAttentionRejects: + """Out-of-domain inputs must be rejected up front.""" + def test_rejects_fp32(self): op = _get_op() q, k, v = _make_qkv(1, 4, 4, 32, 32, torch.float32) @@ -220,14 +228,18 @@ def test_batch_size_1_vs_n(self): assert torch.equal(alone, in_batch), f"drift at batch_size={batch}" def test_different_positions_in_batch(self): + # Non-causal: every position attends to the same full window. Copy the + # same query row content into every position, then require + # bitwise-identical output wherever it sits. (Causal windows differ + # per position, so a causal sweep would compare different reductions.) dtype = torch.bfloat16 - baseline = None - for pos in range(16): - row = self._run_row(2, 8, 2, 64, 64, dtype, pos=pos, seed=11) - if baseline is None: - baseline = row - else: - assert torch.equal(baseline, row), f"drift at position={pos}" + op = _get_op() + q, k, v = _make_qkv(2, 8, 2, 64, 64, dtype, seed=11) + q[0, :, :, :] = q[0, :, 0:1, :] # same row content at every position + out = op(q, k, v, causal=False) + baseline = out[0, 0, 0, :].clone() + for pos in range(1, 16): + assert torch.equal(baseline, out[0, 0, pos, :]), f"drift at position={pos}" def test_block_striding(self): # 2 * 8 * 128 = 2048 work items > MAX_BLOCKS (512): rows are strided From d71b7cfeba0e91acc55d5543c4362ce845c5b65a Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Wed, 19 Aug 2026 21:18:59 +0800 Subject: [PATCH 5/5] fix(lint): black/isort formatting per pre-commit CI --- benchmarks/benchmark_attention.py | 4 +--- rl_engine/_C_npu.pyi | 1 - rl_engine/kernels/ops/ascend/attention/__init__.py | 4 +--- tests/test_attention_ascend.py | 2 +- 4 files changed, 3 insertions(+), 8 deletions(-) diff --git a/benchmarks/benchmark_attention.py b/benchmarks/benchmark_attention.py index 0aebec22..99705a52 100644 --- a/benchmarks/benchmark_attention.py +++ b/benchmarks/benchmark_attention.py @@ -65,9 +65,7 @@ def _maybe_cuda_op(): and hasattr(_C, "deterministic_attention_forward") ): return None - from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( - DeterministicAttentionOp, - ) + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionOp return DeterministicAttentionOp() diff --git a/rl_engine/_C_npu.pyi b/rl_engine/_C_npu.pyi index 6ced124a..bbaaad69 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,7 +8,6 @@ def batch_invariant_logp_ascend( target: torch.Tensor, ignore_index: int, ) -> list[torch.Tensor]: ... - def deterministic_attention_ascend( q: torch.Tensor, k: torch.Tensor, diff --git a/rl_engine/kernels/ops/ascend/attention/__init__.py b/rl_engine/kernels/ops/ascend/attention/__init__.py index 07b76408..f2cdca31 100644 --- a/rl_engine/kernels/ops/ascend/attention/__init__.py +++ b/rl_engine/kernels/ops/ascend/attention/__init__.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from rl_engine.kernels.ops.ascend.attention.deterministic_attn import ( - DeterministicAttentionAscendOp, -) +from rl_engine.kernels.ops.ascend.attention.deterministic_attn import DeterministicAttentionAscendOp __all__ = ["DeterministicAttentionAscendOp"] diff --git a/tests/test_attention_ascend.py b/tests/test_attention_ascend.py index 1d1f614e..086f59da 100644 --- a/tests/test_attention_ascend.py +++ b/tests/test_attention_ascend.py @@ -38,8 +38,8 @@ def _ascend_kernel_available() -> bool: return False try: from rl_engine.kernels.ops.ascend.attention.deterministic_attn import ( - _C_npu, _NPU_EXT_AVAILABLE, + _C_npu, ) except Exception: return False