fix(models): better support for mixed-precision compressed-tensors NVFP4

python/freetoken/models/config.py

@@ -91,6 +91,66 @@ def detect_compressed_tensors_nvfp4(hf_config: Any) -> bool:
     return saw_nvfp4
 
 
+def _nvfp4_global_reciprocal(hf_config: Any) -> bool:
+    """Whether the routed-experts' per-tensor global scale needs to be reciprocated
+    at load time.
+
+    llm-compressor stores the QUANT-side scale (the multiplier baked into the local
+    fp8 scales before the FP4 cast); the dequant kernel divides by it, so the loader
+    must reciprocate 1/x. Modelopt's NVFP4 stores the DEQUANT-side divisor directly
+    and the kernel multiplies (no transform).
+
+    The ``format: nvfp4-pack-quantized`` string on the config_group is set by BOTH
+    exporters, so it alone can't disambiguate. The ground-truth signal is the on-disk
+    tensor naming: llm-compressor uses ``weight_packed`` + ``weight_global_scale`` +
+    ``input_global_scale``; modelopt uses ``weight`` + ``weight_scale_2`` + ``input_scale``.
+    We probe the safetensors index for one routed-expert tensor and pick the convention
+    from its sibling suffixes.
+
+    Returns True (reciprocate) for llm-compressor naming, False (use as-is) for
+    modelopt naming. Returns False when nothing is found (caller can't decide, so the
+    bank loader keeps the on-disk value verbatim -- the modelopt default)."""
+    # Cheap pre-check: ``quant_method`` should be compressed-tensors for either case.
+    quant = getattr(hf_config, "quantization_config", None)
+    if not quant:
+        return False
+    get = quant.get if isinstance(quant, dict) else (lambda k, d=None: getattr(quant, k, d))
+    if str(get("quant_method") or "").lower() != "compressed-tensors":
+        return False
+
+    # Probe the safetensors index for an expert tensor and inspect its sibling suffixes.
+    # Cheap read (one json parse) and bounds-checked (only the first hit is needed).
+    try:
+        from pathlib import Path as _Path
+        from freetoken.utils.hf import download_hf_weight as _download
+        folder = _Path(_download(hf_config._name_or_path
+                                 if hasattr(hf_config, "_name_or_path") else None))
+        idx_path = folder / "model.safetensors.index.json"
+        if not idx_path.exists():
+            return False  # single-file or no index -- fall back to safe default
+        import json as _json
+        weight_map = _json.loads(idx_path.read_text())["weight_map"]
+    except Exception:
+        return False  # any read error -- don't reciprocate (modelopt default is safer)
+
+    # Find one expert key in any shard; the suffix names reveal the layout.
+    has_packed = False       # llm-compressor writes weight_packed
+    has_scale2 = False       # modelopt writes weight_scale_2
+    for name in weight_map:
+        if ".mlp.experts." in name and ".gate_proj." in name:
+            has_packed = has_packed or name.endswith(".weight_packed")
+            has_scale2 = has_scale2 or name.endswith(".weight_scale_2")
+            if has_packed or has_scale2:
+                break
+    # Disambiguate: llm-compressor uses weight_packed (and weight_global_scale); modelopt
+    # uses weight + weight_scale_2 directly. A safetensors checkpoint that has weight_packed
+    # for routed experts is llm-compressor; one that has weight_scale_2 is modelopt.
+    if has_packed:
+        return True
+    return False
+
+
+
 @dataclass(frozen=True)
 class RotaryConfig:
     head_dim: int
@@ -270,6 +330,11 @@ class ModelConfig:
     expert_quant: str = "none"
     # NVFP4 routed-expert GEMM backend (--nvfp4-backend); injected from EngineConfig.
     nvfp4_backend: str = "triton"
+    # Reciprocate the per-tensor NVFP4 expert global scale at load (llm-compressor NVFP4
+    # stores the QUANT-side scale; the dequant kernel multiplies by the DEQUANT-side value,
+    # so we read 1/x from disk). False for modelopt NVFP4 (dequant-side value is on disk).
+    # See ``_nvfp4_global_reciprocal`` for the heuristic.
+    nvfp4_global_reciprocal: bool = False
     # Block size (out, in) for block-wise weight quantization (fp8_block: (128, 128)).
     weight_block_size: tuple[int, int] | None = None
     # Weight quantization of the *dense* attention / GatedDeltaNet projections (separate

