diff --git a/benchmarks/benchmark_attention.py b/benchmarks/benchmark_attention.py index c9509a19..99705a52 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", "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..7d26b89b --- /dev/null +++ b/csrc/ascend/attention/deterministic_attention_ascend.asc @@ -0,0 +1,526 @@ +// 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 + +#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, 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 { +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, 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_); + 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, 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 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, 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); + } + + // 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. + // + // 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, + 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 maskTile = maskBuf_.Get(); + + // 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(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; + } + 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; + bool anyValid = false; + 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); + if (tileMax > NEG_INF) { + anyValid = true; + } + 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); + 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); + 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 = LSE_INVALID; + 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, 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. + 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}; +} 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..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/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..bbaaad69 100644 --- a/rl_engine/_C_npu.pyi +++ b/rl_engine/_C_npu.pyi @@ -8,3 +8,11 @@ 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..f2cdca31 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/attention/__init__.py @@ -0,0 +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 + +__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..32fc9533 --- /dev/null +++ b/rl_engine/kernels/ops/ascend/attention/deterministic_attn.py @@ -0,0 +1,199 @@ +# 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 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 + +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 + # 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_fp32( + 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``. + Out-of-domain inputs are rejected up front (the registry-level + ``PYTORCH_NATIVE_ATTENTION`` entry covers unavailable-kernel fallback). + """ + + 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 e7114827..c97db639 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" @@ -290,9 +295,22 @@ def __init__(self): "silu": [OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], }, + # Ascend NPU: op types without an entry fall back to their CPU + # candidates (see the runtime override below), so only + # Ascend-accelerated ops are listed. + "npu": { + "batch_invariant_logp": [ + OpBackend.ASCEND_BATCH_INVARIANT_LOGP, + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + ], + "attention": [ + OpBackend.ASCEND_DETERMINISTIC_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], + }, } # Preserve the former CPU fallback behavior for every operator on NPU, - # then override only the operator with an Ascend-specific backend. + # then override only the operators with an Ascend-specific backend. self._priority_map["npu"] = { op_type: candidates.copy() for op_type, candidates in self._priority_map["cpu"].items() } @@ -300,6 +318,10 @@ def __init__(self): OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ] + self._priority_map["npu"]["attention"] = [ + OpBackend.ASCEND_DETERMINISTIC_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ] logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") self._adjust_priority_for_hardware() self._adjust_priority_from_env() diff --git a/scripts/check_operator.py b/scripts/check_operator.py index 9dbca48d..ccf18a28 100644 --- a/scripts/check_operator.py +++ b/scripts/check_operator.py @@ -35,9 +35,24 @@ def _parse_dtype(value: str) -> torch.dtype: raise ValueError(f"unsupported dtype: {value}") +def _npu_available() -> bool: + """torch.npu only exists after torch_npu is imported; probe defensively.""" + try: + import torch_npu # noqa: F401 + except ImportError: + return False + return hasattr(torch, "npu") and torch.npu.is_available() + + def _select_device(value: str) -> torch.device: if value == "auto": - return torch.device("cuda" if torch.cuda.is_available() else "cpu") + if torch.cuda.is_available(): + return torch.device("cuda") + if _npu_available(): + return torch.device("npu") + return torch.device("cpu") + if value == "npu" and not _npu_available(): + raise RuntimeError("--device npu was requested, but no Ascend NPU is available") device = torch.device(value) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("--device cuda was requested, but CUDA is not available") @@ -73,7 +88,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--candidate", default="pytorch", - help="Candidate backend to validate, for example pytorch, cuda, cuda-sm90, triton.", + help="Candidate backend to validate, for example pytorch, cuda, cuda-sm90, triton, " + "ascend.", ) parser.add_argument("--dtype", choices=("fp32", "bf16", "fp16"), default="fp32") parser.add_argument("--device", default="auto") diff --git a/setup.py b/setup.py index 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.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 new file mode 100644 index 00000000..086f59da --- /dev/null +++ b/tests/test_attention_ascend.py @@ -0,0 +1,289 @@ +# 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 ( + _NPU_EXT_AVAILABLE, + _C_npu, + ) + 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): + # 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 + + +# --------------------------------------------------------------------------- +# 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 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_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) + 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): + # 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 + 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 + # 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"