From dbabd024595b4154e37b7f1aa5e6c3718b599b5d Mon Sep 17 00:00:00 2001 From: boqi <33243841+wenlive@users.noreply.github.com> Date: Sun, 11 Oct 2026 17:15:32 +0800 Subject: [PATCH] feat(moe): support gsq q2_0 routed experts --- docs/gsq-q2-experts.md | 76 ++++++++ docs/models.md | 2 + .../layers/quantization/configs/base.py | 54 +++++- .../layers/quantization/moe/__init__.py | 2 +- .../freetoken/layers/quantization/moe/q2_0.py | 49 ++++++ .../freetoken/layers/quantization/scheme.py | 1 + python/freetoken/models/config.py | 1 + python/freetoken/models/qwen4_exp/config.py | 1 + .../freetoken/models/qwen4_exp/q2_source.py | 163 ++++++++++++++++++ python/freetoken/moe/expert_banks.py | 8 +- python/freetoken/moe/expert_pieces.py | 4 + python/freetoken/moe/fused_q2_0.py | 85 +++++++++ python/freetoken/moe/legacy_format.py | 1 + python/freetoken/moe/offload_cache.py | 2 + tests/moe/test_q2_0.py | 118 +++++++++++++ tests/moe/test_q2_source.py | 88 ++++++++++ 16 files changed, 651 insertions(+), 4 deletions(-) create mode 100644 docs/gsq-q2-experts.md create mode 100644 python/freetoken/layers/quantization/moe/q2_0.py create mode 100644 python/freetoken/models/qwen4_exp/q2_source.py create mode 100644 python/freetoken/moe/fused_q2_0.py create mode 100644 tests/moe/test_q2_0.py create mode 100644 tests/moe/test_q2_source.py diff --git a/docs/gsq-q2-experts.md b/docs/gsq-q2-experts.md new file mode 100644 index 0000000000..158daabb79 --- /dev/null +++ b/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": "", + "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. diff --git a/docs/models.md b/docs/models.md index cdadc2dcb1..3f343a959e 100644 --- a/docs/models.md +++ b/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. diff --git a/python/freetoken/layers/quantization/configs/base.py b/python/freetoken/layers/quantization/configs/base.py index a64136103b..bcc9ccac25 100644 --- a/python/freetoken/layers/quantization/configs/base.py +++ b/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) diff --git a/python/freetoken/layers/quantization/moe/__init__.py b/python/freetoken/layers/quantization/moe/__init__.py index a5f2dc6645..486279b375 100644 --- a/python/freetoken/layers/quantization/moe/__init__.py +++ b/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 diff --git a/python/freetoken/layers/quantization/moe/q2_0.py b/python/freetoken/layers/quantization/moe/q2_0.py new file mode 100644 index 0000000000..0b74cae815 --- /dev/null +++ b/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,) diff --git a/python/freetoken/layers/quantization/scheme.py b/python/freetoken/layers/quantization/scheme.py index 89f968afc8..7a5e96f0d0 100644 --- a/python/freetoken/layers/quantization/scheme.py +++ b/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 diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 25bc73f5e4..daa68e2907 100644 --- a/python/freetoken/models/config.py +++ b/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 diff --git a/python/freetoken/models/qwen4_exp/config.py b/python/freetoken/models/qwen4_exp/config.py index e413786266..722305cc78 100644 --- a/python/freetoken/models/qwen4_exp/config.py +++ b/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), ) diff --git a/python/freetoken/models/qwen4_exp/q2_source.py b/python/freetoken/models/qwen4_exp/q2_source.py new file mode 100644 index 0000000000..5020b6644b --- /dev/null +++ b/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 diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 7500174ae4..f6bf7cfb88 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -108,7 +108,13 @@ def build_expert_banks( if dummy: for role, per_layer in banks.items(): for tensor in per_layer: - _dummy_fill(role, tensor) + if method.kind is QuantKind.Q2_0: + tensor.fill_(0x55) + blocks = tensor.reshape(-1, 18) + blocks[:, 0] = 0 + blocks[:, 1] = 0x3c # FP16 scale=1, every code=1 -> exact finite zero + else: + _dummy_fill(role, tensor) for alpha in alphas.values(): alpha.fill_(1.0) if hb is not None and torch.cuda.is_available(): diff --git a/python/freetoken/moe/expert_pieces.py b/python/freetoken/moe/expert_pieces.py index 79f655a951..821f5a253a 100644 --- a/python/freetoken/moe/expert_pieces.py +++ b/python/freetoken/moe/expert_pieces.py @@ -48,6 +48,10 @@ def iter_expert_pieces( ``nvfp4_expert_spec``. The reader is resolved here, before any bank is allocated, so a missing parallel reader raises ``NotImplementedError`` while a serial fallback is still cheap. """ + if kind is QuantKind.Q2_0: + from freetoken.models.qwen4_exp.q2_source import iter_q2_expert_pieces + + return iter_q2_expert_pieces(model_path, config, chunk=chunk) spec = get_model_spec(config.architectures[0]) hook = _model_hook(spec, "iter_expert_pieces") if hook is not None: diff --git a/python/freetoken/moe/fused_q2_0.py b/python/freetoken/moe/fused_q2_0.py new file mode 100644 index 0000000000..0c5672b700 --- /dev/null +++ b/python/freetoken/moe/fused_q2_0.py @@ -0,0 +1,85 @@ +"""Q2_0 weight-only MoE: native blocks, BF16 activations, FP32 accumulation.""" +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _q2_values(W, base, row, k, N: tl.constexpr, K: tl.constexpr): + # Each 64-element block has an FP16 scale and 16 adjacent 2-bit bytes, unlike Q2_K/TQ2_0. + block = base + row[:, None] * (K // 64 * 18) + (k[None, :] // 64) * 18 + mask = (row[:, None] < N) & (k[None, :] < K) + low = tl.load(W + block, mask, 0).to(tl.uint16) + high = tl.load(W + block + 1, mask, 0).to(tl.uint16) + scale = (low | (high << 8)).to(tl.float16, bitcast=True).to(tl.float32) + packed = tl.load(W + block + 2 + k[None, :] % 64 // 4, mask, 0).to(tl.int32) + code = (packed >> ((k[None, :] % 4) * 2)) & 3 + # Round decoded weights to BF16; keep BF16 activations and FP32 accumulation. + return ((code - 1).to(tl.float32) * scale).to(tl.bfloat16).to(tl.float32) + + +@triton.jit +def _gate_up(X, W, IDS, A, H: tl.constexpr, I: tl.constexpr, TOPK: tl.constexpr, + BK: tl.constexpr, BN: tl.constexpr): + route = tl.program_id(0) + token = route // TOPK + expert = tl.load(IDS + route).to(tl.int64) + row = tl.program_id(1) * BN + tl.arange(0, BN) + k = tl.arange(0, BK) + base = expert * (2 * I * (H // 64 * 18)) + xv = tl.load(X + token * H + k, k < H, 0).to(tl.float32) + gate = tl.sum(_q2_values(W, base, row, k, 2 * I, H) * xv[None, :], 1).to(tl.bfloat16).to(tl.float32) + up = tl.sum(_q2_values(W, base, row + I, k, 2 * I, H) * xv[None, :], 1).to(tl.bfloat16).to(tl.float32) + value = gate / (1.0 + tl.exp(-gate)) * up + tl.store(A + route * I + row, value, row < I) + + +@triton.jit +def _down(A, W, IDS, R, H: tl.constexpr, I: tl.constexpr, BK: tl.constexpr, BN: tl.constexpr): + route = tl.program_id(0) + expert = tl.load(IDS + route).to(tl.int64) + row = tl.program_id(1) * BN + tl.arange(0, BN) + k = tl.arange(0, BK) + base = expert * (H * (I // 64 * 18)) + av = tl.load(A + route * I + k, k < I, 0).to(tl.float32) + value = tl.sum(_q2_values(W, base, row, k, H, I) * av[None, :], 1) + tl.store(R + route * H + row, value, row < H) + + +@triton.jit +def _reduce(R, WEIGHTS, Y, H: tl.constexpr, TOPK: tl.constexpr, BT: tl.constexpr, BH: tl.constexpr): + token = tl.program_id(0) + row = tl.program_id(1) * BH + tl.arange(0, BH) + route = tl.arange(0, BT) + w = tl.load(WEIGHTS + token * TOPK + route, route < TOPK, 0).to(tl.float32) + values = tl.load(R + (token * TOPK + route[:, None]) * H + row[None, :], + (route[:, None] < TOPK) & (row[None, :] < H), 0).to(tl.float32) + tl.store(Y + token * H + row, tl.sum(values * w[:, None], 0), row < H) + + +def fused_experts_q2_0(x, gate_up, down, topk_weights, topk_ids): + if x.dtype != torch.bfloat16 or gate_up.dtype != torch.uint8 or down.dtype != torch.uint8: + raise ValueError("Q2_0 requires BF16 activations and native uint8 blocks") + if not all(t.is_contiguous() and t.is_cuda for t in (x, gate_up, down, topk_weights, topk_ids)): + raise ValueError("Q2_0 tensors must be contiguous CUDA tensors") + tokens, hidden = x.shape + intermediate = gate_up.shape[1] // 2 + topk = topk_ids.shape[1] + if (hidden % 64 or intermediate % 64 or gate_up.shape[2] != hidden // 64 * 18 + or down.shape != (gate_up.shape[0], hidden, intermediate // 64 * 18) + or topk_ids.shape[0] != tokens or topk_weights.shape != topk_ids.shape): + raise ValueError("incompatible Q2_0 expert geometry or routing tensors") + active = torch.empty((tokens * topk, intermediate), dtype=x.dtype, device=x.device) + routed = torch.empty((tokens * topk, hidden), dtype=x.dtype, device=x.device) + out = torch.empty_like(x) + _gate_up[(tokens * topk, triton.cdiv(intermediate, 4))]( + x, gate_up, topk_ids, active, hidden, intermediate, topk, + triton.next_power_of_2(hidden), 4, num_warps=8) + _down[(tokens * topk, triton.cdiv(hidden, 4))]( + active, down, topk_ids, routed, hidden, intermediate, + triton.next_power_of_2(intermediate), 4, num_warps=4) + _reduce[(tokens, triton.cdiv(hidden, 128))]( + routed, topk_weights, out, hidden, topk, triton.next_power_of_2(topk), 128, num_warps=4) + return out diff --git a/python/freetoken/moe/legacy_format.py b/python/freetoken/moe/legacy_format.py index 973a3dba52..889908ffd3 100644 --- a/python/freetoken/moe/legacy_format.py +++ b/python/freetoken/moe/legacy_format.py @@ -27,6 +27,7 @@ def legacy_bank_names(quant_format: str) -> dict[str, str]: # (kind, kernel) <-> the quant_format tag LEGACY_FORMAT = { + (QuantKind.Q2_0, "triton"): "q2_0", (QuantKind.NONE, "fused"): "bf16", (QuantKind.FP8_BLOCK, "triton"): "fp8_block", (QuantKind.NVFP4, "triton"): "nvfp4", diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 0969f73829..6300cbc18f 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -46,6 +46,7 @@ # native GGUF Q4_0 experts: packed block bytes per output row, dequantized inside # the borrowed ggml MoE kernels. gate_up [L*E, 2I, H//32*18], down [L*E, H, I//32*18]. "q4_0": ("gate_up", "down"), + "q2_0": ("gate_up", "down"), # native ModelOpt rows for the Triton inline-dequant kernels: packed e2m1 codes + # fp8-e4m3 per-16 block scales + per-output-row fp16 globals (w1/w3 carry distinct # globals, and folding them into the e4m3 block scales would underflow) @@ -91,6 +92,7 @@ + (H // 128) * fp8_block_scale_pad(H // 128, I // 128) ) * 2, "q4_0": lambda H, I: 2 * I * (H // 32) * 18 + H * (I // 32) * 18, + "q2_0": lambda H, I: 3 * I * H // 64 * 18, "nvfp4": lambda H, I: 2 * I * (H // 2 + H // 16 + 2) + H * (I // 2 + I // 16 + 2), "mxfp4": lambda H, I: 2 * I * (H // 2 + H // 32 + 2) + H * (I // 2 + I // 32 + 2), "ds_fp4": lambda H, I: 2 * I * (H // 2 + H // 32) + H * (I // 2 + I // 32), diff --git a/tests/moe/test_q2_0.py b/tests/moe/test_q2_0.py new file mode 100644 index 0000000000..21ca219ce1 --- /dev/null +++ b/tests/moe/test_q2_0.py @@ -0,0 +1,118 @@ +"""External expert overrides must preserve the trunk's independent schemes.""" +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.layers.quantization import QuantConfig, QuantKind +from freetoken.layers.quantization.configs.base import ExpertSourceConfig, NoQuantConfig +from freetoken.layers.quantization.moe.base import MoEConfig +from freetoken.layers.quantization.moe.q2_0 import Q20MoEMethod +from freetoken.layers.quantization.names import NameMap + + +def test_q2_override_preserves_trunk_and_shared_experts(): + quant = ExpertSourceConfig(NoQuantConfig(NameMap(roots=(("model.layers", "model.language_model.layers"),)))) + for name in ("model.layers.0.self_attn.q_proj", "model.layers.0.mlp.shared_expert.gate_up_proj", + "model.layers.0.mlp.shared_expert_gate", "model.layers.0.mlp.gate", "lm_head"): + assert quant.scheme_for(name) is None + scheme = quant.scheme_for("model.layers.0.mlp.experts") + assert scheme.kind is QuantKind.Q2_0 + assert quant.scheme_for_name("model.language_model.layers.0.mlp.experts.511.down_proj") == scheme + + +def test_q2_layout_packing_preserves_expert_rows_without_expansion(): + cfg = MoEConfig(num_experts=3, hidden=64, intermediate=64, top_k=2, strategy="offload") + method = Q20MoEMethod(cfg) + layout = method.layout() + assert sum(torch.empty(spec.shape, dtype=spec.dtype).numel() for spec in layout.values()) == 3 * 64 * 64 // 64 * 18 + pieces = {role: torch.randint(0, 256, (3, 64, 18), dtype=torch.uint8) for role in ("gate", "up", "down")} + out = {role: torch.empty((3, *spec.shape), dtype=spec.dtype) for role, spec in layout.items()} + method.pack(pieces, out) + torch.testing.assert_close(out["gate_up"][:, :64], pieces["gate"], rtol=0, atol=0) + torch.testing.assert_close(out["gate_up"][:, 64:], pieces["up"], rtol=0, atol=0) + torch.testing.assert_close(out["down"], pieces["down"], rtol=0, atol=0) + assert method.cpu_format is None + for changes in ({"decode_target": "cpu"}, {"strategy": "resident"}, {"hidden": 65}, + {"dtype": torch.float16}, {"apply_router_weight_on_input": True}): + with pytest.raises(RuntimeError): + Q20MoEMethod(replace(cfg, **changes)) + + +def test_q2_source_is_explicit_and_architecture_scoped(): + config = SimpleNamespace(model_type="qwen4_exp", freetoken_expert_source={"format": "q2_0"}) + assert QuantConfig.from_hf(config).scheme_for("model.layers.0.mlp.experts").kind is QuantKind.Q2_0 + config.model_type = "qwen3_5_moe" + with pytest.raises(ValueError): + QuantConfig.from_hf(config) + + +def _quantized_weights(experts, rows, columns): + scales = (torch.rand(experts, rows, columns // 64, device='cuda') * .06 - .03).half() + scales[0, 0, 0] = 0 + codes = torch.randint(0, 4, (experts, rows, columns), device='cuda', dtype=torch.uint8) + blocks = torch.empty((experts, rows, columns // 64, 18), dtype=torch.uint8, device='cuda') + blocks[..., :2] = scales.contiguous().view(torch.uint8).reshape(experts, rows, columns // 64, 2) + fields = codes.reshape(experts, rows, columns // 64, 16, 4) + blocks[..., 2:] = fields[..., 0] | fields[..., 1] << 2 | fields[..., 2] << 4 | fields[..., 3] << 6 + # Reference values come from the integer grid, not the kernel's decoder. + reference = ((codes.float() - 1) * scales.float().repeat_interleave(64, -1)).bfloat16() + return blocks.reshape(experts, rows, columns // 64 * 18), reference + + +def _reference(x, gate_up, down, ids, weights): + output = torch.zeros(x.shape, dtype=torch.float32, device='cuda') + for token in range(x.shape[0]): + for route in range(ids.shape[1]): + expert = ids[token, route] + gate, up = (gate_up[expert].float() @ x[token].float()).bfloat16().float().chunk(2) + active = (torch.nn.functional.silu(gate) * up).bfloat16().float() + result = (down[expert].float() @ active).bfloat16().float() + output[token] += result * weights[token, route] + return output.bfloat16() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires CUDA') +@pytest.mark.parametrize('tokens,topk,remap', [(1, 1, False), (7, 3, False), (19, 3, False), (5, 3, True)]) +def test_q2_vector_matches_reference_with_signed_scales_and_cache_slots(tokens, topk, remap): + from freetoken.moe.fused_q2_0 import fused_experts_q2_0 + + torch.manual_seed(719) + gate_up, gu_ref = _quantized_weights(8, 128, 128) + down, down_ref = _quantized_weights(8, 128, 64) + x = torch.randn(tokens, 128, device='cuda', dtype=torch.bfloat16) + ids = torch.randint(0, 8, (tokens, topk), device='cuda', dtype=torch.int32) + weights = torch.softmax(torch.randn(tokens, topk, device='cuda'), -1) + expected = _reference(x, gu_ref, down_ref, ids, weights) + if remap: + permutation = torch.randperm(8, device='cuda') + gate_up, down = gate_up[permutation], down[permutation] + ids = permutation.argsort()[ids.long()].int() + actual = fused_experts_q2_0(x, gate_up, down, weights, ids) + torch.testing.assert_close(actual, expected, atol=.002, rtol=.02) + assert torch.isfinite(actual).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires CUDA') +def test_q2_vector_cuda_graph_replay_reads_updated_activations(): + from freetoken.moe.fused_q2_0 import fused_experts_q2_0 + + torch.manual_seed(723) + gate_up, gu_ref = _quantized_weights(8, 128, 128) + down, down_ref = _quantized_weights(8, 128, 64) + x = torch.randn(1, 128, device='cuda', dtype=torch.bfloat16) + ids = torch.tensor([[5, 1, 5]], device='cuda', dtype=torch.int32) + weights = torch.tensor([[.25, .5, .25]], device='cuda') + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + fused_experts_q2_0(x, gate_up, down, weights, ids) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = fused_experts_q2_0(x, gate_up, down, weights, ids) + x.add_(.25) + graph.replay() + torch.testing.assert_close(actual, _reference(x, gu_ref, down_ref, ids, weights), atol=.002, rtol=.02) diff --git a/tests/moe/test_q2_source.py b/tests/moe/test_q2_source.py new file mode 100644 index 0000000000..e469dcc6bf --- /dev/null +++ b/tests/moe/test_q2_source.py @@ -0,0 +1,88 @@ +"""Exercise a real miniature GSQ Q2_0 GGUF, including expert order and corruption.""" +import hashlib +import struct +from types import SimpleNamespace + +import pytest + +from freetoken.models.qwen4_exp.q2_source import iter_q2_expert_pieces, read_q2_directory + + +def _string(value): + data = value.encode() + return struct.pack('