fix(gemma4): load mixed compressed-tensors nvfp4

python/freetoken/models/gemma4/weight.py

@@ -6,12 +6,15 @@
 import safetensors
 import torch
 from freetoken.distributed import get_tp_info
+from freetoken.layers.quantization import QuantKind, get_quant_config
 from freetoken.models.config import FullAttentionGroupConfig
 from freetoken.models.loader import (
     MergeRule,
     ShardReader,
+    ct_nvfp4_fuse,
     drop_page_cache,
     iter_weight_files,
+    nvfp4_parts_ct,
 )
 from freetoken.models.nvfp4_banks import (
     Nvfp4ExpertSourceSpec,
@@ -39,6 +42,15 @@
     layer_to_bank=lambda layer, config: layer,  # every layer is MoE (no dense prefix)
     desc="Gemma4 NVFP4 experts",
 )
+_NVFP4_EXPERT_KEY_TEMPLATE = (
+    r"^model\.language_model\.layers\.(?P<layer>\d+)\.experts\.(?P<expert>\d+)\."
+    r"(?P<proj>gate_proj|up_proj|down_proj)\.(?P<kind>{kinds})$"
+)
+_NVFP4_BANK_KINDS = {
+    "weight": "weight",
+    "weight_scale": "weight_scale",
+    "weight_global": "weight_scale_2",
+}
 _LAYER_INDEX_PATTERN = re.compile(r"layers\.(\d+)\.")
 _LAYER_FF_PREFIX_PATTERN = re.compile(r"^(model\.layers\.\d+)\.")
 _MERGE_RULES = {
@@ -63,6 +75,17 @@
 # The scales are consumed with their .weight.
 _NVFP4_DENSE_SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale")
 _NVFP4_DENSE_MLP_RE = re.compile(r"\.mlp\.(gate_proj|up_proj|down_proj)\.weight$")
+_CT_NVFP4_FUSIONS = {
+    ".feed_forward.shared_mlp.gate_up_proj": (
+        ".feed_forward.shared_mlp.gate_proj",
+        ".feed_forward.shared_mlp.up_proj",
+    ),
+}
+_CT_NVFP4_SCALE_SUFFIXES = (
+    ".weight_scale",
+    ".weight_global_scale",
+    ".input_global_scale",
+)
 
 
 def _nvfp4_dense_parts(reader: ShardReader, raw_base: str):
@@ -156,6 +179,10 @@ def iter_weights(
     include_vision: bool = True,
 ) -> Iterator[tuple[str, torch.Tensor]]:
     def rename_key(raw_name: str) -> str | None:
+        # Static KV quantizer metadata is not part of FreeToken's runtime KV
+        # cache state (same rule used by the generic compressed-tensors reader).
+        if raw_name.endswith((".k_scale", ".v_scale", ".q_scale", ".prob_scale")):
+            return None
         prefix = "model.language_model."
         if raw_name.startswith(prefix):
             return _rename_language_key(raw_name)
@@ -165,8 +192,11 @@ def rename_key(raw_name: str) -> str | None:
 
     def merge_info(key: str) -> tuple[str, MergeRule] | None:
         for suffix, rule in _MERGE_RULES.items():
-            if key.endswith(suffix + ".weight") or key.endswith(suffix):
-                return key.replace(suffix, rule.fused_suffix), rule
+            # Quantized tensors carry roles after the projection name, e.g.
+            # ``q_proj.weight_scale``.  Every role must follow the same fusion as
+            # its weight or the model asks for a fused tensor that was never emitted.
+            if key.endswith(suffix) or suffix + "." in key:
+                return key.replace(suffix, rule.fused_suffix, 1), rule
         return None
 
     config = parse_config(cached_load_hf_config(model_path))
@@ -200,8 +230,19 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None:
                     if _NVFP4_EXPERT_RE.search(raw_name):
                         continue
 
-                    # NVFP4 dense-MLP scales are consumed with their .weight (below), never yielded.
-                    if raw_name.endswith(_NVFP4_DENSE_SCALE_SUFFIXES):
+                    # NVFP4 scales are consumed with their packed weight.  Do not
+                    # blanket-drop ``weight_scale``: FP8 attention projections use
+                    # that same suffix and must be fused q/k/v-wise.
+                    ct_scale_base = next(
+                        (raw_name[: -len(s)] for s in _CT_NVFP4_SCALE_SUFFIXES if raw_name.endswith(s)),
+                        None,
+                    )
+                    if ct_scale_base is not None and reader.has(ct_scale_base + ".weight_packed"):
+                        continue
+                    if (
+                        config.dense_quant == "nvfp4"
+                        and raw_name.endswith(_NVFP4_DENSE_SCALE_SUFFIXES)
+                    ):
                         continue
 
                     is_vision = name.startswith(("vision_tower.", "embed_vision.", "vision_embedder."))
@@ -215,6 +256,25 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None:
                     if not is_expert and not include_non_moe:
                         continue
 
+                    # llm-compressor / compressed-tensors NVFP4.  The stored names
+                    # differ from ModelOpt but map to the same FreeToken buffers.
+                    if raw_name.endswith(".weight_packed"):
+                        raw_base = raw_name[: -len(".weight_packed")]
+                        base = name[: -len(".weight_packed")]
+                        parts = nvfp4_parts_ct(reader, raw_base)
+                        emitted = ct_nvfp4_fuse(base, parts, gateup_buf, _CT_NVFP4_FUSIONS)
+                        if emitted is None:
+                            w, s, g, a = parts
+                            emitted = [
+                                (base + ".weight", w),
+                                (base + ".weight_scale", s),
+                                (base + ".weight_global", g),
+                            ]
+                            if a is not None:
+                                emitted.append((base + ".input_scale", a))
+                        yield from emitted
+                        continue
+
                     # Native W4A16 NVFP4 dense MLP: the .weight is FP4-packed and carries block + per-tensor scales.
                     # The weight_scale_2 sibling guard is defense-in-depth beyond config.dense_quant -- the sibling MoE checkpoint's bf16 shared_mlp has no such sibling, so it falls through to the bf16 path.
                     if (
@@ -230,6 +290,15 @@ def merge_info(key: str) -> tuple[str, MergeRule] | None:
                         continue
 
                     tensor = f.get_tensor(raw_name)
+                    # compressed-tensors stores FP8 per-channel scales as
+                    # ``[out_features, 1]`` (often bf16); FreeToken's FP8 linear
+                    # buffer is a flat fp32 vector.
+                    if (
+                        raw_name.endswith(".weight_scale")
+                        and tensor.ndim == 2
+                        and tensor.shape[1] == 1
+                    ):
+                        tensor = tensor.reshape(-1).to(torch.float32)
                     if is_vision or is_expert:
                         yield name, tensor
                         continue
@@ -296,7 +365,23 @@ def _expert_name(raw_name: str) -> str | None:
 
 
 def nvfp4_expert_spec(model_path: str, config):
-    return _NVFP4_SOURCE_SPEC
+    quant = get_quant_config()
+    stored = quant.stored_tensors(QuantKind.NVFP4)
+    kind_map = {stored[role].name: kind for role, kind in _NVFP4_BANK_KINDS.items()}
+    # Preserve the original ModelOpt descriptor exactly; compressed-tensors uses
+    # weight_packed / weight_global_scale and reciprocal global scales.
+    if set(kind_map) == {"weight", "weight_scale", "weight_scale_2"}:
+        return _NVFP4_SOURCE_SPEC
+    return Nvfp4ExpertSourceSpec(
+        key_pattern=re.compile(
+            _NVFP4_EXPERT_KEY_TEMPLATE.format(kinds="|".join(map(re.escape, kind_map)))
+        ),
+        proj_to_role={"gate_proj": "gate", "up_proj": "up", "down_proj": "down"},
+        layer_to_bank=lambda layer, config: layer,
+        desc=f"Gemma4 NVFP4 experts ({quant.dialect})",
+        kind_map=kind_map,
+        global_reciprocal=stored["weight_global"].reciprocal,
+    )
 
 
 __all__ = [

tests/models/test_gemma4_compressed_tensors.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import save_file
+
+from freetoken.distributed import set_tp_info, try_get_tp_info
+from freetoken.layers.quantization import set_quant_config
+from freetoken.layers.quantization.configs.compressed_tensors import CompressedTensorsConfig
+from freetoken.layers.quantization.configs.modelopt import ModelOptConfig
+from freetoken.models.gemma4.weight import iter_weights, nvfp4_expert_spec
+
+
+@pytest.fixture(scope="session", autouse=True)
+def _tp_info():
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+
+
+def _nvfp4(base: str, rows: int, cols: int, global_scale: float = 2.0):
+    return {
+        base + ".weight_packed": torch.randint(
+            0, 256, (rows, cols // 2), dtype=torch.uint8
+        ),
+        base + ".weight_scale": torch.ones(
+            rows, cols // 16, dtype=torch.float8_e4m3fn
+        ),
+        base + ".weight_global_scale": torch.tensor([global_scale]),
+        base + ".input_global_scale": torch.tensor([2.0]),
+    }
+
+
+def test_mixed_compressed_tensors_dense_weights(tmp_path, monkeypatch):
+    """Gemma 4 Unsloth exports mix FP8 attention with compressed-tensors
+    NVFP4 shared MLPs; both roles must land in the buffers the model builds."""
+    h, q, kv, intermediate = 16, 8, 4, 32
+    layer = "model.language_model.layers.0"
+    attn = layer + ".self_attn"
+    mlp = layer + ".mlp"
+    raw = {
+        attn + ".q_proj.weight": torch.ones(q, h, dtype=torch.float8_e4m3fn),
+        attn + ".q_proj.weight_scale": torch.full((q, 1), 2.0, dtype=torch.bfloat16),
+        attn + ".k_proj.weight": torch.ones(kv, h, dtype=torch.float8_e4m3fn),
+        attn + ".k_proj.weight_scale": torch.full((kv, 1), 3.0, dtype=torch.bfloat16),
+        attn + ".v_proj.weight": torch.ones(kv, h, dtype=torch.float8_e4m3fn),
+        attn + ".v_proj.weight_scale": torch.full((kv, 1), 4.0, dtype=torch.bfloat16),
+        attn + ".o_proj.weight": torch.ones(h, q, dtype=torch.float8_e4m3fn),
+        attn + ".o_proj.weight_scale": torch.full((h, 1), 5.0, dtype=torch.bfloat16),
+        attn + ".k_scale": torch.tensor(1.0),
+        attn + ".v_scale": torch.tensor(1.0),
+    }
+    raw |= _nvfp4(mlp + ".gate_proj", intermediate, h)
+    raw |= _nvfp4(mlp + ".up_proj", intermediate, h)
+    raw |= _nvfp4(mlp + ".down_proj", h, intermediate, global_scale=4.0)
+    save_file(raw, tmp_path / "model.safetensors")
+
+    import freetoken.models.gemma4.weight as weight
+
+    config = SimpleNamespace(
+        num_layers=1,
+        dense_quant="none",
+        attention_group_for_layer=lambda _layer: None,
+    )
+    monkeypatch.setattr(weight, "cached_load_hf_config", lambda _path: object())
+    monkeypatch.setattr(weight, "parse_config", lambda _hf: config)
+
+    loaded = dict(
+        iter_weights(
+            str(tmp_path),
+            torch.device("cpu"),
+            include_moe_experts=False,
+            include_non_moe=True,
+            include_vision=False,
+        )
+    )
+
+    qkv = "model.layers.0.self_attn.qkv_proj"
+    assert loaded[qkv + ".weight"].shape == (q + kv + kv, h)
+    assert loaded[qkv + ".weight_scale"].shape == (q + kv + kv,)
+    assert loaded[qkv + ".weight_scale"].dtype is torch.float32
+    assert torch.equal(
+        loaded[qkv + ".weight_scale"],
+        torch.tensor([2.0] * q + [3.0] * kv + [4.0] * kv),
+    )
+    assert loaded["model.layers.0.self_attn.o_proj.weight_scale"].shape == (h,)
+    assert not any(name.endswith((".k_scale", ".v_scale")) for name in loaded)
+
+    gate_up = "model.layers.0.feed_forward.shared_mlp.gate_up_proj"
+    assert loaded[gate_up + ".weight"].shape == (2 * intermediate, h // 2)
+    assert loaded[gate_up + ".weight_scale"].shape == (2 * intermediate, h // 16)
+    assert loaded[gate_up + ".weight_global"].shape == (2 * intermediate,)
+    assert loaded[gate_up + ".weight_global"][0].item() == pytest.approx(0.5)
+    assert loaded[gate_up + ".input_scale"].item() == pytest.approx(0.5)
+
+    down = "model.layers.0.feed_forward.shared_mlp.down_proj"
+    assert loaded[down + ".weight"].shape == (h, intermediate // 2)
+    assert loaded[down + ".weight_scale"].shape == (h, intermediate // 16)
+    assert loaded[down + ".weight_global"][0].item() == pytest.approx(0.25)
+    assert not any("weight_packed" in name or "global_scale" in name for name in loaded)
+
+
+def test_nvfp4_expert_spec_uses_checkpoint_dialect():
+    ct = CompressedTensorsConfig(
+        {
+            "config_groups": {
+                "group_0": {
+                    "targets": ["Linear"],
+                    "weights": {
+                        "num_bits": 4,
+                        "type": "float",
+                        "strategy": "tensor_group",
+                        "group_size": 16,
+                    },
+                }
+            }
+        }
+    )
+    set_quant_config(ct)
+    spec = nvfp4_expert_spec("unused", object())
+    match = spec.key_pattern.match(
+        "model.language_model.layers.2.experts.7.gate_proj.weight_packed"
+    )
+    assert match and match.group("kind") == "weight_packed"
+    assert spec.kind_map == {
+        "weight_packed": "weight",
+        "weight_scale": "weight_scale",
+        "weight_global_scale": "weight_scale_2",
+    }
+    assert spec.global_reciprocal
+    assert spec.key_pattern.match(
+        "model.language_model.layers.2.experts.7.gate_proj.input_global_scale"
+    ) is None
+
+    set_quant_config(ModelOptConfig({"quant_algo": "NVFP4"}))
+    modelopt = nvfp4_expert_spec("unused", object())
+    assert modelopt.kind_map is None
+    assert not modelopt.global_reciprocal
+    assert modelopt.key_pattern.match(
+        "model.language_model.layers.2.experts.7.gate_proj.weight"
+    )