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

csrc/ascend/batch_invariant_logp_ascend.asc

@@ -308,9 +308,5 @@ std::vector<torch::Tensor> batch_invariant_logp_ascend_forward(torch::Tensor log
     return {logp, lse};
 }
 
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
-{
-    m.def("batch_invariant_logp_ascend",
-          &batch_invariant_logp_ascend_forward,
-          "Batch-invariant selected-token log-probability (Ascend C forward)");
-}
+// The PYBIND11_MODULE for rl_engine._C_npu lives in npu_module.cpp so that
+// every Ascend op shares one compiled module.

csrc/ascend/npu_module.cpp

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

docs/operators/linear-logp.md

@@ -39,6 +39,7 @@ logp.sum().backward()  # gradients flow into hidden, lm_head_weight, bias
 | --- | --- | --- |
 | CUDA SM90 (Hopper) | `FusedLinearLogpSM90Op` | TMA-streamed, Double Buffering, tensor-core forward (`mma.sync.m16n8k16`), online softmax in smem; chunked backward. Compiles for `sm_90a`; validated fp32-accurate on H100. Falls back to Triton/native for fp32/fp16 inputs or hidden dims not divisible by 32. |
 | CUDA / ROCm (Triton) | `TritonLinearLogpOp` | Triton online-softmax forward; Liger-style chunked backward (cuBLAS matmuls, deterministic). Phase 1. |
+| Ascend NPU | `FusedLinearLogpAscendOp` | Batch-invariant Ascend C forward mirroring the SM90 reduction contract: ascending vocab-row scan with the online rescale chain, per-row fp32 dots over a fixed D-tile order, `min(zt - lse, 0)` clamp; the shared chunked backward. Output fp32. |
 | PyTorch native | `NativeLinearLogpOp` | Naive `F.linear` + `log_softmax` + `gather` reference; CPU / Triton-less fallback. |
 
 The SM90 backend (`csrc/cuda/fused_linear_logp_sm90.cu`) streams hidden/weight
@@ -168,10 +169,14 @@ For 4-GPU tensor-parallel validation, use
 ## Implementation Files
 
 - `rl_engine/kernels/ops/triton/loss/linear_logp.py`
-- `rl_engine/kernels/ops/pytorch/loss/linear_logp.py`
-- `rl_engine/kernels/ops/cuda/loss/linear_logp.py` (SM90 wrapper + chunked backward)
-- `csrc/cuda/fused_linear_logp_sm90.cu`, `csrc/ops.cpp`, `setup.py` (SM90 kernel + build)
+- `rl_engine/kernels/ops/pytorch/loss/linear_logp.py` — native reference, chunked backward, TP helpers
+- `rl_engine/kernels/ops/cuda/loss/linear_logp.py` — CUDA fused implementation (SM90 wrapper + chunked backward)
+- `rl_engine/kernels/ops/ascend/loss/linear_logp.py` — Ascend deterministic op
+- `csrc/cuda/fused_linear_logp_sm90.cu`, `csrc/ops.cpp`, `setup.py` — SM90 kernel + build
+- `csrc/ascend/fused_linear_logp_ascend.asc` — Ascend C forward kernel
+- `csrc/ascend/npu_module.cpp` — shared pybind entry for `rl_engine._C_npu`
 - `rl_engine/kernels/registry.py`
 - `tests/test_linear_logp.py`
+- `tests/test_linear_logp_ascend.py` — Ascend correctness + batch-invariance tests
 - `benchmarks/benchmark_linear_logp.py`
 - `docs/design/fused-linear-logp.md`

rl_engine/_C_npu.pyi

@@ -8,3 +8,9 @@ def batch_invariant_logp_ascend(
     target: torch.Tensor,
     ignore_index: int,
 ) -> list[torch.Tensor]: ...
+def fused_linear_logp_ascend(
+    hidden: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    target: torch.Tensor,
+) -> torch.Tensor: ...

rl_engine/kernels/gtest/operator_specs.py

@@ -115,6 +115,7 @@ def _load_object(path: str) -> Any:
             "pytorch": "rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp",
             "triton": "rl_engine.kernels.ops.triton.loss.linear_logp.TritonLinearLogpOp",
             "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.linear_logp.FusedLinearLogpSM90Op",
