feat(qwen3_5_moe): support compressed-tensors NVFP4 experts with block-fp8 dense side

python/freetoken/models/nvfp4_banks.py

@@ -320,8 +320,272 @@ def _load(sink) -> int:
     }
 
 
+
+# --------------------------------------------------------------------------------------
+# compressed-tensors NVFP4 experts (llm-compressor layout): weight_packed (uint8) +
+# weight_scale (fp8-e4m3 block) + scalar weight_global_scale. The dequant global is the
+# reciprocal of the stored quant-side global (vLLM inverts it identically), broadcast
+# per output row. Same native source banks as the modelopt path; single-file (no index)
+# checkpoints are supported. [mixed-ct block-fp8 dense + nvfp4 experts]
+# --------------------------------------------------------------------------------------
+
+
+def _checkpoint_weight_map(folder: str) -> dict[str, str]:
+    """name -> shard basename, from model.safetensors.index.json when present, else from
+    the single-file safetensors header(s)."""
+    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"]
+    import glob
+    import struct
+
+    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 nm in hdr:
+            if nm != "__metadata__":
+                weight_map[nm] = shard
+    return weight_map
+
+
+def _ct_banks_tuple(_hb):
+    return (
+        [b.tensor for b in _hb["gate_up_packed"]],
+        [b.tensor for b in _hb["gate_up_scale"]],
+        [b.tensor for b in _hb["gate_up_global"]],
+        [b.tensor for b in _hb["down_packed"]],
+        [b.tensor for b in _hb["down_scale"]],
+        [b.tensor for b in _hb["down_global"]],
+    )
+
+
+def _ct_place(tensor, kind, role, bank_layer_id, expert, I, banks, global_val):
+    """Place one compressed-tensors NVFP4 expert tensor into the native source banks."""
+    gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global = banks
+    if kind == "weight_packed":
+        if role == "gate":
+            gate_up_packed[bank_layer_id][expert, :I] = tensor
+        elif role == "up":
+            gate_up_packed[bank_layer_id][expert, I:] = tensor
+        else:
+            down_packed[bank_layer_id][expert] = tensor
+    else:  # weight_scale (fp8-e4m3 block); the global is a per-projection scalar -> fill rows
+        if role == "gate":
+            gate_up_scale[bank_layer_id][expert, :I] = tensor
+            gate_up_global[bank_layer_id][expert, :I].fill_(global_val)
+        elif role == "up":
+            gate_up_scale[bank_layer_id][expert, I:] = tensor
+            gate_up_global[bank_layer_id][expert, I:].fill_(global_val)
+        else:
+            down_scale[bank_layer_id][expert] = tensor
+            down_global[bank_layer_id][expert].fill_(global_val)
+
+
+def _ct_sort_shards(weight_map, spec):
+    """Bucket matched expert tensors: weight_global_scale -> globals pass; weight_packed /
+    weight_scale -> bulk pass. Returns (weight_shards, global_names_by_shard)."""
+    weight_shards: dict[str, list[tuple[str, re.Match[str], int]]] = collections.defaultdict(list)
+    global_shards: dict[str, list[str]] = collections.defaultdict(list)
+    for name, shard in weight_map.items():
+        match = spec.key_pattern.match(name)
+        if match is None:
+            continue
+        # Bounds checking against num_layers happens in the parallel builder (which has the
+        # config); the serial build trusts the contiguous 0..L-1 Qwen3.5 MoE layer ids.
+        bank_layer = spec.layer_to_bank(int(match.group("layer")), None)
+        if bank_layer is None:
+            continue
+        kind = match.group("kind")
+        if kind == "weight_global_scale":
+            global_shards[shard].append(name)
+        elif kind in ("weight_packed", "weight_scale"):
+            weight_shards[shard].append((name, match, bank_layer))
+        else:
+            raise ValueError(f"{spec.desc}: unknown CT expert tensor kind {kind!r}")
+    return weight_shards, global_shards
+
+
+def _ct_load_globals(folder, global_shards, spec, drop_page_cache):
+    """Per-(layer, expert, proj) dequant global = 1/weight_global_scale, as python floats."""
+    globals_map: dict[tuple, float] = {}
+    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 in global_shards[shard]:
+                m = spec.key_pattern.match(name)
+                wg = f.get_tensor(name).reshape(-1).to(torch.float32)
+                assert wg.numel() == 1, (
+                    f"{spec.desc}: expected scalar weight_global_scale, got {tuple(wg.shape)}"
+                )
+                globals_map[(int(m.group("layer")), int(m.group("expert")), m.group("proj"))] = (
+                    (1.0 / wg[0]).to(torch.float16).item()
+                )
+        drop_page_cache(path)
+    return globals_map
+
+
+def load_nvfp4_ct_expert_source_banks(
+    model_path: str,
+    config,
+    spec: Nvfp4ExpertSourceSpec,
+    *,
+    drop_page_cache: DropPageCache,
+    primary: bool,
+    layer_sink=None,
+) -> dict[str, list[torch.Tensor]]:
+    """Serial per-shard build of the native NVFP4 source banks from a compressed-tensors
+    (llm-compressor) NVFP4 expert checkpoint."""
+    folder = download_hf_weight(model_path)
+    weight_map = _checkpoint_weight_map(folder)
+
+    E = config.num_experts
+    H = config.hidden_size
+    I = config.moe_intermediate_size
+    num_layers = _num_moe_layers(config)
+
+    for shard in sorted(set(weight_map.values())):
+        drop_page_cache(os.path.join(folder, shard))
+
+    weight_shards, global_shards = _ct_sort_shards(weight_map, spec)
+    globals_map = _ct_load_globals(folder, global_shards, spec, drop_page_cache)
+
+    _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I)
+    banks = _ct_banks_tuple(_hb)
+
+    from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline
+
+    def _load(sink) -> int:
+        tracker = LayerCompletionTracker(E * 6, _hb, sink)
+        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_id in weight_shards[shard]:
+                    layer = int(match.group("layer"))
+                    expert = int(match.group("expert"))
+                    proj = match.group("proj")
+                    role = spec.proj_to_role[proj]
+                    kind = match.group("kind")
+                    g = globals_map[(layer, expert, proj)]
+                    _ct_place(f.get_tensor(name), kind, role, bank_layer_id, expert, I, banks, g)
+                    tracker.note(bank_layer_id)
+                    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 * E * 6
+    assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}"
+    gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global = banks
+    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,
+    }
+
+
+def load_nvfp4_ct_expert_source_banks_parallel(
+    model_path: str,
+    config,
+    spec: Nvfp4ExpertSourceSpec,
+    *,
+    drop_page_cache: DropPageCache,
+    primary: bool,
+    workers: int = 8,
+    chunk: int = 8 << 20,
+    layer_sink=None,
+) -> dict[str, list[torch.Tensor]]:
+    """Parallel counterpart of load_nvfp4_ct_expert_source_banks (byte-for-byte same
+    placement): bulk weight_packed/weight_scale via the chunked O_DIRECT reader; tiny
+    globals stay serial."""
+    from freetoken.models.weight import iter_expert_tensors_parallel
+
+    folder = download_hf_weight(model_path)
+    weight_map = _checkpoint_weight_map(folder)
+
+    E = config.num_experts
+    H = config.hidden_size
+    I = config.moe_intermediate_size
+    num_layers = _num_moe_layers(config)
+
+    weight_info: dict[str, tuple[re.Match[str], int]] = {}
+    global_shards: dict[str, list[str]] = collections.defaultdict(list)
+    for name, shard in weight_map.items():
+        match = spec.key_pattern.match(name)
+        if match is None:
+            continue
+        bank_layer = spec.layer_to_bank(int(match.group("layer")), config)
+        if bank_layer is None:
+            continue
+        _bank_layer(spec, int(match.group("layer")), config)  # bounds check
+        kind = match.group("kind")
+        if kind == "weight_global_scale":
+            global_shards[shard].append(name)
+        elif kind in ("weight_packed", "weight_scale"):
+            weight_info[name] = (match, bank_layer)
+        else:
+            raise ValueError(f"{spec.desc}: unknown CT expert tensor kind {kind!r}")
+
+    globals_map = _ct_load_globals(folder, global_shards, spec, drop_page_cache)
+
+    _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I)
+    banks = _ct_banks_tuple(_hb)
+
+    from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline
+
+    def _load(sink) -> int:
+        tracker = LayerCompletionTracker(E * 6, _hb, sink)
+        placed = 0
+        for name, tensor in iter_expert_tensors_parallel(
+            folder, lambda n: n in weight_info, workers=workers, chunk=chunk
+        ):
+            match, bank_layer_id = weight_info[name]
+            layer = int(match.group("layer"))
+            expert = int(match.group("expert"))
+            proj = match.group("proj")
+            role = spec.proj_to_role[proj]
+            kind = match.group("kind")
+            g = globals_map[(layer, expert, proj)]
+            _ct_place(tensor, kind, role, bank_layer_id, expert, I, banks, g)
+            tracker.note(bank_layer_id)
+            placed += 1
+        return placed
+
+    if layer_sink is not None:
+        placed = _load(layer_sink)
+    else:
+        with PinPipeline() as pins:
+            placed = _load(pins)
+
+    expected = num_layers * E * 6
+    assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}"
+    gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global = banks
+    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_ct_expert_source_banks",
+    "load_nvfp4_ct_expert_source_banks_parallel",
 ]

