fix(load): GGUFs with quantized MoE router gates failed to load

CHANGELOG.md

@@ -6,6 +6,12 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+### Fixed
+
+- GGUFs that quantize the MoE router gate (some community DeepSeek quants;
+  llama.cpp's own quantize leaves it F32) now load: small quantized tensors
+  on raw-array modules are dequantized to f32 at load instead of erroring.
+
 ## [0.2.2] - 2026-08-06
 
 ### Fixed

gmlx/loader.py

@@ -41,7 +41,11 @@
 from .quantized_sdpa_fix import install_quantized_sdpa_mask_fix
 from .rope_batch_fix import install_rope_batch_fix
 from .rotating_cache_fix import install_rotating_cache_fix
-from .modules import KQuantEmbedding, install_kquant_modules
+from .modules import (
+    KQuantEmbedding,
+    dequantize_unattachable_leaves,
+    install_kquant_modules,
+)
 from .populate import (
     maybe_populate_for_load,
     start_populate,
@@ -2799,6 +2803,11 @@ def _install_and_load(
             )
 
     # 6. swap leaves with kquant equivalents.
+    dequant = dequantize_unattachable_leaves(model, hf_weights, hf_kquant_meta)
+    if dequant:
+        log(f"[install] dequantized {len(dequant)} raw-array leaves to f32: "
+            + ", ".join(dequant[:3])
+            + (f" (+{len(dequant) - 3} more)" if len(dequant) > 3 else ""))
     n_replaced = install_kquant_modules(
         model, hf_kquant_meta, native_fp_wire=native_fp_wire)
     log(f"[install] replaced {n_replaced} leaves with kquant modules")
@@ -3241,6 +3250,11 @@ def load_model(
 
     # 6. swap leaves with kquant equivalents.
     loadlog.stage("installing quantized weights")
+    dequant = dequantize_unattachable_leaves(model, hf_weights, hf_kquant_meta)
+    if dequant:
+        _log(f"[install] dequantized {len(dequant)} raw-array leaves to f32: "
+             + ", ".join(dequant[:3])
+             + (f" (+{len(dequant) - 3} more)" if len(dequant) > 3 else ""))
     n_replaced = install_kquant_modules(
         model, hf_kquant_meta, native_fp_wire=native_fp_wire)
     _log(f"[install] replaced {n_replaced} leaves with kquant modules")

gmlx/modules.py

@@ -1409,6 +1409,66 @@ def _wrap(path: str, module):
     return len(wrapped)
 
 
+def dequantize_unattachable_leaves(model: nn.Module,
+                                   hf_weights: dict,
+                                   hf_kquant_meta: dict[str, str],
+                                   max_bytes: int = 64 * 1024 * 1024) -> list[str]:
+    """Dequantize codec'd weights on modules the installer cannot swap.
+
+    Some modules hold their weight as a raw array used inline in the
+    forward (deepseek-family MoEGate routers), so there is no leaf to
+    swap a kquant module into. Dequantize those wire bytes to f32 at
+    load and drop the codec entry. Tensors over ``max_bytes`` (f32
+    size) are left alone so ``install_kquant_modules`` fails loud on
+    them instead of silently materializing gigabytes of float.
+
+    llama.cpp's quantize never codecs router gates, so its GGUFs never
+    hit this path; so far only an antirez DeepSeek-V4-Flash dspark
+    drafter GGUF does. If quantized routers become a pattern in model
+    GGUFs, this approach should change to quantized execution.
+
+    Returns the handled ``path (codec)`` strings for load logging.
+    """
+    import mlx_kquant as kq
+
+    _switch_linear_types, _ = switch_layer_types()
+    attachable = [nn.Linear, nn.Embedding, KQuantLinear, KQuantEmbedding,
+                  KQuantSwitchLinear, KQuantMultiLinear, NativeFPSwitchLinear]
+    attachable.extend(_switch_linear_types)
+    if MultiLinear is not None:
+        attachable.append(MultiLinear)
+    attachable = tuple(attachable)
+    known_codecs = set(kq.codecs())
+    handled: list[str] = []
+
+    def _visit(path: str, module):
+        weight_key = f"{path}.weight"
+        codec = hf_kquant_meta.get(weight_key)
+        if codec is None or isinstance(module, attachable):
+            return module
+        target = getattr(module, "weight", None)
+        if not isinstance(target, mx.array) or codec not in known_codecs:
+            return module
+        if target.size * 4 > max_bytes:
+            return module
+        scales_key = f"{path}.scales"
+        scales = hf_weights.get(scales_key)
+        if scales is None:
+            scales = mx.zeros((1,), dtype=mx.uint8)
+        else:
+            del hf_weights[scales_key]
+        deq = kq.dequantize(hf_weights[weight_key], scales, codec, mx.float32)
+        del hf_weights[weight_key]
+        hf_weights[weight_key] = deq.reshape(target.shape)
+        del hf_kquant_meta[weight_key]
+        handled.append(f"{path} ({codec})")
+        return module
+
+    leaves = model.leaf_modules()
+    tree_map_with_path(_visit, leaves, is_leaf=nn.Module.is_module)
+    return handled
+
+
 def install_kquant_modules(model: nn.Module,
                            hf_kquant_meta: dict[str, str],
                            native_fp_wire: bool = False) -> int:

tests/test_dequant_unattachable.py

@@ -0,0 +1,71 @@
+"""Codec'd tensors on raw-array carrier modules dequantize to f32 at load.
+
+Some community GGUF quantizers codec the MoE router gate (llama.cpp's own
+quantize leaves it F32). The deepseek-family MoEGate holds that weight as a
+raw array, so no KQuant* module can attach; the loader reconstructs the float
+tensor instead of failing. Large unattachable tensors keep failing loud.
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx_kquant as kq
+import pytest
+
+from gmlx.modules import dequantize_unattachable_leaves, install_kquant_modules
+
+
+class RawGate(nn.Module):
+    """MoEGate-style carrier: bare weight array, no Linear."""
+
+    def __init__(self, n_experts, dims):
+        super().__init__()
+        self.weight = mx.zeros((n_experts, dims))
+
+
+class Block(nn.Module):
+    def __init__(self, n_experts, dims):
+        super().__init__()
+        self.gate = RawGate(n_experts, dims)
+        self.proj = nn.Linear(dims, dims, bias=False)
+
+
+def _quantized_block():
+    mx.random.seed(0)
+    model = Block(8, 256)
+    w = mx.random.normal((8, 256))
+    wq, scales = kq.quantize(w, "q8_0")
+    hf_weights = {"gate.weight": wq, "gate.scales": scales}
+    meta = {"gate.weight": "q8_0", "proj.weight": "q8_0"}
+    return model, wq, scales, hf_weights, meta
+
+
+def test_small_raw_leaf_dequantized():
+    model, wq, scales, hf_weights, meta = _quantized_block()
+    handled = dequantize_unattachable_leaves(model, hf_weights, meta)
+
+    assert handled == ["gate (q8_0)"]
+    # Gate is now a plain f32 tensor with the scales sidecar dropped; the
+    # Linear leaf stays codec'd for install_kquant_modules.
+    assert "gate.weight" not in meta
+    assert meta == {"proj.weight": "q8_0"}
+    assert "gate.scales" not in hf_weights
+    deq = hf_weights["gate.weight"]
+    assert deq.shape == (8, 256)
+    assert deq.dtype == mx.float32
+    ref = kq.dequantize(wq, scales, "q8_0", mx.float32).reshape(8, 256)
+    assert mx.array_equal(deq, ref)
+
+    n = install_kquant_modules(model, meta)
+    assert n == 1  # proj swapped; gate left as the raw float carrier
+
+
+def test_large_raw_leaf_still_fails_loud():
+    model, _, _, hf_weights, meta = _quantized_block()
+    handled = dequantize_unattachable_leaves(
+        model, hf_weights, meta, max_bytes=1024)
+
+    assert handled == []
+    assert meta["gate.weight"] == "q8_0"
+    assert "gate.scales" in hf_weights
+    with pytest.raises(ValueError, match="no recognized module class"):
+        install_kquant_modules(model, meta)