python/freetoken/models/nvfp4_banks.py

@@ -130,9 +130,13 @@ def load_nvfp4_expert_source_banks(
             raise ValueError(f"{spec.desc}: unknown NVFP4 expert projection {proj!r}")
         kind = _canon_kind(spec, match.group("kind"))
         if kind == "weight_scale_2":
+            # modelopt's ``weight_scale_2`` and llm-compressor's ``weight_global_scale``
+            # (aliased via spec.kind_map) both carry the per-row global.
             global_shards[shard].append((name, match, bank_layer))
         elif kind in {"weight", "weight_scale"}:
             weight_shards[shard].append((name, match, bank_layer))
+        elif kind in {"input_scale", "input_global_scale"}:
+            pass  # activation scale (W4A8 / W8A8); consumed with its .weight
         else:
             raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}")
 
@@ -254,9 +258,13 @@ def load_nvfp4_expert_source_banks_parallel(
             continue
         kind = _canon_kind(spec, match.group("kind"))
         if kind == "weight_scale_2":
+            # modelopt names the per-row global ``weight_scale_2``; llm-compressor names it
+            # ``weight_global_scale`` (after ``_canon_kind`` maps via spec.kind_map, or alias).
             global_names_by_shard[shard].append(name)
         elif kind in {"weight", "weight_scale"}:
             weight_info[name] = (match, bank_layer)
+        elif kind in {"input_scale", "input_global_scale"}:
+            pass  # activation scale (W4A8 / W8A8); consumed with its .weight, not a bank tensor
         else:
             raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}")
 

python/freetoken/models/qwen3_5_moe/config.py

@@ -7,7 +7,9 @@
     LinearGatedDeltaGroupConfig,
     ModelConfig,
     RotaryConfig,
+    _nvfp4_global_reciprocal,
     detect_compressed_tensors_nvfp4,
+    detect_expert_quant,
 )
 
 
@@ -39,6 +41,15 @@ def _fp8_block_quant(hf_config: Any) -> tuple[str, tuple[int, int] | None]:
     return "none", None
 
 
+def _has_moe_experts(hf_config: Any) -> bool:
+    """True if the model declares any routed experts (``text_config.num_experts > 0``).
+    Used to disambiguate dense compressed-tensors NVFP4 (Qwen3.6-27B) from MoE mixed-precision
+    compressed-tensors (Ornith: FP8 attn/shared_expert + NVFP4 experts) so per-group quant
+    detection isn't overridden by a blanket "everything is NVFP4" assumption."""
+    text = getattr(hf_config, "text_config", hf_config)
+    return (getattr(text, "num_experts", 0) or 0) > 0
+
+
 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``
@@ -60,6 +71,12 @@ def _expert_quant(hf_config: Any) -> str:
                     return "nvfp4"
                 if "fp8" in expert_algo:
                     return "fp8"
+    if str(get("quant_method") or "").lower() == "compressed-tensors":
+        # llm-compressor: the routed experts' group carries ``format: nvfp4-pack-quantized``.
+        # ``detect_expert_quant`` (top-level) already returns "nvfp4" for this -- mirror it
+        # here so the offload expert bank activates.
+        if detect_expert_quant(hf_config) == "nvfp4":
+            return "nvfp4"
     return "none"
 
 