python/freetoken/models/qwen3_5_moe/config.py

@@ -68,6 +68,30 @@ def _expert_quant(hf_config: Any) -> str:
 _compressed_tensors_nvfp4 = detect_compressed_tensors_nvfp4
 
 
+def _compressed_tensors_mixed_fp8block(hf_config: Any) -> bool:
+    """compressed-tensors mixed-precision checkpoint: routed experts are NVFP4 (4-bit,
+    group_size 16, strategy tensor_group) while the dense projections (self_attn
+    q/k/v/o, linear_attn in_proj_*/out_proj, mlp.shared_expert gate/up/down) are 128x128
+    block FP8 (8-bit, strategy block, W8A8). Distinct from pure compressed-tensors NVFP4
+    (detect_compressed_tensors_nvfp4), where the dense side is packed FP4 as well."""
+    get = _quant_accessor(hf_config)
+    if get is None:
+        return False
+    if str(get("quant_method") or "").lower() != "compressed-tensors":
+        return False
+    saw_fp8_block = saw_nvfp4 = False
+    for gspec in (get("config_groups") or {}).values():
+        w = (gspec or {}).get("weights") or {}
+        bits = int(w.get("num_bits", 0) or 0)
+        strategy = str(w.get("strategy", "")).lower()
+        group_size = int(w.get("group_size", 0) or 0)
+        if bits == 8 and strategy == "block":
+            saw_fp8_block = True
+        if bits == 4 and group_size == 16 and strategy == "tensor_group":
+            saw_nvfp4 = True
+    return saw_fp8_block and saw_nvfp4
+
+
 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
