feat(ascend): add prefix-shared attention Ascend C kernel

csrc/ascend/batch_invariant_logp_ascend.asc

@@ -307,10 +307,3 @@ std::vector<torch::Tensor> batch_invariant_logp_ascend_forward(torch::Tensor log
     }
     return {logp, lse};
 }
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
-{
-    m.def("batch_invariant_logp_ascend",
-          &batch_invariant_logp_ascend_forward,
-          "Batch-invariant selected-token log-probability (Ascend C forward)");
-}

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 <vector>
+
+#include <torch/extension.h>
+
+std::vector<torch::Tensor> batch_invariant_logp_ascend_forward(
+    torch::Tensor logits, torch::Tensor target, int64_t ignore_index);
+
+torch::Tensor 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)");
+}

rl_engine/_C_npu.pyi

@@ -8,3 +8,10 @@ def batch_invariant_logp_ascend(
     target: torch.Tensor,
     ignore_index: int,
 ) -> list[torch.Tensor]: ...
+
+
+def prefix_shared_attention_ascend(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+) -> torch.Tensor: ...

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,
         "cp_attention": _make_cp_attention_inputs,
         "logp": _make_logp_inputs,
         "linear_logp": _make_linear_logp_inputs,
@@ -59,6 +60,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}",
         "cp_attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}xcp2",
         "logp": f"{batch}x{seq}x{vocab}",
         "linear_logp": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}",
@@ -177,6 +180,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_cp_attention_inputs(
     args: argparse.Namespace, dtype: torch.dtype, device: torch.device
 ) -> dict[str, Any]:

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"
+            ),
+        },
+    ),
     "cp_attention": OperatorSpec(
         name="cp_attention",
         op_class="attention",
@@ -243,6 +263,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

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"]

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}")

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