feat(qwen4_exp): load block-fp8 dense projections natively

Author: jason-fxz

python/freetoken/models/qwen4_exp/weight.py

@@ -2,7 +2,7 @@
 
 Three separate paths, because the checkpoint's three weight classes live in different places:
 
-* :func:`iter_weights` -- every dense (non-expert) tensor, with the ``model.language_model.`` prefix stripped and fused where the model expects one buffer. See ``_FUSIONS``.
+* :func:`iter_weights` -- every dense (non-expert) tensor, with the ``model.language_model.`` prefix stripped and fused where the model expects one buffer. See ``_DenseFuser``.
 * :func:`load_ple_table` -- the 47.7 GiB FP8 n-gram table, 128 checkpoint shards concatenated into one pinned :class:`HostBank`.
 * :func:`nvfp4_expert_spec` -- how the routed NVFP4 experts are named, for the offload cache's expert reader.
 
@@ -25,8 +25,10 @@
 from freetoken.models.nvfp4_banks import (
     Nvfp4ExpertSourceSpec,
 )
+from freetoken.layers.quantization import get_quant_config
+from freetoken.models.register import get_model_spec
 from freetoken.moe.host_banks import HostBank, read_range_into
-from freetoken.utils import download_hf_weight
+from freetoken.utils import cached_load_hf_config, download_hf_weight
 from freetoken.utils.progress import byte_bar
 from tqdm import tqdm
 
@@ -70,32 +72,13 @@
     ".self_attn.indexer.k_layernorm.weight",
 )
 
-# Fused projections: concat the checkpoint parts along dim 0 in this exact order. A nonzero pad
-# rounds the merged row count up; the model splits the result back with the same sizes.
-_FUSIONS: dict[str, tuple[tuple[str, ...], int]] = {
-    # q carries the output gate, so its half is twice the attention width: [2*qo | kv | kv].
-    ".self_attn.qkv_proj.weight": ((
-        ".self_attn.q_proj.weight", ".self_attn.k_proj.weight", ".self_attn.v_proj.weight",
-    ), 0),
-    ".linear_attn.in_proj.weight": ((
-        ".linear_attn.in_proj_qkv.weight", ".linear_attn.in_proj_z.weight",
-        ".linear_attn.in_proj_b.weight", ".linear_attn.in_proj_a.weight",
-    ), 0),
-    ".mlp.shared_expert.gate_up_proj.weight": ((
-        ".mlp.shared_expert.gate_proj.weight", ".mlp.shared_expert.up_proj.weight",
-    ), 0),
-    # HC mix reads the low-rank down projection and the injection logits from one GEMM; vLLM
-    # pads the merged output to a multiple of 16 rows for cuBLAS (hyperconnection.py pad_size).
-    # The top-level hyper_connection_mixer has no injection and so never fuses.
-    ".attn_hyper_connection.input_mix_weight_down_block_inject.weight": ((
-        ".attn_hyper_connection.input_mix_weight_down.weight",
-        ".attn_hyper_connection.block_inject_weight.weight",
-    ), 16),
-    ".mlp_hyper_connection.input_mix_weight_down_block_inject.weight": ((
-        ".mlp_hyper_connection.input_mix_weight_down.weight",
-        ".mlp_hyper_connection.block_inject_weight.weight",
-    ), 16),
-}
+# The per-layer HC mix reads the low-rank down projection and the injection logits from one GEMM; vLLM pads the merged rows to a multiple of 16 for cuBLAS (hyperconnection.py pad_size).
+# The top-level hyper_connection_mixer has no injection and never fuses.
+_PAD_TO = {"input_mix_weight_down_block_inject": 16}
+_HC_WITH_INJECT = (".attn_hyper_connection", ".mlp_hyper_connection")
+_KIND_SUFFIXES = (".weight_scale_inv", ".weight")
+_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)
+_ELEM_DTYPES = {"e4m3": torch.float8_e4m3fn}
 
 
 def _rename(raw_name: str) -> str | None:
