feat(gguf): serve Qwen2, Qwen3 and Qwen3 MoE GGUF checkpoints

python/freetoken/kernel/aot_models.py

@@ -142,7 +142,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         moe_intermediate_size=768,
         expert_formats=("bf16",),
         aliases=("Qwen/Qwen3-30B-A3B-Thinking-2507",),
-        arch_aliases=("Qwen3VLMoeForConditionalGeneration",),  # Qwen3-VL-30B-A3B: same text tower
+        arch_aliases=("Qwen3VLMoeForConditionalGeneration", "Qwen3MoeGGUFForCausalLM"),  # Qwen3-VL-30B-A3B: same text tower
     ),
     AotModel(
         name="Qwen/Qwen3.5-35B-A3B",
@@ -152,7 +152,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         top_k=8,
         moe_intermediate_size=512,
         expert_formats=("bf16",),
-        arch_aliases=("Qwen3_5MoeForCausalLM",),  # text-only release of the same tower
+        arch_aliases=("Qwen3_5MoeForCausalLM", "Qwen3_5MoeGGUFForCausalLM"),  # text-only release of the same tower
     ),
     AotModel(
         name="Qwen/Qwen3.5-35B-A3B-FP8",
@@ -348,7 +348,7 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         hidden_size=5120,
         kv_groups=((4, 256),),
         aliases=("Qwen/Qwen3.6-27B-FP8", "nvidia/Qwen3.6-27B-NVFP4"),
-        arch_aliases=("Qwen3_5ForCausalLM",),  # text-only release of the same tower
+        arch_aliases=("Qwen3_5ForCausalLM", "Qwen3_5GGUFForCausalLM"),  # text-only release of the same tower
     ),
     AotModel(
         name="google/gemma-4-12B-it",
@@ -381,13 +381,14 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
         architecture="Qwen2ForCausalLM",
         hidden_size=3584,
         kv_groups=((4, 128),),
+        arch_aliases=("Qwen2GGUFForCausalLM",),
     ),
     AotModel(
         name="Qwen/Qwen3-8B",
         architecture="Qwen3ForCausalLM",
         hidden_size=4096,
         kv_groups=((8, 128),),
-        arch_aliases=("Qwen3VLForConditionalGeneration",),  # Qwen3-VL-8B: same text tower
+        arch_aliases=("Qwen3VLForConditionalGeneration", "Qwen3GGUFForCausalLM"),  # Qwen3-VL-8B: same text tower
     ),
     AotModel(
         name="mistralai/Mistral-7B-Instruct-v0.3",

python/freetoken/models/gguf/config.py

@@ -18,6 +18,9 @@
 # reuses the model classes but a GGUF parse_config / iter_weights).
 GGUF_ARCH_TO_REGISTRY: dict[str, str] = {
     "gemma4": "Gemma4GGUFForCausalLM",
+    "qwen2": "Qwen2GGUFForCausalLM",
+    "qwen3": "Qwen3GGUFForCausalLM",
+    "qwen3moe": "Qwen3MoeGGUFForCausalLM",
     "qwen35": "Qwen3_5GGUFForCausalLM",
     "qwen35moe": "Qwen3_5MoeGGUFForCausalLM",
     "qwen4exp": "Qwen4ExpGGUFForCausalLM",

python/freetoken/models/qwen2/__init__.py

@@ -1,5 +1,6 @@
 from .config import parse_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import Qwen2ForCausalLM
 from .weight import iter_weights
 
-__all__ = ["Qwen2ForCausalLM", "parse_config", "iter_weights"]
+__all__ = ["Qwen2ForCausalLM", "parse_config", "iter_weights", "iter_gguf_weights", "parse_gguf_config", "to_gguf"]

python/freetoken/models/qwen2/gguf.py

@@ -0,0 +1,144 @@
+"""Qwen2 / Qwen2.5 from a llama.cpp ``qwen2`` GGUF: config from the GGUF metadata, tensors by name.
+
+The llama.cpp converter (conversion/qwen.py) stores these tensors unchanged: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones.
+"""
+
+from __future__ import annotations
+
+import re
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_ARCH = "qwen2"
+_LAYER_LEAVES = {
+    "self_attn.q_proj": "attn_q",
+    "self_attn.k_proj": "attn_k",
+    "self_attn.v_proj": "attn_v",
+    "self_attn.o_proj": "attn_output",
+    "mlp.gate_proj": "ffn_gate",
+    "mlp.up_proj": "ffn_up",
+    "mlp.down_proj": "ffn_down",
+}
+_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj")
+_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"}
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return (_TOP_LEAVES[name],)
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    blk, leaf = f"blk.{m.group(1)}", m.group(2)
+    if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors:
+        return (f"{blk}.attn_qkv",)
+    tensor = _LAYER_LEAVES.get(leaf)
+    return (f"{blk}.{tensor}",) if tensor is not None else ()
+
+
+def _check_supported(shim: GgufConfigShim) -> None:
+    if any(k.startswith("prism.hadamard.") for k in shim.metadata):
+        raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale")))
+    if scaled:
+        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+
+
+def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
+    """The HF config fields ``parse_config`` reads, from the GGUF metadata."""
+    md = shim.metadata
+
+    def kv(key: str):
+        return md[f"{_ARCH}.{key}"]
+
+    _check_supported(shim)
+    rope = {"rope_theta": float(kv("rope.freq_base")), "rope_type": "default"}
+    scaling = md.get(f"{_ARCH}.rope.scaling.type", "none")
+    if scaling == "yarn":
+        rope.update(
+            rope_type="yarn",
+            factor=float(kv("rope.scaling.factor")),
+            original_max_position_embeddings=int(kv("rope.scaling.original_context_length")),
+        )
+    elif scaling != "none":
+        raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported")
+    # Qwen2.5 files carry no key_length; parse_config then derives head_dim from the hidden size
+    key_length = md.get(f"{_ARCH}.attention.key_length")
+    return SimpleNamespace(
+        model_type="qwen2",
+        architectures=list(shim.architectures),
+        hidden_act="silu",
+        hidden_size=int(kv("embedding_length")),
+        num_hidden_layers=int(kv("block_count")),
+        num_attention_heads=int(kv("attention.head_count")),
+        num_key_value_heads=int(kv("attention.head_count_kv")),
+        head_dim=int(key_length) if key_length is not None else None,
+        rope_scaling=rope,
+        max_position_embeddings=int(kv("context_length")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        intermediate_size=int(kv("feed_forward_length")),
+    )
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    return parse_config(_hf_config(shim))
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The state dict from the GGUF."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    hf = _hf_config(cached_load_hf_config(model_path))
+    index = gguf_tensor_index(model_path)
+    hidden = hf.hidden_size
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    if not hf.tie_word_embeddings:
+        yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    yield "model.norm.weight", float_tensor(t("output_norm"))
+
+    for layer in range(hf.num_hidden_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm"))
+        yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm"))
+        sa = f"{pre}.self_attn"
+        # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix
+        qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+        yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden)
+        bias = index.get(f"{blk}.attn_qkv.bias")
+        yield f"{sa}.qkv_proj.bias", float_tensor(bias) if bias is not None else torch.cat([float_tensor(t(f"{blk}.attn_{p}.bias")) for p in ("q", "k", "v")])
+        out = t(f"{blk}.attn_output")
+        yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+        mlp = f"{pre}.mlp"
+        gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down"))
+        yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden)
+        yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0]))
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]

