fix(models): support llm-compressor NVFP4 MoE export variants

python/freetoken/models/config.py

@@ -69,6 +69,11 @@ def detect_compressed_tensors_nvfp4(hf_config: Any) -> bool:
     if str(get("quant_method") or "").lower() != "compressed-tensors":
         return False
     groups = get("config_groups") or {}
+    if not groups:
+        # No config_groups: some llm-compressor exports (e.g. the AEON / Kwaipilot Qwen3.6
+        # MoE NVFP4 builds) carry only ``format: nvfp4-pack-quantized`` (+ a ``recipe``
+        # string). Gate on the exact format string, like the config_groups branch below.
+        return str(get("format") or "").lower() == "nvfp4-pack-quantized"
     # Verdicts are collected across ALL groups before returning: an early return on
     # the first NVFP4 group would accept a mixed {nvfp4, mxfp4} checkpoint (and the
     # error would depend on the groups' key order).
@@ -277,6 +282,12 @@ class ModelConfig:
     # scale and runs a W8A16 kernel (modelopt MIXED_PRECISION); "none" leaves them bf16
     # (dequant-at-load for any other dense quant, e.g. NVFP4 shared_expert/lm_head).
     attn_quant: str = "none"
+    # Quantization of the GatedDeltaNet's ``out_proj`` only (qwen3_5_moe). Independent of
+    # ``attn_quant`` because llm-compressor checkpoints can quantize the full-attention
+    # projections while leaving the whole GDN bf16 (their ``ignore`` list names the
+    # ``linear_attn.*`` modules). "nvfp4" keeps ``out_proj`` packed (W4A16); "fp8_pertensor"
+    # keeps it per-tensor fp8 (modelopt MIXED_PRECISION); "none" -> bf16.
+    gdn_quant: str = "none"
     # Weight quantization of the *dense* NVFP4 MLP projections -- the shared expert, and dense
     # (non-MoE) MLP layers -- which NVFP4 checkpoints store as packed FP4 like the routed
     # experts. "nvfp4" keeps them packed and runs the W4A16 dense kernels (quartering their

python/freetoken/models/nvfp4_banks.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import collections
+import glob
 import json
 import os
 import re
+import struct
 from dataclasses import dataclass
 from typing import Callable
 
@@ -28,6 +30,10 @@ class Nvfp4ExpertSourceSpec:
     # The checkpoint stores the QUANT-side global scale (local fp8 scales were
     # multiplied by it before the cast); the banks keep its reciprocal.
     global_reciprocal: bool = False
+    # True when the experts are stacked per layer (``experts.gate_up_proj`` U8
+    # [E*rows, cols] + a per-layer scalar global) instead of per expert; the stacked
+    # loader reshapes each bank tensor to [E, ...] and broadcasts the scalar global.
+    stacked: bool = False
 
 
 def _canon_kind(spec: "Nvfp4ExpertSourceSpec", kind: str) -> str:
@@ -78,6 +84,24 @@ def _alloc_nvfp4_host_banks(num_layers: int, E: int, H: int, I: int):
     }, num_layers)
 
 
