feat(ws1): NativeLMHeadOp pure-PyTorch ground-truth reference + numerical contract tests

docs/.nav.yml

@@ -15,6 +15,7 @@ nav:
       - operators/fused-logp.md
       - operators/linear-logp.md
       - operators/grpo-loss.md
+      - operators/lm_head.md
       - operators/ratio-kl.md
       - operators/sampling.md
       - operators/embedding.md

docs/operators/README.md

@@ -22,6 +22,7 @@ Every operator page should include:
 - [Fused LogP](fused-logp.md)
 - [Fused Linear LogP](linear-logp.md)
 - [GRPO Loss](grpo-loss.md)
+- [LM Head](lm_head.md)
 - [Policy Ratio + KL Penalty](ratio-kl.md)
 - [Sampling](sampling.md)
 - [Token Embedding](embedding.md)

docs/operators/lm_head.md

@@ -0,0 +1,127 @@
+# LM Head
+
+The lm_head operator projects hidden states back to vocabulary logits — the final
+layer of the Qwen3/Llama stack. It is a **WS1 ground-truth reference** (issue #108):
+a pure-PyTorch definition of the "correct answer" that downstream fused CUDA/Triton
+kernels are validated against.
+
+- **LM Head** (`NativeLMHeadOp`): `out = hidden @ weight.t() (+ bias)`.
+
+For Qwen3-8B the weight is the output projection `[vocab=151936, hidden=4096]` in the
+HF `nn.Linear` `[out, in]` convention, so it is transposed internally. It is
+**independent** from the embedding table (`tie_word_embeddings=false`) — the two
+weights are not shared — and Qwen3 has **no bias** (`bias=None`).
+
+## Entry Point
+```python
+from rl_engine.kernels.registry import kernel_registry
+
+lm_head = kernel_registry.get_op("lm_head")
+
+logits = lm_head(hidden, weight)          # [B, S, hidden], [vocab, hidden] -> [B, S, vocab]
+logits = lm_head(hidden, weight, bias=b)  # optional [vocab] bias
+```
+
+The op exposes the WS1 dual-path contract:
+
+- `forward(...)` — projects in the input dtype, returns the input dtype (Axis-B accuracy
+  candidate / dtype-behavior path).
+- `forward_fp32(...)` — upcasts to fp32, accumulates in fp32, returns fp32 (the
+  ground-truth golden path). The matmul runs with autocast disabled and CUDA TF32
+  turned off, so it stays a true fp32 reference regardless of the caller's ambient
+  precision context (the global `allow_tf32` flag is saved and restored around it).
+
+## Backends
+
+| Backend | Wrapper | Native symbol | Status |
+| --- | --- | --- | --- |
+| PyTorch fallback | `NativeLMHeadOp` | None | fp32 ground-truth reference; CPU and any GPU. |
+| CUDA / ROCm / Triton | — | — | Planned: downstream fused kernels validate against this reference. |
+
+## Tensor Contract
+
+| Argument | Shape | Dtype | Requirements |
+| --- | --- | --- | --- |
+| `hidden` | `[B, S, hidden]` (any leading dims) | float (fp16/bf16/fp32) | Hidden states (Qwen3-8B `hidden=4096`). |
+| `weight` | `[vocab, hidden]` | float | Output projection in HF `[out, in]` layout; transposed internally. Qwen3-8B `[151936, 4096]`. |
+| `bias` | `[vocab]` or `None` | float | Optional; Qwen3 has none (`None`). |
+| output | `hidden.shape[:-1] + (vocab,)` | `forward`: hidden dtype · `forward_fp32`: float32 | Logits. |
+
+Output dtype follows `hidden`. Pure function — no randomness, no in-place mutation,
+device/dtype follow the inputs.
+
+> **Difference from the bare `matmul` op**: lm_head takes the weight in HF `[out, in]`
+> layout and transposes it internally (`weight.t()`); the `matmul` op computes a bare
+> `a @ b` with no transpose. Do not use them interchangeably.
+
+## Dispatch Behavior
+
+`kernel_registry.get_op("lm_head")` resolves through the `OpBackend` priority map. On
+`cuda` / `rocm` / `cpu` the only registered backend today is the PyTorch native op
+(`PYTORCH_NATIVE_LM_HEAD`), so every device dispatches to the fp32 reference. When fused
+kernels land, they are prepended to the priority list and the native op becomes the fallback.
+
+## Accuracy
+
+Reference semantics (`forward_fp32`):
+
+```python
+out = hidden.float() @ weight.float().t()
+if bias is not None:
+    out = out + bias.float()
+```
+
+- **Ground truth**: `forward_fp32` accumulates in and returns fp32, with autocast and
+  CUDA TF32 disabled so it is a true fp32 reference even if the caller has TF32 or
+  autocast enabled.
+- **Dtype path**: `forward` runs the projection in the input dtype. Because this is a
+  reduction over `hidden`, low-precision accumulation **drifts** from the fp32 reference
+  (unlike the lossless embedding gather). Unlike `forward_fp32`, this path intentionally
+  follows the ambient precision context (it is the dtype-behavior path): with an fp32
+  input it is bitwise-equal to the ground truth **when ambient TF32/autocast is off**, but
+  on a TF32-enabled GPU it tracks real hardware behavior and may drift. bf16/fp16 are
+  always checked with a tolerance.
+- **Axis-B — accuracy tolerance**: measured as max absolute error relative to the output
+  peak magnitude. On a SMALL load point with the real `hidden=4096` reduction length,
+  bf16 drifts ~0.3–0.4% of peak and fp16 ~0.05%. Elementwise `rtol` is not used: many
+  logits are near zero while the accumulated error tracks the reduction length, not the
+  output value.
+- **Axis-A — batch invariance**: a row's logits are independent of the rest of the batch,
+  so the output is bitwise-identical regardless of batch size or padding (`torch.equal`,
+  `atol=0`) — **provided the reduction order is fixed**. Multi-threaded CPU GEMM splits
+  the `hidden` reduction across threads by the `M = batch*seq` dimension, which silently
+  breaks bitwise batch invariance for large `hidden`; the tests pin a single thread to fix
+  the order. On GPU, cuBLAS likewise splits K by `M`, so a bitwise batch-invariant GEMM is
+  a downstream kernel concern, not a free property of `torch.matmul`.
+
+## Performance Notes
+
+Reference operator — no fused kernel or benchmark yet. Downstream fused kernels carry their
+own benchmarks and are measured against this reference for correctness.
+
+## Tests
+
+```bash
+python -m pytest tests/test_lm_head.py -v
+```
+
+Covers: fp32 correctness vs naive matmul (bitwise, with ambient TF32 pinned off),
+`forward_fp32` precision-context safety (true fp32 under ambient autocast + restores the
+global TF32 flag on CPU; numerically beats a TF32 matmul on GPU), bf16/fp16 dtype-path
+accuracy (relative-to-peak tolerance, with `bias`), output shape, bias semantics, Axis-A
+batch invariance (slice + padding, single-thread reduction, all dtypes), input purity,
+gradient flow to `hidden`/`weight` (closed-form check), registry dispatch, and a GPU-only
+smoke test at the real Qwen3-8B dims (`vocab=151936, hidden=4096`) that skips when CUDA or
+GPU memory is unavailable.
+
+## Implementation Files
+
+- `rl_engine/kernels/ops/pytorch/linear/lm_head.py`
+- `rl_engine/kernels/registry.py`
+- `tests/test_lm_head.py`
+
+## Known Limitations
+
+- PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work).
+- Axis-A bitwise batch invariance holds only with a fixed reduction order (single-thread on
+  CPU); a batch-invariant GEMM on GPU is a downstream concern.

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

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Optional
+
+import torch
+
+
+class NativeLMHeadOp:
+    """
+    Pure PyTorch native language-model-head reference.
+    out = hidden @ weight.t() (+ bias)
+
+    Projects hidden states back to vocabulary logits -- the final layer of the
+    Qwen3/Llama stack. For Qwen3-8B the weight is the output projection
+    ``[vocab=151936, hidden=4096]`` in the HF ``nn.Linear`` ``[out, in]``
+    convention, so it is transposed internally (``weight.t()``). This is the
+    one difference from the bare ``matmul`` op (which computes ``a @ b`` with no
+    transpose) -- do not use them interchangeably. The lm_head weight is
+    *independent* from the embedding table (``tie_word_embeddings=false``), and
+    Qwen3 has no bias (pass ``bias=None``).
+
+    Unlike embedding (a lossless row gather), this is a reduction over the
+    ``hidden`` dimension: the low-precision ``forward`` path accumulates in the
+    input dtype and therefore drifts from the fp32 ``forward_fp32`` ground
+    truth. Axis-B accuracy uses a tolerance (``torch.allclose``), not bitwise
+    equality. Axis-A batch invariance still holds bitwise within a single dtype
+    (each output row reduces over ``hidden`` independently of the batch).
+    """
+
+    def __init__(self) -> None:
+        """No state; the op is a pure function over (hidden, weight, bias)."""
+
+    def __call__(
+        self,
+        hidden: torch.Tensor,
+        weight: torch.Tensor,
+        *,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Alias for ``forward`` so the op is callable like a module."""
+        return self.forward(hidden, weight, bias=bias)
+
+    def forward(
+        self,
+        hidden: torch.Tensor,
+        weight: torch.Tensor,
+        *,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Canonical entry: project in the input dtype, output the input dtype.
+        This is the dtype-behavior path used as the Axis-B accuracy candidate.
+        """
+        return self._lm_head(
+            hidden, weight, bias, compute_dtype=hidden.dtype, output_dtype=hidden.dtype
+        )
+
+    def forward_fp32(
+        self,
+        hidden: torch.Tensor,
+        weight: torch.Tensor,
+        *,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Ground truth: upcast to fp32, accumulate in fp32, force fp32 output.
+
+        The matmul is wrapped to disable autocast and TF32 so this stays a true
+        fp32 reference regardless of the caller's ambient precision context.
+        """
+        return self._lm_head(
+            hidden,
+            weight,
+            bias,
+            compute_dtype=torch.float32,
+            output_dtype=torch.float32,
+            strict_fp32=True,
+        )
+
+    # ------------------------------------------------------------------ #
+    # Helpers
+    # ------------------------------------------------------------------ #
+    @staticmethod
+    def _lm_head(
+        hidden: torch.Tensor,
+        weight: torch.Tensor,
+        bias: Optional[torch.Tensor],
+        *,
+        compute_dtype: torch.dtype,
+        output_dtype: torch.dtype,
+        strict_fp32: bool = False,
+    ) -> torch.Tensor:
+        """Core matmul: cast to ``compute_dtype``, project, optionally add bias, cast out.
+
+        When ``strict_fp32`` is set, the matmul runs with autocast disabled and
+        CUDA TF32 turned off so the fp32 reference path is not silently downcast
+        by the caller's ambient autocast/TF32 settings.
+        """
+        h = hidden.to(compute_dtype)
+        w = weight.to(compute_dtype)
+        # [..., hidden] @ [hidden, vocab] -> [..., vocab]; weight is [vocab, hidden] (HF [out, in]).
+        if strict_fp32:
+            with NativeLMHeadOp._strict_fp32_matmul(h.device.type):
+                out = h @ w.t()
+        else:
+            out = h @ w.t()
+        if bias is not None:
+            out = out + bias.to(compute_dtype)
+        return out.to(output_dtype)
+
+    @staticmethod
+    @contextmanager
+    def _strict_fp32_matmul(device_type: str):
+        """Disable autocast and TF32 for a true fp32 matmul, restoring state after."""
+        prev_tf32 = torch.backends.cuda.matmul.allow_tf32
+        torch.backends.cuda.matmul.allow_tf32 = False
+        try:
+            with torch.autocast(device_type=device_type, enabled=False):
+                yield
+        finally:
+            torch.backends.cuda.matmul.allow_tf32 = prev_tf32

rl_engine/kernels/registry.py

@@ -56,6 +56,8 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta):
     PYTORCH_NATIVE_SILU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp"
     PYTORCH_NATIVE_SWIGLU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp"
 
+    # WS1 pure-PyTorch ground-truth linear ops
+    PYTORCH_NATIVE_LM_HEAD = "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp"
     # WS1 pure-PyTorch ground-truth embedding ops
     PYTORCH_NATIVE_EMBEDDING = "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp"
 
@@ -94,6 +96,7 @@ def __init__(self):
                 "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS],
                 "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP],
                 "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL],
+                "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD],
                 "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING],
                 "silu": [OpBackend.PYTORCH_NATIVE_SILU],
                 "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU],
@@ -109,6 +112,7 @@ def __init__(self):
                 "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS],
                 "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP],
                 "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL],
+                "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD],
                 "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING],
                 "silu": [OpBackend.PYTORCH_NATIVE_SILU],
                 "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU],
@@ -119,6 +123,7 @@ def __init__(self):
                 "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS],
                 "linear_logp": [OpBackend.PYTORCH_LINEAR_LOGP],
                 "ratio_kl": [OpBackend.PYTORCH_RATIO_KL],
+                "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD],
                 "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING],
                 "silu": [OpBackend.PYTORCH_NATIVE_SILU],
                 "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU],