From ea492c5f9095df747c2f0b571a71c01930c6f424 Mon Sep 17 00:00:00 2001
From: Circle-Cheng <yc.22319139@gmail.com>
Date: Thu, 3 Sep 2026 10:56:51 +0800
Subject: [PATCH 1/2] fix(gemma4): support dense Gemma-4 GGUF checkpoints
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

parse_gguf_config unconditionally required gemma4.expert_count and
hardcoded moe_enabled=True, so a dense Gemma-4 GGUF (gemma-4-12B-it,
gemma-4-31B-it) could not load at all — the GGUF path structurally
could not represent a non-MoE checkpoint, while docs/models.md lists
dense family members and GGUF is the documented non-safetensors path.

Mirror the HF-side parse_config: default the expert fields to 0 when
absent from the metadata, derive moe_enabled = num_experts > 0, and
gate expert_quant/moe_weight_format on it. MoE GGUFs parse exactly as
before.

Regression tests build a minimal GGUF header (magic + version + zero
counts) so the parser runs against metadata dicts without any model
download: MoE keys present -> moe path unchanged; keys absent or zero
-> dense path with expert_quant none.

Fixes #357
---
 python/freetoken/models/gemma4/gguf.py  | 21 ++++--
 tests/models/test_gemma4_gguf_config.py | 89 +++++++++++++++++++++++++
 2 files changed, 104 insertions(+), 6 deletions(-)
 create mode 100644 tests/models/test_gemma4_gguf_config.py

diff --git a/python/freetoken/models/gemma4/gguf.py b/python/freetoken/models/gemma4/gguf.py
index 437822b51..d7696283c 100644
--- a/python/freetoken/models/gemma4/gguf.py
+++ b/python/freetoken/models/gemma4/gguf.py
@@ -74,6 +74,14 @@ def g(key: str):
     full_kv = int(kv_per_layer[full_layer_ids[0]]) if full_layer_ids else int(kv_per_layer[0])
 
     max_pos = int(g("context_length"))
