From 8e94860d070c529a8e407141616e78ae1660187c Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 25 Aug 2026 16:49:57 +0800 Subject: [PATCH] feat(ascend): add prefix-shared attention Ascend C kernel - csrc/ascend/attention/prefix_shared_attention_ascend.asc: Ascend C prefix-shared fused attention forward (bf16, fp32 online softmax over fixed 64-key tiles, one block per (bs, g, 64-row query block), no split-K -> batch-invariant); D=128, non-causal, same surface as the CUDA prefix_shared_attention op - rl_engine/kernels/ops/ascend/attention/prefix_shared_attn.py: PrefixSharedAttentionAscendOp wrapper (op(q, k, v) -> out) - gtest: prefix_shared_attention spec (pytorch/cuda/ascend candidates) with GtestPrefixSharedAttentionOp gold; check_operator.py learns --device npu - pybind consolidated in csrc/ascend/ops_npu.asc (single PYBIND11_MODULE; batch_invariant_logp_ascend.asc only drops its module block) - tests/test_prefix_shared_attention_ascend.py: correctness, batch invariance, validation - setup.py: recursive .asc glob --- .../prefix_shared_attention_ascend.asc | 391 ++++++++++++++++++ csrc/ascend/batch_invariant_logp_ascend.asc | 7 - csrc/ascend/ops_npu.asc | 25 ++ rl_engine/kernels/gtest/operator_inputs.py | 17 + rl_engine/kernels/gtest/operator_specs.py | 40 ++ .../kernels/ops/ascend/attention/__init__.py | 8 + .../ascend/attention/prefix_shared_attn.py | 119 ++++++ scripts/check_operator.py | 11 + setup.py | 2 +- tests/test_prefix_shared_attention_ascend.py | 250 +++++++++++ tests/test_ws1_gtest_gpu.py | 1 + 11 files changed, 863 insertions(+), 8 deletions(-) create mode 100644 csrc/ascend/attention/prefix_shared_attention_ascend.asc create mode 100644 csrc/ascend/ops_npu.asc create mode 100644 rl_engine/kernels/ops/ascend/attention/__init__.py create mode 100644 rl_engine/kernels/ops/ascend/attention/prefix_shared_attn.py create mode 100644 tests/test_prefix_shared_attention_ascend.py diff --git a/csrc/ascend/attention/prefix_shared_attention_ascend.asc b/csrc/ascend/attention/prefix_shared_attention_ascend.asc new file mode 100644 index 00000000..6b9e292b --- /dev/null +++ b/csrc/ascend/attention/prefix_shared_attention_ascend.asc @@ -0,0 +1,391 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Prefix-shared fused attention, Ascend C (CANN) forward kernel. +// +// out = softmax(Q K^T * scale) @ V +// +// Mirrors the CUDA kernel in csrc/cuda/attention/prefix_shared_attention.cu +// (the GRPO decode workload): every one of the G query groups attends over +// the same shared prompt-prefix key/value sequence, so K and V are stored +// once per batch instead of once per group. +// - layout : q [bs, G, len_q, D], k/v [bs, len_kv, D] contiguous, D = 128 +// - numerics: bf16 in/out; fp32 online-softmax accumulation (per-row max / +// sum-exp rescale per key tile), the same flash-style single +// pass as the CUDA kernel; no causal mask, no key-padding mask +// (same surface as the CUDA op). +// - blocking: each (bs, g, 64-row query block) is processed end-to-end by +// one AI-core block over fixed 64-key tiles. The per-row +// reduction order depends only on len_kv, so row outputs are +// batch-invariant: they never depend on batch size, batch +// position, or how many blocks were launched (items are strided +// across blocks). +// +// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by +// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu. +// Python bindings live in csrc/ascend/ops_npu.asc (single PYBIND11_MODULE). + +#include +#include + +#include "kernel_operator.h" + +#include + +#include "torch_npu/csrc/core/npu/NPUStream.h" + +namespace { + +// Fixed head dimension (matches the CUDA op gate). +constexpr uint32_t HEAD_DIM = 128; +// Query rows per block (CUDA BLOCK_Q). +constexpr uint32_t BLOCK_Q = 64; +// Keys per tile (CUDA BLOCK_KV). Fixed for all runs; this is what makes the +// per-row reduction order batch-invariant. +constexpr uint32_t TILE_N = 64; +// 1/sqrt(HEAD_DIM); matches the CUDA kernel's rsqrtf(dim) softmax scale. +constexpr float SCALE = 0.088388348f; +// 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 KernelPrefixSharedAttention { +public: + __aicore__ inline KernelPrefixSharedAttention(AscendC::TPipe* pipe) : pipe_(pipe) {} + + __aicore__ inline void Init(GM_ADDR q, + GM_ADDR k, + GM_ADDR v, + GM_ADDR out, + int64_t bs, + int64_t G, + int64_t lenQ, + int64_t lenKv) + { + bs_ = bs; + G_ = G; + lenQ_ = lenQ; + lenKv_ = lenKv; + qGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(q)); + kGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(k)); + vGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(v)); + outGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(out)); + + // UB budget stays well under 192 KB: + // q/k/v/acc fp32 tiles 4 x 32 KB + bf16 staging 16 KB + // + prod/work/scores/state/scalar ~2 KB. + pipe_->InitBuffer(qBufF_, BLOCK_Q * HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(kBufF_, TILE_N * HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(vBufF_, TILE_N * HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(accBufF_, BLOCK_Q * HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(bufT_, BLOCK_Q * HEAD_DIM * sizeof(T)); + pipe_->InitBuffer(prodBufF_, HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(workBufF_, HEAD_DIM * sizeof(float)); + pipe_->InitBuffer(scoresBuf_, TILE_N * sizeof(float)); + pipe_->InitBuffer(rowMaxBuf_, BLOCK_Q * sizeof(float)); + pipe_->InitBuffer(rowSumExpBuf_, BLOCK_Q * sizeof(float)); + // 64 B: floats [0,8) hold the per-row online-softmax rescale exp. + 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); + eventVMTE2_ = pipe_->FetchEventID(AscendC::HardEvent::V_MTE2); + eventVMTE3_ = pipe_->FetchEventID(AscendC::HardEvent::V_MTE3); + eventMTE3S_ = pipe_->FetchEventID(AscendC::HardEvent::MTE3_S); + } + + __aicore__ inline void Process() + { + const int64_t qBlocks = (lenQ_ + BLOCK_Q - 1) / BLOCK_Q; + const int64_t items = bs_ * G_ * qBlocks; + for (int64_t item = AscendC::GetBlockIdx(); item < items; + item += AscendC::GetBlockNum()) { + const int64_t qb = item % qBlocks; + const int64_t g = (item / qBlocks) % G_; + const int64_t b = item / (qBlocks * G_); + ProcessBlock(b, g, qb); + } + } + +private: + __aicore__ inline void LoadQTile(int64_t b, int64_t g, int64_t rowStart, uint32_t numRows) + { + const int64_t offset = ((b * G_ + g) * lenQ_ + rowStart) * HEAD_DIM; + AscendC::LocalTensor qT = bufT_.Get(); + AscendC::DataCopyExtParams cp{ + 1, static_cast(numRows * 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, + numRows * HEAD_DIM); + } + + __aicore__ inline void LoadKTile(int64_t b, int64_t start, uint32_t count) + { + // bufT_ staging may hold data the vector pipe is still reading (the + // Q cast or the previous tile's V cast). + AscendC::SetFlag(eventVMTE2_); + AscendC::WaitFlag(eventVMTE2_); + const int64_t offset = (b * lenKv_ + start) * HEAD_DIM; + AscendC::LocalTensor kT = bufT_.Get(); + 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_); + AscendC::WaitFlag(eventMTE2S_); + AscendC::Cast(kBufF_.Get(), kT, AscendC::RoundMode::CAST_NONE, + count * HEAD_DIM); + } + + __aicore__ inline void LoadVTile(int64_t b, int64_t start, uint32_t count) + { + // bufT_ staging may hold data the vector pipe is still reading (the + // K cast of this tile). + AscendC::SetFlag(eventVMTE2_); + AscendC::WaitFlag(eventVMTE2_); + const int64_t offset = (b * lenKv_ + start) * HEAD_DIM; + AscendC::LocalTensor vT = bufT_.Get(); + 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_); + AscendC::WaitFlag(eventMTE2S_); + AscendC::Cast(vBufF_.Get(), vT, AscendC::RoundMode::CAST_NONE, + count * HEAD_DIM); + } + + // Online-softmax step for query row r over one 64-key tile: recompute the + // row max, rescale the running (sum-exp, accumulator), then fold the tile + // into both with P = exp(scores - mNew) and P . V. + __aicore__ inline void ProcessRowTile(uint32_t r, uint32_t count) + { + AscendC::LocalTensor qRow = qBufF_.Get()[r * HEAD_DIM]; + AscendC::LocalTensor accRow = accBufF_.Get()[r * HEAD_DIM]; + AscendC::LocalTensor kTile = kBufF_.Get(); + AscendC::LocalTensor vTile = vBufF_.Get(); + AscendC::LocalTensor scores = scoresBuf_.Get(); + AscendC::LocalTensor prod = prodBufF_.Get(); + AscendC::LocalTensor work = workBufF_.Get(); + AscendC::LocalTensor scalar = scalarBuf_.Get(); + AscendC::LocalTensor rowMax = rowMaxBuf_.Get(); + AscendC::LocalTensor rowSumExp = rowSumExpBuf_.Get(); + + // scores[j] = SCALE * (q_r . k_j) for j in [0, count); past-the-end + // lanes become NEG_INF so they drop out of max/exp exactly. + for (uint32_t j = 0; j < count; ++j) { + AscendC::Mul(prod, qRow, kTile[j * HEAD_DIM], HEAD_DIM); + AscendC::ReduceSum(scalar, prod, work, HEAD_DIM); + WaitVector(); + scores.SetValue(j, scalar.GetValue(0) * SCALE); + } + for (uint32_t j = count; j < TILE_N; ++j) { + scores.SetValue(j, NEG_INF); + } + + AscendC::SetFlag(eventSV_); + AscendC::WaitFlag(eventSV_); + AscendC::ReduceMax(scalar, scores, work, TILE_N, false); + WaitVector(); + const float mOld = rowMax.GetValue(r); + const float tileMax = scalar.GetValue(0); + const float mNew = mOld > tileMax ? mOld : tileMax; + + // rescale = exp(mOld - mNew), computed through the vector Exp (aicore + // scalar code has no expf); 0 on the first tile, where mOld == NEG_INF. + scalar.SetValue(0, mOld - mNew); + AscendC::SetFlag(eventSV_); + AscendC::WaitFlag(eventSV_); + AscendC::Exp(scalar, scalar, 8); + WaitVector(); + const float rescale = scalar.GetValue(0); + + // p[j] = exp(scores[j] - mNew); past-the-end lanes stay exp(-inf) = 0. + AscendC::Adds(scores, scores, -mNew, TILE_N); + AscendC::Exp(scores, scores, TILE_N); + AscendC::ReduceSum(scalar, scores, work, TILE_N); + WaitVector(); + const float tileSumExp = scalar.GetValue(0); + + rowMax.SetValue(r, mNew); + rowSumExp.SetValue(r, rowSumExp.GetValue(r) * rescale + tileSumExp); + + // acc_r *= rescale, then acc_r += sum_j p[j] * v_j. p[j] == 0 exactly + // on past-the-end lanes, and adding a zero term is a no-op, so the + // skip below is exact. + AscendC::Muls(accRow, accRow, rescale, HEAD_DIM); + 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(accRow, accRow, prod, HEAD_DIM); + } + } + + __aicore__ inline void ProcessBlock(int64_t b, int64_t g, int64_t qb) + { + const int64_t rowStart = qb * BLOCK_Q; + const int64_t remaining = lenQ_ - rowStart; + const uint32_t numRows = static_cast(remaining < BLOCK_Q ? remaining : BLOCK_Q); + + LoadQTile(b, g, rowStart, numRows); + + // Real zeroing: the buffers start as uninitialized UB and + // 0 * inf == NaN in the online-softmax rescale. + AscendC::Duplicate(accBufF_.Get(), 0.0f, numRows * HEAD_DIM); + AscendC::Duplicate(rowMaxBuf_.Get(), NEG_INF, numRows); + AscendC::Duplicate(rowSumExpBuf_.Get(), 0.0f, numRows); + + const int64_t numTiles = (lenKv_ + TILE_N - 1) / TILE_N; + for (int64_t tile = 0; tile < numTiles; ++tile) { + const int64_t start = tile * TILE_N; + const uint32_t count = TileCount(start); + LoadKTile(b, start, count); + LoadVTile(b, start, count); + for (uint32_t r = 0; r < numRows; ++r) { + ProcessRowTile(r, count); + } + } + WriteOutputs(b, g, rowStart, numRows); + } + + __aicore__ inline void WriteOutputs(int64_t b, int64_t g, int64_t rowStart, uint32_t numRows) + { + AscendC::LocalTensor acc = accBufF_.Get(); + AscendC::LocalTensor rowSumExp = rowSumExpBuf_.Get(); + AscendC::LocalTensor outT = bufT_.Get(); + + // out_r = acc_r / sumExp_r, then one bf16 cast and one copy-out for + // the whole contiguous row block. sumExp > 0 always (every block + // attends over at least one real key; len_kv >= 1 is host-checked), + // the guard just mirrors the CUDA op's defensive division. + for (uint32_t r = 0; r < numRows; ++r) { + const float sumExp = rowSumExp.GetValue(r); + const float invDenom = (sumExp > 0.0f) ? (1.0f / sumExp) : 0.0f; + AscendC::Muls(acc[r * HEAD_DIM], acc[r * HEAD_DIM], invDenom, HEAD_DIM); + } + + AscendC::Cast(outT, acc, AscendC::RoundMode::CAST_ROUND, numRows * HEAD_DIM); + AscendC::SetFlag(eventVMTE3_); + AscendC::WaitFlag(eventVMTE3_); + const int64_t offset = ((b * G_ + g) * lenQ_ + rowStart) * HEAD_DIM; + AscendC::DataCopyExtParams outCp{ + 1, static_cast(numRows * HEAD_DIM * sizeof(T)), 0, 0, 0}; + // UB -> GM has no pad-params overload (CANN 8.5.1): the row block is + // always contiguous, so no padding is needed anyway. + AscendC::DataCopyPad(outGm_[offset], outT, outCp); + // Drain MTE3 before the next block 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-out. + 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 = lenKv_ - start; + return static_cast(remaining < TILE_N ? remaining : TILE_N); + } + + AscendC::TPipe* pipe_; + AscendC::GlobalTensor qGm_; + AscendC::GlobalTensor kGm_; + AscendC::GlobalTensor vGm_; + AscendC::GlobalTensor outGm_; + AscendC::TBuf qBufF_; + AscendC::TBuf kBufF_; + AscendC::TBuf vBufF_; + AscendC::TBuf accBufF_; + AscendC::TBuf bufT_; + AscendC::TBuf prodBufF_; + AscendC::TBuf workBufF_; + AscendC::TBuf scoresBuf_; + AscendC::TBuf rowMaxBuf_; + AscendC::TBuf rowSumExpBuf_; + AscendC::TBuf scalarBuf_; + AscendC::TEventID eventVS_; + AscendC::TEventID eventSV_; + AscendC::TEventID eventMTE2S_; + AscendC::TEventID eventVMTE2_; + AscendC::TEventID eventVMTE3_; + AscendC::TEventID eventMTE3S_; + int64_t bs_; + int64_t G_; + int64_t lenQ_; + int64_t lenKv_; +}; + +} // namespace + +extern "C" __global__ __vector__ void prefix_shared_attention_ascend_kernel_bf16( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR out, + int64_t bs, int64_t G, int64_t lenQ, int64_t lenKv) +{ + AscendC::TPipe pipe; + KernelPrefixSharedAttention op(&pipe); + op.Init(q, k, v, out, bs, G, lenQ, lenKv); + op.Process(); +} + +torch::Tensor prefix_shared_attention_ascend_forward( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v) +{ + 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() == 3 && v.dim() == 3, + "q must be 4-D [bs, G, len_q, D]; k and v must be 3-D [bs, len_kv, 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 && + k.scalar_type() == at::kBFloat16 && v.scalar_type() == at::kBFloat16, + "prefix-shared attention requires bf16 (matches the CUDA op)"); + TORCH_CHECK(q.size(3) == HEAD_DIM && k.size(2) == HEAD_DIM && v.size(2) == 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/v must share the same key length"); + TORCH_CHECK(q.size(2) >= 1 && k.size(1) >= 1, "len_q and len_kv must be positive"); + + const int64_t bs = q.size(0); + const int64_t G = q.size(1); + const int64_t lenQ = q.size(2); + const int64_t lenKv = k.size(1); + + torch::Tensor out = at::empty({bs, G, lenQ, HEAD_DIM}, q.options()); + + // stream(true): flush the task queue before launch so the kernel cannot + // overtake earlier NPU work; the output was allocated with at::empty (no + // queued initializer) for the same reason. + auto aclStream = c10_npu::getCurrentNPUStream().stream(true); + const int64_t qBlocks = (lenQ + BLOCK_Q - 1) / BLOCK_Q; + const int64_t items = bs * G * qBlocks; + const uint32_t blockNum = static_cast(std::min(items, MAX_BLOCKS)); + prefix_shared_attention_ascend_kernel_bf16<<>>( + reinterpret_cast(q.mutable_data_ptr()), + reinterpret_cast(k.mutable_data_ptr()), + reinterpret_cast(v.mutable_data_ptr()), + reinterpret_cast(out.mutable_data_ptr()), + bs, G, lenQ, lenKv); + return out; +} diff --git a/csrc/ascend/batch_invariant_logp_ascend.asc b/csrc/ascend/batch_invariant_logp_ascend.asc index dead4cbe..49139240 100644 --- a/csrc/ascend/batch_invariant_logp_ascend.asc +++ b/csrc/ascend/batch_invariant_logp_ascend.asc @@ -307,10 +307,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..14ba6416 --- /dev/null +++ b/csrc/ascend/ops_npu.asc @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +// Single Python module entry point for the unified Ascend C extension. +// Keep PYBIND11_MODULE in one translation unit as new .asc kernels are added. + +#include + +#include + +std::vector batch_invariant_logp_ascend_forward( + torch::Tensor logits, torch::Tensor target, int64_t ignore_index); + +torch::Tensor prefix_shared_attention_ascend_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v); + +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("prefix_shared_attention_ascend", + &prefix_shared_attention_ascend_forward, + "Prefix-shared fused attention (Ascend C forward)"); +} diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index f37fce36..31ac5608 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -31,6 +31,7 @@ def make_operator_inputs( "matmul": _make_matmul_inputs, "det_gemm": _make_det_gemm_inputs, "attention": _make_attention_inputs, + "prefix_shared_attention": _make_prefix_shared_attention_inputs, "logp": _make_logp_inputs, "linear_logp": _make_linear_logp_inputs, "batch_invariant_logp": _make_batch_invariant_logp_inputs, @@ -58,6 +59,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "matmul": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "det_gemm": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", + "prefix_shared_attention": f"{batch}x{_arg_int(args, 'n_heads', DEFAULT_N_HEADS)}" + f"x{seq}x{DEFAULT_HEAD_DIM}", "logp": f"{batch}x{seq}x{vocab}", "linear_logp": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "batch_invariant_logp": f"{batch}x{seq}x{vocab}", @@ -175,6 +178,20 @@ def _make_attention_inputs( return inputs +def _make_prefix_shared_attention_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + """Prefix-shared layout: k/v are 3-D [B, Skv, D] shared by all G groups.""" + batch, seq = _batch_seq(args) + skv = _arg_int(args, "skv", seq) + n_groups = _arg_int(args, "n_heads", DEFAULT_N_HEADS) + return { + "q": _floating_tensor((batch, n_groups, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0), + "k": _floating_tensor((batch, skv, DEFAULT_HEAD_DIM), args, dtype, device, 1), + "v": _floating_tensor((batch, skv, DEFAULT_HEAD_DIM), args, dtype, device, 2), + } + + def _make_logp_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index bde87edb..589d3861 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -76,6 +76,26 @@ def _load_object(path: str) -> Any: }, grad_input_names=("q", "k", "v"), ), + # GRPO decode: every G group attends over one shared K/V sequence + # ([B, Skv, D] instead of [B, Hkv, Skv, D]). Forward-only (no backward), + # same surface as the CUDA PrefixSharedAttentionOp. + "prefix_shared_attention": OperatorSpec( + name="prefix_shared_attention", + op_class="attention", + gold_path="rl_engine.kernels.gtest.operator_specs.GtestPrefixSharedAttentionOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.gtest.operator_specs.GtestPrefixSharedAttentionOp", + "cuda": ( + "rl_engine.kernels.ops.cuda.attention.prefix_shared_attn." + "PrefixSharedAttentionOp" + ), + "ascend": ( + "rl_engine.kernels.ops.ascend.attention.prefix_shared_attn." + "PrefixSharedAttentionAscendOp" + ), + }, + ), "logp": OperatorSpec( name="logp", op_class="logprob", @@ -227,6 +247,26 @@ def forward_fp32(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: return packed +class GtestPrefixSharedAttentionOp: + """gtest view of the prefix-shared layout: expand the shared K/V over the + G groups and reuse the standard fp32 attention reference (non-causal, + default scale), which is the gold for the CUDA/Ascend prefix-shared ops. + """ + + op_class = "attention" + + def __init__(self) -> None: + from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + self._op = NativeAttentionOp() + + def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + return self._op(q, k.unsqueeze(1), v.unsqueeze(1), causal=False) + + def forward_fp32(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + return self._op.forward_fp32(q, k.unsqueeze(1), v.unsqueeze(1), causal=False) + + class _LogpSM90CandidateAdapter: def __init__(self, candidate: Any) -> None: self._candidate = candidate 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..c4ba16df --- /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.prefix_shared_attn import ( + PrefixSharedAttentionAscendOp, +) + +__all__ = ["PrefixSharedAttentionAscendOp"] diff --git a/rl_engine/kernels/ops/ascend/attention/prefix_shared_attn.py b/rl_engine/kernels/ops/ascend/attention/prefix_shared_attn.py new file mode 100644 index 00000000..2f8e0cd2 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/attention/prefix_shared_attn.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Ascend NPU prefix-shared fused attention (GRPO decode workload). + +Port of the CUDA op in rl_engine/kernels/ops/cuda/attention/prefix_shared_attn.py: +in GRPO, the G generated responses share the exact same prompt-prefix KV cache, +so K/V are stored once per batch and broadcast across all G query groups. + +Forward: softmax(Q K^T * scale) @ V on an Ascend C kernel +(`_C_npu.prefix_shared_attention_ascend`) with fp32 online-softmax +accumulation. bf16 in/out, non-causal, no key-padding mask, head dim fixed at +128 -- the same surface as the CUDA `PrefixSharedAttentionOp`, which is +forward-only and so is this port. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from rl_engine.utils.logger import logger + +_C_npu: Any = None +try: + from rl_engine import _C_npu + + _NPU_EXT_AVAILABLE = True +except ImportError: # pragma: no cover - Ascend extension not built + _NPU_EXT_AVAILABLE = False + +_HEAD_DIM = 128 + + +class PrefixSharedAttentionAscendOp: + """Prefix-shared softmax attention on Ascend NPU. + + q [bs, G, len_q, D] attends over a single shared k/v sequence + [bs, len_kv, D] that every G group reuses. Mirrors the CUDA + ``PrefixSharedAttentionOp`` surface (``op(q, k, v) -> out``). + """ + + def __init__(self) -> None: + if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "prefix_shared_attention_ascend"): + raise RuntimeError( + "prefix_shared_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.prefix_shared_attention_ascend kernel." + ) + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: + """ + Prefix-shared attention forward pass. + + Args: + q: Query tensor of shape [bs, G, len_q, head_dim] + k: Shared Key tensor of shape [bs, len_kv, head_dim] + v: Shared Value tensor of shape [bs, len_kv, head_dim] + + Returns: + Output tensor of shape [bs, G, len_q, head_dim] + """ + return self.forward(q, k, v) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: + self._validate_inputs(q, k, v) + return _C_npu.prefix_shared_attention_ascend( + q.contiguous(), k.contiguous(), v.contiguous() + ) + + @staticmethod + def _validate_inputs(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + if q.dim() != 4 or k.dim() != 3 or v.dim() != 3: + raise ValueError( + f"q must be 4-D [B, G, Sq, D] and k/v 3-D [B, Skv, D], got " + f"q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}" + ) + b, g, sq, d = q.shape + skv = k.shape[1] + if k.shape[0] != b or v.shape[0] != b: + raise ValueError("batch size mismatch between q/k/v") + if k.shape[2] != d or v.shape[2] != d: + raise ValueError( + f"k/v head dim mismatch: k={tuple(k.shape)}, v={tuple(v.shape)}, " + f"expected D={d}" + ) + if v.shape[1] != skv: + raise ValueError( + f"k/v key length mismatch: k={tuple(k.shape)}, v={tuple(v.shape)}" + ) + if d != _HEAD_DIM: + raise ValueError(f"head dim D must be {_HEAD_DIM}, got {d}") + if q.dtype != torch.bfloat16 or k.dtype != torch.bfloat16 or v.dtype != torch.bfloat16: + raise ValueError( + f"only BF16 is supported (matches the CUDA op), got " + f"q={q.dtype}, k={k.dtype}, v={v.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 sq < 1 or skv < 1: + raise ValueError(f"Sq and Skv must be positive, got Sq={sq}, Skv={skv}") + if g < 1: + raise ValueError(f"G must be positive, got G={g}") diff --git a/scripts/check_operator.py b/scripts/check_operator.py index 9dbca48d..499e205d 100644 --- a/scripts/check_operator.py +++ b/scripts/check_operator.py @@ -41,6 +41,17 @@ def _select_device(value: str) -> torch.device: 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") + if device.type == "npu": + # torch.npu only exists after torch_npu is imported (mirrors the + # defensive probe in rl_engine.platforms.device). + try: + import torch_npu # noqa: F401 + except ImportError as exc: + raise RuntimeError( + "--device npu was requested, but torch_npu is not installed" + ) from exc + if not torch.npu.is_available(): + raise RuntimeError("--device npu was requested, but no NPU is available") return device 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_prefix_shared_attention_ascend.py b/tests/test_prefix_shared_attention_ascend.py new file mode 100644 index 00000000..e9c6301d --- /dev/null +++ b/tests/test_prefix_shared_attention_ascend.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for the Ascend NPU prefix-shared fused attention. + +Validates the same properties as the CUDA prefix-shared op: +1. **Correctness** - output matches the ``NativeAttentionOp.forward_fp32`` + ground truth (full softmax over the shared K/V, no causal mask) 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 (bs, g, 64-row block) item is processed end-to-end by one block). +""" + +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} +_RTOL = {torch.bfloat16: 2.0e-2} + + +def _npu_available() -> bool: + try: + import torch_npu # noqa: F401 + except ImportError: + return False + return hasattr(torch, "npu") and torch.npu.is_available() + + +def _ascend_kernel_available() -> bool: + if not _npu_available(): + return False + try: + from rl_engine.kernels.ops.ascend.attention.prefix_shared_attn import ( + _C_npu, + _NPU_EXT_AVAILABLE, + ) + except Exception: + return False + return _NPU_EXT_AVAILABLE and hasattr(_C_npu, "prefix_shared_attention_ascend") + + +requires_ascend = pytest.mark.skipif( + not _ascend_kernel_available(), + reason="prefix_shared_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.prefix_shared_attn import ( + PrefixSharedAttentionAscendOp, + ) + + return PrefixSharedAttentionAscendOp() + + +def _gold(q, k, v): + """fp32 ground truth: softmax(Q K^T / sqrt(D)) V over the shared K/V.""" + return NativeAttentionOp().forward_fp32( + q, + k.unsqueeze(1), # [bs, 1, Skv, D]: every G group shares the same KV head + v.unsqueeze(1), + causal=False, + scale=1.0 / math.sqrt(_D), + ) + + +def _make_qkv(batch, groups, sq, skv, dtype, seed=0): + generator = torch.Generator(device="cpu").manual_seed(seed) + q = torch.randn(batch, groups, sq, _D, dtype=dtype, generator=generator).to("npu") + k = torch.randn(batch, skv, _D, dtype=dtype, generator=generator).to("npu") + v = torch.randn(batch, skv, _D, dtype=dtype, generator=generator).to("npu") + return q, k, v + + +# --------------------------------------------------------------------------- +# Correctness +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendPrefixSharedAttentionCorrectness: + def test_basic(self): + op = _get_op() + q, k, v = _make_qkv(2, 4, 65, 130, torch.bfloat16) # ragged Sq/Skv + out = op(q, k, v) + gold = _gold(q, k, v) + assert out.dtype == torch.bfloat16 + assert out.shape == (2, 4, 65, _D) + assert torch.allclose( + out.float(), gold, atol=_ATOL[torch.bfloat16], rtol=_RTOL[torch.bfloat16] + ) + + def test_exact_tiles(self): + op = _get_op() + q, k, v = _make_qkv(2, 8, 64, 64, torch.bfloat16) # one Q block, one KV tile + out = op(q, k, v) + gold = _gold(q, k, v) + assert torch.allclose( + out.float(), gold, atol=_ATOL[torch.bfloat16], rtol=_RTOL[torch.bfloat16] + ) + + def test_decode_window(self): + op = _get_op() + q, k, v = _make_qkv(1, 16, 1, 512, torch.bfloat16) # Sq << Skv, 8 KV tiles + out = op(q, k, v) + gold = _gold(q, k, v) + assert torch.allclose( + out.float(), gold, atol=_ATOL[torch.bfloat16], rtol=_RTOL[torch.bfloat16] + ) + + def test_long_prefix_multi_tile(self): + op = _get_op() + q, k, v = _make_qkv(1, 2, 32, 1024, torch.bfloat16) # 16 KV tiles + out = op(q, k, v) + gold = _gold(q, k, v) + assert torch.allclose( + out.float(), gold, atol=_ATOL[torch.bfloat16], rtol=_RTOL[torch.bfloat16] + ) + + def test_shared_kv_across_groups(self): + # Every G group attends over the exact same K/V; a G-sweep must match + # the per-group reference (and be bitwise equal across G for equal q). + op = _get_op() + q, k, v = _make_qkv(1, 4, 32, 96, torch.bfloat16) + q[0, 1, :, :] = q[0, 0, :, :] # force identical query rows in 2 groups + out = op(q, k, v) + assert torch.equal(out[0, 0], out[0, 1]) + gold = _gold(q, k, v) + assert torch.allclose( + out.float(), gold, atol=_ATOL[torch.bfloat16], rtol=_RTOL[torch.bfloat16] + ) + + +# --------------------------------------------------------------------------- +# Batch invariance +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendPrefixSharedAttentionBatchInvariance: + def test_batch_size_1_vs_n(self): + # The same (b=0, g=0, row 0) content computed alone vs embedded in a + # larger batch must be bitwise identical. (Content is copied in + # explicitly: CPU bf16 randn consumes two fp32 draws per element, so + # same-seed tensors of different batch sizes would not line up.) + dtype = torch.bfloat16 + op = _get_op() + q1, k1, v1 = _make_qkv(1, 4, 64, 64, dtype, seed=7) + alone = op(q1, k1, v1)[0, 0, 0, :].clone() + for batch in (2, 4, 8): + q, k, v = _make_qkv(batch, 4, 64, 64, dtype, seed=7) + q[0, 0, 0, :] = q1[0, 0, 0, :] + k[0, :, :] = k1[0, :, :] + v[0, :, :] = v1[0, :, :] + in_batch = op(q, k, v)[0, 0, 0, :].clone() + assert torch.equal(alone, in_batch), f"drift at batch_size={batch}" + + def test_different_positions_in_batch(self): + # One fixed query row embedded at several positions (spanning both + # 64-row query blocks) must give the bitwise-identical output at each. + dtype = torch.bfloat16 + op = _get_op() + q1, k1, v1 = _make_qkv(1, 1, 1, 64, dtype, seed=11) + baseline = op(q1, k1, v1)[0, 0, 0, :].clone() + q, k, v = _make_qkv(2, 4, 128, 64, dtype, seed=11) + k[0, :, :] = k1[0, :, :] + v[0, :, :] = v1[0, :, :] + for pos in (0, 63, 64, 127): + q[0, 0, pos, :] = q1[0, 0, 0, :] + out = op(q, k, v) + for pos in (0, 63, 64, 127): + assert torch.equal(baseline, out[0, 0, pos, :]), f"drift at position={pos}" + + def test_block_striding(self): + # The strided run below has 8 * 16 * (320/64) = 640 work items > + # MAX_BLOCKS (512), so items are strided across blocks; numerics must + # not depend on block assignment. The 1-item run and the strided run + # must give the bitwise-identical rows for the same content. + dtype = torch.bfloat16 + op = _get_op() + small_q, small_k, small_v = _make_qkv(1, 1, 64, 64, dtype, seed=3) + small = op(small_q, small_k, small_v) + big_q, big_k, big_v = _make_qkv(8, 16, 320, 64, dtype, seed=3) + big_q[0, 0, :64, :] = small_q[0, 0, :, :] + big_k[0, :, :] = small_k[0, :, :] + big_v[0, :, :] = small_v[0, :, :] + big = op(big_q, big_k, big_v) + assert torch.equal(big[0, 0, :64, :], small[0, 0, :, :]) + + def test_repeated_runs_deterministic(self): + dtype = torch.bfloat16 + q, k, v = _make_qkv(2, 4, 128, 128, dtype, seed=5) + op = _get_op() + first = op(q, k, v) + for _ in range(3): + again = op(q, k, v) + assert torch.equal(first, again) + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + + +@requires_ascend +class TestAscendPrefixSharedAttentionValidation: + def test_rejects_fp32(self): + op = _get_op() + q, k, v = _make_qkv(1, 4, 32, 32, torch.float32) + with pytest.raises(ValueError, match="only BF16"): + op(q, k, v) + + def test_rejects_fp16(self): + op = _get_op() + q, k, v = _make_qkv(1, 4, 32, 32, torch.float16) + with pytest.raises(ValueError, match="only 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.bfloat16) + k = torch.randn(1, 32, 64, device="npu", dtype=torch.bfloat16) + v = torch.randn(1, 32, 64, device="npu", dtype=torch.bfloat16) + with pytest.raises(ValueError, match="head dim D must be 128"): + op(q, k, v) + + def test_rejects_4d_kv(self): + op = _get_op() + q = torch.randn(1, 4, 32, 128, device="npu", dtype=torch.bfloat16) + k = torch.randn(1, 1, 32, 128, device="npu", dtype=torch.bfloat16) + v = torch.randn(1, 1, 32, 128, device="npu", dtype=torch.bfloat16) + with pytest.raises(ValueError, match="k/v 3-D"): + op(q, k, v) + + def test_rejects_kv_length_mismatch(self): + op = _get_op() + q = torch.randn(1, 4, 32, 128, device="npu", dtype=torch.bfloat16) + k = torch.randn(1, 32, 128, device="npu", dtype=torch.bfloat16) + v = torch.randn(1, 64, 128, device="npu", dtype=torch.bfloat16) + with pytest.raises(ValueError, match="key length mismatch"): + op(q, k, v) diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index ae4efa51..2f4ec864 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -48,6 +48,7 @@ def test_all_ws1_single_ops_are_registered(): "swiglu", "pack", "linear_logp", + "prefix_shared_attention", } <= names