feat(moe): support gsq q2_0 routed experts

docs/gsq-q2-experts.md

@@ -0,0 +1,76 @@
+# GSQ Q2_0 routed experts
+
+Qwen3.8-Flash-Next can use an HF checkpoint for the trunk, shared experts,
+tokenizer and PLE, with a separate GSQ Q2_0 GGUF file for routed experts.
+This path reads the complete expert set in bounded chunks and keeps native
+packed bytes in the host bank. It does not load a complete GGUF model or
+support IQ3_XXS, Q2_K or TQ2_0.
+
+GSQ Q2_0 (GGML type 42) stores 64 weights in 18 bytes: one little-endian
+FP16 scale and sixteen bytes of adjacent 2-bit codes. The decoded value is
+`(code - 1) * scale`. The GPU kernel rounds weights and activations to BF16,
+then uses FP32 accumulation and output-side routing weights.
+
+The initial implementation requires TP=1, SiLU, BF16 activations and pure
+GPU offload. CPU, hybrid and resident expert execution are unsupported.
+It overrides routed-expert quantization only; the HF dialect still decides
+all trunk and shared-expert schemes. It uses the existing offload cache and
+CUDA graph path, without new cache policy or prefill optimization.
+
+## Source contract
+
+Provide `freetoken_expert_source` in the HF config or through `--hf-overrides`:
+
+```json
+{
+  "freetoken_expert_source": {
+    "format": "q2_0",
+    "scope": "routed_experts_only",
+    "path": "experts-q2.gguf",
+    "sha256": "<verified 64-character lowercase SHA256>",
+    "size": 0,
+    "verified_mtime_ns": 0,
+    "layers": 48,
+    "experts_per_layer": 512,
+    "hidden": 2560,
+    "intermediate": 640
+  }
+}
+```
+
+The numbers above illustrate this model's geometry. Set `size` and
+`verified_mtime_ns` from the actual file after verifying its complete SHA256
+against a trusted checksum. A relative path resolves against the HF model
+directory. Keep the source immutable after verification. Runtime checks the
+receipt's size/mtime and the complete GGUF expert directory, including tensor
+types, geometry, alignment, bounds and overlap; it does not rehash the entire
+file on each load. Invalid directories are rejected before bank allocation.
+
+Use an HF index containing only non-expert tensors. Verify those tensors and
+tokenizer independently against the intended HF revision; this feature does
+not convert, download or certify a mixed checkpoint. Do not substitute pruned
+experts or assume the GGUF trunk is interchangeable with the HF trunk.
+
+An illustrative launch after preparing and verifying both sources is:
+
+```bash
+ft serve --model /models/verified-hf-trunk --text-model-only \
+  --dtype bfloat16 --moe-strategy offload --expert-load serial \
+  --ple-backend disk --hf-overrides "$(cat expert-overrides.json)"
+```
+
+Choose expert cache and KV capacity for the actual device. Expert host bytes
+are `3 * layers * experts * hidden * intermediate / 64 * 18`; the illustrated
+geometry requires 31.640625 GiB for the complete compressed bank, plus other
+runtime memory. This launch is documentation, not an end-to-end test claim.
+
+## Tests
+
+```bash
+PYTHONPATH=python python -m pytest tests/moe/test_q2_0.py tests/moe/test_q2_source.py -q
+```
+
+Tests use a real miniature GGUF and an independent Torch numerical reference.
+They cover expert order, incomplete/corrupt directories, native bank packing,
+signed and zero scales, slot remapping and CUDA graph replay. GPU cases skip
+when CUDA is unavailable.

docs/models.md

