[WS1] Close embedding and LM-head invariance coverage

docs/operators/lm_head.md

@@ -1,18 +1,20 @@
 # 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
+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 for 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)`.
+- **LM Head** (`NativeLMHeadOp`): mathematically `out = hidden @ weight.t() (+ bias)`.
+  The native reference implements this as row-wise fixed-K GEMV projections so the
+  reference path is batch-invariant.
 
 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`).
+HF `nn.Linear` `[out, in]` convention. It is independent from the embedding table
+(`tie_word_embeddings=false`), and Qwen3 has no bias (`bias=None`).
 
 ## Entry Point
+
 ```python
 from rl_engine.kernels.registry import kernel_registry
 
@@ -24,104 +26,75 @@ 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).
+- `forward(...)` projects in the input dtype and returns the input dtype.
+- `forward_fp32(...)` upcasts to fp32, accumulates in fp32, and returns fp32. The
+  fixed-K projection runs with autocast disabled and CUDA TF32 turned off, so it
+  stays a true fp32 reference regardless of the caller's ambient precision context.
 
 ## 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. |
+| CUDA / ROCm / Triton | N/A | N/A | 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.
+| `hidden` | `[B, S, hidden]` or any leading dims | fp16/bf16/fp32 | Hidden states. |
+| `weight` | `[vocab, hidden]` | fp16/bf16/fp32 | Output projection in HF `[out, in]` layout. |
+| `bias` | `[vocab]` or `None` | fp16/bf16/fp32 | Optional; Qwen3 uses `None`. |
+| output | `hidden.shape[:-1] + (vocab,)` | `forward`: hidden dtype; `forward_fp32`: fp32 | Logits. |
 
-## 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.
+Output dtype follows `hidden`. The op is pure: no randomness and no in-place mutation.
 
 ## Accuracy
 
 Reference semantics (`forward_fp32`):
 
 ```python
-out = hidden.float() @ weight.float().t()
+flat_hidden = hidden.float().reshape(-1, hidden.size(-1))
+if flat_hidden.size(0) == 0:
+    flat_out = flat_hidden @ weight.float().t()
+else:
+    rows = [torch.mv(weight.float(), row) for row in flat_hidden]
+    flat_out = torch.stack(rows)
+out = flat_out.reshape(*hidden.shape[:-1], weight.size(0))
 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.
+  CUDA TF32 disabled.
 - **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.
+  reduction over `hidden`, low-precision accumulation drifts from the fp32 reference and
+  is checked with tolerance.
+- **Axis-A batch invariance**: a row's logits are bitwise-identical regardless of batch
+  size or padding. The native reference enforces this by flattening leading dimensions
+  and projecting each row through the same GEMV-shaped K reduction instead of relying on
+  batched GEMM, whose reduction tree can change with `M = batch * seq`.
+
+## Dispatch Behavior
+
+`kernel_registry.get_op("lm_head")` resolves through the `OpBackend` priority map. On
+`cuda`, `rocm`, and `cpu`, the only registered backend today is the PyTorch native op
+(`PYTORCH_NATIVE_LM_HEAD`), so every device dispatches to this reference. When fused
+kernels land, they should preserve the same batch-invariant contract.
 
 ## 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.
+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.
 
 ## 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/executors/deepspeed_trainer.py

@@ -246,7 +246,9 @@ def _next_published_weight_version(self, consumed_weight_version: int) -> int:
         self._latest_published_weight_version = published
         return published
 
-    def _export_zero3_full_state_model(self) -> tuple[torch.nn.Module, Mapping[str, Any]]:
+    def _export_zero3_full_state_model(
+        self,
+    ) -> tuple[torch.nn.Module, Mapping[str, Any]]:
         model = getattr(self.engine, "module", self.model)
         rank = self._engine_rank()
         if rank != 0:
