feat(qwen3_5_moe): keep a natively-fp8 lm_head native instead of failing to load

python/freetoken/kernel/triton/fp8_pertensor_linear.py

@@ -423,8 +423,45 @@ def __init__(self, in_features: int, output_sizes: list[int], has_bias: bool = F
         super().__init__(in_features, sum(output_sizes), has_bias)
 
 
+class Fp8LMHead(Fp8PerTensorLinear):
+    """FP8 (W8A16) LM head: ``Fp8PerTensorLinear`` with ``ParallelLMHead``'s prefill
+    behaviour at TP=1 -- slice to the last token per sequence, then the W8A16 GEMV/GEMM
+    instead of a bf16 ``F.linear`` over the dequantized weight. Weight buffers, the
+    ``input_scale`` dance and the uniform-scale/segment precompute are all inherited; the
+    prefill slice is the only new code.
+
+    Exists because a checkpoint can store ``lm_head`` as fp8 natively and no head class could
+    consume that: ``Nvfp4LMHead`` is FP4, and glm_moe_dsa/glm5_next's fp8 heads subclass the
+    bf16 ``ParallelLMHead`` and quantize at load. So a natively-fp8 head had to be
+    dequantized at load. For unsloth/Qwen3.6-35B-A3B-NVFP4-Fast that matrix is
+    ``[248320, 2048]``: 0.474 GiB kept native versus 0.947 GiB as bf16. On a 16 GiB card the
+    difference is not only decode traffic, it is the headroom an 8k prefill needs -- the
+    dequantized head left 118 MiB free and the GDN chunked-prefill workspace then OOM'd.
+
+    ``weight_scale`` is per output row ``[vocab]``, which covers both compressed-tensors
+    strategies: "channel" ships one scalar per row, "tensor" ships one for the whole matrix
+    and the loader broadcasts it. TP=1 and untied embeddings, same as ``Nvfp4LMHead``."""
+
+    def __init__(self, num_embeddings: int, embedding_dim: int):
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        super().__init__(in_features=embedding_dim, out_features=num_embeddings)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        from freetoken.core import get_global_ctx
+
+        batch = get_global_ctx().batch
+        if batch.is_prefill:
+            # Only the last position of each sequence produces logits, so slicing here keeps
+            # the [vocab, hidden] GEMM at M=batch rather than M=prompt_tokens.
+            indices = batch.attn_metadata.get_last_indices(batch.size)
+            x = x[indices].contiguous()
+        return super().forward(x)
+
+
 __all__ = [
     "FP8",
+    "Fp8LMHead",
     "Fp8PerTensorLinear",
     "Fp8PerTensorColMerged",
     "fp8_pertensor_linear",

python/freetoken/models/qwen3_5_moe/config.py

@@ -71,7 +71,9 @@ def _expert_quant(hf_config: Any) -> str:
 def _lm_head_quant(hf_config: Any) -> str:
     """Whether the checkpoint stores ``lm_head`` as NVFP4. modelopt MIXED_PRECISION lists it in
     the per-layer ``quantized_layers`` map (``W4A16_NVFP4``); pure-NVFP4 checkpoints have no
-    per-layer map and leave lm_head bf16. Returns ``"nvfp4"`` or ``"none"``."""
+    per-layer map and leave lm_head bf16. llm-compressor mixed-precision instead puts
+    ``lm_head`` in an fp8 ``config_groups`` entry. Returns ``"nvfp4"``, ``"fp8"``
+    or ``"none"``."""
     get = _quant_accessor(hf_config)
     if get is None:
         return "none"
@@ -82,6 +84,23 @@ def _lm_head_quant(hf_config: Any) -> str:
         if name == "lm_head" or name.endswith(".lm_head"):
             if "fp4" in str((spec or {}).get("quant_algo", "")).lower():
                 return "nvfp4"
+    # compressed-tensors ``ignore`` wins over a group's ``targets``. A head listed there is
+    # bf16 on disk however broadly the group matches, so claiming it as fp8 would build an
+    # Fp8LMHead for a bf16 weight and fail the load on the dtype check.
+    if any("lm_head" in str(x) for x in (get("ignore") or [])):
+        return "none"
+    # llm-compressor mixed-precision puts lm_head in the fp8 group (unsloth NVFP4-Fast), with
+    # no per-layer `quantized_layers` map at all. Keeping it fp8 rather than dequantizing is
+    # worth 0.473 GiB on a [248320, 2048] head -- decode traffic, and the headroom an 8k
+    # prefill needs. Gated on the same weights geometry _attn_quant uses.
+    for g in (get("config_groups") or {}).values():
+        if not g or not any("lm_head" in t for t in (g.get("targets") or [])):
+            continue
+        w = g.get("weights") or {}
+        if str(w.get("type", "")).lower() != "float" or int(w.get("num_bits", 0) or 0) != 8:
+            continue
+        if w.get("group_size") is None and w.get("strategy") in ("tensor", "channel"):
+            return "fp8"
     return "none"
 
 

python/freetoken/models/qwen3_5_moe/model.py

@@ -100,6 +100,15 @@ def __init__(self, config: ModelConfig):
             self.lm_head = Nvfp4LMHead(
                 num_embeddings=config.vocab_size, embedding_dim=config.hidden_size
             )
+        elif getattr(config, "lm_head_quant", "none") == "fp8":
+            # checkpoint stores the (untied) lm_head as fp8 (llm-compressor mixed-precision):
+            # keep it native (W8A16). Halves this matrix versus dequantizing to bf16.
+            from freetoken.kernel.triton.fp8_pertensor_linear import Fp8LMHead
+
+            assert not config.tie_word_embeddings, "fp8 lm_head assumes untied embeddings"
+            self.lm_head = Fp8LMHead(
+                num_embeddings=config.vocab_size, embedding_dim=config.hidden_size
+            )
         else:
             self.lm_head = ParallelLMHead(
                 num_embeddings=config.vocab_size,

python/freetoken/models/qwen3_5_moe/weight.py

@@ -207,6 +207,7 @@ def iter_weights(
             include_non_moe=include_non_moe, include_moe_experts=include_moe_experts,
             dense_nvfp4=config.dense_quant == "nvfp4",
             lmhead_nvfp4=config.lm_head_quant == "nvfp4",
+            lmhead_fp8=config.lm_head_quant == "fp8",
         )
         return
     tp_info = get_tp_info()
@@ -311,9 +312,30 @@ def iter_weights(
 }
 
 
-def _per_row_scale(scalar: torch.Tensor, rows: int) -> torch.Tensor:
-    """Per-tensor scalar -> per-output-row fp32 vector ``[rows]`` (exact broadcast)."""
-    return scalar.reshape(1).to(torch.float32).expand(rows)
+def _per_row_scale(scale: torch.Tensor, rows: int) -> torch.Tensor:
+    """FP8 weight scale -> per-output-row fp32 vector ``[rows]``, which is what
+    ``Fp8PerTensorLinear.weight_scale`` wants whichever on-disk shape arrives:
+
+    * **per-tensor** (``strategy: "tensor"``) -- one scalar for the whole weight, broadcast
+      across the rows. Exact.
+    * **per-channel** (``strategy: "channel"``) -- ``[rows, 1]``, already one scalar per row,
+      so the reshape alone is enough. This case is why the per-tensor ``reshape(1)`` cannot be
+      unconditional: on a per-channel checkpoint it raises ``RuntimeError: shape '[1]' is
+      invalid for input of size <rows>``.
+
+    Any other element count is raised on rather than broadcast: a silently mis-shaped scale
+    would be applied to the wrong output rows and produce plausible garbage.
+    """
+    flat = scale.reshape(-1).to(torch.float32)
+    if flat.numel() == 1:
+        return flat.expand(rows)
+    if flat.numel() != rows:
+        raise ValueError(
+            f"fp8 weight_scale has {flat.numel()} elements for a weight with {rows} output "
+            f"rows (shape {tuple(scale.shape)}); expected either 1 (per-tensor) or {rows} "
+            "(per-channel)"
+        )
+    return flat
 
 
 def _pt_fp8_fuse(base: str, weight: torch.Tensor, scalar: torch.Tensor,
@@ -426,7 +448,7 @@ def _dense_nvfp4_emit(
 
 def _iter_weights_attn_fp8(
     model_path: str, device: torch.device, *, include_non_moe: bool, include_moe_experts: bool,
-    dense_nvfp4: bool = False, lmhead_nvfp4: bool = False,
+    dense_nvfp4: bool = False, lmhead_nvfp4: bool = False, lmhead_fp8: bool = False,
 ) -> Iterator[tuple[str, torch.Tensor]]:
     """Dense pass for the modelopt MIXED_PRECISION Qwen3.5 checkpoint.
 
@@ -474,7 +496,14 @@ def _iter_weights_attn_fp8(
                     raw_base = raw_name[: -len(".weight")]
                     has_s2 = raw_base + ".weight_scale_2" in keyset
                     has_s = raw_base + ".weight_scale" in keyset
-                    if has_s and not has_s2:  # per-tensor FP8 dense projection
+                    # An fp8 lm_head is only safe to keep native if the model built an
+                    # Fp8LMHead for it (lm_head_quant == "fp8"). The bf16 ParallelLMHead has
+                    # no weight_scale buffer, so emitting fp8 into it fails the load with
+                    # "Unexpected keys in state_dict: ['lm_head.weight_scale']". Keep it
+                    # native only when the model asked for an fp8 head; otherwise fall
+                    # through and dequantize.
+                    is_lmhead = base == "lm_head" or base.endswith(".lm_head")
+                    if has_s and not has_s2 and (lmhead_fp8 or not is_lmhead):
                         w = f.get_tensor(raw_name)  # fp8-e4m3, kept verbatim
                         sc = f.get_tensor(raw_base + ".weight_scale")
                         # modelopt's calibrated activation scale: kept (not dropped with the

tests/models/test_qwen3_5_moe_fp8_lm_head.py

@@ -0,0 +1,99 @@
+"""qwen3_5_moe fp8 ``lm_head``: detection, layer choice, and the loader gate.
+
+A compressed-tensors mixed-precision checkpoint can put ``lm_head`` in its fp8 group
+(unsloth/Qwen3.6-35B-A3B-NVFP4-Fast). Three things have to agree for that to load: the config
+detector reports ``"fp8"``, the model builds an ``Fp8LMHead``, and the dense loader keeps the
+weight fp8 instead of dequantizing it. The loader gate is the subtle one -- emitting fp8 into
+the bf16 ``ParallelLMHead`` fails with an unexpected ``lm_head.weight_scale`` key.
+
+Runs off trimmed copies of the real ``quantization_config`` blocks.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from freetoken.kernel.triton.fp8_pertensor_linear import Fp8LMHead, Fp8PerTensorLinear
+from freetoken.models.qwen3_5_moe.config import _lm_head_quant
+
+
+def _quant_config(*, lm_head_in_targets: bool, strategy: str = "channel") -> dict:
+    """Trimmed from unsloth/Qwen3.6-35B-A3B-NVFP4-Fast: fp8 group + NVFP4 expert group."""
+    targets = [".self_attn.q_proj", ".linear_attn.in_proj_qkv"]
+    if lm_head_in_targets:
+        targets = targets + ["lm_head"]
+    return {
+        "quant_method": "compressed-tensors",
+        "format": "mixed-precision",
+        "config_groups": {
+            "group_0": {
+                "targets": targets,
+                "weights": {"num_bits": 8, "type": "float",
+                            "strategy": strategy, "group_size": None},
+            },
+            "group_1": {
+                "targets": ["re:.*mlp.experts.*"],
+                "weights": {"num_bits": 4, "type": "float",
+                            "strategy": "tensor_group", "group_size": 16},
+            },
+        },
+    }
+
+
+class _Cfg:
+    def __init__(self, quantization_config):
+        self.quantization_config = quantization_config
+
+
+@pytest.mark.parametrize("strategy", ["tensor", "channel"])
+def test_lm_head_in_the_fp8_group_is_detected(strategy):
+    cfg = _Cfg(_quant_config(lm_head_in_targets=True, strategy=strategy))
+    assert _lm_head_quant(cfg) == "fp8"
+
+
+def test_lm_head_outside_the_fp8_group_is_not_claimed():
+    """Only the group that targets ``lm_head`` may decide the head's dtype; a checkpoint whose
+    fp8 group covers the attention projections alone must leave the head bf16."""
+    cfg = _Cfg(_quant_config(lm_head_in_targets=False))
+    assert _lm_head_quant(cfg) == "none"
+
+
+def test_an_ignored_lm_head_is_not_claimed_even_when_a_group_targets_it():
+    """compressed-tensors ``ignore`` wins over ``targets``. This is the common llm-compressor
+    shape: a group matches broadly, then ``ignore`` carves modules back out. Such a head is
+    bf16 on disk, so claiming it as fp8 would fail the load on the dtype check."""
+    cfg = _Cfg(_quant_config(lm_head_in_targets=True))
+    cfg.quantization_config["ignore"] = ["re:.*lm_head"]
+    assert _lm_head_quant(cfg) == "none"
+
+
+def test_an_unrelated_ignore_entry_does_not_suppress_the_fp8_head():
+    cfg = _Cfg(_quant_config(lm_head_in_targets=True))
+    cfg.quantization_config["ignore"] = ["re:.*mlp.experts.*", "model.embed_tokens"]
+    assert _lm_head_quant(cfg) == "fp8"
+
+
+def test_a_4bit_lm_head_group_is_not_read_as_fp8():
+    cfg = _Cfg(_quant_config(lm_head_in_targets=True))
+    cfg.quantization_config["config_groups"]["group_0"]["weights"]["num_bits"] = 4
+    assert _lm_head_quant(cfg) == "none"
+
+
+def test_a_grouped_fp8_lm_head_is_not_read_as_per_tensor():
+    """Fp8LMHead consumes a per-output-row scale; a block/group scale is a different layout
+    and must not be routed here."""
+    cfg = _Cfg(_quant_config(lm_head_in_targets=True, strategy="group"))
+    cfg.quantization_config["config_groups"]["group_0"]["weights"]["group_size"] = 128
+    assert _lm_head_quant(cfg) == "none"
+
+
+def test_fp8_lm_head_buffers_match_the_per_output_row_contract():
+    head = Fp8LMHead(num_embeddings=64, embedding_dim=8)
+    assert isinstance(head, Fp8PerTensorLinear)
+    assert head.weight.shape == (64, 8)
+    assert head.weight.dtype == torch.float8_e4m3fn
+    # Per output row, fp32 -- covers "channel" directly and "tensor" once broadcast.
+    assert head.weight_scale.shape == (64,)
+    assert head.weight_scale.dtype == torch.float32
+    assert head.bias is None

tests/models/test_qwen3_5_moe_weight.py

@@ -0,0 +1,54 @@
+"""qwen3_5_moe dense weight-loader helpers: the fp8 ``weight_scale`` shape contract.
+
+``_per_row_scale`` is what turns whatever ``weight_scale`` a compressed-tensors or modelopt
+checkpoint puts on disk into the per-output-row fp32 vector ``Fp8PerTensorLinear.weight_scale``
+declares. Both on-disk granularities are exercised, plus the refusal that keeps a mis-shaped
+scale from being applied to the wrong output rows.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from freetoken.models.qwen3_5_moe.weight import _per_row_scale
+
+_ROWS = 7
+
+
+@pytest.mark.parametrize("shape", [(), (1,), (1, 1)])
+def test_per_tensor_scale_broadcasts_to_every_row(shape):
+    """``strategy: "tensor"`` -- one scalar for the whole weight."""
+    out = _per_row_scale(torch.full(shape, 0.25), _ROWS)
+    assert out.shape == (_ROWS,)
+    assert out.dtype == torch.float32
+    assert out.tolist() == [0.25] * _ROWS
+
+
+@pytest.mark.parametrize("shape", [(_ROWS, 1), (1, _ROWS), (_ROWS,)])
+def test_per_channel_scale_keeps_row_order(shape):
+    """``strategy: "channel"`` -- already one scalar per row; order must survive the reshape.
+
+    Order is asserted with distinct values rather than a set/sum: permuting the scales would
+    keep every aggregate identical while silently scaling each output row by another row's
+    factor.
+    """
+    values = [1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0]
+    out = _per_row_scale(torch.tensor(values).reshape(shape), _ROWS)
+    assert out.shape == (_ROWS,)
+    assert out.dtype == torch.float32
+    assert out.tolist() == values
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64])
+def test_scale_is_promoted_to_fp32_from_any_storage_dtype(dtype):
+    out = _per_row_scale(torch.ones(_ROWS, 1, dtype=dtype), _ROWS)
+    assert out.dtype == torch.float32
+
+
+@pytest.mark.parametrize("shape", [(3, 1), (_ROWS + 1, 1), (2, 3)])
+def test_a_mismatched_scale_raises_instead_of_broadcasting(shape):
+    """Neither 1 nor ``rows`` elements: raise. Broadcasting row 0 over every output row would
+    load without error and serve fluent, wrong tokens."""
+    with pytest.raises(ValueError, match="expected either 1"):
+        _per_row_scale(torch.ones(shape), _ROWS)