@@ -115,26 +98,95 @@ def _rename(raw_name: str) -> str | None:
     return raw_name
 
 
-def _try_fuse(
-    name: str, tensor: torch.Tensor, buf: dict[str, dict[int, torch.Tensor]]
-) -> tuple[str, torch.Tensor] | tuple[()] | None:
-    """Buffer a fusion part; return the merged ``(name, tensor)`` once all parts arrive, ``()`` while incomplete, ``None`` if ``name`` is not a fusion part."""
-    for fused_suffix, (parts, pad_to) in _FUSIONS.items():
-        for idx, part in enumerate(parts):
-            if not name.endswith(part):
-                continue
-            key = name[: -len(part)] + fused_suffix
-            slots = buf.setdefault(key, {})
-            slots[idx] = tensor
-            if len(slots) < len(parts):
-                return ()
-            del buf[key]
-            rows = [slots[i] for i in range(len(parts))]
-            pad = (-sum(t.shape[0] for t in rows)) % pad_to if pad_to else 0
-            if pad:
-                rows.append(torch.zeros(pad, *rows[0].shape[1:], dtype=rows[0].dtype, device=rows[0].device))
-            return key, torch.cat(rows, dim=0)
-    return None
+def _split_kind(name: str) -> tuple[str, str]:
+    """``name`` -> ``(module, kind)``; kind is "" for tensors that are neither a weight nor a block scale."""
+    for suffix in _KIND_SUFFIXES:
+        if name.endswith(suffix):
+            return name[: -len(suffix)], suffix
+    return name, ""
+
+
+class _DenseFuser:
+    """Concatenates checkpoint projection parts into the model's merged buffers, per kind (weight / block scale).
+
+    The part table is the family's packed_modules_mapping. The QuantConfig picks the GDN in_proj layout and validates each part against the scheme the model built its buffer from.
+    """
+
+    def __init__(self, quant, packed: tuple[tuple[str, tuple[str, ...]], ...]) -> None:
+        self.quant = quant
+        self.groups = {fused: parts for fused, parts in packed if fused != "experts"}  # experts: bank reader
+        self.by_part: dict[str, list[tuple[str, int]]] = {}
+        for fused, parts in self.groups.items():
+            for idx, part in enumerate(parts):
+                self.by_part.setdefault(part, []).append((fused, idx))
+        self.buf: dict[tuple[str, str], dict[int, torch.Tensor]] = {}
+
+    def scheme(self, module: str):
+        return None if self.quant is None else self.quant.scheme_for(module)
+
+    def _target(self, parent: str, leaf: str) -> tuple[str, int] | None:
+        candidates = self.by_part.get(leaf)
+        if not candidates:
+            return None
+        if len(candidates) > 1:
+            # GDN: quantized checkpoints split qkv|z from the bf16 b|a; same test as gdn.py
+            split = self.scheme(f"{parent}.in_proj_qkvz") is not None
+            keep = {"in_proj_qkvz", "in_proj_ba"} if split else {"in_proj"}
+            candidates = [c for c in candidates if c[0] in keep]
+            if not candidates:
+                raise ValueError(f"{parent}.{leaf}: no merged projection for the {'split' if split else 'fused'} GDN layout")
+        fused, idx = candidates[0]
+        if fused in _PAD_TO and not parent.endswith(_HC_WITH_INJECT):
+            return None
+        return f"{parent}.{fused}", idx
+
+    def check(self, module: str, name: str, tensor: torch.Tensor) -> None:
+        """``tensor`` (checkpoint key ``name``) must match the scheme the model built ``module`` from."""
+        scheme = self.scheme(module)
+        if name.endswith(".weight_scale_inv"):
+            if scheme is None or not scheme.has("weight_scale_inv"):
+                raise ValueError(f"{name}: {module} has no block scale in the checkpoint's quant config ({scheme})")
+            return
+        is_fp8 = tensor.dtype in _FP8_DTYPES
+        if scheme is None:
+            if is_fp8:
+                raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} unquantized")
+            return
+        expected = _ELEM_DTYPES.get(scheme.weight.elem)
+        if expected is not None and tensor.dtype is not expected:
+            raise ValueError(f"{name} is {tensor.dtype} but the checkpoint's quant config declares {module} {scheme}")
+        rows, cols = (scheme.weight.group or (1, 1))
+        if rows > 1 and tensor.shape[0] % rows or cols > 1 and tensor.shape[1] % cols:
+            raise ValueError(f"{name}: {tuple(tensor.shape)} is not a multiple of the {rows}x{cols} scale block of {module}")
+
+    def check_unfused(self, name: str, tensor: torch.Tensor) -> None:
+        module, kind = _split_kind(name)
+        if kind == ".weight_scale_inv" or (kind == ".weight" and tensor.dtype in _FP8_DTYPES):
+            self.check(module, name, tensor)
+
+    def fuse(self, name: str, tensor: torch.Tensor) -> list[tuple[str, torch.Tensor]] | None:
+        """Buffer a part; return the merged ``[(name, tensor)]`` once its kind is complete, ``[]`` while incomplete, ``None`` if ``name`` is not a part."""
+        module, kind = _split_kind(name)
+        if not kind:
+            return None
+        parent, _, leaf = module.rpartition(".")
+        hit = self._target(parent, leaf)
+        if hit is None:
+            return None
+        fused, idx = hit
+        self.check(fused, name, tensor)
+        slots = self.buf.setdefault((fused, kind), {})
+        slots[idx] = tensor
+        parts = self.groups[fused.rpartition(".")[2]]
+        if len(slots) < len(parts):
+            return []
+        del self.buf[(fused, kind)]
+        rows = [slots[i] for i in range(len(parts))]
+        pad_to = _PAD_TO.get(fused.rpartition(".")[2], 0) if kind == ".weight" else 0
+        pad = (-sum(t.shape[0] for t in rows)) % pad_to if pad_to else 0
+        if pad:
+            rows.append(torch.zeros(pad, *rows[0].shape[1:], dtype=rows[0].dtype, device=rows[0].device))
+        return [(fused + kind, torch.cat(rows, dim=0))]
 
 
 def iter_weights(
@@ -146,24 +198,19 @@ def iter_weights(
 ) -> Iterator[tuple[str, torch.Tensor]]:
     """Yield the dense (non-expert) weights, prefix-stripped and fused to the model's buffers.
 
-    Keys keep the checkpoint's module names below the stripped prefix, so the emitted set is the
-    model's state dict minus the routed experts. Nothing here is quantized: every release's skip
-    list (modelopt ``ignore``, fp8 ``modules_to_not_convert``) covers everything except those experts,
-    so attention, GDN, HC, PLE, the shared expert and lm_head are all plain bf16 (the n-gram hash
-    constants stay int64). Fusions:
-    attention q|k|v -> ``qkv_proj``, GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, shared-expert
-    gate|up -> ``gate_up_proj``, and each per-layer HC's ``input_mix_weight_down`` |
-    ``block_inject_weight`` -> a zero-padded ``input_mix_weight_down_block_inject``.
-
-    ``include_moe_experts`` is accepted for the loader contract but never yields anything: the
-    routed experts are NVFP4 and always come from the offload cache's expert reader.
+    Keys keep the checkpoint's module names below the stripped prefix, so the emitted set is the model's state dict minus the routed experts.
+    A dense projection is bf16 or 128x128 block-fp8 (``.weight`` e4m3 + ``.weight_scale_inv``) as the checkpoint's QuantConfig says: the official releases skip everything but the routed experts, the community NVFP4-FP8 requants quantize the attention / GDN projections.
+    Fusions, per kind: attention q|k|v -> ``qkv_proj``; GDN ``in_proj_{qkv,z,b,a}`` -> ``in_proj``, or ``in_proj_qkvz`` + bf16 ``in_proj_ba`` when qkv|z is quantized; shared-expert gate|up -> ``gate_up_proj``; each per-layer HC's ``input_mix_weight_down`` | ``block_inject_weight`` -> a zero-padded ``input_mix_weight_down_block_inject``.
+    ``include_moe_experts`` is accepted for the loader contract but never yields anything: the routed experts are NVFP4 and always come from the offload cache's expert reader.
     """
     if get_tp_info().size > 1:
         raise NotImplementedError("qwen4_exp weight loading supports TP=1 only")
     if not include_non_moe:
         return
 
-    fuse_buf: dict[str, dict[int, torch.Tensor]] = {}
+    hf_config = cached_load_hf_config(model_path)
+    spec = get_model_spec(hf_config.architectures[0])
+    fuser = _DenseFuser(get_quant_config(), spec.packed_modules_mapping)
     for file in tqdm(
         iter_weight_files(model_path),
         desc="Loading weights",
@@ -175,14 +222,14 @@ def iter_weights(
                 if name is None:
                     continue
                 tensor = f.get_tensor(raw_name)
-                fused = _try_fuse(name, tensor, fuse_buf)
-                if fused is not None:
-                    if fused != ():  # () means buffered, not yet complete
-                        yield fused
-                    continue
-                yield name, tensor
-
-    assert not fuse_buf, f"Incomplete projection fusions: {sorted(fuse_buf)}"
+                fused = fuser.fuse(name, tensor)
+                if fused is None:
+                    fuser.check_unfused(name, tensor)
+                    yield name, tensor
+                else:
+                    yield from fused
+
+    assert not fuser.buf, f"Incomplete projection fusions: {sorted(k[0] + k[1] for k in fuser.buf)}"
 
 
 # ======================================================================================

tests/models/qwen4_exp/common.py

@@ -255,3 +255,118 @@ def spy(self, index, md, slot):
 
     monkeypatch.setattr(QSASparseAttnBackend, "_select", spy)
     return seen
+
+
+LM = "model.language_model"
+
+# quantization_config of each released checkpoint, trimmed to the entries parse_config and the reader look at
+
+# RadixArk/Qwen3.8-Flash-Next-NVFP4: modelopt NVFP4 everywhere except the ignore list
+RADIXARK_NVFP4 = {
+    "quant_algo": "NVFP4",
+    "quant_method": "modelopt",
+    "ignore": [
+        "model.embed_tokens",
+        "mtp.*",
+        "model.mtp.*",
+        "*.self_attn.*",
+        "*.linear_attn.*",
+        "*.mlp.gate*",
+        "*.mlp.shared_expert.*",
+        "*.mlp.shared_expert_gate*",
+        "*hyper_connection*",
+        "*.ple.*",
+        "model.visual.*",
+        "model.language_model.embed_tokens",
+        "lm_head",
+    ],
+}
+
+# nvidia/Qwen3.8-Flash-Next-NVFP4: modelopt MIXED_PRECISION, the per-module algo sits in quantized_layers
+NVIDIA_NVFP4 = {
+    "quant_algo": "MIXED_PRECISION",
+    "quant_method": "modelopt",
+    "quantized_layers": {
+        **{f"model.language_model.layers.{i}.mlp.experts": {"quant_algo": "NVFP4", "group_size": 16} for i in range(48)},
+        "model.language_model.layers.1.ple.ple_embedding.ngram_embedding": {"quant_algo": "FP8"},
+        "mtp.layers.0.mlp.experts": {"quant_algo": "FP8_PB_WO", "group_size": 128},
+    },
+    "ignore": ["lm_head", "model.language_model.embed_tokens", "model.language_model.layers.0.mlp.shared_expert*", "model.visual*"],
+}
+
+# Qwen/Qwen3.8-Flash-Next-FP8: 128x128 block-fp8 experts, everything else listed in modules_to_not_convert
+QWEN_FP8 = {
+    "quant_method": "fp8",
+    "activation_scheme": "dynamic",
+    "weight_per_tensor": False,
+    "act_per_tensor": False,
+    "weight_block_size": [128, 128],
+    "modules_to_not_convert": [
+        "lm_head",
+        "model.language_model.embed_tokens",
+        "model.language_model.hyper_connection_mixer.input_mix_weight_up",
+        "model.language_model.layers.0.linear_attn.in_proj_qkv",
+        "model.language_model.layers.3.self_attn.q_proj",
+        "model.language_model.layers.3.mlp.gate",
+        "model.language_model.layers.3.mlp.shared_expert.gate_proj",
+    ],
+    "modules_to_convert": ["ple.ple_embedding.ngram_embedding"],
+}
+
+
+def mixed_precision_quant(gdn_layers, attn_layers, moe_layers) -> dict:
+    """modelopt MIXED_PRECISION with NVFP4 routed experts and FP8_PB_WO attention / GDN projections, ignore list as in lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8."""
+    return {
+        "quant_method": "modelopt",
+        "quant_algo": "MIXED_PRECISION",
+        "quantized_layers": {
+            **{f"{LM}.layers.{i}.mlp.experts": {"quant_algo": "NVFP4", "group_size": 16} for i in moe_layers},
+            **{f"{LM}.layers.{i}.linear_attn.{p}": {"quant_algo": "FP8_PB_WO", "group_size": 128}
+               for i in gdn_layers for p in ("in_proj_qkv", "in_proj_z", "out_proj")},
+            **{f"{LM}.layers.{i}.self_attn.{p}_proj": {"quant_algo": "FP8_PB_WO", "group_size": 128}
+               for i in attn_layers for p in "qkvo"},
+        },
+        "ignore": [
+            "model.embed_tokens", "mtp.*", "model.mtp.*", "*.mlp.gate*", "*.mlp.shared_expert.*",
+            "*.mlp.shared_expert_gate*", "*hyper_connection*", "*.ple.*", "model.visual.*",
+            "model.language_model.embed_tokens", "lm_head", "*.self_attn.indexer*",
+        ],
+    }
+
+
+# lovedheart/Qwen3.8-Flash-Next-NVFP4-FP8, trimmed to layers 0 (GDN) and 3 (attention)
+LOVEDHEART_NVFP4_FP8 = mixed_precision_quant(gdn_layers=(0,), attn_layers=(3,), moe_layers=(0, 3))
+
+
+def install_quant_config(model_path: str) -> None:
+    """Install ``model_path``'s QuantConfig process-wide, as EngineConfig does before the reader runs."""
+    from freetoken.layers.quantization import set_quant_config
+    from freetoken.models.register import checkpoint_quant_config, get_model_spec
+    from freetoken.utils import cached_load_hf_config
+
+    hf = cached_load_hf_config(model_path)
+    set_quant_config(checkpoint_quant_config(model_path, hf, get_model_spec(hf.architectures[0])))
+
+
+def meta_state_dict(model_path: str) -> dict[str, torch.Tensor]:
+    """State dict of the model the engine builds for ``model_path`` (experts offloaded), on the meta device."""
+    from freetoken.engine.config import EngineConfig
+    from freetoken.engine.engine import _decode_target
+    from freetoken.layers import rotary
+    from freetoken.models import create_model
+    from freetoken.utils.torch_utils import torch_dtype
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    config = EngineConfig(model_path=model_path, tp_info=try_get_tp_info(), dtype=torch.bfloat16, moe_strategy="offload")
+    object.__setattr__(config.model_config, "moe_strategy", "offload")
+    object.__setattr__(config.model_config, "decode_target", _decode_target(config))
+    saved = rotary._ROPE_DEVICE
+    rotary.set_rope_device(torch.device("cpu"))  # get_rope refuses to build on meta
+    rotary.get_rope.cache_clear()
+    try:
+        with torch.device("meta"), torch_dtype(torch.bfloat16):
+            return create_model(config.model_config).state_dict()
+    finally:
+        rotary.set_rope_device(saved)
+        rotary.get_rope.cache_clear()  # the cpu rope must not leak into the GPU tests' cache