@@ -120,3 +120,5 @@ unset on PCIe GPUs: wider grids add no bandwidth there.
   "reasoning_effort": 37}`, and the checkpoint's encoder validates it. Runtime window-cache
   controls use 128-token pages.
 - Qwen3.8-Flash-Next keeps a 47.7 GiB PLE n-gram table pinned in host RAM.
+  [GSQ Q2_0 routed experts](gsq-q2-experts.md) can use a separate verified expert
+  source with an HF trunk; this does not enable full GSQ GGUF model loading.

python/freetoken/layers/quantization/configs/base.py

@@ -108,10 +108,10 @@ def from_hf(
         ``hf_config`` comes from ``cached_load_hf_config``, which folds an old ModelOpt ``hf_quant_config.json`` into ``quantization_config``."""
         q = quantization_config_of(hf_config)
         if q is None:
-            return NoQuantConfig(name_map, unquantized)
+            return _expert_override(hf_config, NoQuantConfig(name_map, unquantized))
         for cls in dialects():
             if cls.claims(q):
-                return cls(q, hf_config, name_map=name_map, unquantized=unquantized)
+                return _expert_override(hf_config, cls(q, hf_config, name_map=name_map, unquantized=unquantized))
         raise NotImplementedError(f"quantization method {q.get('quant_method')!r} is not supported")
 
 
@@ -128,3 +128,53 @@ def claims(cls, q: dict[str, Any]) -> bool:
 
     def scheme_for_name(self, name: str) -> QuantScheme | None:
         return None
+
+
+class ExpertSourceConfig(QuantConfig):
+    """Independent routed-expert storage; all other schemes come from the HF dialect."""
+
+    dialect = "expert_source"
+    STORAGE = {}
+
+    def __init__(self, base: QuantConfig):
+        super().__init__(base.name_map)
+        self.base = base
+
+    def scheme_for_name(self, name: str) -> QuantScheme | None:
+        from ..names import is_routed_expert
+        from ..scheme import WeightDesc
+
+        # Override routed experts only; preserve the HF trunk and shared-expert schemes.
+        if is_routed_expert(name):
+            return QuantScheme(QuantKind.Q2_0, WeightDesc("q2_0", (1, 64), "f16"), {"weight"})
+        return self.base.scheme_for_name(name)
+
+    def scheme_for(self, prefix: str) -> QuantScheme | None:
+        from ..names import is_routed_expert
+
+        if is_routed_expert(prefix):
+            return self.scheme_for_name(prefix)
+        return self.base.scheme_for(prefix)
+
+    def stored_tensors(self, kind: QuantKind) -> dict[str, Stored]:
+        if kind is QuantKind.Q2_0:
+            raise ValueError("Q2_0 experts are read from the explicit GGUF expert source")
+        return self.base.stored_tensors(kind)
+
+    def layer_config(self, layer, layer_kind, scheme):
+        if scheme is not None and scheme.kind is QuantKind.Q2_0:
+            if layer_kind is not LayerKind.MOE:
+                raise ValueError("Q2_0 expert override may only bind a routed MoE layer")
+            return super().layer_config(layer, layer_kind, scheme)
+        return self.base.layer_config(layer, layer_kind, scheme)
+
+
+def _expert_override(hf_config, base: QuantConfig) -> QuantConfig:
+    source = cfg_get(hf_config, "freetoken_expert_source")
+    if source is None:
+        return base
+    if not isinstance(source, dict) or source.get("format") != "q2_0":
+        raise ValueError("unsupported freetoken_expert_source contract")
+    if cfg_get(hf_config, "model_type") != "qwen4_exp":
+        raise ValueError("Q2_0 external experts currently require qwen4_exp")
+    return ExpertSourceConfig(base)

python/freetoken/layers/quantization/moe/__init__.py

@@ -1,5 +1,5 @@
 from .base import BankSpec, ExpertView, MoEConfig, MoEKernel, MoEMethod
-from . import fp8_block, mxfp4, mxfp8, nvfp4, unquantized
+from . import fp8_block, mxfp4, mxfp8, nvfp4, unquantized, q2_0
 from .fp8_block import Fp8BlockMoEMethod
 from .mxfp4 import Mxfp4MoEMethod
 from .mxfp8 import Mxfp8MoEMethod

python/freetoken/layers/quantization/moe/q2_0.py

@@ -0,0 +1,49 @@
+"""GSQ Q2_0 routed experts, kept in their native 64-element / 18-byte blocks."""
+from __future__ import annotations
+
+import torch
+
+from ..registry import LayerKind, register_method
+from ..scheme import QuantKind
+from .base import BankSpec, MoEConfig, MoEKernel, MoEMethod
+
+
+class TritonQ20MoEKernel(MoEKernel):
+    name = "triton"
+
+    def unusable_reason(self, cfg: MoEConfig):
+        reason = self._common_reject(cfg, tp_ok=False, cpu_ok=False, plain_silu_only=True)
+        if reason:
+            return reason
+        if cfg.interleaved or cfg.apply_router_weight_on_input:
+            return "Q2_0 requires contiguous gate/up halves and output-side routing weights"
+        if cfg.hidden % 64 or cfg.intermediate % 64:
+            return "Q2_0 input dimensions must be divisible by 64"
+        if cfg.strategy != "offload":
+            return "Q2_0 currently requires pure GPU offload"
+        if cfg.dtype != torch.bfloat16:
+            return "Q2_0 mixed-precision path currently requires BF16 activations"
+        return None
+
+    def layout(self, cfg):
+        return {
+            "gate_up": BankSpec((2 * cfg.intermediate, cfg.hidden // 64 * 18), torch.uint8),
+            "down": BankSpec((cfg.hidden, cfg.intermediate // 64 * 18), torch.uint8),
+        }
+
+    def pack(self, pieces, cfg, out):
+        # Copy packed bytes directly; do not expand a full expert layer to BF16.
+        out["gate_up"][:, :cfg.intermediate].copy_(pieces["gate"])
+        out["gate_up"][:, cfg.intermediate:].copy_(pieces["up"])
+        out["down"].copy_(pieces["down"])
+        return {}
+
+    def apply(self, layer, x, topk_weights, topk_ids, view, *, is_prefill):
+        from freetoken.moe.fused_q2_0 import fused_experts_q2_0
+
+        return fused_experts_q2_0(x, view.tensors["gate_up"], view.tensors["down"], topk_weights, topk_ids)
+
+
+@register_method(QuantKind.Q2_0, LayerKind.MOE)
+class Q20MoEMethod(MoEMethod):
+    candidates = (TritonQ20MoEKernel,)

python/freetoken/layers/quantization/scheme.py

@@ -20,6 +20,7 @@ class QuantKind(Enum):
     MXFP8 = "mxfp8"
     NVFP4 = "nvfp4"
     MXFP4 = "mxfp4"
+    Q2_0 = "q2_0"
 
     def __str__(self) -> str:
         return self.value

python/freetoken/models/config.py

@@ -289,6 +289,7 @@ class ModelConfig:
     # "fp8_block" is DeepSeek-V3-style 128x128 block-fp8 (weight fp8-e4m3 +
     # weight_scale_inv per block), also applied to the dense projections.
     expert_quant: str = "none"
+    expert_source: dict | None = None
     # Block size (out, in) for block-wise weight quantization (fp8_block: (128, 128)).
     weight_block_size: tuple[int, int] | None = None
     # the checkpoint's quant kind for the dense attention / GatedDeltaNet projections, detected by the family's parse_config for its reader

python/freetoken/models/qwen4_exp/config.py

@@ -256,6 +256,7 @@ def parse_config(hf_config: Any) -> ModelConfig:
         image_token_id=getattr(hf_config, "image_token_id", None),
         attention_groups=groups,
         expert_quant=expert_quant,
+        expert_source=getattr(hf_config, "freetoken_expert_source", None),
         qwen4_args=qwen4_args,
         slot_states=ple_slot_states(qwen4_args),
     )

python/freetoken/models/qwen4_exp/q2_source.py

@@ -0,0 +1,163 @@
+"""Expert-only GGUF v3 source; deliberately independent of GGUF trunk/PLE readers."""
+from __future__ import annotations
+
+import math
+from pathlib import Path
+import re
+import struct
+
+import torch
+
+GGML_GSQ_Q2_0 = 42
+Q2_BLOCK_ELEMENTS = 64
+Q2_BLOCK_BYTES = 18
+_FORMATS = {0: "B", 1: "b", 2: "H", 3: "h", 4: "I", 5: "i", 6: "f", 7: "?", 10: "Q", 11: "q", 12: "d"}
+
+
+def _read(stream, count):
+    data = stream.read(count)
+    if len(data) != count:
+        raise ValueError("truncated GGUF expert source")
+    return data
+
+
+def _unpack(stream, fmt):
+    return struct.unpack("<" + fmt, _read(stream, struct.calcsize("<" + fmt)))
+
+
+def _string(stream):
+    size, = _unpack(stream, "Q")
+    if size > 16 << 20:
+        raise ValueError("oversized GGUF string")
+    return _read(stream, size).decode("utf-8")
+
+
+def _value(stream, kind=None, *, keep=True):
+    if kind is None:
+        kind, = _unpack(stream, "I")
+    if kind == 8:
+        return _string(stream)
+    if kind == 9:
+        elem, count = _unpack(stream, "IQ")
+        if count > 10_000_000:
+            raise ValueError("oversized GGUF metadata array")
+        if not keep and elem in _FORMATS:
+            stream.seek(count * struct.calcsize("<" + _FORMATS[elem]), 1)
+            return None
+        values = [_value(stream, elem, keep=keep) for _ in range(count)]
+        return values if keep else None
+    if kind not in _FORMATS:
+        raise ValueError(f"unsupported GGUF metadata type {kind}")
+    return _unpack(stream, _FORMATS[kind])[0]
+
+
+def read_q2_directory(path: Path, *, layers, experts, hidden, intermediate):
+    """Validate the complete expert directory against the caller's model geometry."""
+    if (layers <= 0 or experts <= 0 or hidden <= 0 or intermediate <= 0
+            or hidden % Q2_BLOCK_ELEMENTS or intermediate % Q2_BLOCK_ELEMENTS):
+        raise ValueError("invalid GSQ Q2_0 model geometry")
+    metadata, directory = {}, {}
+    with path.open("rb") as stream:
+        magic, version, tensors, values = _unpack(stream, "IIQQ")
+        if magic != 0x46554747 or version != 3 or tensors > 100_000 or values > 100_000:
+            raise ValueError("invalid GGUF v3 header")
+        for _ in range(values):
+            key = _string(stream)
+            value = _value(stream, keep=not key.startswith("tokenizer."))
+            if not key.startswith("tokenizer."):
+                metadata[key] = value
+        for _ in range(tensors):
+            name = _string(stream)
+            dims, = _unpack(stream, "I")
+            if not 1 <= dims <= 4:
+                raise ValueError("invalid GGUF tensor rank")
+            shape = _unpack(stream, "Q" * dims)
+            kind, offset = _unpack(stream, "IQ")
+            if name in directory:
+                raise ValueError(f"duplicate GGUF tensor {name}")
+            directory[name] = (shape, kind, offset)
+        alignment = int(metadata.get("general.alignment", 32))
+        if alignment <= 0 or alignment & (alignment - 1):
+            raise ValueError("invalid GGUF alignment")
+        data_start = (stream.tell() + alignment - 1) // alignment * alignment
+    required = {"general.architecture": "qwen4exp", "qwen4exp.block_count": layers,
+                "qwen4exp.expert_count": experts, "qwen4exp.embedding_length": hidden,
+                "qwen4exp.expert_feed_forward_length": intermediate}
+    for key, expected in required.items():
+        if metadata.get(key) != expected:
+            raise ValueError(f"Q2_0 geometry mismatch: {key}={metadata.get(key)!r}, expected {expected!r}")
+    selected, spans = {}, []
+    source_size = path.stat().st_size
+    for layer in range(layers):
+        for role in ("gate", "up", "down"):
+            name = f"blk.{layer}.ffn_{role}_exps.weight"
+            expected = (intermediate, hidden, experts) if role == "down" else (hidden, intermediate, experts)
+            shape, kind, offset = directory.get(name, (None, None, None))
+            if shape != expected or kind != GGML_GSQ_Q2_0:
+                raise ValueError(f"missing or incompatible Q2_0 expert tensor {name}: {shape}, type {kind}")
+            size = math.prod(shape) // Q2_BLOCK_ELEMENTS * Q2_BLOCK_BYTES
+            start = data_start + offset
+            if start < data_start or start + size > source_size or offset % alignment:
+                raise ValueError(f"out-of-bounds or misaligned expert tensor {name}")
+            selected[layer, role] = (start, size)
+            spans.append((start, start + size))
+    spans.sort()
+    if any(a[1] > b[0] for a, b in zip(spans, spans[1:])):
+        raise ValueError("overlapping GGUF expert tensors")
+    actual = {name for name in directory if re.fullmatch(r"blk\.\d+\.ffn_(gate|up|down)_exps\.weight", name)}
+    if len(actual) != layers * 3:
+        raise ValueError("unexpected routed-expert tensor count")
+    return selected
+
+
+def iter_q2_expert_pieces(model_path, config, *, chunk=8 << 20):
+    """Check the verified source identity, then stream every expert without rehashing on startup."""
+    source = config.expert_source
+    if (not isinstance(source, dict) or source.get("format") != "q2_0"
+            or source.get("scope") != "routed_experts_only"):
+        raise ValueError("Q2_0 requires an explicit expert source")
+    E, H, I = config.num_experts, config.hidden_size, config.moe_intermediate_size
+    geometry = {"layers": config.num_layers, "experts_per_layer": E,
+                "hidden": H, "intermediate": I,
+                "expert_bytes": 3 * config.num_layers * E * H * I // Q2_BLOCK_ELEMENTS * Q2_BLOCK_BYTES}
+    if any(source.get(key) != value for key, value in geometry.items()):
+        raise ValueError("Q2_0 source contract does not match model geometry")
+    digest = source.get("sha256")
+    if not isinstance(digest, str) or re.fullmatch(r"[0-9a-f]{64}", digest) is None:
+        raise ValueError("Q2_0 source needs a verified SHA256 receipt")
+    path = Path(source["path"])
+    if not path.is_absolute():
+        path = Path(model_path) / path
+    stat = path.stat()
+    if (stat.st_size != source.get("size")
+            or source.get("verified_mtime_ns") != stat.st_mtime_ns):
+        raise ValueError("Q2_0 source identity changed; reverify the complete artifact")
+    directory = read_q2_directory(path, layers=config.num_layers, experts=config.num_experts,
+                                  hidden=config.hidden_size, intermediate=config.moe_intermediate_size)
+    # Reject an invalid directory before build_expert_banks reserves a large host bank.
+    return _read_expert_chunks(path, directory, config.num_layers, E, H, I, chunk)
+
+
+def _read_expert_chunks(path, directory, layers, E, H, I, chunk):
+    per_expert = 3 * H * I // 64 * 18
+    count = max(1, min(E, chunk // per_expert))
+    # Read bounded expert chunks; retain every layer and expert.
+    with path.open("rb", buffering=0) as stream:
+        for layer in range(layers):
+            for e0 in range(0, E, count):
+                e1 = min(E, e0 + count)
+                pieces = {}
+                for role in ("gate", "up", "down"):
+                    rows, width = (H, I // 64 * 18) if role == "down" else (I, H // 64 * 18)
+                    piece = torch.empty((e1 - e0, rows, width), dtype=torch.uint8)
+                    dst = memoryview(piece.numpy()).cast("B")
+                    start, total = directory[layer, role]
+                    stream.seek(start + e0 * (total // E))
+                    done = 0
+                    while done < len(dst):
+                        got = stream.readinto(dst[done:])
+                        if not got:
+                            raise ValueError("short Q2_0 expert read")
+                        done += got
+                    pieces[role] = piece
+                yield layer, e0, e1, pieces