@@ -164,9 +188,14 @@ def parse_config(hf_config: Any) -> ModelConfig:
         else {k: v for k, v in rope_params.items() if not isinstance(v, (list, dict))}
     )
 
+    mixed_ct = _compressed_tensors_mixed_fp8block(hf_config)
     expert_quant, weight_block_size = _fp8_block_quant(hf_config)
     if expert_quant == "none":
         expert_quant = _expert_quant(hf_config)  # nvfp4 / mixed-precision modelopt
+    if mixed_ct:
+        # Routed experts are packed NVFP4 (offload banks); dense side is 128x128 block FP8.
+        expert_quant = "nvfp4"
+        weight_block_size = (128, 128)
     # Dense attention/GDN quant is independent of the routed experts (block-fp8 already
     # quantizes both, so only probe for per-tensor FP8 when experts aren't block-fp8).
     attn_quant = "none" if expert_quant == "fp8_block" else _attn_quant(hf_config)
@@ -182,10 +211,16 @@ def parse_config(hf_config: Any) -> ModelConfig:
     # 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).
-    if _compressed_tensors_nvfp4(hf_config):
+    if _compressed_tensors_nvfp4(hf_config) and not mixed_ct:
         attn_quant = "nvfp4"
         dense_quant = "nvfp4"
         lm_head_quant = "none"
+    if mixed_ct:
+        # Dense attn/GDN projections and the shared expert stay native 128x128 block FP8
+        # (W8A8); lm_head / embeddings / norms are bf16.
+        attn_quant = "fp8_block"
+        dense_quant = "fp8_block"
+        lm_head_quant = "none"
 
     # 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.

python/freetoken/models/qwen3_5_moe/gdn.py

@@ -74,7 +74,7 @@ def __init__(
         # qkv|z carry a weight scale (block-fp8 weight_scale_inv, or per-tensor FP8
         # weight_scale); b|a stay bf16. Both quant modes therefore split the four-way
         # fusion into an fp8 qkvz GEMM + a bf16 ba GEMM (matches sglang/vLLM).
-        self._block_fp8 = expert_quant == "fp8_block"
+        self._block_fp8 = expert_quant == "fp8_block" or attn_quant == "fp8_block"
         self._pertensor_fp8 = attn_quant == "fp8_pertensor"
         self._fp8 = self._block_fp8 or self._pertensor_fp8
 

python/freetoken/models/qwen3_5_moe/moe.py

@@ -22,7 +22,9 @@ class _SharedExpert(BaseOP):
     """Always-present shared SwiGLU expert of width ``shared_expert_intermediate_size``."""
 
     def __init__(self, config: ModelConfig, hidden_size: int, intermediate_size: int):
-        if getattr(config, "expert_quant", "none") == "fp8_block":
+        if getattr(config, "expert_quant", "none") == "fp8_block" or getattr(
+            config, "dense_quant", "none"
+        ) == "fp8_block":
             self.gate_up_proj = Fp8BlockColMerged(
                 hidden_size, [intermediate_size, intermediate_size], has_bias=False
             )