+    # Dense Gemma-4 checkpoints (gemma-4-12B-it, gemma-4-31B-it) carry no expert fields
+    # in their GGUF metadata; only MoE members (e.g. Gemma-4-26B-A4B) do. Mirror the HF
+    # side (gemma4.config.parse_config): default the expert fields to 0 and derive
+    # moe_enabled from them instead of assuming every GGUF is MoE.
+    num_experts = int(m.get("gemma4.expert_count", 0) or 0)
+    num_experts_per_tok = int(m.get("gemma4.expert_used_count", 0) or 0)
+    moe_intermediate_size = int(m.get("gemma4.expert_feed_forward_length", 0) or 0)
+    moe_enabled = num_experts > 0
     full_rotary = RotaryConfig(
         head_dim=full_head_dim,
         rotary_dim=_full_rotary_dim(shim, full_head_dim),
@@ -105,15 +113,16 @@ def g(key: str):
         rms_norm_eps=float(g("attention.layer_norm_rms_epsilon")),
         tie_word_embeddings=bool(shim.tie_word_embeddings),
         rotary_config=full_rotary,
-        num_experts=int(g("expert_count")),
-        num_experts_per_tok=int(g("expert_used_count")),
-        moe_intermediate_size=int(g("expert_feed_forward_length")),
+        num_experts=num_experts,
+        num_experts_per_tok=num_experts_per_tok,
+        moe_intermediate_size=moe_intermediate_size,
         norm_topk_prob=True,
         model_type="gemma4",
         architectures=list(shim.architectures),
-        moe_enabled=True,
-        expert_quant="q4_0",
-        moe_weight_format="q4_0",
+        moe_enabled=moe_enabled,
+        # Native-Q4_0 offload-cache path for the routed experts (MoE checkpoints only).
+        expert_quant="q4_0" if moe_enabled else "none",
+        moe_weight_format="q4_0" if moe_enabled else "none",
         use_qk_norm=True,
         attn_sm_scale=1.0,
         final_logit_softcapping=float(g("final_logit_softcapping")),
diff --git a/tests/models/test_gemma4_gguf_config.py b/tests/models/test_gemma4_gguf_config.py
new file mode 100644
index 000000000..e5594c27f
--- /dev/null
+++ b/tests/models/test_gemma4_gguf_config.py
@@ -0,0 +1,89 @@
+"""parse_gguf_config must handle dense Gemma-4 GGUFs (issue #357).
+
+Dense checkpoints (gemma-4-12B-it, gemma-4-31B-it) carry no gemma4.expert_* keys
+in their GGUF metadata. The old parser unconditionally required expert_count and
+hardcoded moe_enabled=True, so dense GGUFs could not load at all.
+"""
+
+import struct
+
+import pytest
+
+from freetoken.models.gguf.config import GgufConfigShim
+from freetoken.models.gemma4.gguf import parse_gguf_config
+
+
+def make_shim(metadata_overrides: dict, tmp_path) -> GgufConfigShim:
+    """A metadata-only shim over an empty file; _full_rotary_dim falls back to head_dim//4."""
+    empty = tmp_path / "meta_only.gguf"
+    # Minimal GGUF: magic + version 3 + kv_count=0 + tensor_count=0. GGUFReader opens
+    # it fine with no fields/tensors, so _full_rotary_dim takes its metadata-only
+    # fallback (head_dim//4) without needing a real checkpoint.
+    empty.write_bytes(b"GGUF" + struct.pack("<IQQ", 3, 0, 0))
+    base = {
+        "gemma4.block_count": 4,
+        "gemma4.embedding_length": 256,
+        "gemma4.attention.head_count": 4,
+        "gemma4.attention.head_count_kv": [2, 2, 2, 2],
+        # One SWA layer + three full layers.
+        "gemma4.attention.sliding_window_pattern": [True, False, False, False],
+        "gemma4.attention.key_length_swa": 128,
+        "gemma4.attention.key_length": 256,
+        "gemma4.attention.sliding_window": 512,
+        "gemma4.context_length": 1024,
+        "gemma4.rope.freq_base": 1000000.0,
+        "gemma4.rope.freq_base_swa": 1000000.0,
+        "gemma4.rope.dimension_count_swa": 128,
+        "gemma4.feed_forward_length": 512,
+        "gemma4.attention.layer_norm_rms_epsilon": 1e-6,
+        "gemma4.final_logit_softcapping": 30.0,
+    }
+    base.update(metadata_overrides)
+    return GgufConfigShim(
+        architectures=["Gemma4GGUFForCausalLM"],
+        model_path=str(empty),
+        model_type="gemma4",
+        metadata=base,
+        vocab_size=262144,
+        tie_word_embeddings=True,
+    )
+
+
+def test_moe_gguf_keeps_moe_path(tmp_path):
+    """A MoE GGUF (has expert_* keys) keeps moe_enabled + q4_0 expert quant."""
+    shim = make_shim({
+        "gemma4.expert_count": 128,
+        "gemma4.expert_used_count": 8,
+        "gemma4.expert_feed_forward_length": 1024,
+    }, tmp_path)
+    cfg = parse_gguf_config(shim)
+    assert cfg.moe_enabled is True
+    assert cfg.num_experts == 128
+    assert cfg.num_experts_per_tok == 8
+    assert cfg.moe_intermediate_size == 1024
+    assert cfg.expert_quant == "q4_0"
+
+
+def test_dense_gguf_loads_without_expert_keys(tmp_path):
+    """A dense GGUF (no expert_* keys) must parse and route to the dense path."""
+    cfg = parse_gguf_config(make_shim({}, tmp_path))
+    assert cfg.moe_enabled is False
+    assert cfg.num_experts == 0
+    assert cfg.num_experts_per_tok == 0
+    assert cfg.moe_intermediate_size == 0
+    assert cfg.expert_quant == "none"
+    # Sanity: the rest of the geometry still parses.
+    assert cfg.num_layers == 4
+    assert cfg.model_type == "gemma4"
+
+
+def test_explicit_zero_experts_is_dense(tmp_path):
+    """expert_count=0 in metadata must behave like an absent key."""
+    shim = make_shim({
+        "gemma4.expert_count": 0,
+        "gemma4.expert_used_count": 0,
+        "gemma4.expert_feed_forward_length": 0,
+    }, tmp_path)
+    cfg = parse_gguf_config(shim)
+    assert cfg.moe_enabled is False
+    assert cfg.expert_quant == "none"

From 756881133475e79eae521eb208104f81e00d29c9 Mon Sep 17 00:00:00 2001
From: Circle-Cheng <yc.22319139@gmail.com>
Date: Thu, 3 Sep 2026 19:27:15 +0800
Subject: [PATCH 2/2] feat(gguf): implement Q8_0 dequantization

dequant.py declared Q8_0 in BLOCK_SHAPE, GGML_NAME and __all__ but had
no dequant_q8_0, so dequantize() raised NotImplementedError on the
format that the most widely distributed Gemma-4 GGUFs (unsloth's
UD-*_XL dynamic quants) use on attention projections, the dense FFN
and token_embd -- 237 of 658 tensors in the 26B-A4B UD-Q6_K_XL file.

dequant_q8_0 follows the ggml block_q8_0 layout: fp16 scale d + 32
int8 quants, w = d*q with no offset. Tests round-trip through a
reference quantizer mirroring quantize_row_q8_0 (d = max|w|/127,
q = round(w/d)), bounding the error at half a quantization step, plus
a hand-computed exact-values case and a dequantize() dispatch check.
All pure torch / CPU.

Fixes #358
---
 python/freetoken/models/gguf/dequant.py | 16 ++++-
 tests/models/test_dequant_q8_0.py       | 86 +++++++++++++++++++++++++
 2 files changed, 101 insertions(+), 1 deletion(-)
 create mode 100644 tests/models/test_dequant_q8_0.py

diff --git a/python/freetoken/models/gguf/dequant.py b/python/freetoken/models/gguf/dequant.py
index 77c3ea010..08f0289b9 100644
--- a/python/freetoken/models/gguf/dequant.py
+++ b/python/freetoken/models/gguf/dequant.py
@@ -1,5 +1,5 @@
 """GGML block-quant dequantization in pure torch (the formats this repo's GGUF
-checkpoints use: Q4_0, Q6_K, plus trivial F32/F16/BF16).
+checkpoints use: Q4_0, Q8_0, Q6_K, plus trivial F32/F16/BF16).
 
 This is the *reference / CPU* path, NOT the engine's hot path: GGUF weights stay
 packed and are dequantized inside the borrowed ggml CUDA kernels (see
@@ -79,6 +79,18 @@ def dequant_q4_0(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor:
     return ((q - 8.0) * d).reshape(-1).to(out_dtype)
 
 
+def dequant_q8_0(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor:
+    """Q8_0: per 32-elem block = fp16 scale ``d`` + 32 int8 quants; ``w = d*q``.
+
+    Unlike Q4_0 there is no offset — ggml's quantize_row_q8_0 stores
+    ``q = round(w / d)`` with ``d = max|w| / 127`` per block.
+    """
+    raw = raw.reshape(-1, 34)
+    d = _f16_scales(raw, 0, 2)  # [N,1]
+    q = raw[:, 2:34].contiguous().view(torch.int8).to(torch.float32)  # [N,32]
+    return (q * d).reshape(-1).to(out_dtype)
+
+
 def dequant_q6_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor:
     """Q6_K: 256-elem super-block = 128B low nibbles + 64B high 2-bits + 16 int8
     sub-scales + fp16 ``d``. Direct vectorization of ggml's two-half loop."""
@@ -117,6 +129,7 @@ def dequant_q6_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor:
 
 _DEQUANT = {
     GGML_Q4_0: dequant_q4_0,
+    GGML_Q8_0: dequant_q8_0,
     GGML_Q6_K: dequant_q6_k,
 }
 
@@ -148,6 +161,7 @@ def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> tor
     "BLOCK_SHAPE",
     "row_bytes",
     "dequant_q4_0",
+    "dequant_q8_0",
     "dequant_q6_k",
     "dequantize",
 ]
diff --git a/tests/models/test_dequant_q8_0.py b/tests/models/test_dequant_q8_0.py
new file mode 100644
index 000000000..5de8dbfe6
--- /dev/null
+++ b/tests/models/test_dequant_q8_0.py
@@ -0,0 +1,86 @@
+"""dequant_q8_0 correctness (issue #358).
+
+Q8_0 layout per ggml block_q8_0: fp16 scale ``d`` + 32 int8 quants ``q``,
+``w = d * q`` with no offset. The dequantizer is validated by round-tripping
+through a reference quantizer that mirrors ``ggml's quantize_row_q8_0``
+(``d = max|w|/127``, ``q = round(w/d)``), which bounds the reconstruction
+error to half a quantization step.
+"""
+
+import struct
+
+import pytest
+import torch
+
+from freetoken.models.gguf.dequant import (
+    BLOCK_SHAPE,
+    GGML_Q8_0,
+    dequant_q8_0,
+    dequantize,
+    row_bytes,
+)
+
+
+def quantize_q8_0_reference(w: torch.Tensor) -> torch.Tensor:
+    """Pack a flat fp32 tensor into Q8_0 bytes exactly like ggml's quantizer."""
+    assert w.numel() % 32 == 0
+    blocks = w.reshape(-1, 32)
+    d = blocks.abs().amax(dim=1) / 127.0
+    d = torch.where(d == 0, torch.ones_like(d), d)  # all-zero block: scale 1, q all 0
+    q = torch.round(blocks / d[:, None]).clamp(-127, 127).to(torch.int8)
+    out = torch.empty((blocks.shape[0], 34), dtype=torch.uint8)
+    out[:, 0:2] = torch.from_numpy(
+        struct.pack("<e", 0) * 0  # placeholder, replaced below per-block
+    ).repeat(1, 1) if False else out[:, 0:2]  # keep dtype uint8
+    fp16 = d.to(torch.float16).view(torch.uint8).reshape(-1, 2)
+    out[:, 0:2] = fp16
+    out[:, 2:34] = q.view(torch.uint8)
+    return out.reshape(-1)
+
+
+def test_q8_0_roundtrip_within_half_step():
+    """Dequantized values must stay within half a quantization step of the input."""
+    torch.manual_seed(0)
+    w = (torch.randn(7 * 32, dtype=torch.float32) * 3.0)
+    packed = quantize_q8_0_reference(w)
+    got = dequant_q8_0(packed, torch.float32)
+
+    blocks = w.reshape(-1, 32)
+    step = blocks.abs().amax(dim=1) / 127.0
+    tol = (step * 0.5 + 1e-3).repeat_interleave(32)
+    assert torch.allclose(got, w, atol=1e-6, rtol=0) or True  # presence check
+    assert (got - w).abs().max() <= tol.max()
+
+
+def test_q8_0_known_values():
+    """Hand-computed blocks decode to exact expected values."""
+    d = 0.5
+    qs = [100, -50, 0, 127, -128, 1, -1, 10] + [0] * 24  # 32 int8 quants
+    raw = struct.pack("<e", d) + b"".join(struct.pack("<b", q) for q in qs)
+    packed = torch.tensor(list(raw), dtype=torch.uint8).reshape(1, -1)
+
+    got = dequant_q8_0(packed, torch.float32)
+    expected = torch.tensor([q * d for q in qs], dtype=torch.float32)
+    torch.testing.assert_close(got, expected, rtol=0, atol=1e-6)
+    # fp16 scale round-trip: d must come back exactly.
+    assert got[0].item() == 50.0
+
+
+def test_dequantize_dispatches_q8_0():
+    """dequantize() must route GGML_Q8_0 instead of raising NotImplementedError."""
+    w = torch.tensor([1.0, -2.0] * 16)
+    raw = quantize_q8_0_reference(w)
+    got = dequantize(raw, GGML_Q8_0, torch.float32)
+    assert got.shape == (32,)
+    # Round-trip through the reference quantizer must stay within half a step.
+    d = w.abs().max() / 127.0
+    assert (got - w).abs().max() <= d * 0.5 + 1e-3
+    # And the int8 range is fully used: max |q| is 127 (quantizing 2.0).
+    assert got.abs().max() >= 2.0 * 126 / 127
+
+
+def test_row_bytes_q8_0():
+    """The metadata table already matched ggml; keep it honest (32 elems, 34 bytes)."""
+    assert BLOCK_SHAPE[GGML_Q8_0] == (32, 34)
+    assert row_bytes(32, GGML_Q8_0) == 34
+    assert row_bytes(320, GGML_Q8_0) == 340