+            "ascend": "rl_engine.kernels.ops.ascend.loss.linear_logp.FusedLinearLogpAscendOp",
         },
         grad_input_names=("hidden", "lm_head_weight"),
     ),

rl_engine/kernels/ops/ascend/loss/__init__.py

@@ -1,2 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright (c) 2026 RL-Kernel Contributors
+
+from . import batch_invariant_logp  # noqa: F401
+from . import linear_logp  # noqa: F401

rl_engine/kernels/ops/ascend/loss/linear_logp.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+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
+
+_SUPPORTED_DTYPES = {torch.float32, torch.float16, torch.bfloat16}
+
+
+class _FusedLinearLogpAscendFunction(torch.autograd.Function):
+    """Autograd bridge for the Ascend fused linear log-prob forward.
+
+    The backward is the shared Liger-style chunked formula from
+    ``rl_engine.kernels.ops.pytorch.loss.linear_logp.chunked_linear_logp_backward``
+    (the same formula the CUDA SM90 op falls back to), so gradients follow
+    the CUDA op's portable backward exactly.
+    """
+
+    @staticmethod
+    def forward(ctx, hidden, lm_head_weight, target_ids):
+        hidden_2d = hidden.reshape(-1, hidden.size(-1)).contiguous()
+        weight = lm_head_weight.contiguous()
+        target_1d = (
+            target_ids.reshape(-1).to(device=hidden_2d.device, dtype=torch.long).contiguous()
+        )
+        output = _C_npu.fused_linear_logp_ascend(hidden_2d, weight, None, target_1d)
+        ctx.save_for_backward(hidden_2d, weight, target_1d)
+        ctx.lead_shape = hidden.shape[:-1]
+        ctx.hidden_dtype = hidden.dtype
+        ctx.weight_dtype = lm_head_weight.dtype
+        return output.reshape(hidden.shape[:-1])
+
+    @staticmethod
+    def backward(ctx, grad_logp):
+        from rl_engine.kernels.ops.pytorch.loss.linear_logp import chunked_linear_logp_backward
+
+        hidden_2d, weight, target_1d = ctx.saved_tensors
+        grad_hidden, grad_weight, _ = chunked_linear_logp_backward(
+            grad_logp,
+            hidden_2d,
+            weight,
+            target_1d,
+            hidden_2d,  # bias placeholder; has_bias=False
+            has_bias=False,
+            lead_shape=ctx.lead_shape,
+            hidden_dtype=ctx.hidden_dtype,
+            weight_dtype=ctx.weight_dtype,
+            bias_dtype=None,
+        )
+        return grad_hidden, grad_weight, None
+
+
+class FusedLinearLogpAscendOp:
+    """Batch-invariant fused linear log-prob for Ascend NPU.
+
+    Computes ``log_softmax(hidden @ W^T + b)[target]`` without materializing
+    the ``[N, V]`` logits. The Ascend C forward mirrors the SM90 kernel's
+    reduction contract: ascending vocab-row scan with the online rescale
+    chain, per-row fp32 dots over a fixed D-tile order, final
+    ``min(zt - lse, 0)`` clamp; the output is fp32.
+    """
+
+    is_fused_logp = True
+    is_batch_invariant = True
+
+    def __init__(self) -> None:
+        if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "fused_linear_logp_ascend"):
+            raise RuntimeError(
+                "fused_linear_logp_ascend is not compiled into the extension. Rebuild with "
+                "KERNEL_ALIGN_FORCE_ASCEND=1 on an Ascend NPU host: 'pip install -e .'"
+            )
+        self.op = _C_npu.fused_linear_logp_ascend
+        logger.info("Successfully linked to precompiled _C_npu.fused_linear_logp_ascend kernel.")
+
+    def __call__(
+        self,
+        hidden: torch.Tensor,
+        lm_head_weight: torch.Tensor,
+        target_ids: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        *,
+        tp_group: Any = None,
+        vocab_start_index: int = 0,
+        global_vocab_size: Optional[int] = None,
+    ) -> torch.Tensor:
+        return self.apply(
+            hidden,
+            lm_head_weight,
+            target_ids,
+            bias,
+            tp_group=tp_group,
+            vocab_start_index=vocab_start_index,
+            global_vocab_size=global_vocab_size,
+        )
+
+    def apply(
+        self,
+        hidden: torch.Tensor,
+        lm_head_weight: torch.Tensor,
+        target_ids: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+        *,
+        tp_group: Any = None,
+        vocab_start_index: int = 0,
+        global_vocab_size: Optional[int] = None,
+    ) -> torch.Tensor:
+        from rl_engine.kernels.ops.pytorch.loss.linear_logp import (
+            NativeLinearLogpOp,
+            should_use_tensor_parallel_linear_logp,
+        )
+
+        if lm_head_weight.size(-1) != hidden.size(-1):
+            raise ValueError(
+                f"hidden dim {hidden.size(-1)} must match lm_head_weight dim "
+                f"{lm_head_weight.size(-1)}"
+            )
+        if lm_head_weight.device != hidden.device:
+            raise ValueError(
+                f"lm_head_weight device {lm_head_weight.device} must match hidden "
+                f"device {hidden.device}"
+            )
+        if hidden.shape[:-1] != target_ids.shape:
+            raise ValueError(
+                f"hidden leading shape {tuple(hidden.shape[:-1])} must match "
+                f"target_ids shape {tuple(target_ids.shape)}"
+            )
+        # Tensor-parallel and bias paths are not covered by the Ascend forward;
+        # delegate to the native reference (same fallback as the CUDA op).
+        if (
+            should_use_tensor_parallel_linear_logp(
+                tp_group,
+                int(vocab_start_index),
+                global_vocab_size,
+                lm_head_weight.size(0),
+            )
+            or bias is not None
+        ):
+            return NativeLinearLogpOp().apply(
+                hidden,
+                lm_head_weight,
+                target_ids,
+                bias,
+                tp_group=tp_group,
+                vocab_start_index=vocab_start_index,
+                global_vocab_size=global_vocab_size,
+            )
+        if not self._ascend_supported(hidden, lm_head_weight):
+            return NativeLinearLogpOp().apply(hidden, lm_head_weight, target_ids)
+        return _FusedLinearLogpAscendFunction.apply(hidden, lm_head_weight, target_ids)
+
+    @staticmethod
+    def _ascend_supported(hidden: torch.Tensor, lm_head_weight: torch.Tensor) -> bool:
+        return (
+            hidden.device.type == "npu"
+            and lm_head_weight.device.type == "npu"
+            and hidden.is_contiguous()
+            and lm_head_weight.is_contiguous()
+            and hidden.dtype in _SUPPORTED_DTYPES
+            and lm_head_weight.dtype == hidden.dtype
+        )