@@ -439,7 +441,7 @@ def _linear_logp_op_for_device(device: torch.device | str) -> Any:
     resolved = torch.device(device)
     if resolved.type == "cpu":
         return NativeLinearLogpOp()
-    return kernel_registry.get_op("linear_logp")
+    return kernel_registry.get_op("linear_logp", device=resolved)
 
 
 def _unwrap_training_model(engine: Any, fallback_model: torch.nn.Module) -> torch.nn.Module:

rl_engine/kernels/gtest/operator_inputs.py

@@ -55,8 +55,8 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str:
         "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}",
         "silu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}",
         "swiglu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}",
-        "embedding": f"{batch}x{seq}x{vocab}x{DEFAULT_HIDDEN}",
-        "lm_head": f"{batch}x{seq}x{vocab}",
+        "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}",
+        "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}",
         "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}",
     }
     try:
@@ -169,9 +169,10 @@ def _make_embedding_inputs(
 ) -> dict[str, Any]:
     batch, seq = _batch_seq(args)
     vocab = _arg_int(args, "vocab", DEFAULT_VOCAB)
+    hidden_dim = _normalized_dim(args)
     return {
         "token_ids": _token_ids((batch, seq), vocab, args, device),
-        "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 0),
+        "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 0),
     }
 
 
@@ -180,9 +181,10 @@ def _make_lm_head_inputs(
 ) -> dict[str, Any]:
     batch, seq = _batch_seq(args)
     vocab = _arg_int(args, "vocab", DEFAULT_VOCAB)
+    hidden_dim = _normalized_dim(args)
     return {
-        "hidden": _floating_tensor((batch, seq, DEFAULT_HIDDEN), args, dtype, device, 0),
-        "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 1),
+        "hidden": _floating_tensor((batch, seq, hidden_dim), args, dtype, device, 0),
+        "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 1),
         "bias": None,
     }
 

rl_engine/kernels/gtest/operator_specs.py

@@ -57,6 +57,26 @@ def _load_object(path: str) -> Any:
         },
         grad_input_names=("hidden", "lm_head_weight"),
     ),
+    "embedding": OperatorSpec(
+        name="embedding",
+        op_class="elementwise",
+        gold_path="rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp",
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp",
+        },
+        grad_input_names=("weight",),
+    ),
+    "lm_head": OperatorSpec(
+        name="lm_head",
+        op_class="reduction",
+        gold_path="rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp",
+        gold_method="forward_fp32",
+        candidate_paths={
+            "pytorch": "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp",
+        },
+        grad_input_names=("hidden", "weight"),
+    ),
 }
 
 

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

@@ -28,7 +28,9 @@ class NativeLMHeadOp:
     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).
+    because each flattened output row is projected through the same GEMV-shaped
+    reduction over ``hidden``. The reduction path therefore depends on K and V,
+    not on how many other rows share the batch.
     """
 
     def __init__(self) -> None:
@@ -101,16 +103,25 @@ def _lm_head(
         """
         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()
+                out = NativeLMHeadOp._fixed_k_projection(h, w)
         else:
-            out = h @ w.t()
+            out = NativeLMHeadOp._fixed_k_projection(h, w)
         if bias is not None:
             out = out + bias.to(compute_dtype)
         return out.to(output_dtype)
 
+    @staticmethod
+    def _fixed_k_projection(hidden: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+        flat_hidden = hidden.reshape(-1, hidden.size(-1))
+        if flat_hidden.size(0) == 0:
+            flat_out = flat_hidden @ weight.t()
+            return flat_out.reshape(*hidden.shape[:-1], weight.size(0))
+        rows = [torch.mv(weight, row) for row in flat_hidden]
+        flat_out = torch.stack(rows, dim=0)
+        return flat_out.reshape(*hidden.shape[:-1], weight.size(0))
+
     @staticmethod
     @contextmanager
     def _strict_fp32_matmul(device_type: str):

rl_engine/kernels/registry.py

@@ -6,6 +6,8 @@
 from enum import Enum, EnumMeta
 from typing import Any, Dict, Optional, Set, Type
 
+import torch
+
 from rl_engine.platforms.device import device_ctx
 from rl_engine.utils.logger import logger
 
@@ -163,11 +165,18 @@ def __init__(self):
                     OpBackend.CUDA_DETERMINISTIC_LOGP,
                     OpBackend.PYTORCH_NATIVE,
                 ],