python/freetoken/models/qwen2/model.py

@@ -46,6 +46,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [

python/freetoken/models/qwen3/__init__.py

@@ -1,5 +1,6 @@
 from .config import parse_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import Qwen3ForCausalLM
 from .weight import iter_weights
 
-__all__ = ["Qwen3ForCausalLM", "parse_config", "iter_weights"]
+__all__ = ["Qwen3ForCausalLM", "parse_config", "iter_weights", "iter_gguf_weights", "parse_gguf_config", "to_gguf"]

python/freetoken/models/qwen3/gguf.py

@@ -0,0 +1,143 @@
+"""Qwen3 dense from a llama.cpp ``qwen3`` GGUF: config from the GGUF metadata, tensors by name.
+
+The llama.cpp converter (conversion/qwen.py) stores these tensors unchanged: both sides run NeoX rope, so q/k carry no rope permute, and the RMSNorm weights are the HF ones.
+"""
+
+from __future__ import annotations
+
+import re
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Iterator
+
+import torch
+
+from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.tensors import float_tensor, fused, matrix
+
+from .config import parse_config
+
+if TYPE_CHECKING:
+    from freetoken.models.config import ModelConfig
+    from freetoken.models.gguf.config import GgufConfigShim
+
+_ARCH = "qwen3"
+_LAYER_LEAVES = {
+    "self_attn.q_proj": "attn_q",
+    "self_attn.k_proj": "attn_k",
+    "self_attn.v_proj": "attn_v",
+    "self_attn.o_proj": "attn_output",
+    "mlp.gate_proj": "ffn_gate",
+    "mlp.up_proj": "ffn_up",
+    "mlp.down_proj": "ffn_down",
+}
+_QKV_LEAVES = ("self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj")
+_TOP_LEAVES = {"model.embed_tokens": "token_embd", "lm_head": "output"}
+_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$")
+
+
+def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]:
+    """Checkpoint module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for."""
+    if name in _TOP_LEAVES:
+        return (_TOP_LEAVES[name],)
+    m = _LAYER_RE.match(name)
+    if m is None:
+        return ()
+    blk, leaf = f"blk.{m.group(1)}", m.group(2)
+    if leaf in _QKV_LEAVES and f"{blk}.attn_qkv.weight" in tensors:
+        return (f"{blk}.attn_qkv",)
+    tensor = _LAYER_LEAVES.get(leaf)
+    return (f"{blk}.{tensor}",) if tensor is not None else ()
+
+
+def _check_supported(shim: GgufConfigShim) -> None:
+    if any(k.startswith("prism.hadamard.") for k in shim.metadata):
+        raise NotImplementedError(f"{shim.model_path}: Hadamard-rotated (prism) GGUF weights are not supported")
+    scaled = sorted(n for n in gguf_tensor_index(shim.model_path) if n.endswith((".scale", ".input_scale")))
+    if scaled:
+        raise NotImplementedError(f"{shim.model_path}: GGUF tensors with per-tensor scales are not supported (e.g. {scaled[0]})")
+
+
+def _hf_config(shim: GgufConfigShim) -> SimpleNamespace:
+    """The HF config fields ``parse_config`` reads, from the GGUF metadata."""
+    md = shim.metadata
+
+    def kv(key: str):
+        return md[f"{_ARCH}.{key}"]
+
+    _check_supported(shim)
+    rope = {"rope_theta": float(kv("rope.freq_base")), "rope_type": "default"}
+    scaling = md.get(f"{_ARCH}.rope.scaling.type", "none")
+    if scaling == "yarn":
+        rope.update(
+            rope_type="yarn",
+            factor=float(kv("rope.scaling.factor")),
+            original_max_position_embeddings=int(kv("rope.scaling.original_context_length")),
+        )
+    elif scaling != "none":
+        raise NotImplementedError(f"{shim.model_path}: rope scaling {scaling!r} is not supported")
+    key_length = md.get(f"{_ARCH}.attention.key_length")
+    return SimpleNamespace(
+        model_type="qwen3",
+        architectures=list(shim.architectures),
+        hidden_act="silu",
+        hidden_size=int(kv("embedding_length")),
+        num_hidden_layers=int(kv("block_count")),
+        num_attention_heads=int(kv("attention.head_count")),
+        num_key_value_heads=int(kv("attention.head_count_kv")),
+        head_dim=int(key_length) if key_length is not None else None,
+        rope_scaling=rope,
+        max_position_embeddings=int(kv("context_length")),
+        rms_norm_eps=float(kv("attention.layer_norm_rms_epsilon")),
+        vocab_size=int(shim.vocab_size),
+        tie_word_embeddings=bool(shim.tie_word_embeddings),
+        intermediate_size=int(kv("feed_forward_length")),
+    )
+
+
+def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
+    return parse_config(_hf_config(shim))
+
+
+def iter_gguf_weights(
+    model_path: str,
+    device: torch.device,
+    *,
+    include_moe_experts: bool,
+    include_non_moe: bool,
+) -> Iterator[tuple[str, torch.Tensor]]:
+    """The state dict from the GGUF."""
+    if not include_non_moe:
+        return
+    from freetoken.utils import cached_load_hf_config
+
+    hf = _hf_config(cached_load_hf_config(model_path))
+    index = gguf_tensor_index(model_path)
+    hidden = hf.hidden_size
+
+    def t(name: str) -> GgufTensorInfo:
+        return index[name if name in index else f"{name}.weight"]
+
+    yield from fused("model.embed_tokens", [t("token_embd")], [matrix(t("token_embd"))], hidden)
+    if not hf.tie_word_embeddings:
+        yield from fused("lm_head", [t("output")], [matrix(t("output"))], hidden)
+    yield "model.norm.weight", float_tensor(t("output_norm"))
+
+    for layer in range(hf.num_hidden_layers):
+        blk, pre = f"blk.{layer}", f"model.layers.{layer}"
+        yield f"{pre}.input_layernorm.weight", float_tensor(t(f"{blk}.attn_norm"))
+        yield f"{pre}.post_attention_layernorm.weight", float_tensor(t(f"{blk}.ffn_norm"))
+        sa = f"{pre}.self_attn"
+        # convert_hf_to_gguf.py --fuse-qkv stores q, k and v as one attn_qkv matrix
+        qkv = [t(f"{blk}.attn_qkv")] if f"{blk}.attn_qkv.weight" in index else [t(f"{blk}.attn_{p}") for p in ("q", "k", "v")]
+        yield from fused(f"{sa}.qkv_proj", qkv, [matrix(p) for p in qkv], hidden)
+        yield f"{sa}.q_norm.weight", float_tensor(t(f"{blk}.attn_q_norm"))
+        yield f"{sa}.k_norm.weight", float_tensor(t(f"{blk}.attn_k_norm"))
+        out = t(f"{blk}.attn_output")
+        yield from fused(f"{sa}.o_proj", [out], [matrix(out)], int(out.ne[0]))
+        mlp = f"{pre}.mlp"
+        gate, up, down = (t(f"{blk}.ffn_{p}") for p in ("gate", "up", "down"))
+        yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden)
+        yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0]))
+
+
+__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]

python/freetoken/models/qwen3/model.py

@@ -46,6 +46,8 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
             num_embeddings=config.vocab_size,
             embedding_dim=config.hidden_size,
+            quant_config=config.quant,
+            prefix=f"{prefix}.embed_tokens",
         )
         self.layers = OPList(
             [

python/freetoken/models/qwen3_moe/__init__.py

@@ -1,4 +1,5 @@
 from .config import parse_config
+from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
 from .model import Qwen3MoeForCausalLM
 from .weight import iter_weights, iter_weights_parallel
 
@@ -7,4 +8,7 @@
     "parse_config",
     "iter_weights",
     "iter_weights_parallel",
+    "iter_gguf_weights",
+    "parse_gguf_config",
+    "to_gguf",
 ]