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/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("= 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 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("