-                "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN],
+                "attn": [
+                    OpBackend.FLASH_ATTN,
+                    OpBackend.TRITON_GENERIC,
+                    OpBackend.PYTORCH_ATTN,
+                ],
                 "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION],
                 "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN],
                 "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS],
-                "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP],
+                "linear_logp": [
+                    OpBackend.TRITON_LINEAR_LOGP,
+                    OpBackend.PYTORCH_LINEAR_LOGP,
+                ],
                 "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL],
                 "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM],
                 "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD],
@@ -179,7 +188,11 @@ def __init__(self):
                 "rope": [OpBackend.PYTORCH_NATIVE_ROPE],
             },
             "rocm": {
-                "logp": [OpBackend.ROCM_AITER, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE],
+                "logp": [
+                    OpBackend.ROCM_AITER,
+                    OpBackend.TRITON_GENERIC,
+                    OpBackend.PYTORCH_NATIVE,
+                ],
                 "logp_deterministic": [OpBackend.PYTORCH_NATIVE],
                 "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE],
                 "attn": [
@@ -191,7 +204,10 @@ def __init__(self):
                 "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN],
                 "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS],
                 "rope": [OpBackend.PYTORCH_NATIVE_ROPE],
-                "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP],
+                "linear_logp": [
+                    OpBackend.TRITON_LINEAR_LOGP,
+                    OpBackend.PYTORCH_LINEAR_LOGP,
+                ],
                 "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL],
                 "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL],
                 "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM],
@@ -237,7 +253,11 @@ def _adjust_priority_from_env(self):
                 OpBackend.ROCM_FLASH_ATTN,
                 OpBackend.TRITON_GENERIC,
             ]
-        elif rocm_attn_backend and rocm_attn_backend not in {"native", "pytorch", "sdpa"}:
+        elif rocm_attn_backend and rocm_attn_backend not in {
+            "native",
+            "pytorch",
+            "sdpa",
+        }:
             logger.warning(
                 "Unknown RL_KERNEL_ROCM_ATTN_BACKEND=%s; using default ROCm attention priority.",
                 rocm_attn_backend,
@@ -281,16 +301,11 @@ def _adjust_priority_for_hardware(self):
         except Exception as e:
             logger.warning(f"Failed to probe device capability: {e}")
 
-    def get_op(self, op_type: str) -> Any:
+    def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any:
         """Core distribution logic: Automatically select the best operator
         based on hardware and priority.
         """
-        if device_ctx.is_rocm:
-            platform = "rocm"
-        elif device_ctx.device_type == "cuda":
-            platform = "cuda"
-        else:
-            platform = "cpu"
+        platform = self._platform_for_device(device)
         candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE])
 
         for backend in candidates:
@@ -314,6 +329,21 @@ def get_op(self, op_type: str) -> Any:
 
         raise RuntimeError(f"No functional backend found for {op_type} on {platform}")
 
+    def _platform_for_device(self, device: torch.device | str | None) -> str:
+        if device is None:
+            if device_ctx.is_rocm:
+                return "rocm"
+            if device_ctx.device_type == "cuda":
+                return "cuda"
+            return "cpu"
+
+        resolved = torch.device(device)
+        if resolved.type == "cuda":
+            return "rocm" if torch.version.hip is not None else "cuda"
+        if resolved.type in self._priority_map:
+            return resolved.type
+        return "cpu"
+
     def _load_backend(self, backend: OpBackend) -> Optional[Type]:
         """Dynamic loading technique: Import modules only when needed
         and check environment dependencies.