rl_engine/kernels/registry.py

@@ -107,6 +107,9 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta):
     ASCEND_BATCH_INVARIANT_LOGP = (
         "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAscendOp"
     )
+    ASCEND_FUSED_LINEAR_LOGP = (
+        "rl_engine.kernels.ops.ascend.loss.linear_logp.FusedLinearLogpAscendOp"
+    )
     # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3)
     PYTORCH_VOCAB_PARALLEL_LOGP = (
         "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp"
@@ -622,6 +625,10 @@ def __init__(self):
             OpBackend.ASCEND_BATCH_INVARIANT_LOGP,
             OpBackend.PYTORCH_BATCH_INVARIANT_LOGP,
         ]
+        self._priority_map["npu"]["linear_logp"] = [
+            OpBackend.ASCEND_FUSED_LINEAR_LOGP,
+            OpBackend.PYTORCH_LINEAR_LOGP,
+        ]
         logger.info(f"KernelRegistry initialized for {device_ctx.device_type}")
         self._adjust_priority_for_hardware()
         self._adjust_priority_from_env()

rl_engine/tests/test_dispatch.py

@@ -166,6 +166,10 @@ def fake_load_backend(backend):
         OpBackend.ASCEND_BATCH_INVARIANT_LOGP,
         OpBackend.PYTORCH_BATCH_INVARIANT_LOGP,
     ]
+    assert registry._priority_map["npu"]["linear_logp"] == [
+        OpBackend.ASCEND_FUSED_LINEAR_LOGP,
+        OpBackend.PYTORCH_LINEAR_LOGP,
+    ]
 
 
 def test_npu_available_handles_runtime_failure(monkeypatch):

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