@@ -108,17 +125,31 @@ def _attn_quant(hf_config: Any) -> str:
     """Per-tensor FP8 on the *dense* attention/GDN projections. The modelopt
     ``MIXED_PRECISION`` checkpoints tag ``self_attn.{q,k,v,o}_proj`` and
     ``linear_attn.{in_proj_qkv,in_proj_z,out_proj}`` with ``quant_algo`` ``FP8`` (fp8-e4m3
-    weight + a scalar ``weight_scale``; W8A16). Returns ``"fp8_pertensor"`` when present,
-    else ``"none"`` (NVFP4 dense weights -- shared_expert/lm_head -- stay dequant-at-load)."""
+    weight + a scalar ``weight_scale``; W8A16). llm-compressor mixed-precision checkpoints
+    (Ornith) use a per-tensor-FP8 ``config_groups`` entry targeting the same projections.
+    Returns ``"fp8_pertensor"`` when present, else ``"none"``."""
     get = _quant_accessor(hf_config)
     if get is None:
         return "none"
     layers = get("quantized_layers") or {}
-    if not isinstance(layers, dict):
-        return "none"
-    for name, spec in layers.items():
-        algo = str((spec or {}).get("quant_algo", "")).lower()
-        if algo == "fp8" and (".self_attn." in name or ".linear_attn." in name):
+    if isinstance(layers, dict):
+        for name, spec in layers.items():
+            algo = str((spec or {}).get("quant_algo", "")).lower()
+            if algo == "fp8" and (".self_attn." in name or ".linear_attn." in name):
+                return "fp8_pertensor"
+    # llm-compressor path: a config_group with format "float-quantized" and per-tensor
+    # strategy (``group_size is None and strategy == "tensor"``) targeting ``.self_attn.``
+    # or ``.linear_attn.`` is per-tensor FP8 (W8A16). Targets are regex strings
+    # (``re:.*\.self_attn\.``), so the substring check uses the unescaped leaf path.
+    for g in (get("config_groups") or {}).values():
+        if not g:
+            continue
+        if not any((".self_attn" in t or "linear_attn" 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") == "tensor":
             return "fp8_pertensor"
     return "none"
 
@@ -170,19 +201,36 @@ def parse_config(hf_config: Any) -> ModelConfig:
     # 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)
-    # NVFP4 checkpoints store the dense MLP projections (shared_expert; dense non-MoE MLP) as
-    # packed FP4 exactly like the routed experts -- independent of whether attention is FP8
-    # (mixed) or bf16 (pure NVFP4). Keep them native FP4 (W4A16) whenever the experts are
-    # NVFP4. The lm_head is detected separately (only the mixed checkpoint quantizes it).
-    # MoE-NVFP4 keeps the shared_expert dense MLP native FP4 (expert_quant=="nvfp4"); a dense
-    # (non-MoE) modelopt checkpoint instead tags the bare .mlp.{gate,up,down}_proj as NVFP4.
-    dense_quant = "nvfp4" if expert_quant == "nvfp4" else _dense_mlp_quant(hf_config)
+    # Default: shared expert / dense MLP match the routed experts' quant (NVFP4 when experts
+    # are NVFP4). The modelopt dense (non-MoE) checkpoint sets a per-projection override via
+    # ``quantized_layers`` that ``_dense_mlp_quant`` picks up. llm-compressor mixed-precision
+    # MoE (Ornith: NVFP4 experts + FP8 attn + FP8 shared_expert) is the one true exception:
+    # its config_groups put the shared_expert in the FP8 group, so we explicitly read that
+    # and override dense_quant to "none" -- the FP8 path is selected via attn_quant.
+    if expert_quant == "nvfp4":
+        mlp_q = _dense_mlp_quant(hf_config)
+        dense_quant = mlp_q if mlp_q != "none" else "nvfp4"
+        if _compressed_tensors_nvfp4(hf_config) and _has_moe_experts(hf_config):
+            quant = getattr(hf_config, "quantization_config", None)
+            get = quant.get if isinstance(quant, dict) else (lambda k, d=None: getattr(quant, k, d))
+            for g in (get("config_groups") or {}).values():
+                if not g:
+                    continue
+                if any("shared_expert" in t for t in (g.get("targets") or [])):
+                    if str(g.get("format") or "").lower() == "float-quantized":
+                        dense_quant = "none"
+                    break
+    else:
+        dense_quant = _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).
-    if _compressed_tensors_nvfp4(hf_config):
+    # compressed-tensors NVFP4 *dense* (Qwen3.6-27B-NVFP4): every Linear is W4A16 NVFP4
+    # (q/k/v/o, GDN out_proj, dense MLP). GDN in_proj_*, lm_head, norms stay bf16. Wire the
+    # shared W4A16 kernels (attn_quant=="nvfp4" routes the attention/GDN linears through
+    # them too). MoE compressed-tensors checkpoints (Ornith: FP8 attn + FP8 shared_expert
+    # + NVFP4 experts) skip the override -- the per-group detectors above already chose
+    # the right per-tensor quant for each block.
+    if _compressed_tensors_nvfp4(hf_config) and not _has_moe_experts(hf_config):
         attn_quant = "nvfp4"
         dense_quant = "nvfp4"
         lm_head_quant = "none"
@@ -257,6 +305,7 @@ def parse_config(hf_config: Any) -> ModelConfig:
         attn_quant=attn_quant,
         dense_quant=dense_quant,
         lm_head_quant=lm_head_quant,
+        nvfp4_global_reciprocal=_nvfp4_global_reciprocal(hf_config),
     )
 
 

python/freetoken/models/qwen3_5_moe/moe.py

@@ -3,43 +3,41 @@
 from typing import TYPE_CHECKING
 
 import torch
-from freetoken.layers import (
-    BaseOP,
-    LinearColParallelMerged,
-    LinearReplicated,
-    LinearRowParallel,
-    make_moe_layer,
-    silu_and_mul,
-)
-
-from freetoken.kernel.triton.fp8_block_linear import Fp8BlockColMerged, Fp8BlockLinear
+from freetoken.layers import BaseOP, LinearReplicated, make_moe_layer, silu_and_mul
 
 if TYPE_CHECKING:
     from freetoken.models.config import ModelConfig
 
 
 class _SharedExpert(BaseOP):
-    """Always-present shared SwiGLU expert of width ``shared_expert_intermediate_size``."""
+    """Always-present shared SwiGLU expert of width ``shared_expert_intermediate_size``.
+    The quant type follows ``dense_quant`` (NVFP4 W4A16 for modelopt/shared_expert NVFP4
+    and dense Qwen3.6-27B; per-tensor FP8 W8A8 for llm-compressor mixed-precision Ornith),
+    then ``attn_quant`` (per-tensor FP8 when shared_expert is in the FP8 group), then
+    ``expert_quant`` (block FP8), else bf16. The dispatch lives in models/quant_linear.py
+    so every dense col-merged linear (attn, GDN, shared_expert) picks the same kernel."""
 
     def __init__(self, config: ModelConfig, hidden_size: int, intermediate_size: int):
-        if getattr(config, "expert_quant", "none") == "fp8_block":
-            self.gate_up_proj = Fp8BlockColMerged(
-                hidden_size, [intermediate_size, intermediate_size], has_bias=False
-            )
-            self.down_proj = Fp8BlockLinear(intermediate_size, hidden_size, has_bias=False)
-        elif getattr(config, "dense_quant", "none") == "nvfp4":
-            # NVFP4 checkpoint: keep the shared expert's NVFP4 weights native (W4A16).
-            from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged, Nvfp4DenseLinear
-
-            self.gate_up_proj = Nvfp4DenseColMerged(
-                hidden_size, [intermediate_size, intermediate_size], has_bias=False
-            )
-            self.down_proj = Nvfp4DenseLinear(intermediate_size, hidden_size, has_bias=False)
-        else:
-            self.gate_up_proj = LinearColParallelMerged(
-                hidden_size, [intermediate_size, intermediate_size], has_bias=False
-            )
-            self.down_proj = LinearRowParallel(intermediate_size, hidden_size, has_bias=False)
+        from freetoken.models.quant_linear import make_col_merged_quant, make_replicated_quant
+
+        # ``dense_quant`` wins for the NVFP4 path (shared_expert NVFP4 in modelopt mixed
+        # precision and in dense Qwen3.6-27B); ``attn_quant`` for FP8 paths.
+        # The factory's pre-NVFP4 check (``expert_quant == "fp8_block"``) still fires for
+        # block-fp8 models; per-tensor FP8 takes the attn_quant path; everything else bf16.
+        dense_quant = getattr(config, "dense_quant", "none")
+        attn_quant = getattr(config, "attn_quant", "none")
+        expert_quant = getattr(config, "expert_quant", "none")
+        # For NVFP4 dense, the factory's check is on ``attn_quant``; promote ``dense_quant``
+        # to ``attn_quant`` for the call so a single dispatch covers both. Block-fp8 and
+        # per-tensor-fp8 are unchanged.
+        fused_attn = "nvfp4" if dense_quant == "nvfp4" else attn_quant
+        self.gate_up_proj = make_col_merged_quant(
+            expert_quant, fused_attn, hidden_size,
+            [intermediate_size, intermediate_size], has_bias=False,
+        )
+        self.down_proj = make_replicated_quant(
+            expert_quant, fused_attn, intermediate_size, hidden_size, has_bias=False,
+        )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.down_proj.forward(silu_and_mul(self.gate_up_proj.forward(x)))