qwen3_5_moe: load mixed-precision compressed-tensors checkpoints (unsloth NVFP4+FP8)

docs/models.md

@@ -11,6 +11,7 @@ for them; other checkpoints of the same architectures work too.
 | GLM-4.7 | [nvidia/GLM-4.7-NVFP4](https://huggingface.co/nvidia/GLM-4.7-NVFP4) |
 | Qwen3.6 / Qwen3.5 MoE | [Qwen/Qwen3.6-35B-A3B](https://huggingface.co/Qwen/Qwen3.6-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-35B-A3B-FP8)), [nvidia/Qwen3.6-35B-A3B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-35B-A3B-NVFP4), [Qwen/Qwen3.5-35B-A3B](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) ([-FP8](https://huggingface.co/Qwen/Qwen3.5-35B-A3B-FP8)) |
 | Qwen3.6 dense | [Qwen/Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) ([-FP8](https://huggingface.co/Qwen/Qwen3.6-27B-FP8)), [nvidia/Qwen3.6-27B-NVFP4](https://huggingface.co/nvidia/Qwen3.6-27B-NVFP4) |
+| Qwen3.8 dense | [unsloth/Qwen3.8-27B-NVFP4](https://huggingface.co/unsloth/Qwen3.8-27B-NVFP4) (mixed NVFP4+FP8 export) |
 | Qwen3-MoE | [Qwen/Qwen3-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) |
 | gpt-oss | [openai/gpt-oss-120b](https://huggingface.co/openai/gpt-oss-120b), [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) |
 | Gemma-4 | [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it), [nvidia/Gemma-4-26B-A4B-NVFP4](https://huggingface.co/nvidia/Gemma-4-26B-A4B-NVFP4), [google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it), [nvidia/Gemma-4-31B-IT-NVFP4](https://huggingface.co/nvidia/Gemma-4-31B-IT-NVFP4) .. |
@@ -38,3 +39,6 @@ for them; other checkpoints of the same architectures work too.
 - DeepSeek-V4 checkpoints must keep the `inference/config.json` subdir — the
   authoritative model args are read from there.
 - Multimodal checkpoints are served text-only.
+- unsloth's dense Qwen3.x NVFP4 exports are mixed-precision per module (NVFP4 MLP
+  layers, FP8 attention/GDN/lm_head, bf16 residual parts); each linear loads in the
+  storage the checkpoint actually uses.

python/freetoken/models/config.py

@@ -243,6 +243,11 @@ class ModelConfig:
     # it bf16. Separate from dense_quant because only some NVFP4 checkpoints quantize lm_head
     # (modelopt MIXED_PRECISION does; pure NVFP4 leaves it bf16).
     lm_head_quant: str = "none"
+    # Per-layer override of the dense-MLP storage for mixed compressed-tensors exports
+    # (unsloth: most layers packed NVFP4, a few FP8 layers). Keys are layer ids; the value
+    # replaces dense_quant for that layer's dense MLP at construction ("nvfp4" native W4A16
+    # / "bf16" dequant-at-load). None = every dense-MLP layer follows dense_quant.
+    dense_mlp_storage: dict[int, str] | None = None
     shared_expert_intermediate_size: int = 0
     use_qk_norm: bool = False
     # ----- DeepSeek/GLM-style MoE extensions (default keeps other models intact) -----

python/freetoken/models/loader.py

@@ -182,6 +182,9 @@ def files(self) -> list[str]:
     def names_in(self, file: str) -> list[str]:
         return [name for name, shard in self._map.items() if shard == file]
 
+    def has(self, name: str) -> bool:
+        return name in self._map
+
     def get_tensor(self, name: str) -> torch.Tensor:
         import safetensors
 

python/freetoken/models/qwen3_5_moe/config.py

@@ -68,6 +68,77 @@ def _expert_quant(hf_config: Any) -> str:
 _compressed_tensors_nvfp4 = detect_compressed_tensors_nvfp4
 
 
+def _compressed_linear_storage(
+    hf_config: Any,
+) -> tuple[str, str, dict[int, str] | None, bool]:
+    """Per-module dense-linear storage sniffed from the checkpoint's
+    ``model.safetensors.index.json`` weight_map:
+    ``(attention storage, dense-MLP fallback, per-layer dense-MLP overrides, lm_head
+    is fp8)``.
+
+    The attention storage is ``"nvfp4"`` (packed q/k/v/o + GDN out_proj linears),
+    ``"fp8"`` (weight + scale siblings, e.g. unsloth's mixed export), or ``"none"``
+    (bf16). The fallback is ``"nvfp4"`` when at least one dense-MLP layer is packed,
+    else ``"none"``. The override map gives the storage of dense-MLP layers that are
+    not packed NVFP4 (``"fp8"`` keeps them native W8A16; ``"bf16"`` dequantizes at
+    load); ``None`` means every dense-MLP layer follows the fallback. When the index
+    is unavailable (a hub id before download, or a single-file checkpoint) assume
+    the official dense-NVFP4 layout (all packed, no overrides, bf16 lm_head).
+    """
+    import json
+    import os
+    import re
+
+    path = getattr(hf_config, "name_or_path", None)
+    if not isinstance(path, str) or not os.path.isdir(path):
+        return "nvfp4", "nvfp4", None, False
+    index = os.path.join(path, "model.safetensors.index.json")
+    if not os.path.isfile(index):
+        return "nvfp4", "nvfp4", None, False
+    with open(index, encoding="utf-8") as f:
+        weight_map = json.load(f).get("weight_map", {})
+    keys = set(weight_map)
+    attn_suffix_re = re.compile(r"\.(?:self_attn\.(?:q|k|v|o)_proj|linear_attn\.out_proj)\.")
+    attn_bases = set()
+    for k in keys:
+        if attn_suffix_re.search(k) and k.endswith((".weight", ".weight_packed")):
+            attn_bases.add(k.rsplit(".", 1)[0])
+    attn_packed = any(b + ".weight_packed" in keys for b in attn_bases)
+    attn_fp8 = any(b + ".weight_scale" in keys for b in attn_bases)
+    attn = "nvfp4" if attn_packed else ("fp8" if attn_fp8 else "none")
+    # Per-layer dense-MLP storage (order-independent): packed -> nvfp4, any scale
+    # sibling -> fp8 (native W8A16), bare weight -> bf16. lm_head: fp8 when it ships a
+    # scale sibling.
+    mlp_key_re = re.compile(
+        r"\.layers\.(\d+)\.mlp\.(?:gate|up|down)_proj\.(weight_packed|weight_scale|weight)$"
+    )
+    packed_layers: set[int] = set()
+    scaled_layers: set[int] = set()
+    plain_layers: set[int] = set()
+    for k in keys:
+        mm = mlp_key_re.search(k)
+        if mm:
+            layer = int(mm.group(1))
+            suffix = mm.group(2)
+            if suffix == "weight_packed":
+                packed_layers.add(layer)
+            elif suffix == "weight_scale":
+                scaled_layers.add(layer)
+            else:
+                plain_layers.add(layer)
+    lmhead_fp8 = any(k == "lm_head.weight_scale" or k.endswith(".lm_head.weight_scale") for k in keys)
+    # No mlp.{gate,up,down}_proj keys at all: routed-expert MoE naming (shared_expert /
+    # experts) -- keep the native assumption with no per-layer overrides.
+    if not (packed_layers or scaled_layers or plain_layers):
+        return attn, "nvfp4", None, lmhead_fp8
+    all_layers = packed_layers | scaled_layers | plain_layers
+    overrides = {
+        layer: ("fp8" if layer in scaled_layers else "bf16")
+        for layer in all_layers - packed_layers
+    }
+    return attn, "nvfp4" if packed_layers else "none", overrides or None, lmhead_fp8
+
+
 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
@@ -182,10 +253,21 @@ 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).
+    # Mixed compressed-tensors exports (unsloth's dynamic per-module quant) store the
+    # attention projections as weight-only FP8 (per-row scale) and some dense-MLP layers
+    # as FP8: keep each group in its native storage (W8A16 attention, W4A16 packed MLP
+    # layers, bf16 dequant for the rest) instead of dequantizing everything to bf16.
+    dense_mlp_storage: dict[int, str] | None = None
     if _compressed_tensors_nvfp4(hf_config):
-        attn_quant = "nvfp4"
-        dense_quant = "nvfp4"
-        lm_head_quant = "none"
+        attn_storage, dense_quant, dense_mlp_storage, lmhead_fp8 = _compressed_linear_storage(hf_config)
+        attn_quant = (
+            "fp8_pertensor" if attn_storage == "fp8"
+            else "nvfp4" if attn_storage == "nvfp4"
+            else "none"
+        )
+        # unsloth's mixed export stores lm_head as weight-only FP8 (per-row scale): keep
+        # it native W8A16 (halves the ~2.5 GB bf16 lm_head); official layouts are bf16.
+        lm_head_quant = "fp8_pertensor" if lmhead_fp8 else "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.
@@ -257,6 +339,7 @@ def parse_config(hf_config: Any) -> ModelConfig:
         attn_quant=attn_quant,
         dense_quant=dense_quant,
         lm_head_quant=lm_head_quant,
+        dense_mlp_storage=dense_mlp_storage,
     )
 
 

python/freetoken/models/qwen3_5_moe/model.py

@@ -49,7 +49,7 @@ def __init__(self, config: ModelConfig, layer_id: int):
             self.self_attn = Qwen3_5Attention(config, layer_id)
         # Dense variants (num_experts==0, e.g. Qwen3.6-27B) use a plain SwiGLU MLP instead of
         # the routed MoE block; both expose ``forward(hidden)->hidden`` and the same key prefix.
-        self.mlp = Qwen3_5MoE(config, layer_id) if config.moe_enabled else Qwen3_5DenseMLP(config)
+        self.mlp = Qwen3_5MoE(config, layer_id) if config.moe_enabled else Qwen3_5DenseMLP(config, layer_id)
         self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
@@ -100,6 +100,16 @@ def __init__(self, config: ModelConfig):
             self.lm_head = Nvfp4LMHead(
                 num_embeddings=config.vocab_size, embedding_dim=config.hidden_size
             )
+        elif getattr(config, "lm_head_quant", "none") == "fp8_pertensor":
+            # unsloth's mixed export stores lm_head as weight-only FP8 (per-row scale):
+            # keep it native W8A16 -- the bf16 dequant of this ~2.5 GB matrix would both
+            # double the memory and be the single largest load-time materialization.
+            from freetoken.kernel.triton.fp8_pertensor_linear import Fp8PerTensorLinear
+
+            assert not config.tie_word_embeddings, "FP8 lm_head assumes untied embeddings"
+            self.lm_head = Fp8PerTensorLinear(
+                in_features=config.hidden_size, out_features=config.vocab_size
+            )
         else:
             self.lm_head = ParallelLMHead(
                 num_embeddings=config.vocab_size,

python/freetoken/models/qwen3_5_moe/moe.py

@@ -21,13 +21,42 @@
 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):
+    def __init__(self, config: ModelConfig, hidden_size: int, intermediate_size: int,
+                 layer_id: int | None = None):
+        # Mixed compressed-tensors exports store some dense-MLP layers in a different
+        # precision than dense_quant (unsloth: most layers packed NVFP4, a few FP8):
+        # the per-layer override (dense_mlp_storage) wins when it covers this layer.
+        storage = "bf16"
+        overrides = getattr(config, "dense_mlp_storage", None)
+        if overrides is not None and layer_id is not None and layer_id in overrides:
+            storage = overrides[layer_id]
+        elif getattr(config, "dense_quant", "none") == "nvfp4":
+            storage = "nvfp4"
+        # storage == "fp8": mixed exports (unsloth) store this dense-MLP layer weight-only
+        # FP8 (per-row scale) -- keep it native W8A16, the same classes the fp8-split GDN
+        # uses, instead of the bf16 dequant (halves this layer's memory).
+        if storage == "fp8":
+            from freetoken.kernel.triton.fp8_pertensor_linear import (
+                Fp8PerTensorColMerged,
+                Fp8PerTensorLinear,
+            )
+
+            self.gate_up_proj = Fp8PerTensorColMerged(
+                hidden_size, [intermediate_size, intermediate_size], has_bias=False
+            )
+            # down_proj: input = intermediate, output = hidden (in_features, out_features)
+            # -> weight [hidden, intermediate], matching the other branches
+            # (Nvfp4DenseLinear / LinearRowParallel both take (intermediate, hidden)).
+            self.down_proj = Fp8PerTensorLinear(
+                intermediate_size, hidden_size, has_bias=False
+            )
+            return
         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":
+        elif storage == "nvfp4":
             # NVFP4 checkpoint: keep the shared expert's NVFP4 weights native (W4A16).
             from freetoken.kernel.triton.nvfp4_linear import Nvfp4DenseColMerged, Nvfp4DenseLinear
 
@@ -48,12 +77,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 class Qwen3_5DenseMLP(_SharedExpert):
     """Dense (non-MoE) SwiGLU MLP for dense Qwen3.x checkpoints (e.g. 27B): ``gate_up_proj``
     (fused gate|up) + ``down_proj`` at full ``intermediate_size``. Same structure (and quant
-    dispatch) as the shared expert -- NVFP4 (W4A16) when ``dense_quant=="nvfp4"``, else bf16 --
-    so it reuses ``_SharedExpert`` directly and keeps the state-dict keys flat
+    dispatch) as the shared expert -- NVFP4 (W4A16) when ``dense_quant=="nvfp4"`` (or the
+    per-layer ``dense_mlp_storage`` override says so), else bf16 -- so it reuses
+    ``_SharedExpert`` directly and keeps the state-dict keys flat
     (``...layers.N.mlp.{gate_up_proj,down_proj}``)."""
 
-    def __init__(self, config: ModelConfig):
-        super().__init__(config, config.hidden_size, config.intermediate_size)
+    def __init__(self, config: ModelConfig, layer_id: int | None = None):
+        super().__init__(config, config.hidden_size, config.intermediate_size, layer_id)
 
 
 class Qwen3_5MoE(BaseOP):
@@ -73,7 +103,7 @@ def __init__(self, config: ModelConfig, layer_id: int | None = None):
         )
         self.gate = LinearReplicated(config.hidden_size, config.num_experts, has_bias=False)
         self.shared_expert = _SharedExpert(
-            config, config.hidden_size, config.shared_expert_intermediate_size
+            config, config.hidden_size, config.shared_expert_intermediate_size, layer_id
         )
         self.shared_expert_gate = LinearReplicated(config.hidden_size, 1, has_bias=False)