[WS1][Ascend] [Qwen3-8b] LM head ops

csrc/ascend/npu_module.cpp

@@ -32,6 +32,11 @@ torch::Tensor embedding_ascend_forward(torch::Tensor token_ids,
 torch::Tensor fused_logp_ascend_forward(torch::Tensor logits,
                                         torch::Tensor target);
 
+torch::Tensor lm_head_ascend_forward(torch::Tensor hidden,
+                                     torch::Tensor weight,
+                                     torch::optional<torch::Tensor> bias,
+                                     bool output_fp32);
+
 torch::Tensor rmsnorm_ascend_forward(torch::Tensor x,
                                      torch::Tensor weight,
                                      torch::Tensor rstd);
@@ -78,4 +83,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
     m.def("fused_logp_ascend",
           &fused_logp_ascend_forward,
           "Batch-invariant fused selected-token log-probability (Ascend C forward)");
+    m.def("lm_head_ascend",
+          &lm_head_ascend_forward,
+          "Batch-invariant LM-head projection (Ascend C forward)");
 }

docs/operators/lm_head.md

@@ -37,6 +37,7 @@ The op exposes the WS1 dual-path contract:
 | --- | --- | --- | --- |
 | PyTorch fallback | `NativeLMHeadOp` | None | fp32 ground-truth reference; CPU and any GPU. |
 | CUDA SM90 (H200/Hopper) | `SM90LMHeadOp` | `_C.lm_head_sm90_forward` | Single-card batch-invariant forward backend; no Split-K; bf16 backward uses deterministic GEMM. |
+| Ascend NPU | `AscendLMHeadOp` | `_C_npu.lm_head_ascend` | Single-card batch-invariant forward backend: one output element per block, full K reduction in fp32 over a fixed tile order; fp32-formula VJP backward. |
 | ROCm / Triton | N/A | N/A | Falls back to the PyTorch native reference. |
 
 ## Tensor Contract
@@ -89,6 +90,21 @@ hidden row and vocab weight row independently, so large-vocab projections are ex
 to be memory-bandwidth bound compared with a tiled GEMM. This path exists to preserve a
 fixed hidden-dimension accumulation order for the WS1/H200 correctness gate.
 
+On `npu` the priority is:
+
+1. `ASCEND_LM_HEAD` — `AscendLMHeadOp` (batch-invariant Ascend C forward; fp32/bf16/fp16).
+2. `PYTORCH_NATIVE_LM_HEAD` — `NativeLMHeadOp` (fallback).
+
+The Ascend kernel implements the same structure as the SM90 CUDA kernel: one output
+element per block iteration, the full hidden-dimension reduction inside that block over
+a fixed tile order (products -> per-tile sum -> sequential scalar accumulation), bias
+added in fp32, final cast to the output dtype. There is no Split-K, so a row's logits
+depend only on H — never on N or block assignment — and are bitwise identical across
+batch sizes, row positions, and block assignments on the NPU. The per-tile sums use the
+Ascend vector unit's fixed hardware tree instead of CUDA's warp-shuffle tree, so the
+comparison against the PyTorch reference (torch.mv) is tolerance-based per the
+reduction contract, not bitwise.
+
 For bf16 H200 training, `SM90LMHeadOp.backward` routes `dhidden` through
 `_C.det_gemm_da` and `dweight` through `_C.det_gemm_db` (`hidden.T @ dlogits`, transposed
 back to the HF `[vocab, hidden]` layout). The wrapper fails fast if those deterministic
@@ -97,18 +113,25 @@ GEMM symbols are missing instead of silently falling back to cuBLAS for bf16 gra
 ## Tests
 
 ```bash
-python -m pytest tests/test_lm_head.py -v
+python -m pytest tests/test_lm_head.py tests/test_lm_head_ascend.py -v
 ```
 
 Covers fp32 correctness vs the fixed-K reference, precision-context safety, bf16/fp16
 accuracy, output shape, bias semantics, Axis-A batch invariance, input purity, gradient
 flow to `hidden` and `weight`, registry dispatch, and a GPU-only smoke test at the real
-Qwen3-8B dimensions.
+Qwen3-8B dimensions. The Ascend suite adds: contract-tolerance correctness vs the
+reference (fp32/bf16/fp16, with and without bias), fp32-formula VJP backward, bitwise
+batch invariance (batch sizes 1 vs {2,4,16,300}, row positions, multi-tile H=10000,
+repeated runs), and NPU registry dispatch.
 
 ## Implementation Files
 
 - `rl_engine/kernels/ops/pytorch/linear/lm_head.py`
 - `rl_engine/kernels/ops/cuda/linear/lm_head.py`
 - `csrc/cuda/embedding_lm_head_sm90.cu`
+- `rl_engine/kernels/ops/ascend/linear/lm_head.py` — Ascend deterministic op
+- `csrc/ascend/lm_head_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_lm_head.py`
+- `tests/test_lm_head_ascend.py`

rl_engine/_C_npu.pyi

@@ -63,3 +63,9 @@ def fused_logp_ascend(
     logits: torch.Tensor,
     target: torch.Tensor,
 ) -> torch.Tensor: ...
+def lm_head_ascend(
+    hidden: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    output_fp32: bool,
+) -> torch.Tensor: ...

rl_engine/kernels/gtest/operator_specs.py

@@ -166,6 +166,7 @@ def _load_object(path: str) -> Any:
             "pytorch": "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp",
             "triton": "rl_engine.kernels.ops.triton.linear.lm_head.TritonLMHeadOp",
             "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp",
+            "ascend": "rl_engine.kernels.ops.ascend.linear.lm_head.AscendLMHeadOp",
         },
         grad_input_names=("hidden", "weight"),
     ),

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

@@ -2,3 +2,4 @@
 # Copyright (c) 2026 RL-Kernel Contributors
 
 from . import embedding  # noqa: F401
+from . import lm_head  # noqa: F401

rl_engine/kernels/ops/ascend/linear/lm_head.py

@@ -0,0 +1,159 @@
+# 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.kernels.ops.backward_runtime import record_backward
+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 _AscendLMHeadFunction(torch.autograd.Function):
+    """Autograd bridge for the Ascend batch-invariant LM-head forward.
+
+    The VJP is the standard linear formula computed in fp32 on the NPU
+    (``grad_hidden = grad @ W``, ``grad_weight = grad^T @ H``, bias = the
+    fixed-order row sum), then cast back to the input dtypes. The CUDA op
+    uses its declared deterministic GEMM for the backward; on NPU the
+    plain torch matmuls are the deterministic-in-practice equivalent, and
+    the gtest compares gradients at the reduction contract tolerance.
+    """
+
+    @staticmethod
+    def forward(ctx, hidden, weight, bias, output_fp32: bool):
+        bias_to_save = (
+            bias if bias is not None else torch.empty(0, device=hidden.device, dtype=hidden.dtype)
+        )
+        ctx.save_for_backward(hidden, weight, bias_to_save)
+        ctx.has_bias = bias is not None
+        ctx.output_fp32 = bool(output_fp32)
+        return _C_npu.lm_head_ascend(hidden, weight.contiguous(), bias, bool(output_fp32))
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor):
+        hidden, weight, bias = ctx.saved_tensors
+        grad_2d = grad_output.reshape(-1, weight.size(0)).contiguous().float()
+        hidden_2d = hidden.reshape(-1, hidden.size(-1)).contiguous().float()
+        weight_f = weight.contiguous().float()
+        grad_hidden = grad_weight = grad_bias = None
+        if ctx.needs_input_grad[0]:
+            grad_hidden = grad_2d @ weight_f
+            grad_hidden = grad_hidden.reshape_as(hidden).to(hidden.dtype)
+        if ctx.needs_input_grad[1]:
+            grad_weight = (grad_2d.t() @ hidden_2d).to(weight.dtype)
+        if ctx.has_bias and ctx.needs_input_grad[2]:
+            rows = grad_output.reshape(-1, weight.size(0)).float()
+            acc = torch.zeros((rows.shape[1],), device=rows.device, dtype=torch.float32)
+            for index in range(rows.shape[0]):
+                acc = acc + rows[index]
+            grad_bias = acc.to(bias.dtype)
+        record_backward(
+            "lm_head",
+            kernel_id="rl_engine.kernels.ops.ascend.linear.lm_head._AscendLMHeadFunction",
+            impl="ascend_fp32_matmul_vjp",
+            family="ascend",
+        )
+        return grad_hidden, grad_weight, grad_bias, None
+
+
+class AscendLMHeadOp:
+    """Single-card batch-invariant Ascend LM-head op.
+
+    The Ascend C forward mirrors the SM90 CUDA kernel's structure: one
+    output element per block iteration, full hidden-dimension fp32 reduction
+    inside that block over a fixed tile order, bias added in fp32, final
+    cast to the output dtype. There is no Split-K and no algorithm
+    selection, so a row's logits do not depend on batch layout.
+    """
+
+    op_class = "reduction"
+    is_batch_invariant = True
+    backward_impl = "ascend_fp32_matmul_vjp"
+
+    def __init__(self) -> None:
+        if not _NPU_EXT_AVAILABLE or not hasattr(_C_npu, "lm_head_ascend"):
+            raise RuntimeError(
+                "lm_head_ascend is not compiled into the extension. "
+                "Rebuild on an Ascend NPU host with KERNEL_ALIGN_FORCE_ASCEND=1."
+            )
+        logger.info("Successfully linked to precompiled _C_npu.lm_head_ascend kernel.")
+
+    def __call__(
+        self,
+        hidden: torch.Tensor,
+        weight: torch.Tensor,
+        *,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return self.forward(hidden, weight, bias=bias)
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        weight: torch.Tensor,
+        *,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if not self._can_use_ascend(hidden, weight, bias):
+            raise RuntimeError(
+                "AscendLMHeadOp requires Ascend NPU bf16/fp16/fp32 inputs; "
+                "Native/Triton fallback is forbidden"
+            )
+        return _AscendLMHeadFunction.apply(hidden, weight, bias, False)
+
+    def forward_fp32(
+        self,
+        hidden: torch.Tensor,
+        weight: torch.Tensor,
+        *,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if not self._can_use_ascend(hidden, weight, bias):
+            raise RuntimeError(
+                "AscendLMHeadOp requires Ascend NPU bf16/fp16/fp32 inputs; "
+                "Native/Triton fallback is forbidden"
+            )
+        return _AscendLMHeadFunction.apply(hidden, weight, bias, True)
+
+    def parameter_vjp_contributions_fp32(self, *, hidden, weight, grad_output, bias=None):
+        del weight, bias
+        rows_h = hidden.reshape(-1, hidden.size(-1)).float()
+        rows_g = grad_output.reshape(-1, grad_output.size(-1)).float()
+        return {"weight": rows_g[:, :, None] * rows_h[:, None, :]}
+
+    @staticmethod
+    def _can_use_ascend(
+        hidden: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor],
+    ) -> bool:
+        bias_ok = bias is None or (
+            bias.device.type == "npu"
+            and bias.device == hidden.device
+            and bias.dim() == 1
+            and bias.dtype in _SUPPORTED_DTYPES
+        )
+        return (
+            hidden.device.type == "npu"
+            and weight.device.type == "npu"
+            and hidden.device == weight.device
+            and hidden.dim() >= 2
+            and weight.dim() == 2
+            and hidden.size(-1) == weight.size(1)
+            and hidden.dtype in _SUPPORTED_DTYPES
+            and weight.dtype in _SUPPORTED_DTYPES
+            and bias_ok
+        )

rl_engine/kernels/registry.py

@@ -128,6 +128,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta):
     ASCEND_RMS_NORM = "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp"
     ASCEND_EMBEDDING = "rl_engine.kernels.ops.ascend.linear.embedding.AscendEmbeddingOp"
     ASCEND_FUSED_LOGP = "rl_engine.kernels.ops.ascend.loss.logp.FusedLogpAscendOp"
+    ASCEND_LM_HEAD = "rl_engine.kernels.ops.ascend.linear.lm_head.AscendLMHeadOp"
     # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3)
     PYTORCH_VOCAB_PARALLEL_LOGP = (
         "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp"
@@ -732,6 +733,10 @@ def __init__(self):
             OpBackend.ASCEND_FUSED_LOGP,
             OpBackend.PYTORCH_NATIVE,
         ]
+        self._priority_map["npu"]["lm_head"] = [
+            OpBackend.ASCEND_LM_HEAD,
+            OpBackend.PYTORCH_NATIVE_LM_HEAD,
+        ]
         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

@@ -182,6 +182,10 @@ def fake_load_backend(backend):
         OpBackend.ASCEND_FUSED_LOGP,
         OpBackend.PYTORCH_NATIVE,
     ]
+    assert registry._priority_map["npu"]["lm_head"] == [
+        OpBackend.ASCEND_LM_HEAD,
+        OpBackend.PYTORCH_NATIVE_LM_HEAD,
+    ]
 
 
 def test_npu_available_handles_runtime_failure(monkeypatch):