+def _weight_map(folder: str) -> dict[str, str]:
+    """name -> shard (basename) from the index, or from each safetensors header when the
+    checkpoint ships a single shard without an index (llm-compressor single-file exports)."""
+    index = os.path.join(folder, "model.safetensors.index.json")
+    if os.path.exists(index):
+        with open(index, encoding="utf-8") as f:
+            return json.load(f)["weight_map"]
+    weight_map: dict[str, str] = {}
+    for shard in sorted(os.path.basename(p) for p in glob.glob(os.path.join(folder, "*.safetensors"))):
+        with open(os.path.join(folder, shard), "rb") as fh:
+            n = struct.unpack("<Q", fh.read(8))[0]
+            hdr = json.loads(fh.read(n))
+        for name in hdr:
+            if name != "__metadata__":
+                weight_map[name] = shard
+    return weight_map
+
+
 def load_nvfp4_expert_source_banks(
     model_path: str,
     config,
@@ -103,9 +127,7 @@ def load_nvfp4_expert_source_banks(
     until then (the caller owns that tradeoff).
     """
     folder = download_hf_weight(model_path)
-    index_path = os.path.join(folder, "model.safetensors.index.json")
-    with open(index_path, encoding="utf-8") as f:
-        weight_map = json.load(f)["weight_map"]
+    weight_map = _weight_map(folder)
 
     E = config.num_experts
     H = config.hidden_size
@@ -235,8 +257,7 @@ def load_nvfp4_expert_source_banks_parallel(
     from freetoken.models.weight import iter_expert_tensors_parallel
 
     folder = download_hf_weight(model_path)
-    with open(os.path.join(folder, "model.safetensors.index.json"), encoding="utf-8") as f:
-        weight_map = json.load(f)["weight_map"]
+    weight_map = _weight_map(folder)
 
     E = config.num_experts
     H = config.hidden_size
@@ -336,8 +357,119 @@ def _load(sink) -> int:
     }
 
 
+def load_nvfp4_stacked_expert_sources(
+    model_path: str,
+    config,
+    spec: Nvfp4ExpertSourceSpec,
+    *,
+    drop_page_cache: DropPageCache,
+    primary: bool,
+    layer_sink=None,
+) -> dict[str, list[torch.Tensor]]:
+    """Build the 6 native NVFP4 source banks for a STACKED (per-layer) expert layout.
+
+    llm-compressor can store the routed experts as one packed tensor per layer instead of
+    per expert: ``...experts.gate_up_proj.weight_packed`` U8 [E*rows, cols] (rows are
+    expert-major, so the bank tensor just reshapes to [E, rows, cols]) plus ONE
+    layer-global ``weight_global_scale`` scalar (reciprocated at ingest). Placement and
+    the resulting 6-bank dict are identical to :func:`load_nvfp4_expert_source_banks`
+    (which is why the marlin/b12x repack and the offload cache never notice the
+    difference). ``layer_sink``: see :func:`load_nvfp4_expert_source_banks`."""
+    folder = download_hf_weight(model_path)
+    weight_map = _weight_map(folder)
+
+    E = config.num_experts
+    H = config.hidden_size
+    I = config.moe_intermediate_size
+    num_layers = _num_moe_layers(config)
+
+    weight_shards: dict[str, list[tuple[str, re.Match[str], int]]] = collections.defaultdict(list)
+    global_shards: dict[str, list[tuple[str, int, str]]] = collections.defaultdict(list)
+    for name, shard in weight_map.items():
+        match = spec.key_pattern.match(name)
+        if match is None:
+            continue
+        bank_layer = _bank_layer(spec, int(match.group("layer")), config)
+        if bank_layer is None:
+            continue
+        kind = _canon_kind(spec, match.group("kind"))
+        if kind == "weight_scale_2":
+            global_shards[shard].append((name, bank_layer, match.group("proj")))
+        elif kind in {"weight", "weight_scale"}:
+            weight_shards[shard].append((name, match, bank_layer))
+        else:
+            raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}")
+
+    # Pass 1: the per-layer scalar globals (reciprocal at ingest). gate_up and down carry
+    # their own global, so key by (bank_layer, proj).
+    globals_map: dict[tuple[int, str], torch.Tensor] = {}
+    for shard in sorted(global_shards):
+        path = os.path.join(folder, shard)
+        drop_page_cache(path)
+        with safetensors.safe_open(path, framework="pt", device="cpu") as f:
+            for name, bank_layer, proj in global_shards[shard]:
+                globals_map[(bank_layer, proj)] = _ingest_global(spec, f.get_tensor(name))
+        drop_page_cache(path)
+
+    _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I)  # unpinned; pinned after fill
+    gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]]
+    gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]]
+    gate_up_global = [b.tensor for b in _hb["gate_up_global"]]
+    down_packed = [b.tensor for b in _hb["down_packed"]]
+    down_scale = [b.tensor for b in _hb["down_scale"]]
+    down_global = [b.tensor for b in _hb["down_global"]]
+
+    from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline
+
+    def _load(sink) -> int:
+        tracker = LayerCompletionTracker(4, _hb, sink)  # 2 gate_up + 2 down per layer
+        placed = 0
+        for shard in tqdm(sorted(weight_shards), desc=f"Loading {spec.desc}", disable=not primary):
+            path = os.path.join(folder, shard)
+            with safetensors.safe_open(path, framework="pt", device="cpu") as f:
+                for name, match, bank_layer in weight_shards[shard]:
+                    proj = match.group("proj")
+                    kind = _canon_kind(spec, match.group("kind"))
+                    tensor = f.get_tensor(name)
+                    if kind == "weight":
+                        if proj == "gate_up_proj":
+                            gate_up_packed[bank_layer].copy_(tensor.view(E, 2 * I, H // 2))
+                        else:
+                            down_packed[bank_layer].copy_(tensor.view(E, H, I // 2))
+                    else:
+                        g = globals_map[(bank_layer, proj)]
+                        if proj == "gate_up_proj":
+                            gate_up_scale[bank_layer].copy_(tensor.view(E, 2 * I, H // 16))
+                            gate_up_global[bank_layer].fill_(g.item())
+                        else:
+                            down_scale[bank_layer].copy_(tensor.view(E, H, I // 16))
+                            down_global[bank_layer].fill_(g.item())
+                    tracker.note(bank_layer)
+                    placed += 1
+            drop_page_cache(path)
+        return placed
+
+    if layer_sink is not None:
+        placed = _load(layer_sink)
+    else:
+        with PinPipeline() as pins:
+            placed = _load(pins)
+
+    expected = num_layers * 4
+    assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}"
+    return {
+        "gate_up_packed": gate_up_packed,
+        "gate_up_scale": gate_up_scale,
+        "gate_up_global": gate_up_global,
+        "down_packed": down_packed,
+        "down_scale": down_scale,
+        "down_global": down_global,
+    }
+
+
 __all__ = [
     "Nvfp4ExpertSourceSpec",
     "load_nvfp4_expert_source_banks",
     "load_nvfp4_expert_source_banks_parallel",
+    "load_nvfp4_stacked_expert_sources",
 ]

python/freetoken/models/qwen3_5_moe/config.py

@@ -39,16 +39,110 @@ def _fp8_block_quant(hf_config: Any) -> tuple[str, tuple[int, int] | None]:
     return "none", None
 
 
+def _ct_expert_groups_nvfp4(hf_config: Any) -> bool:
+    """compressed-tensors MoE checkpoint: are the *routed experts* NVFP4? Mirrors the
+    mixed-precision scan in ``models.config.detect_expert_quant``: groups whose
+    ``targets`` name the experts decide (a generic ``["Linear"]`` group falls back to
+    covering everything). Exports without ``config_groups`` (format-only, e.g.
+    doth4580/Kwaipilot-KAT-Coder-V2.5-Dev-NVFP4-MIXED) ride the top-level format."""
+    get = _quant_accessor(hf_config)
+    if get is None:
+        return False
+    groups = get("config_groups") or {}
+    if not groups:
+        return str(get("format") or "").lower() == "nvfp4-pack-quantized"
+    groups = [g or {} for g in (groups.values() if isinstance(groups, dict) else [])]
+    expert_groups = [g for g in groups if any("experts" in str(t) for t in (g.get("targets") or []))]
+    for g in expert_groups or groups:
+        w = (g or {}).get("weights") or {}
+        if int(w.get("num_bits", 0) or 0) != 4 or str(w.get("type", "")).lower() != "float":
+            continue
+        if int(w.get("group_size", 0) or 0) == 16 and str(w.get("strategy", "")).lower() == "tensor_group":
+            return True
+    return False
+
+
+def _ct_recipe(hf_config: Any) -> dict[str, str]:
+    """Parse llm-compressor's ``recipe`` string (``all,gdn:fp8,-router``) into a map of
+    module-key -> value. Tokens: ``name:quant`` (per-module override), ``-name`` (skip),
+    bare ``name`` (quantize). Unknown keys are ignored; absence means the default."""
+    get = _quant_accessor(hf_config)
+    if get is None:
+        return {}
+    out: dict[str, str] = {}
+    for tok in str(get("recipe") or "").split(","):
+        tok = tok.strip()
+        if not tok:
+            continue
+        if tok.startswith("-"):
+            out.setdefault(tok[1:], "skip")
+        elif ":" in tok:
+            key, _, val = tok.partition(":")
+            out.setdefault(key, val.lower())
+        else:
+            out.setdefault(tok, "quant")
+    return out
+
+
+def _ct_ignored(hf_config: Any, probe: str) -> bool:
+    """Is the canonical module path ``probe`` excluded from quantization by the export's
+    ``ignore`` list (exact or ``re:`` entries) or a ``recipe`` skip token (``-lm_head``)?"""
+    get = _quant_accessor(hf_config)
+    if get is None:
+        return False
+    import re
+
+    for entry in get("ignore") or []:
+        e = str(entry)
+        if e.startswith("re:"):
+            try:
+                if re.search(e[3:], probe):
+                    return True
+            except re.error:
+                continue
+        elif e == probe:
+            return True
+    return _ct_recipe(hf_config).get(probe.split(".")[-1]) == "skip"
+
+
+def _ct_gdn_nvfp4(hf_config: Any) -> bool:
+    """Is the GDN ``out_proj`` stored NVFP4 (as opposed to per-tensor fp8 / bf16)?
+
+    The full-attention projections can be NVFP4 while the whole GDN is left bf16 (the
+    AEON Qwen3.6-35B-A3B ``ignore`` list) or per-tensor fp8 (a ``gdn:fp8`` recipe, e.g.
+    Kwaipilot-KAT-Coder-V2.5). Returns False in both cases; True when the GDN rides the
+    default NVFP4 format."""
+    if _ct_ignored(hf_config, "model.language_model.layers.0.linear_attn.out_proj"):
+        return False
+    return _ct_recipe(hf_config).get("gdn", "nvfp4") in ("nvfp4", "fp4", "quant", "all")
+
+
+def _ct_linear_attn_ignored(hf_config: Any) -> bool:
+    """Backward-compatible wrapper: the GDN out_proj is "ignored" exactly when it is not
+    NVFP4 (the ignore-list or ``recipe`` override left it fp8/bf16)."""
+    return not _ct_gdn_nvfp4(hf_config)
+
+
 def _expert_quant(hf_config: Any) -> str:
     """Quantization format of the *routed* experts (the only weights served from the
     offload cache). The nvidia/modelopt checkpoints are either plain NVFP4 (``quant_algo``
     ``NVFP4``) or ``MIXED_PRECISION`` (per-layer ``quantized_layers`` map); in the mixed
-    case the routed experts carry their own ``W4A16_NVFP4``/``FP8`` algo. Dense quantized
-    weights (attention/shared-expert/lm_head) are handled separately by dequant-at-load."""
+    case the routed experts carry their own ``W4A16_NVFP4``/``FP8`` algo. llm-compressor
+    (compressed-tensors) MoE checkpoints keep the routed experts NVFP4 in the offload
+    cache. Dense quantized weights (attention/shared-expert/lm_head) are handled
+    separately by dequant-at-load."""
     get = _quant_accessor(hf_config)
     if get is None:
         return "none"
     algo = str(get("quant_algo") or get("quant_method") or "").lower()
+    if algo == "compressed-tensors":
+        # Dense exports (e.g. Qwen3.6-27B, no routed experts) stay "none" -- their
+        # weights ride the dense compressed-tensors reader. An MoE export (e.g. the
+        # Qwen3.6-35B-A3B NVFP4 builds) keeps its experts NVFP4 in the offload cache.
+        text = getattr(hf_config, "text_config", hf_config)
+        if int(getattr(text, "num_experts", 0) or 0) > 0 and _ct_expert_groups_nvfp4(hf_config):
+            return "nvfp4"
+        return "none"
     if "fp4" in algo:
         return "nvfp4"
     if "mixed" in algo:
@@ -179,13 +273,18 @@ def parse_config(hf_config: Any) -> ModelConfig:
     dense_quant = "nvfp4" if expert_quant == "nvfp4" else _dense_mlp_quant(hf_config)
     lm_head_quant = _lm_head_quant(hf_config)
 
-    # compressed-tensors NVFP4 (dense Qwen3.6-27B): the attention (q/k/v/o, GDN out_proj) AND
-    # the dense MLP are W4A16 NVFP4; GDN in_proj_*, lm_head, norms stay bf16. Wire the shared
-    # W4A16 kernels (attn_quant=="nvfp4" routes the attention/GDN linears through them too).
+    # compressed-tensors NVFP4: the attention (q/k/v/o) and dense MLP are W4A16 NVFP4.
+    # The GDN out_proj follows only when the export actually quantized it (its
+    # ``ignore``/``recipe`` may leave the GDN bf16 or fp8). lm_head is NVFP4 unless the
+    # export skipped it (dense Qwen3.6-27B and the AEON MoE builds put it in ``ignore``).
     if _compressed_tensors_nvfp4(hf_config):
         attn_quant = "nvfp4"
         dense_quant = "nvfp4"
-        lm_head_quant = "none"
+        lm_head_quant = "none" if _ct_ignored(hf_config, "lm_head") else "nvfp4"
+    # The GDN's out_proj follows attn_quant EXCEPT when the export left it non-NVFP4
+    # (the AEON ``ignore`` list keeps the GDN bf16; a ``gdn:fp8`` recipe keeps it
+    # per-tensor fp8); the full-attention q/k/v/o are quantized either way.
+    gdn_quant = "none" if (attn_quant == "nvfp4" and not _ct_gdn_nvfp4(hf_config)) else attn_quant
 
     # Dense variants (e.g. Qwen3.6-27B) report num_experts==0: route the decoder MLP through
     # the dense Qwen3_5DenseMLP instead of the MoE block.
@@ -255,6 +354,7 @@ def parse_config(hf_config: Any) -> ModelConfig:
         expert_quant=expert_quant,
         weight_block_size=weight_block_size,
         attn_quant=attn_quant,
+        gdn_quant=gdn_quant,
         dense_quant=dense_quant,
         lm_head_quant=lm_head_quant,
     )

python/freetoken/models/qwen3_5_moe/gdn.py

@@ -53,7 +53,7 @@ class Qwen3_5GatedDeltaNet(BaseOP):
     def __init__(
         self, hidden_size, num_k_heads, num_v_heads, head_k_dim, head_v_dim,
         conv_kernel_size, rms_norm_eps, layer_id, expert_quant: str = "none",
-        attn_quant: str = "none",
+        attn_quant: str = "none", gdn_quant: str | None = None,
     ):
         self.layer_id = layer_id
         # The fla chunk/decode kernels read+write the recurrent state and the per-chunk h as
@@ -100,9 +100,11 @@ def __init__(
         self.norm = _GatedRMSNorm(head_v_dim, eps=rms_norm_eps)
         # out_proj follows the checkpoint quant: block-fp8 / per-tensor-fp8 / compressed-tensors
         # NVFP4 (W4A16) / bf16. in_proj_* stay bf16 in every mode (above), so a compressed-tensors
-        # NVFP4 checkpoint (attn_quant=="nvfp4") only makes out_proj native FP4.
+        # NVFP4 checkpoint only makes out_proj native FP4 -- and only when the export actually
+        # quantized it (its ``ignore`` list may leave the whole GDN bf16, see config.gdn_quant).
         self.out_proj = make_replicated_quant(
-            expert_quant, attn_quant, self.value_dim, hidden_size, has_bias=False
+            expert_quant, attn_quant if gdn_quant is None else gdn_quant,
+            self.value_dim, hidden_size, has_bias=False,
         )
 
     def _gate_params(self, a: torch.Tensor, b: torch.Tensor):

python/freetoken/models/qwen3_5_moe/model.py

@@ -44,6 +44,7 @@ def __init__(self, config: ModelConfig, layer_id: int):
                 layer_id=layer_id,
                 expert_quant=config.expert_quant,
                 attn_quant=config.attn_quant,
+                gdn_quant=config.gdn_quant,
             )
         else:
             self.self_attn = Qwen3_5Attention(config, layer_id)