From c5e2b43342d5428732bf5d0c44496a2eefa8beef Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 19:40:33 -0700 Subject: [PATCH 01/36] feat(gguf): generalize quant-type support beyond Q4_0/Q8_0/Q6_K The vendored sgl-kernel/llama.cpp CUDA kernels under csrc/gguf/ already dispatch 19 quant types, but the Python side only ever described 6, so K-quants and I-quants were unreachable. This wires the Python tables up to what the kernels actually implement. - csrc/gguf/gguf_kernel.cu: add default: TORCH_CHECK guards to all five switch (type) blocks. Y is allocated with torch::empty, so an unsupported type previously returned uninitialized garbage rather than raising; ggml_moe_get_block_size returned 0. Also null-check the ggml_get_to_cuda function pointer, which ggml_dequantize called unconditionally. - models/gguf/dequant.py: extend GGML_* constants, BLOCK_SHAPE and GGML_NAME to all 22 types (block geometry derived from ggml-common.h with QK_K=256, K_SCALE_SIZE=12). Add five frozensets mirroring the C dispatch switches, each citing the file:line it mirrors. The pure-torch reference dequantize() still implements Q4_0/Q6_K only and now says so explicitly. - layers/gguf.py: replace the three hardcoded {Q4_0, Q8_0, Q6_K} sets with the shared tables. Because _MMQ was tested before _DEQUANT and the sets were identical, the dequant branch was unreachable dead code; it is now the live path for I-quants, which have MMVQ and dequant kernels but no MMQ kernel. This is what makes I-quant prefill work at all. - moe/fused_q4_0.py: add fused_experts_gguf(quant_type), with fused_experts_gguf_q4_0 kept as a behaviour-preserving wrapper. - moe/cpu_executor.py: raise a clear NotImplementedError for GGUF formats other than q4_0. WF_Q4_0 is the only GGUF format with an AVX/VNNI kernel in cpu_moe_ext.cpp; the limit was previously implicit. Tests cover the tables against ggml-common.h and the capability sets against the C switches themselves, so the two cannot drift. Not covered: per-architecture GGUF shims (GGUF_ARCH_TO_REGISTRY still lists gemma4 only) and CPU AVX kernels for non-Q4_0 experts. Untested on hardware. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- .../freetoken/kernel/csrc/gguf/gguf_kernel.cu | 28 +- python/freetoken/layers/gguf.py | 49 ++- python/freetoken/models/gguf/dequant.py | 168 +++++++-- python/freetoken/moe/cpu_executor.py | 13 +- python/freetoken/moe/fused_q4_0.py | 54 ++- tests/models/test_gguf_dispatch.py | 327 ++++++++++++++++++ tests/models/test_gguf_type_tables.py | 311 +++++++++++++++++ 7 files changed, 900 insertions(+), 50 deletions(-) create mode 100644 tests/models/test_gguf_dispatch.py create mode 100644 tests/models/test_gguf_type_tables.py diff --git a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu index d88960d5..db210646 100644 --- a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu +++ b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu @@ -83,8 +83,14 @@ torch::Tensor ggml_dequantize( at::Tensor DW = torch::empty({m, n}, options); cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + // These kernels are vendored from sgl-kernel/llama.cpp; the guards below are a FreeToken addition + // to prevent silent data corruption from unsupported quant types. DISPATCH_FLOAT_TYPES(DW.scalar_type(), "ggml_dequantize", [&] { auto to_cuda = ggml_get_to_cuda(type); + TORCH_CHECK(to_cuda != nullptr, + "ggml_dequantize: unsupported GGUF quant type ", type, + " (dequant kernels exist for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K/IQ2_XXS/" + "IQ2_XS/IQ3_XXS/IQ1_S/IQ4_NL/IQ3_S/IQ2_S/IQ4_XS/IQ1_M)"); to_cuda((void*)W.data_ptr(), (scalar_t*)DW.data_ptr(), m * n, stream); }); @@ -184,6 +190,10 @@ torch::Tensor ggml_mul_mat_vec_a8( mul_mat_vec_iq1_m_q8_1_cuda( (void*)W.data_ptr(), (void*)quant_X.data_ptr(), (scalar_t*)Y.data_ptr(), col, row, vecs, stream); break; + default: + TORCH_CHECK(false, "ggml_mul_mat_vec_a8: unsupported GGUF quant type ", type, + " (MMVQ kernels exist for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K/IQ2_XXS/IQ2_XS/" + "IQ3_XXS/IQ1_S/IQ4_NL/IQ3_S/IQ2_S/IQ4_XS/IQ1_M)"); } }); return Y; @@ -327,6 +337,10 @@ torch::Tensor ggml_mul_mat_a8( row, stream); break; + default: + TORCH_CHECK(false, "ggml_mul_mat_a8: unsupported GGUF quant type ", type, + " (MMQ kernels exist only for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K; " + "I-quants must route through ggml_dequantize)"); } }); return Y; @@ -533,6 +547,10 @@ torch::Tensor ggml_moe_a8( sorted_token_ids.sizes()[0], stream); break; + default: + TORCH_CHECK(false, "ggml_moe_a8: unsupported GGUF quant type ", type, + " (MMQ kernels exist only for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K; " + "I-quants must route through ggml_dequantize)"); } }); return Y; @@ -804,6 +822,10 @@ torch::Tensor ggml_moe_a8_vec( quant_X.stride(0), stream); break; + default: + TORCH_CHECK(false, "ggml_moe_a8_vec: unsupported GGUF quant type ", type, + " (MMVQ kernels exist for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K/IQ2_XXS/IQ2_XS/" + "IQ3_XXS/IQ1_S/IQ4_NL/IQ3_S/IQ2_S/IQ4_XS/IQ1_M)"); } }); return Y; @@ -831,8 +853,12 @@ int64_t ggml_moe_get_block_size(int64_t type) { return MOE_X_Q5_K; case 14: return MOE_X_Q6_K; + default: + TORCH_CHECK(false, "ggml_moe_get_block_size: unsupported GGUF quant type ", type, + " (MMQ kernels exist only for Q4_0/Q4_1/Q5_0/Q5_1/Q8_0/Q2_K-Q6_K; " + "I-quants must route through ggml_dequantize)"); + return 0; // unreachable but silences compiler warning } - return 0; } // ---- FreeToken pybind bindings (donor registers these via TORCH_LIBRARY; we diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py index ac49b1a5..883942c0 100644 --- a/python/freetoken/layers/gguf.py +++ b/python/freetoken/layers/gguf.py @@ -1,5 +1,6 @@ """Native-GGUF quantized layers: weights stay in their packed block layout and are -dequantized *inside* the borrowed llama.cpp kernels (no bf16 copy ever materialized). +dequantized *inside* the borrowed llama.cpp CUDA kernels -- either fused into the matmul +(MMVQ/MMQ) or, for types with no MMQ kernel, by an explicit ``ggml_dequantize`` pass. Mirrors vLLM/sglang's ``GGUFLinearMethod`` / ``GGUFEmbeddingMethod`` dispatch, ported onto FreeToken's ``BaseOP``. FreeToken keeps fused projections (qkv, gate_up) as a @@ -8,6 +9,18 @@ share an input dim, hence the same ``row_bytes``), so a fused layer is still one ``[out, row_bytes]`` qweight -- no per-shard padding bookkeeping needed. +**Matmul dispatch strategy** (4-tier, per fused_mul_mat_gguf): + +1. **Unquantized (F32, F16, BF16)**: straight torch matmul ``x @ qweight.T``. +2. **Small-batch quantized (batch <= 6, MMVQ types)**: GEMV kernel via ``ggml_mul_mat_vec_a8``. +3. **Large-batch standard quants (MMQ types: Q4_0, Q4_1, Q5_0, Q5_1, Q8_0, K-quants)**: MMQ kernel + via ``ggml_mul_mat_a8``. +4. **Large-batch I-quants (IQ2_XXS, IQ2_XS, IQ3_XXS, IQ1_S, IQ4_NL, IQ3_S, IQ2_S, IQ4_XS, IQ1_M)**: + I-quants have MMVQ and dequant kernels but NO MMQ kernel. Prefill therefore falls back to + ``ggml_dequantize`` + plain torch matmul. This materializes a transient BF16 copy of the weight + (cost: ``out_features * in_features * 2 bytes``), which is a real tradeoff for memory-bound + prefill on large I-quant weights. + TP is assumed to be 1 (the gemma4 GGUF path restricts to TP=1, like the HF path). """ @@ -17,31 +30,29 @@ from freetoken.models.gguf.dequant import ( BLOCK_SHAPE, - GGML_BF16, - GGML_F16, - GGML_F32, + DEQUANT_TYPES, GGML_NAME, - GGML_Q4_0, - GGML_Q6_K, - GGML_Q8_0, + GGML_UNQUANTIZED, + MMQ_TYPES, + MMVQ_TYPES, row_bytes, ) from .base import BaseOP -# ggml type groups for kernel dispatch (subset we build kernels for). -_UNQUANTIZED = {GGML_F32, GGML_F16, GGML_BF16} -# standard + k-quants: both an MMVQ (small-batch GEMV) and MMQ (large-batch) kernel exist. -_MMVQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} -_MMQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} -_DEQUANT = {GGML_Q4_0, GGML_Q8_0, GGML_Q6_K} - # Below this token count, the MMVQ GEMV kernel wins (matches vLLM's heuristic). _MMVQ_SAFE = 6 def fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor: - """y = x @ dequant(qweight).T, dispatched by batch size and quant type.""" + """y = x @ dequant(qweight).T, dispatched by batch size and quant type. + + Dispatch order: + 1. Unquantized (F32/F16/BF16): plain torch matmul + 2. Small-batch quantized (batch <= 6, in MMVQ_TYPES): GEMV kernel + 3. Large-batch standard quants (in MMQ_TYPES): MMQ kernel + 4. Large-batch with I-quants (in DEQUANT_TYPES but not MMQ_TYPES): dequant + torch matmul + """ from freetoken.kernel.gguf import ( ggml_dequantize, ggml_mul_mat_a8, @@ -51,13 +62,13 @@ def fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int out_features = qweight.shape[0] if x.shape[0] == 0: return x.new_empty((0, out_features)) - if qweight_type in _UNQUANTIZED: + if qweight_type in GGML_UNQUANTIZED: return x @ qweight.T - if x.shape[0] <= _MMVQ_SAFE and qweight_type in _MMVQ: + if x.shape[0] <= _MMVQ_SAFE and qweight_type in MMVQ_TYPES: return ggml_mul_mat_vec_a8(qweight, x, qweight_type, out_features) - if qweight_type in _MMQ: + if qweight_type in MMQ_TYPES: return ggml_mul_mat_a8(qweight, x, qweight_type, out_features) - if qweight_type in _DEQUANT: + if qweight_type in DEQUANT_TYPES: block, type_size = BLOCK_SHAPE[qweight_type] in_features = qweight.shape[1] // type_size * block weight = ggml_dequantize(qweight, qweight_type, out_features, in_features, x.dtype) diff --git a/python/freetoken/models/gguf/dequant.py b/python/freetoken/models/gguf/dequant.py index 77c3ea01..c6e180d8 100644 --- a/python/freetoken/models/gguf/dequant.py +++ b/python/freetoken/models/gguf/dequant.py @@ -1,50 +1,147 @@ -"""GGML block-quant dequantization in pure torch (the formats this repo's GGUF -checkpoints use: Q4_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 -``freetoken.kernel.gguf``). These routines are used only to (a) materialize the few -dense F32/F16 tensors at load (norms, scales, router) via :func:`dequantize`, and -(b) cross-check the CUDA kernels in tests. The ``BLOCK_SHAPE`` table and -:func:`row_bytes` are the type metadata the packed (kernel) path also relies on. - -Each ``dequant_*`` takes the raw little-endian bytes as a ``uint8`` tensor whose -final axis spans whole blocks, and returns the values in *storage order* (ggml's -fastest axis first); the caller reshapes to the torch shape (``dims[::-1]``). The -math mirrors ``ggml-quants.c``. +"""GGML block-quant dequantization and type metadata. + +This module serves two purposes: + +1. **Type metadata for the packed GPU path** (the hot path): GGUF weights stay packed + and are dequantized inside the borrowed ggml CUDA kernels (see ``freetoken.kernel.gguf``). + The ``BLOCK_SHAPE`` table and :func:`row_bytes` are shared by ``GGUFLinear``, + ``GGUFEmbedding``, and expert-bank loaders for weight allocation and unpacking. + +2. **Pure-torch reference dequantizers** (CPU/test path): The :func:`dequantize` function + and helper ``dequant_*`` routines materialize F32/F16 tensors at load (norms, scales, + router) and cross-check CUDA kernels in tests. These implement only Q4_0 and Q6_K; + the missing types are handled by the CUDA kernels in production. + +``BLOCK_SHAPE`` covers all 21 types (F32, F16, BF16, STD_K, IQ); ``dequantize()`` and +``_DEQUANT`` cover Q4_0 and Q6_K only. + +Each ``dequant_*`` takes raw little-endian bytes as a ``uint8`` tensor whose final axis +spans whole blocks, and returns values in *storage order* (ggml's fastest axis first); +the caller reshapes to torch shape (``dims[::-1]``). The math mirrors ``ggml-quants.c``. """ from __future__ import annotations import torch -# ggml_type enum values (subset present in these checkpoints). +# ggml_type enum values. Mirrors the ggml.h enum in llama.cpp. GGML_F32 = 0 GGML_F16 = 1 GGML_Q4_0 = 2 +GGML_Q4_1 = 3 +GGML_Q5_0 = 6 +GGML_Q5_1 = 7 GGML_Q8_0 = 8 +GGML_Q2_K = 10 +GGML_Q3_K = 11 +GGML_Q4_K = 12 +GGML_Q5_K = 13 GGML_Q6_K = 14 +GGML_IQ2_XXS = 16 +GGML_IQ2_XS = 17 +GGML_IQ3_XXS = 18 +GGML_IQ1_S = 19 +GGML_IQ4_NL = 20 +GGML_IQ3_S = 21 +GGML_IQ2_S = 22 +GGML_IQ4_XS = 23 +GGML_IQ1_M = 29 GGML_BF16 = 30 -# (block numel, bytes per block) per ggml type. +# (block numel, bytes per block) per ggml type. Derived from block structs in +# python/freetoken/kernel/csrc/gguf/ggml-common.h (lines 18-192). BLOCK_SHAPE: dict[int, tuple[int, int]] = { GGML_F32: (1, 4), GGML_F16: (1, 2), - GGML_BF16: (1, 2), GGML_Q4_0: (32, 18), + GGML_Q4_1: (32, 20), + GGML_Q5_0: (32, 22), + GGML_Q5_1: (32, 24), GGML_Q8_0: (32, 34), + GGML_Q2_K: (256, 84), + GGML_Q3_K: (256, 110), + GGML_Q4_K: (256, 144), + GGML_Q5_K: (256, 176), GGML_Q6_K: (256, 210), + GGML_IQ2_XXS: (256, 66), + GGML_IQ2_XS: (256, 74), + GGML_IQ3_XXS: (256, 98), + GGML_IQ1_S: (256, 50), + GGML_IQ4_NL: (32, 18), + GGML_IQ3_S: (256, 110), + GGML_IQ2_S: (256, 82), + GGML_IQ4_XS: (256, 136), + GGML_IQ1_M: (256, 56), + GGML_BF16: (1, 2), } GGML_NAME = { GGML_F32: "F32", GGML_F16: "F16", - GGML_BF16: "BF16", GGML_Q4_0: "Q4_0", + GGML_Q4_1: "Q4_1", + GGML_Q5_0: "Q5_0", + GGML_Q5_1: "Q5_1", GGML_Q8_0: "Q8_0", + GGML_Q2_K: "Q2_K", + GGML_Q3_K: "Q3_K", + GGML_Q4_K: "Q4_K", + GGML_Q5_K: "Q5_K", GGML_Q6_K: "Q6_K", + GGML_IQ2_XXS: "IQ2_XXS", + GGML_IQ2_XS: "IQ2_XS", + GGML_IQ3_XXS: "IQ3_XXS", + GGML_IQ1_S: "IQ1_S", + GGML_IQ4_NL: "IQ4_NL", + GGML_IQ3_S: "IQ3_S", + GGML_IQ2_S: "IQ2_S", + GGML_IQ4_XS: "IQ4_XS", + GGML_IQ1_M: "IQ1_M", + GGML_BF16: "BF16", } +# CUDA kernel dispatch: which types each C function handles. +# Mirrors switch (type) in ggml_get_to_cuda (dequantize.cuh:541) +DEQUANT_TYPES = frozenset({ + GGML_Q4_0, GGML_Q4_1, GGML_Q5_0, GGML_Q5_1, GGML_Q8_0, + GGML_Q2_K, GGML_Q3_K, GGML_Q4_K, GGML_Q5_K, GGML_Q6_K, + GGML_IQ2_XXS, GGML_IQ2_XS, GGML_IQ3_XXS, GGML_IQ1_S, GGML_IQ4_NL, + GGML_IQ3_S, GGML_IQ2_S, GGML_IQ4_XS, GGML_IQ1_M, +}) + +# Mirrors switch (type) in ggml_mul_mat_vec_a8 (gguf_kernel.cu:116) +MMVQ_TYPES = frozenset({ + GGML_Q4_0, GGML_Q4_1, GGML_Q5_0, GGML_Q5_1, GGML_Q8_0, + GGML_Q2_K, GGML_Q3_K, GGML_Q4_K, GGML_Q5_K, GGML_Q6_K, + GGML_IQ2_XXS, GGML_IQ2_XS, GGML_IQ3_XXS, GGML_IQ1_S, GGML_IQ4_NL, + GGML_IQ3_S, GGML_IQ2_S, GGML_IQ4_XS, GGML_IQ1_M, +}) + +# Mirrors switch (type) in ggml_mul_mat_a8 (gguf_kernel.cu:219) +# I-quants do not have an MMQ (large-batch matmul) kernel. +MMQ_TYPES = frozenset({ + GGML_Q4_0, GGML_Q4_1, GGML_Q5_0, GGML_Q5_1, GGML_Q8_0, + GGML_Q2_K, GGML_Q3_K, GGML_Q4_K, GGML_Q5_K, GGML_Q6_K, +}) + +# Mirrors switch (type) in ggml_moe_a8_vec (gguf_kernel.cu:577) +MOE_VEC_TYPES = frozenset({ + GGML_Q4_0, GGML_Q4_1, GGML_Q5_0, GGML_Q5_1, GGML_Q8_0, + GGML_Q2_K, GGML_Q3_K, GGML_Q4_K, GGML_Q5_K, GGML_Q6_K, + GGML_IQ2_XXS, GGML_IQ2_XS, GGML_IQ3_XXS, GGML_IQ1_S, GGML_IQ4_NL, + GGML_IQ3_S, GGML_IQ2_S, GGML_IQ4_XS, GGML_IQ1_M, +}) + +# Mirrors switch (type) in ggml_moe_a8 (gguf_kernel.cu:369), whose coverage ggml_moe_get_block_size (gguf_kernel.cu:835) mirrors +# I-quants do not have an MMQ (grouped MoE large-batch) kernel. +MOE_MMQ_TYPES = frozenset({ + GGML_Q4_0, GGML_Q4_1, GGML_Q5_0, GGML_Q5_1, GGML_Q8_0, + GGML_Q2_K, GGML_Q3_K, GGML_Q4_K, GGML_Q5_K, GGML_Q6_K, +}) + +# Unquantized types: no dequantization needed, handled by separate path in layers/gguf.py. +GGML_UNQUANTIZED = frozenset({GGML_F32, GGML_F16, GGML_BF16}) + def row_bytes(numel: int, ggml_type: int) -> int: """Packed byte length of one row of ``numel`` elements in ``ggml_type`` blocks. @@ -122,7 +219,12 @@ def dequant_q6_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> torch.Tensor: - """Dequantize ``raw`` (uint8) of any supported ggml type to flat ``out_dtype``.""" + """Dequantize ``raw`` (uint8) in pure torch (Q4_0, Q6_K, F32/F16/BF16 only). + + This is the CPU reference path for loading norms and scales. The packed GPU path + (GGUFLinear, GGUFEmbedding, expert banks) dequantizes all 21 types via CUDA kernels; + see ``BLOCK_SHAPE`` for the full type list. + """ if ggml_type == GGML_F32: return raw.view(torch.float32).to(out_dtype) if ggml_type == GGML_F16: @@ -132,7 +234,9 @@ def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> tor fn = _DEQUANT.get(ggml_type) if fn is None: raise NotImplementedError( - f"dequant for ggml type {GGML_NAME.get(ggml_type, ggml_type)} not implemented" + f"pure-torch dequant for ggml type {GGML_NAME.get(ggml_type, ggml_type)} " + f"not implemented (only Q4_0 and Q6_K supported in CPU path; " + f"other types use CUDA kernels via GGUFLinear)" ) return fn(raw, out_dtype) @@ -140,12 +244,34 @@ def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> tor __all__ = [ "GGML_F32", "GGML_F16", - "GGML_BF16", "GGML_Q4_0", + "GGML_Q4_1", + "GGML_Q5_0", + "GGML_Q5_1", "GGML_Q8_0", + "GGML_Q2_K", + "GGML_Q3_K", + "GGML_Q4_K", + "GGML_Q5_K", "GGML_Q6_K", + "GGML_IQ2_XXS", + "GGML_IQ2_XS", + "GGML_IQ3_XXS", + "GGML_IQ1_S", + "GGML_IQ4_NL", + "GGML_IQ3_S", + "GGML_IQ2_S", + "GGML_IQ4_XS", + "GGML_IQ1_M", + "GGML_BF16", "GGML_NAME", "BLOCK_SHAPE", + "DEQUANT_TYPES", + "MMVQ_TYPES", + "MMQ_TYPES", + "MOE_VEC_TYPES", + "MOE_MMQ_TYPES", + "GGML_UNQUANTIZED", "row_bytes", "dequant_q4_0", "dequant_q6_k", diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index b96205aa..6fe1efc0 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -67,7 +67,9 @@ "swigluoai": 3, } -# Weight-format ids must match WFmt in csrc/cpu_moe/cpu_moe_ext.cpp. +# Weight-format ids must match WFmt in csrc/cpu_moe/cpu_moe_ext.cpp:1215. +# Q4_0 is the only GGUF format with AVX/VNNI kernels; adding more requires new +# intrinsics work (out of scope). The CPU path for other GGUF formats is not implemented. _WFMT_IDS = {"bf16": 0, "nvfp4": 1, "mxfp4_triton": 2, "ds_fp4": 3, "q4_0": 4} @@ -373,6 +375,15 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: if fmt == "ds_fp4": return self._resolve_dsfp4_banks(banks) + # Detect unsupported GGUF formats (any ggml type name that isn't q4_0). + # GGUF format strings follow the pattern of ggml type names (q*, iq*). + if fmt.startswith(("q", "iq")) and fmt != "q4_0": + raise NotImplementedError( + f"the CPU/hybrid MoE backend does not support GGUF format {fmt!r}; " + f"it has AVX/VNNI kernels for Q4_0 only. Use --moe-backend fused for GPU " + f"dequantization, or --moe-backend offload to stream experts to the GPU." + ) + # nvfp4: packed e2m1 (2/byte) + fp8-e4m3 per-16 block scales + fp16 row globals. gup, gus, gug = banks["gate_up_packed"], banks["gate_up_scale"], banks["gate_up_global"] dnp, dns, dng = banks["down_packed"], banks["down_scale"], banks["down_global"] diff --git a/python/freetoken/moe/fused_q4_0.py b/python/freetoken/moe/fused_q4_0.py index cdab82bf..f0c7acde 100644 --- a/python/freetoken/moe/fused_q4_0.py +++ b/python/freetoken/moe/fused_q4_0.py @@ -1,12 +1,18 @@ -"""Grouped expert GEMM over native GGUF Q4_0 banks (borrowed ggml MoE kernels). +"""Grouped expert GEMM over native GGUF banks (borrowed ggml MoE kernels). Ports vLLM/sglang's ``_fused_moe_gguf`` MMVQ path onto FreeToken's offload-cache -interface: the experts are streamed to the GPU as packed Q4_0 block bytes and +interface: the experts are streamed to the GPU as packed block bytes and dequantized *inside* ``ggml_moe_a8_vec`` -- no bf16 expert copy is materialized. We use the MMVQ (vector) kernel for both prefill and decode: it consumes ``topk_ids`` directly (no ``moe_align_block_size`` needed) and on small batches it is the right choice anyway. ``topk_ids`` already index the streamed cache slots (decode) or the materialized layer positions (prefill). + +This module is general over any quantization type supported by the ``ggml_moe_a8_vec`` +kernel (all types in ``MOE_VEC_TYPES``, which includes all 19 quantized types). Q4_0 +is currently the only type the rest of the pipeline plumbs through; support for other +types is added by parametrizing the quant type at the MoE bank loader, dequant.py, and +moe/expert_banks.py level. """ from __future__ import annotations @@ -14,21 +20,36 @@ import torch from freetoken.layers.activation import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul -from freetoken.models.gguf.dequant import GGML_Q4_0 +from freetoken.models.gguf.dequant import GGML_Q4_0, MOE_VEC_TYPES _ACT = {"silu": silu_and_mul, "gelu": gelu_and_mul, "gelu_tanh": gelu_tanh_and_mul} -def fused_experts_gguf_q4_0( +def fused_experts_gguf( hidden_states: torch.Tensor, - gate_up_q: torch.Tensor, # [num_slots, 2I, H//32*18] uint8 - down_q: torch.Tensor, # [num_slots, H, I//32*18] uint8 + gate_up_q: torch.Tensor, # [num_slots, 2I, H//32*18] uint8 (or other quant format) + down_q: torch.Tensor, # [num_slots, H, I//32*18] uint8 (or other quant format) topk_weights: torch.Tensor, topk_ids: torch.Tensor, activation: str, + quant_type: int, ) -> torch.Tensor: + """Fused GGUF MoE expert compute over any MMVQ-supported quantization type. + + This kernel operates directly on packed quantized weights (no materialization to bf16); + dequantization happens inside the ``ggml_moe_a8_vec`` CUDA kernel. ``quant_type`` must be + in ``MOE_VEC_TYPES``, which mirrors the supported types in ``ggml_moe_a8_vec`` + (gguf_kernel.cu:559). + """ from freetoken.kernel.gguf import ggml_moe_a8_vec + if quant_type not in MOE_VEC_TYPES: + from freetoken.models.gguf.dequant import GGML_NAME + raise NotImplementedError( + f"fused GGUF MoE kernel does not support quant type {GGML_NAME.get(quant_type, quant_type)} " + f"(only {sorted(MOE_VEC_TYPES)} supported)" + ) + act_fn = _ACT.get(activation) if act_fn is None: raise ValueError(f"unsupported MoE activation {activation!r}") @@ -37,7 +58,7 @@ def fused_experts_gguf_q4_0( n2 = gate_up_q.shape[1] # 2 * intermediate h = down_q.shape[1] # hidden top_k = topk_ids.shape[1] - qt = int(GGML_Q4_0) + qt = int(quant_type) # gate_up: [num_tokens*top_k, 2I] -> activation -> [num_tokens*top_k, I] gate_up = ggml_moe_a8_vec(hidden_states, gate_up_q, topk_ids, top_k, qt, n2, num_tokens) @@ -50,4 +71,21 @@ def fused_experts_gguf_q4_0( return out.sum(dim=1) -__all__ = ["fused_experts_gguf_q4_0"] +def fused_experts_gguf_q4_0( + hidden_states: torch.Tensor, + gate_up_q: torch.Tensor, # [num_slots, 2I, H//32*18] uint8 + down_q: torch.Tensor, # [num_slots, H, I//32*18] uint8 + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str, +) -> torch.Tensor: + """GGUF Q4_0 MoE (backward-compat wrapper). + + This is a thin wrapper around ``fused_experts_gguf`` that hardcodes the Q4_0 type. + All existing callers use this for now; the general function is available for future + multi-quant pipelines. + """ + return fused_experts_gguf(hidden_states, gate_up_q, down_q, topk_weights, topk_ids, activation, int(GGML_Q4_0)) + + +__all__ = ["fused_experts_gguf", "fused_experts_gguf_q4_0"] diff --git a/tests/models/test_gguf_dispatch.py b/tests/models/test_gguf_dispatch.py new file mode 100644 index 00000000..a54284a4 --- /dev/null +++ b/tests/models/test_gguf_dispatch.py @@ -0,0 +1,327 @@ +"""Unit tests for GGUF linear dispatch routing (fused_mul_mat_gguf). + +Tests the 3-tier dispatch strategy without requiring CUDA or the compiled kernel extension. +Uses monkeypatch to replace the C kernel functions with mocks that record which dispatch +path was taken and return correctly-shaped CPU tensors. + +Dispatch order (from python/freetoken/layers/gguf.py:46-75): +1. Empty input early return +2. Unquantized (F32/F16/BF16) → torch matmul +3. Small batch (<= _MMVQ_SAFE) AND MMVQ_TYPES → ggml_mul_mat_vec_a8 +4. MMQ_TYPES → ggml_mul_mat_a8 +5. DEQUANT_TYPES (but not MMQ_TYPES) → ggml_dequantize + torch matmul +6. Else → NotImplementedError +""" + +from __future__ import annotations + +import sys +from types import ModuleType + +import pytest +import torch + +from freetoken.models.gguf.dequant import ( + GGML_BF16, + GGML_F16, + GGML_F32, + GGML_IQ1_M, + GGML_IQ2_S, + GGML_Q2_K, + GGML_Q4_K, + GGML_Q6_K, + BLOCK_SHAPE, +) +from freetoken.layers.gguf import _MMVQ_SAFE, fused_mul_mat_gguf + + +@pytest.fixture +def mock_kernel_module(monkeypatch): + """Replace freetoken.kernel.gguf with a mock that tracks kernel calls. + + Returns a dict tracking which kernels were called and with what arguments. + """ + call_log = { + "ggml_mul_mat_vec_a8": None, + "ggml_mul_mat_a8": None, + "ggml_dequantize": None, + } + + def make_mmvq_kernel(call_log): + """Mock ggml_mul_mat_vec_a8: GEMV kernel for small batch.""" + def kernel(qweight, x, qweight_type, out_features): + call_log["ggml_mul_mat_vec_a8"] = { + "qweight_shape": qweight.shape, + "x_shape": x.shape, + "qweight_type": qweight_type, + "out_features": out_features, + } + batch_size = x.shape[0] + return torch.randn(batch_size, out_features, dtype=x.dtype) + return kernel + + def make_mmq_kernel(call_log): + """Mock ggml_mul_mat_a8: MMQ kernel for large batch.""" + def kernel(qweight, x, qweight_type, out_features): + call_log["ggml_mul_mat_a8"] = { + "qweight_shape": qweight.shape, + "x_shape": x.shape, + "qweight_type": qweight_type, + "out_features": out_features, + } + batch_size = x.shape[0] + return torch.randn(batch_size, out_features, dtype=x.dtype) + return kernel + + def make_dequant_kernel(call_log): + """Mock ggml_dequantize: materializes weight into BF16.""" + def kernel(qweight, qweight_type, out_features, in_features, out_dtype): + call_log["ggml_dequantize"] = { + "qweight_shape": qweight.shape, + "qweight_type": qweight_type, + "out_features": out_features, + "in_features": in_features, + "out_dtype": out_dtype, + } + return torch.randn(out_features, in_features, dtype=out_dtype) + return kernel + + mock_module = ModuleType("freetoken.kernel.gguf") + mock_module.ggml_mul_mat_vec_a8 = make_mmvq_kernel(call_log) + mock_module.ggml_mul_mat_a8 = make_mmq_kernel(call_log) + mock_module.ggml_dequantize = make_dequant_kernel(call_log) + + monkeypatch.setitem(sys.modules, "freetoken.kernel.gguf", mock_module) + + yield call_log + + # Cleanup: ensure mock is removed so it doesn't interfere with other tests + if "freetoken.kernel.gguf" in sys.modules: + del sys.modules["freetoken.kernel.gguf"] + + +def make_qweight(out_features: int, in_features: int, qweight_type: int) -> torch.Tensor: + """Create a mock qweight tensor in packed format for the given type.""" + block, type_size = BLOCK_SHAPE[qweight_type] + row_bytes_val = in_features // block * type_size + return torch.randint(0, 256, (out_features, row_bytes_val), dtype=torch.uint8) + + +class TestEmptyInput: + """Test early return for zero-row input.""" + + def test_empty_input_short_circuits(self, mock_kernel_module): + """Zero-row input returns shape (0, out_features) with no kernel call.""" + out_features = 4096 + in_features = 4096 + qweight_type = GGML_Q4_K + + x = torch.randn(0, in_features) + qweight = make_qweight(out_features, in_features, qweight_type) + + result = fused_mul_mat_gguf(x, qweight, qweight_type) + + assert result.shape == (0, out_features) + assert mock_kernel_module["ggml_mul_mat_vec_a8"] is None + assert mock_kernel_module["ggml_mul_mat_a8"] is None + assert mock_kernel_module["ggml_dequantize"] is None + + +class TestUnquantized: + """Test unquantized paths (F32, F16, BF16).""" + + @pytest.mark.parametrize("qweight_type,dtype", [ + (GGML_F32, torch.float32), + (GGML_F16, torch.float16), + (GGML_BF16, torch.bfloat16), + ]) + def test_unquantized_never_calls_kernels(self, mock_kernel_module, qweight_type, dtype): + """F32/F16/BF16 go through plain torch matmul, no kernel call.""" + out_features = 4096 + in_features = 4096 + batch_size = 32 + + x = torch.randn(batch_size, in_features, dtype=dtype) + # For unquantized, qweight is stored as-is (1 element per 1-4 bytes) + qweight = torch.randn(out_features, in_features, dtype=dtype) + + result = fused_mul_mat_gguf(x, qweight, qweight_type) + + # Check no kernels were called + assert mock_kernel_module["ggml_mul_mat_vec_a8"] is None + assert mock_kernel_module["ggml_mul_mat_a8"] is None + assert mock_kernel_module["ggml_dequantize"] is None + # Result should match torch matmul + expected = x @ qweight.T + assert result.shape == expected.shape + + +class TestSmallBatchMMVQ: + """Test small-batch dispatch to ggml_mul_mat_vec_a8 (MMVQ).""" + + @pytest.mark.parametrize("qweight_type", [ + GGML_Q2_K, # K-quant + GGML_Q4_K, # K-quant + GGML_Q6_K, # K-quant + GGML_IQ2_S, # I-quant + GGML_IQ1_M, # I-quant + ]) + def test_small_batch_uses_mmvq(self, mock_kernel_module, qweight_type): + """Batch <= _MMVQ_SAFE in MMVQ_TYPES calls ggml_mul_mat_vec_a8.""" + out_features = 4096 + in_features = 4096 + batch_size = 1 # Small batch (well below _MMVQ_SAFE) + + x = torch.randn(batch_size, in_features, dtype=torch.bfloat16) + qweight = make_qweight(out_features, in_features, qweight_type) + + result = fused_mul_mat_gguf(x, qweight, qweight_type) + + # MMVQ kernel should have been called + assert mock_kernel_module["ggml_mul_mat_vec_a8"] is not None + assert mock_kernel_module["ggml_mul_mat_a8"] is None + assert mock_kernel_module["ggml_dequantize"] is None + + call_info = mock_kernel_module["ggml_mul_mat_vec_a8"] + assert call_info["x_shape"] == (batch_size, in_features) + assert call_info["qweight_type"] == qweight_type + assert call_info["out_features"] == out_features + assert result.shape == (batch_size, out_features) + + +class TestLargeBatchStandardQuants: + """Test large-batch K-quants and standard quants dispatch to ggml_mul_mat_a8 (MMQ).""" + + @pytest.mark.parametrize("qweight_type", [ + GGML_Q2_K, + GGML_Q4_K, + GGML_Q6_K, + ]) + def test_kquant_large_batch_takes_mmq(self, mock_kernel_module, qweight_type): + """K-quants at large batch call ggml_mul_mat_a8.""" + out_features = 4096 + in_features = 4096 + batch_size = _MMVQ_SAFE + 1 # Large batch (above threshold) + + x = torch.randn(batch_size, in_features, dtype=torch.bfloat16) + qweight = make_qweight(out_features, in_features, qweight_type) + + result = fused_mul_mat_gguf(x, qweight, qweight_type) + + # MMQ kernel should have been called + assert mock_kernel_module["ggml_mul_mat_a8"] is not None + assert mock_kernel_module["ggml_mul_mat_vec_a8"] is None + assert mock_kernel_module["ggml_dequantize"] is None + + call_info = mock_kernel_module["ggml_mul_mat_a8"] + assert call_info["x_shape"] == (batch_size, in_features) + assert call_info["qweight_type"] == qweight_type + assert call_info["out_features"] == out_features + assert result.shape == (batch_size, out_features) + + +class TestIQuantDispatch: + """Test I-quant dispatch: they have MMVQ but NO MMQ kernels.""" + + @pytest.mark.parametrize("qweight_type", [ + GGML_IQ2_S, # enum=22 + GGML_IQ1_M, # enum=29 + ]) + def test_iquant_small_batch_takes_mmvq(self, mock_kernel_module, qweight_type): + """I-quants with batch <= _MMVQ_SAFE call ggml_mul_mat_vec_a8.""" + out_features = 4096 + in_features = 4096 + batch_size = 1 # Small batch + + x = torch.randn(batch_size, in_features, dtype=torch.bfloat16) + qweight = make_qweight(out_features, in_features, qweight_type) + + result = fused_mul_mat_gguf(x, qweight, qweight_type) + + # MMVQ (GEMV) kernel should have been called + assert mock_kernel_module["ggml_mul_mat_vec_a8"] is not None + assert mock_kernel_module["ggml_mul_mat_a8"] is None + assert mock_kernel_module["ggml_dequantize"] is None + + call_info = mock_kernel_module["ggml_mul_mat_vec_a8"] + assert call_info["qweight_type"] == qweight_type + assert result.shape == (batch_size, out_features) + + @pytest.mark.parametrize("qweight_type", [ + GGML_IQ2_S, # enum=22 + GGML_IQ1_M, # enum=29 + ]) + def test_iquant_large_batch_takes_dequant_path(self, mock_kernel_module, qweight_type): + """I-quants have no MMQ kernel, so large batch falls back to dequant + matmul. + + Rationale: I-quants have MMVQ and dequant kernels but no MMQ kernel. Routing them + to ggml_mul_mat_a8 (which doesn't support them) would return uninitialized memory. + Instead, we dequantize and use plain torch matmul. + """ + out_features = 4096 + in_features = 4096 + batch_size = _MMVQ_SAFE + 1 # Large batch (above threshold) + + x = torch.randn(batch_size, in_features, dtype=torch.bfloat16) + qweight = make_qweight(out_features, in_features, qweight_type) + + result = fused_mul_mat_gguf(x, qweight, qweight_type) + + # Dequant path should have been used + assert mock_kernel_module["ggml_dequantize"] is not None + assert mock_kernel_module["ggml_mul_mat_a8"] is None + assert mock_kernel_module["ggml_mul_mat_vec_a8"] is None + + call_info = mock_kernel_module["ggml_dequantize"] + assert call_info["qweight_type"] == qweight_type + assert call_info["out_features"] == out_features + assert call_info["in_features"] == in_features + assert result.shape == (batch_size, out_features) + + +class TestUnsupportedTypes: + """Test error handling for unsupported types.""" + + def test_unknown_type_raises(self, mock_kernel_module): + """Unsupported type (e.g. 99) raises NotImplementedError with type info.""" + out_features = 4096 + in_features = 4096 + batch_size = 32 + unsupported_type = 99 + + x = torch.randn(batch_size, in_features, dtype=torch.bfloat16) + qweight = torch.randint(0, 256, (out_features, 128), dtype=torch.uint8) + + with pytest.raises(NotImplementedError) as excinfo: + fused_mul_mat_gguf(x, qweight, unsupported_type) + + # Error message should contain the type (either the enum value or a name) + error_msg = str(excinfo.value) + assert "99" in error_msg or "unsupported" in error_msg.lower() + + +class TestMMVQThresholdTracking: + """Test that _MMVQ_SAFE constant is properly used in dispatch.""" + + def test_mmvq_threshold_boundary(self, mock_kernel_module): + """At batch == _MMVQ_SAFE, MMVQ path is taken; at +1, MMQ path is taken.""" + out_features = 4096 + in_features = 4096 + qweight_type = GGML_Q4_K + + # Test at threshold: should use MMVQ + x_at_threshold = torch.randn(_MMVQ_SAFE, in_features, dtype=torch.bfloat16) + qweight = make_qweight(out_features, in_features, qweight_type) + + result = fused_mul_mat_gguf(x_at_threshold, qweight, qweight_type) + assert mock_kernel_module["ggml_mul_mat_vec_a8"] is not None + + # Reset call log + mock_kernel_module["ggml_mul_mat_vec_a8"] = None + mock_kernel_module["ggml_mul_mat_a8"] = None + + # Test above threshold: should use MMQ + x_above_threshold = torch.randn(_MMVQ_SAFE + 1, in_features, dtype=torch.bfloat16) + result = fused_mul_mat_gguf(x_above_threshold, qweight, qweight_type) + assert mock_kernel_module["ggml_mul_mat_a8"] is not None + assert mock_kernel_module["ggml_mul_mat_vec_a8"] is None diff --git a/tests/models/test_gguf_type_tables.py b/tests/models/test_gguf_type_tables.py new file mode 100644 index 00000000..e2684c8a --- /dev/null +++ b/tests/models/test_gguf_type_tables.py @@ -0,0 +1,311 @@ +"""Unit tests for GGML type tables in dequant.py. + +These tests verify that the BLOCK_SHAPE table and capability sets accurately +reflect the runtime behavior of the CUDA kernels. They run without CUDA (no kernel +compilation), checking only the type definitions and switch statement extraction. + +Tests cover: +1. Block struct byte sizes derived from ggml-common.h match BLOCK_SHAPE entries +2. row_bytes() roundtrip consistency for various numerals +3. Logical consistency between type-capability sets (subsets, unions, exclusions) +4. Switch statement cases in gguf_kernel.cu match Python frozenset definitions +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +from freetoken.models.gguf.dequant import ( + BLOCK_SHAPE, + DEQUANT_TYPES, + GGML_UNQUANTIZED, + GGML_NAME, + MMQ_TYPES, + MMVQ_TYPES, + MOE_MMQ_TYPES, + MOE_VEC_TYPES, + row_bytes, +) + + +def _parse_ggml_common_h() -> dict[str, int]: + """Extract block struct definitions from ggml-common.h and compute packed byte sizes. + + Returns a dict mapping quant type names (e.g. "block_q4_0") to their byte sizes. + Substitutes QK_K=256 and K_SCALE_SIZE=12 before computing sizes. + """ + header_path = Path(__file__).parent.parent.parent / "python" / "freetoken" / "kernel" / "csrc" / "gguf" / "ggml-common.h" + with open(header_path) as f: + content = f.read() + + # Extract block struct definitions + # Pattern: typedef struct { ... } block_; + pattern = r"typedef\s+struct\s*\{([^}]+)\}\s*block_(\w+);" + matches = re.findall(pattern, content) + + # Define substitutions for macro constants + macros = { + "QK_K": 256, + "K_SCALE_SIZE": 12, + "QK4_0": 32, + "QK4_1": 32, + "QK5_0": 32, + "QK5_1": 32, + "QK8_0": 32, + "QK8_1": 32, + "QK4_NL": 32, + } + + sizes = {} + + for struct_body, type_name in matches: + # Parse field declarations from the struct body + # Format: []; + field_pattern = r"(\w+)\s+(\w+)(?:\[([^\]]+)\])?;" + fields = re.findall(field_pattern, struct_body) + + total_size = 0 + for field_type, field_name, array_size in fields: + # Compute size of each field + if field_type == "half": + field_size = 2 + elif field_type == "half2": + field_size = 4 + elif field_type == "uint8_t": + field_size = 1 + elif field_type == "uint16_t": + field_size = 2 + elif field_type == "uint32_t": + field_size = 4 + elif field_type == "int8_t": + field_size = 1 + elif field_type == "int32_t": + field_size = 4 + else: + # Unknown type; skip + continue + + # Handle array sizes + if array_size: + # Substitute macros in array size + size_expr = array_size + for macro_name, macro_val in macros.items(): + size_expr = size_expr.replace(macro_name, str(macro_val)) + # Evaluate the expression (handles division, multiplication, etc.) + try: + array_len = eval(size_expr) + except Exception: + # If evaluation fails, try simple substitution + continue + field_size *= array_len + + total_size += field_size + + if total_size > 0: + sizes[type_name] = total_size + + return sizes + + +def _extract_switch_cases(file_path: str, func_name: str) -> set[int]: + """Extract case labels from a switch statement in a C++ file. + + Args: + file_path: Path to the .cu file + func_name: Name of the function containing the switch + + Returns: + Set of case numbers extracted from the switch statement + """ + with open(file_path) as f: + content = f.read() + + # Find the function + func_pattern = rf"(?:torch::Tensor|int64_t|void)\s+{func_name}\s*\([^)]*\)\s*\{{" + func_match = re.search(func_pattern, content) + if not func_match: + return set() + + # Find the switch statement within the function + start_pos = func_match.end() + # Scan forward to find "switch (" + switch_pos = content.find("switch (", start_pos) + if switch_pos == -1 or switch_pos > start_pos + 5000: + return set() + + # Find the opening brace of the switch + brace_pos = content.find("{", switch_pos) + # Find the matching closing brace + brace_count = 1 + end_pos = brace_pos + 1 + while brace_count > 0 and end_pos < len(content): + if content[end_pos] == "{": + brace_count += 1 + elif content[end_pos] == "}": + brace_count -= 1 + end_pos += 1 + + switch_body = content[brace_pos + 1:end_pos - 1] + + # Extract all "case :" labels + case_pattern = r"case\s+(\d+):" + cases = set(int(n) for n in re.findall(case_pattern, switch_body)) + + return cases + + +def test_block_shape_matches_ggml_common(): + """Verify BLOCK_SHAPE table entries match struct sizes from ggml-common.h. + + Parses typedef struct { ... } block_ definitions from the header, + computes packed byte sizes from field declarations (accounting for QK_K=256, + K_SCALE_SIZE=12, and field type sizes), and asserts each matches the + corresponding BLOCK_SHAPE[type_enum][1] entry. + """ + parsed_sizes = _parse_ggml_common_h() + + # Map C struct names to GGML type enums and expected sizes + type_mappings = { + "q4_0": (2, 18), + "q4_1": (3, 20), + "q5_0": (6, 22), + "q5_1": (7, 24), + "q8_0": (8, 34), + "q2_K": (10, 84), + "q3_K": (11, 110), + "q4_K": (12, 144), + "q5_K": (13, 176), + "q6_K": (14, 210), + "iq2_xxs": (16, 66), + "iq2_xs": (17, 74), + "iq3_xxs": (18, 98), + "iq1_s": (19, 50), + "iq4_nl": (20, 18), + "iq3_s": (21, 110), + "iq2_s": (22, 82), + "iq4_xs": (23, 136), + "iq1_m": (29, 56), + } + + for struct_name, (ggml_type, expected_bytes) in type_mappings.items(): + parsed_size = parsed_sizes.get(struct_name) + + # If parsing failed for this type (too complex), accept the hard-coded expectation + if parsed_size is not None: + parsed_size = int(parsed_size) + assert parsed_size == expected_bytes, ( + f"block_{struct_name}: parsed size {parsed_size} != expected {expected_bytes}" + ) + + # Also verify BLOCK_SHAPE matches + assert ggml_type in BLOCK_SHAPE, f"GGML type {ggml_type} not in BLOCK_SHAPE" + block_numel, bytes_per_block = BLOCK_SHAPE[ggml_type] + assert bytes_per_block == expected_bytes, ( + f"BLOCK_SHAPE[{ggml_type}][1]={bytes_per_block} != expected {expected_bytes}" + ) + + +def test_row_bytes_roundtrip(): + """Test row_bytes() consistency for various block counts. + + For each quant type, verify that row_bytes(k * block_numel, type) == k * bytes_per_block + for k in [1, 2, 4, 16], and that non-multiples of block_numel raise AssertionError. + """ + for ggml_type, (block_numel, bytes_per_block) in BLOCK_SHAPE.items(): + # Test valid multiples + for k in [1, 2, 4, 16]: + numel = k * block_numel + expected = k * bytes_per_block + actual = row_bytes(numel, ggml_type) + assert actual == expected, ( + f"row_bytes({numel}, {ggml_type}): got {actual}, expected {expected}" + ) + + # Test that non-multiples raise + if block_numel > 1: + bad_numel = block_numel + 1 + with pytest.raises(AssertionError): + row_bytes(bad_numel, ggml_type) + + +def test_capability_sets_are_consistent(): + """Verify logical consistency and completeness of type-capability sets. + + Checks: + - MMVQ_TYPES == DEQUANT_TYPES (both handle all STD_K and IQ types) + - MMQ_TYPES == MOE_MMQ_TYPES (both handle STD_K only) + - MOE_VEC_TYPES == DEQUANT_TYPES (both handle all STD_K and IQ types) + - MMQ_TYPES is a strict subset of MMVQ_TYPES + - MOE_MMQ_TYPES is a strict subset of MOE_VEC_TYPES + - Every member of every set has a BLOCK_SHAPE entry and GGML_NAME entry + - None of the five sets intersects GGML_UNQUANTIZED (F32, F16, BF16) + """ + # Check set equalities + assert MMVQ_TYPES == DEQUANT_TYPES, ( + f"MMVQ_TYPES {MMVQ_TYPES} != DEQUANT_TYPES {DEQUANT_TYPES}" + ) + assert MMQ_TYPES == MOE_MMQ_TYPES, ( + f"MMQ_TYPES {MMQ_TYPES} != MOE_MMQ_TYPES {MOE_MMQ_TYPES}" + ) + assert MOE_VEC_TYPES == DEQUANT_TYPES, ( + f"MOE_VEC_TYPES {MOE_VEC_TYPES} != DEQUANT_TYPES {DEQUANT_TYPES}" + ) + + # Check subset relationships + assert MMQ_TYPES < MMVQ_TYPES, ( + f"MMQ_TYPES {MMQ_TYPES} is not a strict subset of MMVQ_TYPES {MMVQ_TYPES}" + ) + assert MOE_MMQ_TYPES < MOE_VEC_TYPES, ( + f"MOE_MMQ_TYPES {MOE_MMQ_TYPES} is not a strict subset of MOE_VEC_TYPES {MOE_VEC_TYPES}" + ) + + # Check every member is in BLOCK_SHAPE and GGML_NAME + all_types = DEQUANT_TYPES | MMVQ_TYPES | MMQ_TYPES | MOE_VEC_TYPES | MOE_MMQ_TYPES + for ggml_type in all_types: + assert ggml_type in BLOCK_SHAPE, f"GGML type {ggml_type} not in BLOCK_SHAPE" + assert ggml_type in GGML_NAME, f"GGML type {ggml_type} not in GGML_NAME" + + # Check no intersection with unquantized types + for ggml_type_set in [DEQUANT_TYPES, MMVQ_TYPES, MMQ_TYPES, MOE_VEC_TYPES, MOE_MMQ_TYPES]: + assert ggml_type_set.isdisjoint(GGML_UNQUANTIZED), ( + f"Type set {ggml_type_set} intersects GGML_UNQUANTIZED {GGML_UNQUANTIZED}" + ) + + +def test_capability_sets_match_cuda_switches(): + """Extract switch cases from gguf_kernel.cu and verify against Python sets. + + For each CUDA kernel (ggml_mul_mat_vec_a8, ggml_mul_mat_a8, ggml_moe_a8, + ggml_moe_a8_vec, ggml_moe_get_block_size), extracts the case labels from + the switch(type) block and asserts they match the corresponding Python frozenset. + + This is the critical test that prevents Python tables from drifting from C source. + """ + kernel_path = Path(__file__).parent.parent.parent / "python" / "freetoken" / "kernel" / "csrc" / "gguf" / "gguf_kernel.cu" + + # Extract cases for each kernel + mmvq_cases = _extract_switch_cases(str(kernel_path), "ggml_mul_mat_vec_a8") + mmq_cases = _extract_switch_cases(str(kernel_path), "ggml_mul_mat_a8") + moe_a8_cases = _extract_switch_cases(str(kernel_path), "ggml_moe_a8") + moe_vec_cases = _extract_switch_cases(str(kernel_path), "ggml_moe_a8_vec") + moe_block_size_cases = _extract_switch_cases(str(kernel_path), "ggml_moe_get_block_size") + + # Verify against Python sets + assert mmvq_cases == MMVQ_TYPES, ( + f"ggml_mul_mat_vec_a8 cases {mmvq_cases} != MMVQ_TYPES {MMVQ_TYPES}" + ) + assert mmq_cases == MMQ_TYPES, ( + f"ggml_mul_mat_a8 cases {mmq_cases} != MMQ_TYPES {MMQ_TYPES}" + ) + assert moe_a8_cases == MOE_MMQ_TYPES, ( + f"ggml_moe_a8 cases {moe_a8_cases} != MOE_MMQ_TYPES {MOE_MMQ_TYPES}" + ) + assert moe_vec_cases == MOE_VEC_TYPES, ( + f"ggml_moe_a8_vec cases {moe_vec_cases} != MOE_VEC_TYPES {MOE_VEC_TYPES}" + ) + assert moe_block_size_cases == MOE_MMQ_TYPES, ( + f"ggml_moe_get_block_size cases {moe_block_size_cases} != MOE_MMQ_TYPES {MOE_MMQ_TYPES}" + ) From 6f40e506bde4085b9b79c99e5f821a49ddc3e2f8 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 20:32:29 -0700 Subject: [PATCH 02/36] feat(gguf): qwen35moe config + tensor-name adapter (not yet registered) Builds the FreeToken ModelConfig from llama.cpp qwen35moe GGUF metadata and maps tensor names back to FreeToken parameter names. Derived against vcruz305/Ornith-1.5-35B-A3B-GGUF. The GDN projections are the non-obvious part. llama.cpp's qwen3.5 mapping splits what qwen3next fused -- attn_qkv<-in_proj_qkv, attn_gate<-in_proj_z, ssm_beta<-in_proj_b, ssm_alpha<-in_proj_a -- and FreeToken's HF loader already fuses those same pairs (_PT_FP8_FUSE / _PT_BF16_FUSE in weight.py), so we emit the fused representation and the model code sees one shape regardless of source. Head layout is pinned by arithmetic rather than assumed: attn_qkv's packed output width of 8192 equals q(group_count*state_size) + k(group_count*state_size) + v(time_step_rank*state_size) = 2048+2048+4096, so num_k_heads==ssm.group_count and num_v_heads==ssm.time_step_rank at head_dim==ssm.state_size==128. parse_gguf_config asserts the inner_size relation rather than trusting it. Deliberately NOT added to GGUF_ARCH_TO_REGISTRY yet: iter_gguf_weights is not implemented, and the routed-expert path cannot serve this checkpoint regardless -- Ornith's experts are IQ3_S (gate/up) + Q4_K (down) as separate stacks, while the offload/CPU bank machinery keys on a single fused "q4_0" schema. Registering the arch now would advertise support that does not exist. See the follow-up plan. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/qwen3_5_moe/gguf.py | 248 ++++++++++++++++++++ 1 file changed, 248 insertions(+) create mode 100644 python/freetoken/models/qwen3_5_moe/gguf.py diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py new file mode 100644 index 00000000..e91c838b --- /dev/null +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -0,0 +1,248 @@ +"""Qwen3.5-MoE GGUF adapter: build the FreeToken ``ModelConfig`` and stream weights +from a llama.cpp ``qwen35moe`` checkpoint. + +The geometry is identical to the HF qwen3_5_moe model (hybrid GDN/full attention on a +``full_attention_interval`` stride, 256 routed experts + a shared expert, NextN/MTP head), +so this produces the *same* ``ModelConfig`` as ``qwen3_5_moe.config.parse_config`` -- only +the source is GGUF KV metadata instead of a HF config object. + +Tensor-name mapping is the inverse of llama.cpp's ``gguf-py/gguf/tensor_mapping.py``. +The one non-obvious part is the GDN projections. llama.cpp's *qwen3.5* mapping splits +what qwen3next fused:: + + attn_qkv <- model.layers.{i}.linear_attn.in_proj_qkv + attn_gate <- model.layers.{i}.linear_attn.in_proj_z + ssm_beta <- model.layers.{i}.linear_attn.in_proj_b + ssm_alpha <- model.layers.{i}.linear_attn.in_proj_a + +FreeToken's HF loader already knows how to put those back together -- see ``_PT_FP8_FUSE`` +and ``_PT_BF16_FUSE`` in ``weight.py``, which fuse ``(in_proj_qkv, in_proj_z) -> +in_proj_qkvz`` and ``(in_proj_b, in_proj_a) -> in_proj_ba`` in that order. We emit the +same fused buffers here so the model code sees one representation regardless of source. + +Verified against vcruz305/Ornith-1.5-35B-A3B-GGUF (IQ3_M), whose metadata gives +block_count=41 (40 decoder layers + 1 NextN block), embedding_length=2048, +head_count=16, head_count_kv=2, key_length=value_length=256, expert_count=256, +expert_used_count=8, expert_feed_forward_length=512, full_attention_interval=4, +ssm.conv_kernel=4, ssm.state_size=128, ssm.group_count=16, ssm.time_step_rank=32, +ssm.inner_size=4096. Those are self-consistent: the packed ``attn_qkv`` output width of +8192 is exactly q(16*128) + k(16*128) + v(32*128), i.e. num_k_heads == group_count and +num_v_heads == time_step_rank, both with head_dim == state_size == 128. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Iterator + +import torch + +from freetoken.models.config import ( + FullAttentionGroupConfig, + LinearGatedDeltaGroupConfig, + ModelConfig, + RotaryConfig, +) +from freetoken.models.gguf.dequant import GGML_NAME, dequantize + +if TYPE_CHECKING: + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "qwen35moe" + + +def _kv(shim: "GgufConfigShim", key: str, default: Any = None) -> Any: + """Read ``qwen35moe.`` from the GGUF metadata.""" + val = shim.metadata.get(f"{_ARCH}.{key}", default) + if val is None and default is None: + raise ValueError(f"GGUF {shim.model_path}: missing required key {_ARCH}.{key}") + return val + + +def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + block_count = int(_kv(shim, "block_count")) + # llama.cpp appends the NextN/MTP block to the decoder stack. FreeToken serves + # text-only without speculative decoding, so the MTP block is not a decoder layer. + nextn = int(_kv(shim, "nextn_predict_layers", 0)) + num_layers = block_count - nextn + + hidden_size = int(_kv(shim, "embedding_length")) + num_qo_heads = int(_kv(shim, "attention.head_count")) + num_kv_heads = int(_kv(shim, "attention.head_count_kv")) + head_dim = int(_kv(shim, "attention.key_length")) + rms_eps = float(_kv(shim, "attention.layer_norm_rms_epsilon")) + rope_base = float(_kv(shim, "rope.freq_base")) + rotary_dim = int(_kv(shim, "rope.dimension_count")) + max_pos = int(_kv(shim, "context_length")) + + num_experts = int(_kv(shim, "expert_count", 0)) + experts_per_tok = int(_kv(shim, "expert_used_count", 0)) + moe_inter = int(_kv(shim, "expert_feed_forward_length", 0)) + shared_inter = int(_kv(shim, "expert_shared_feed_forward_length", 0)) + + # GDN geometry. state_size is the per-head dim; group_count is the number of k heads + # and time_step_rank the number of v heads (see module docstring for the arithmetic + # that pins this down against the packed attn_qkv width). + conv_kernel = int(_kv(shim, "ssm.conv_kernel")) + state_size = int(_kv(shim, "ssm.state_size")) + num_k_heads = int(_kv(shim, "ssm.group_count")) + num_v_heads = int(_kv(shim, "ssm.time_step_rank")) + inner_size = int(_kv(shim, "ssm.inner_size")) + if num_v_heads * state_size != inner_size: + raise ValueError( + f"GGUF {shim.model_path}: ssm.time_step_rank({num_v_heads}) * " + f"ssm.state_size({state_size}) != ssm.inner_size({inner_size}); the GDN head " + "layout assumed by this adapter does not hold for this checkpoint" + ) + + # llama.cpp writes the stride, not a per-layer list: layer i is full attention when + # (i + 1) % interval == 0. For Ornith (interval=4, 40 layers) that is 3,7,...,39. + interval = int(_kv(shim, "full_attention_interval")) + full_ids = tuple(i for i in range(num_layers) if (i + 1) % interval == 0) + linear_ids = tuple(i for i in range(num_layers) if i not in set(full_ids)) + + full_rotary = RotaryConfig( + head_dim=head_dim, + rotary_dim=rotary_dim, + max_position=max_pos, + base=rope_base, + scaling=None, + ) + groups = tuple( + sorted( + ( + FullAttentionGroupConfig( + name="full", + layer_ids=full_ids, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + rotary_config=full_rotary, + ), + LinearGatedDeltaGroupConfig( + name="linear", + layer_ids=linear_ids, + num_key_heads=num_k_heads, + num_value_heads=num_v_heads, + key_head_dim=state_size, + value_head_dim=state_size, + conv_kernel_dim=conv_kernel, + output_gate=True, + ), + ), + key=lambda g: g.layer_ids[0] if g.layer_ids else 1 << 30, + ) + ) + + return ModelConfig( + num_layers=num_layers, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + hidden_size=hidden_size, + vocab_size=shim.vocab_size, + intermediate_size=0, # every layer is MoE in qwen35moe + hidden_act="silu", + rms_norm_eps=rms_eps, + tie_word_embeddings=shim.tie_word_embeddings, + rotary_config=full_rotary, + num_experts=num_experts, + num_experts_per_tok=experts_per_tok, + moe_intermediate_size=moe_inter, + shared_expert_intermediate_size=shared_inter, + norm_topk_prob=True, + moe_enabled=num_experts > 0, + use_qk_norm=True, + model_type=_ARCH, + architectures=["Qwen35MoeGGUFForCausalLM"], + vision_config=None, + image_token_id=None, + attention_groups=groups, + expert_quant="gguf", + weight_block_size=None, + attn_quant="gguf", + dense_quant="gguf", + lm_head_quant="gguf", + ) + + +# -------------------------------------------------------------------------------------- +# Tensor-name mapping (inverse of llama.cpp gguf-py/gguf/tensor_mapping.py for qwen3.5) +# -------------------------------------------------------------------------------------- + +# Per-layer 1:1 renames that need no reshaping or fusing. +_LAYER_MAP: dict[str, str] = { + # shared by both layer kinds + "attn_norm.weight": "input_layernorm.weight", + "post_attention_norm.weight": "post_attention_layernorm.weight", + # full-attention layers + "attn_q.weight": "self_attn.q_proj.weight", + "attn_k.weight": "self_attn.k_proj.weight", + "attn_v.weight": "self_attn.v_proj.weight", + "attn_output.weight": "self_attn.o_proj.weight", + "attn_q_norm.weight": "self_attn.q_norm.weight", + "attn_k_norm.weight": "self_attn.k_norm.weight", + # GDN (linear-attention) layers + "ssm_conv1d.weight": "linear_attn.conv1d.weight", + "ssm_norm.weight": "linear_attn.norm.weight", + "ssm_out.weight": "linear_attn.out_proj.weight", + "ssm_a": "linear_attn.A_log", + "ssm_dt.bias": "linear_attn.dt_bias", + # MoE router + shared expert + "ffn_gate_inp.weight": "mlp.gate.weight", + "ffn_gate_inp_shexp.weight": "mlp.shared_expert_gate.weight", + "ffn_gate_shexp.weight": "mlp.shared_expert.gate_proj.weight", + "ffn_up_shexp.weight": "mlp.shared_expert.up_proj.weight", + "ffn_down_shexp.weight": "mlp.shared_expert.down_proj.weight", +} + +# Pairs llama.cpp splits that FreeToken's model code wants fused, in concat order. +# Mirrors _PT_FP8_FUSE / _PT_BF16_FUSE in weight.py. +_FUSE: dict[str, tuple[str, str]] = { + "linear_attn.in_proj_qkvz.weight": ("attn_qkv.weight", "attn_gate.weight"), + "linear_attn.in_proj_ba.weight": ("ssm_beta.weight", "ssm_alpha.weight"), +} + +# Routed-expert stacks: [num_experts, out, in] packed blocks, handled by the offload +# expert-bank loader rather than yielded as ordinary parameters. +_EXPERT_SUFFIXES = ( + "ffn_gate_exps.weight", + "ffn_up_exps.weight", + "ffn_down_exps.weight", +) + +_GLOBAL_MAP: dict[str, str] = { + "token_embd.weight": "model.embed_tokens.weight", + "output_norm.weight": "model.norm.weight", + "output.weight": "lm_head.weight", +} + + +def gguf_name_to_freetoken(name: str, num_layers: int) -> str | None: + """Map one llama.cpp tensor name to its FreeToken parameter name. + + Returns ``None`` for tensors FreeToken does not consume (the NextN/MTP block, and + the routed-expert stacks, which the expert-bank loader reads directly). + """ + if name in _GLOBAL_MAP: + return _GLOBAL_MAP[name] + if not name.startswith("blk."): + return None + _, idx, suffix = name.split(".", 2) + layer = int(idx) + if layer >= num_layers: + return None # the trailing NextN/MTP block: served text-only, no speculation + if suffix.startswith("nextn."): + return None + if suffix in _EXPERT_SUFFIXES: + return None + mapped = _LAYER_MAP.get(suffix) + if mapped is None: + return None + return f"model.layers.{layer}.{mapped}" + + +__all__ = [ + "parse_gguf_config", + "gguf_name_to_freetoken", + "_FUSE", + "_EXPERT_SUFFIXES", +] From a9614af8c1f15f7613e9a3a6c02989f3d9cdd2fd Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 20:53:28 -0700 Subject: [PATCH 03/36] test(gguf): resolve IQ3S_N_SCALE when parsing block_iq3_s from ggml-common.h The header-derived size check dropped block_iq3_s's scales[IQ3S_N_SCALE] field because the macro table lacked that define, yielding 106 bytes and failing against BLOCK_SHAPE's 110. 110 is correct and independently confirmed against a real checkpoint: Ornith-1.5-35B IQ3_M's token_embd.weight is 248320*2048/256*110 = 218,521,600 bytes on disk, matching row_bytes() exactly. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- tests/models/test_gguf_type_tables.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/models/test_gguf_type_tables.py b/tests/models/test_gguf_type_tables.py index e2684c8a..5433657a 100644 --- a/tests/models/test_gguf_type_tables.py +++ b/tests/models/test_gguf_type_tables.py @@ -50,6 +50,11 @@ def _parse_ggml_common_h() -> dict[str, int]: macros = { "QK_K": 256, "K_SCALE_SIZE": 12, + # ggml-common.h: `#define IQ3S_N_SCALE QK_K / 64`. Without it block_iq3_s's + # `scales[IQ3S_N_SCALE]` field parses as 0 and the struct comes out 106 instead + # of 110 -- which the real-checkpoint check disproves: Ornith's IQ3_S + # token_embd.weight is 248320*2048/256*110 == 218,521,600 bytes on disk exactly. + "IQ3S_N_SCALE": 4, "QK4_0": 32, "QK4_1": 32, "QK5_0": 32, From 8f17e8d8903d2d596c7172d978790b2d51e22e7f Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 21:49:26 -0700 Subject: [PATCH 04/36] feat(gguf): serve qwen35moe GGUF checkpoints (Ornith-1.5-35B-A3B) Adds the architecture path plus the mixed-quant machinery llama.cpp checkpoints actually need. GGUFMergedLinear (layers/gguf.py): a merged projection whose parts carry different ggml types. FreeToken fuses projections by concatenating packed rows, which needs a shared row_bytes -- true for gemma4 (all Q4_0) but not for llama.cpp's mixed quants. Ornith's full-attention qkv is q,k IQ3_S + v Q4_K and its GDN in_proj is attn_qkv Q4_K + gate/beta/alpha IQ3_S; 880- and 1152-byte rows cannot interleave. Concatenating the per-part OUTPUTS is identical to the merged GEMM since every part reads the same x. gguf_merged_or_plain() keeps the uniform case on the cheaper single-kernel GGUFLinear. qwen3_5_moe/gguf.py: iter_gguf_weights + convert_qwen35_to_gguf. Maps the GGUF tensor table onto the module tree, fusing per layer in the orders gdn.py and attention.py declare (in_proj = attn_qkv, attn_gate, ssm_beta, ssm_alpha, matching _in_proj_split = [conv_dim, value_dim, n_v, n_v]). A_log/dt_bias stay fp32 as gdn.py requires; conv1d is reshaped to [conv_dim, 1, kernel] for its .squeeze(1). qwen3_5_moe/gguf_experts.py: routed-expert host banks in native packed bytes. Offload path generalized off the q4_0 literal: a "gguf" bank schema, byte accounting from row_bytes, and the (gate_up, down) ggml types threaded ModelConfig -> ExpertBanks -> OffloadMoeCache -> the MoE kernels, because a GGUF's row stride is a property of the file, not the format. fused_experts_gguf takes a separate down_quant_type: the two banks own separate slot pools and may legitimately differ. What is NOT supported, and fails loudly rather than silently: a bank whose ggml type varies BY LAYER. moe_vec.cuh addresses the pool as expert * nrows * (ncols / qk) with no padding allowance, and the pool is one allocation shared by all layers, so mixed strides would read every block at the wrong offset and emit fluent garbage. _gguf_banks raises with the offending layers named. Ornith IQ3_S and IQ3_XXS have uniform banks; IQ3_M, IQ2_M, IQ2_XXS and IQ1_S split ffn_down_exps across two types. Also: _model_setup_override now declines for expert_quant == "gguf". qwen3_5_moe exports an NVFP4 setup_offload_expert_banks that would otherwise hijack the GGUF load and fail deep inside an NVFP4 reader. Not yet run end-to-end; weights have not been loaded on hardware. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/engine/engine.py | 2 + python/freetoken/layers/gguf.py | 140 +++++- python/freetoken/layers/moe.py | 18 + python/freetoken/models/config.py | 14 + python/freetoken/models/gguf/config.py | 4 +- .../freetoken/models/qwen3_5_moe/__init__.py | 8 + python/freetoken/models/qwen3_5_moe/gguf.py | 426 +++++++++++++++++- .../models/qwen3_5_moe/gguf_experts.py | 280 ++++++++++++ python/freetoken/models/register.py | 7 + python/freetoken/moe/expert_banks.py | 83 +++- python/freetoken/moe/fused_q4_0.py | 24 +- python/freetoken/moe/offload_cache.py | 12 + tests/models/test_qwen35moe_gguf.py | 315 +++++++++++++ 13 files changed, 1321 insertions(+), 12 deletions(-) create mode 100644 python/freetoken/models/qwen3_5_moe/gguf_experts.py create mode 100644 tests/models/test_qwen35moe_gguf.py diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 22c4a6c7..6771719d 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -491,6 +491,7 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks) -> tuple[int kv_reserve_tokens=max(config.kv_reserve_tokens, min_reserve), page_size=page_tokens, quant_format=banks.quant_format, + gguf_expert_types=banks.gguf_expert_types, ) def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: @@ -609,6 +610,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: prefill_overlap=config.moe_prefill_overlap, prefill_hit_d2d=config.moe_prefill_hit_d2d, quant_format=banks.quant_format, + gguf_expert_types=banks.gguf_expert_types, decode_target=decode_target, hybrid_max_fetch=config.moe_hybrid_max_fetch, ) diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py index 883942c0..28db4df3 100644 --- a/python/freetoken/layers/gguf.py +++ b/python/freetoken/layers/gguf.py @@ -9,6 +9,16 @@ share an input dim, hence the same ``row_bytes``), so a fused layer is still one ``[out, row_bytes]`` qweight -- no per-shard padding bookkeeping needed. +**Merged vs. plain fused projections**: + +When all output parts share the same quant type (the common case in gemma4), a plain +``GGUFLinear`` with concatenated packed rows is valid and efficient -- one kernel launch +dequantizes and multiplies. When parts use different quant types (as in Ornith's IQ3_M +checkpoint, where qkv_proj mixes IQ3_S and Q4_K), row_bytes differs per part, so torch.cat +would produce garbage. ``GGUFMergedLinear`` instead materializes the output of each part +separately via ``fused_mul_mat_gguf`` and concatenates the results along dim=-1 (equivalent +to the GEMM because all parts read the same input: ``cat([x @ W1.T, x @ W2.T]) == x @ cat([W1, W2], 0).T``). + **Matmul dispatch strategy** (4-tier, per fused_mul_mat_gguf): 1. **Unquantized (F32, F16, BF16)**: straight torch matmul ``x @ qweight.T``. @@ -99,6 +109,92 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return out +class GGUFMergedLinear(BaseOP): + """Merged linear projection with parts that have different quant types. + + Used when fusing output-parallel projections (qkv, gate_up) whose parts use different + quantization types. Unlike GGUFLinear (which concatenates packed rows along dim 0 and + requires all parts to share row_bytes), GGUFMergedLinear materializes the output of + each part separately via fused_mul_mat_gguf, then concatenates the results. + + Mathematically equivalent to a single GEMM, since all parts read the same input x: + cat([x @ W1.T, x @ W2.T]) == x @ cat([W1, W2], 0).T + (source: llama.cpp's iq*_m mixed-quant strategy). + """ + + def __init__( + self, + in_features: int, + output_sizes: list[int], + quant_types: list[int], + has_bias: bool = False, + ): + """Initialize a merged linear projection. + + Args: + in_features: Input feature dimension (shared by all parts). + output_sizes: List of output sizes for each part; must all be > 0. + quant_types: List of GGML quant types, one per part; must match output_sizes length. + has_bias: Whether to allocate a bias term. + + Raises: + ValueError: If output_sizes and quant_types lengths do not match, or if any output_size <= 0. + NotImplementedError: If any quant_type is not supported (not in MMVQ_TYPES or GGML_UNQUANTIZED). + """ + if len(output_sizes) != len(quant_types): + raise ValueError( + f"output_sizes length {len(output_sizes)} != quant_types length {len(quant_types)}" + ) + if not all(o > 0 for o in output_sizes): + raise ValueError(f"all output_sizes must be > 0, got {output_sizes}") + + # Validate each quant type is supported. + for qt in quant_types: + if qt not in MMVQ_TYPES and qt not in GGML_UNQUANTIZED: + raise NotImplementedError( + f"quant type {GGML_NAME.get(qt, qt)} not in MMVQ_TYPES or GGML_UNQUANTIZED" + ) + + self.in_features = in_features + self.output_sizes = output_sizes + self.out_features = sum(output_sizes) + self._quant_types = quant_types + self.part_names = [] + + # Allocate packed weight buffers: one named tensor per part (qweight_0, qweight_1, ...). + # Named (not underscore-prefixed) so they are discovered by state_dict. + for i, (out_size, qt) in enumerate(zip(output_sizes, quant_types)): + name = f"qweight_{i}" + self.part_names.append(name) + setattr( + self, + name, + torch.empty(out_size, row_bytes(in_features, qt), dtype=torch.uint8), + ) + + self.bias = torch.empty(self.out_features) if has_bias else None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass: compute each part's output and concatenate along dim=-1. + + Args: + x: Input tensor of shape [..., in_features]. + + Returns: + Tensor of shape [..., out_features] with parts concatenated along dim=-1. + """ + parts = [] + for name, qt in zip(self.part_names, self._quant_types): + qweight = getattr(self, name) + part_out = fused_mul_mat_gguf(x, qweight, qt) + parts.append(part_out) + + out = torch.cat(parts, dim=-1) + if self.bias is not None: + out = out + self.bias + return out + + class GGUFEmbedding(BaseOP): """Vocab embedding stored as a native GGUF block-quantized table. @@ -136,4 +232,46 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y -__all__ = ["GGUFLinear", "GGUFEmbedding", "fused_mul_mat_gguf"] +def gguf_merged_or_plain( + in_features: int, + output_sizes: list[int], + quant_types: list[int], + has_bias: bool = False, +) -> GGUFLinear | GGUFMergedLinear: + """Choose between GGUFLinear (uniform quant types) and GGUFMergedLinear (mixed types). + + When all output parts share the same quant type (the uniform case, common in gemma4), + return a GGUFLinear with concatenated packed rows -- valid and cheaper since row_bytes + is identical per part (one kernel launch instead of N). + + When quant types differ (the mixed case, produced by llama.cpp's IQ*_M / Q*_K_M), + return a GGUFMergedLinear to avoid torch.cat garbage from misaligned row_bytes. + + Args: + in_features: Input feature dimension. + output_sizes: List of output sizes for each part. + quant_types: List of GGML quant types, one per part. + has_bias: Whether to allocate a bias term. + + Returns: + GGUFLinear if all quant types are identical, else GGUFMergedLinear. + """ + if len(set(quant_types)) == 1: + # Uniform case: all parts use the same quant type. + # Concatenate packed rows (they share row_bytes) into a single [sum(output_sizes), row_bytes] weight. + out_features = sum(output_sizes) + qt = quant_types[0] + lin = GGUFLinear(in_features, out_features, qt, has_bias=has_bias) + return lin + else: + # Mixed case: parts use different quant types. + return GGUFMergedLinear(in_features, output_sizes, quant_types, has_bias=has_bias) + + +__all__ = [ + "GGUFLinear", + "GGUFMergedLinear", + "GGUFEmbedding", + "fused_mul_mat_gguf", + "gguf_merged_or_plain", +] diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index d68d8ded..eec3aa83 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -531,6 +531,24 @@ def _expert_gemm( return fused_experts_gguf_q4_0( hidden_states, gate_up, down, topk_weights, topk_ids, self.activation ) + if fmt == "gguf": + # Same MMVQ grouped GEMV as q4_0, but the two banks carry whatever ggml types + # the checkpoint chose (Ornith IQ3_S: both banks IQ3_S; the attention-side + # projections mix, but expert banks must be uniform -- see ModelConfig + # .gguf_expert_types for why the slot pool cannot hold two strides). + from freetoken.moe.fused_q4_0 import fused_experts_gguf + + gate_up, down = views + types = cache.gguf_expert_types + assert types is not None, ( + "quant_format 'gguf' requires gguf_expert_types on the offload cache " + "(set from ModelConfig.gguf_expert_types at cache construction)" + ) + t_gate_up, t_down = types + return fused_experts_gguf( + hidden_states, gate_up, down, topk_weights, topk_ids, self.activation, + quant_type=t_gate_up, down_quant_type=t_down, + ) if fmt == "mxfp4_triton": # gpt-oss MXFP4 experts (biased, clamped swiglu): transposed split-K GEMV # decode + grouped `_t` prefill. The swiglu scalars live on the layer diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index f6105e1f..09c61151 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -243,6 +243,20 @@ class ModelConfig: # it bf16. Separate from dense_quant because only some NVFP4 checkpoints quantize lm_head # (modelopt MIXED_PRECISION does; pure NVFP4 leaves it bf16). lm_head_quant: str = "none" + # Routed-expert bank ggml types for ``expert_quant == "gguf"``, as + # ``(gate_up_type, down_type)``. Unlike the other quant tags, which name a *format* + # whose geometry is then fixed, a GGUF checkpoint picks a ggml type per tensor, so the + # bank row stride (``row_bytes``) is only knowable from the file. The offload cache + # sizes its slot pool from these and the MoE kernels take the type as an argument. + # + # Both entries must be uniform across layers: the GPU slot pool is one contiguous + # allocation shared by every layer, and moe_vec.cuh addresses it as + # ``expert * nrows * (ncols / qk)`` with no padding allowance -- so two layers with + # different row strides in one pool would read every block at the wrong offset. + # ``load_gguf_expert_sources`` raises when a checkpoint violates that (llama.cpp's + # *_M mixes do: Ornith IQ3_M splits ffn_down_exps across Q4_K and IQ3_S, while its + # IQ3_S/IQ3_XXS siblings are uniform and load fine). + gguf_expert_types: tuple[int, int] | None = None shared_expert_intermediate_size: int = 0 use_qk_norm: bool = False # ----- DeepSeek/GLM-style MoE extensions (default keeps other models intact) ----- diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 63b1a18b..29d7efda 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -15,9 +15,11 @@ from .reader import gguf_architecture, load_gguf_metadata, gguf_tensor_names # GGUF ``general.architecture`` -> FreeToken registry key (a GGUF-specific spec that -# reuses the model classes but a GGUF parse_config / iter_weights). +# reuses the model classes but a GGUF parse_config / iter_weights). The key is the +# value of the GGUF general.architecture metadata key. GGUF_ARCH_TO_REGISTRY: dict[str, str] = { "gemma4": "Gemma4GGUFForCausalLM", + "qwen35moe": "Qwen35MoeGGUFForCausalLM", } diff --git a/python/freetoken/models/qwen3_5_moe/__init__.py b/python/freetoken/models/qwen3_5_moe/__init__.py index 98936e9f..1ee40363 100644 --- a/python/freetoken/models/qwen3_5_moe/__init__.py +++ b/python/freetoken/models/qwen3_5_moe/__init__.py @@ -7,6 +7,10 @@ load_nvfp4_expert_sources_parallel, setup_offload_expert_banks, ) +from .gguf import parse_gguf_config, iter_gguf_weights +# Resolved off this package by freetoken.moe.expert_banks._gguf_banks (the GGUF expert +# layout is architecture-specific, so the provider looks it up via the model registry). +from .gguf_experts import gguf_expert_types, load_gguf_expert_sources __all__ = [ "Qwen3_5MoEForCausalLM", @@ -16,4 +20,8 @@ "load_nvfp4_expert_sources", "load_nvfp4_expert_sources_parallel", "setup_offload_expert_banks", + "parse_gguf_config", + "iter_gguf_weights", + "gguf_expert_types", + "load_gguf_expert_sources", ] diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index e91c838b..8bdc9cb5 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -36,13 +36,23 @@ import torch +# Verify that LinearGatedDeltaGroupConfig is available for isinstance checks +# (imported above in the config module import) + from freetoken.models.config import ( FullAttentionGroupConfig, LinearGatedDeltaGroupConfig, ModelConfig, RotaryConfig, ) -from freetoken.models.gguf.dequant import GGML_NAME, dequantize +from freetoken.models.gguf.dequant import ( + GGML_IQ3_S, + GGML_NAME, + GGML_Q4_K, + GGML_Q6_K, + dequantize, + row_bytes, +) if TYPE_CHECKING: from freetoken.models.gguf.config import GgufConfigShim @@ -58,6 +68,29 @@ def _kv(shim: "GgufConfigShim", key: str, default: Any = None) -> Any: return val +def _uniform_expert_types(model_path: str, num_layers: int) -> tuple[int, int] | None: + """``(gate_up, down)`` ggml types of the routed-expert banks, or None if not uniform. + + The offload slot pool is one allocation per bank shared by every layer, and + ``moe_vec.cuh`` addresses it as ``expert * nrows * (ncols / qk)`` with no padding + allowance -- so a bank whose type varies by layer cannot be served. We return None + rather than raising here because ``parse_gguf_config`` also runs for metadata-only + inspection; ``expert_banks._gguf_banks`` is where the load actually fails, with the + offending layers named. (llama.cpp's *_M mixes hit this: Ornith IQ3_M splits + ffn_down_exps across Q4_K and IQ3_S, while IQ3_S / IQ3_XXS are uniform.) + """ + from .gguf_experts import gguf_expert_types + + try: + types = gguf_expert_types(model_path, num_layers) + except Exception: + return None + gate_up, down = set(types["gate_up"]), set(types["down"]) + if len(gate_up) != 1 or len(down) != 1: + return None + return (next(iter(gate_up)), next(iter(down))) + + def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: block_count = int(_kv(shim, "block_count")) # llama.cpp appends the NextN/MTP block to the decoder stack. FreeToken serves @@ -157,6 +190,7 @@ def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: image_token_id=None, attention_groups=groups, expert_quant="gguf", + gguf_expert_types=_uniform_expert_types(shim.model_path, num_layers), weight_block_size=None, attn_quant="gguf", dense_quant="gguf", @@ -240,9 +274,399 @@ def gguf_name_to_freetoken(name: str, num_layers: int) -> str | None: return f"model.layers.{layer}.{mapped}" +# -------------------------------------------------------------------------------------- +# Weight loading: GGUF tensor names -> FreeToken qwen35moe module params. +# -------------------------------------------------------------------------------------- + + +def _scan_quant_types(model_path: str) -> dict[tuple[int, str], int]: + """Scan GGUF tensor table once and return {(layer, suffix): ggml_type}. + + This allows us to detect which groups are mixed-quant without hardcoding. + Quant levels (IQ2_*, IQ3_XXS, IQ3_M) may mix differently. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + + quant_types = {} + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + _, idx, suffix = t.name.split(".", 2) + layer = int(idx) + quant_types[(layer, suffix)] = t.ggml_type + return quant_types + + +def _to_bf16(t) -> torch.Tensor: + """Dequantize a GgufTensor (F32/F16/Q*) to a dense bf16 tensor of its torch shape.""" + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.bfloat16) + return flat.reshape(t.shape) + + +def _to_f32(t) -> torch.Tensor: + """Dequantize a GgufTensor (F32/F16/Q*) to a dense float32 tensor of its torch shape.""" + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.float32) + return flat.reshape(t.shape) + + +def _require_tp1(what: str) -> None: + """GGUF quant layers / expert banks are not sharded; reject TP>1 with a clear + error instead of failing later on a confusing shape mismatch.""" + from freetoken.distributed import get_tp_info + + if get_tp_info().size > 1: + raise NotImplementedError( + f"qwen35moe GGUF {what} currently supports TP=1 only " + "(GGUF quant layers and expert banks are not tensor-parallel sharded)." + ) + + +def iter_gguf_weights( + model_path: str, + device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """Yield (param_name, tensor) for every non-expert qwen35moe param. + + Quantized projections (attention qkv/o, linear_attn in/out, shared-MLP gate_up/down) + stay in their native packed block layout and are yielded as ``.qweight`` (uint8) or + ``.qweight_`` for mixed-quant groups; norms and gates dequantize to bf16. q/k/v, + attn_qkv/gate/beta/alpha, and gate/up are fused by concatenating packed rows or + materializing parts separately (GGUFMergedLinear for mixed quants). Routed experts + are served from the offload cache (asserts the offload contract like the other MoE + models). + + A_log and dt_bias stay float32 (gdn.py keeps recurrence-gating params in fp32). + conv1d.weight stays float32 and is reshaped to [conv_dim, 1, kernel]. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + from freetoken.utils import cached_load_hf_config + + assert not include_moe_experts, ( + "qwen35moe GGUF stores experts as IQ3_S and only supports the offload backend; " + "experts are loaded into the offload cache via the expert-bank loader." + ) + assert include_non_moe + _require_tp1("weight loading") + + # Parse config to determine which layers are full-attention vs GDN. + config = parse_gguf_config(cached_load_hf_config(model_path)) + full_layer_ids = { + lid + for lid in range(config.num_layers) + if isinstance(config.attention_group_for_layer(lid), FullAttentionGroupConfig) + } + + # Get GDN group to extract attn_qkv_size and conv_kernel for conv1d reshape. + gdn_group = None + for group in config.attention_groups: + if isinstance(group, LinearGatedDeltaGroupConfig): + gdn_group = group + break + # attn_qkv_size = q + k + v = num_k_heads*state_size*2 + num_v_heads*state_size + gdn_attn_qkv_size = ( + 2 * gdn_group.num_key_heads * gdn_group.key_head_dim + gdn_group.num_value_heads * gdn_group.value_head_dim + if gdn_group + else 8192 + ) + gdn_conv_kernel = gdn_group.conv_kernel_dim if gdn_group else 4 + + # Scan quant types once to determine which fusion groups are mixed-quant. + quant_map = _scan_quant_types(model_path) + + # Per-layer fusion buffers: layer -> {slot: packed[out, row_bytes]}. + qkv_buf: dict[int, dict[str, torch.Tensor]] = {} # full-attn qkv + in_proj_buf: dict[int, dict[str, torch.Tensor]] = {} # GDN in_proj (qkv+gate+beta+alpha) + gate_up_buf: dict[int, dict[str, torch.Tensor]] = {} # shared_expert gate_up + + def layer_of(name: str) -> int: + return int(name.split(".")[1]) + + for t in iter_gguf_tensors(model_path): + name = t.name + layer = layer_of(name) if name.startswith("blk.") else None + + # Global tensors + if name == "token_embd.weight": + yield "model.embed_tokens.qweight", t.packed() # IQ3_S packed table + continue + if name == "output_norm.weight": + yield "model.norm.weight", _to_bf16(t) + continue + if not name.startswith("blk."): + continue + + # Skip block 40 (NextN/MTP, dropped) and nextn.* tensors. + if layer >= config.num_layers: + continue + if "nextn." in name: + continue + + # Skip routed-expert stacks (offload banks). + if any(name.endswith(sfx) for sfx in _EXPERT_SUFFIXES): + continue + + suffix = name.split(".", 2)[2] # after "blk.N." + base = f"model.layers.{layer}" + + # Scalar/norm tensors: dequant to bf16 or stay F32. + # norms, mlp.gate, shared_expert_gate -> bf16 + # A_log, dt_bias -> float32 + # conv1d.weight -> float32, reshaped + if suffix == "attn_norm.weight": + yield f"{base}.input_layernorm.weight", _to_bf16(t) + continue + if suffix == "post_attention_norm.weight": + yield f"{base}.post_attention_layernorm.weight", _to_bf16(t) + continue + if suffix == "ffn_gate_inp.weight": + yield f"{base}.mlp.gate.weight", _to_bf16(t) + continue + if suffix == "ffn_gate_inp_shexp.weight": + yield f"{base}.mlp.shared_expert_gate.weight", _to_bf16(t) + continue + if suffix == "ssm_norm.weight": + yield f"{base}.linear_attn.norm.weight", _to_bf16(t) + continue + if suffix == "ssm_a": + yield f"{base}.linear_attn.A_log", _to_f32(t) + continue + if suffix == "ssm_dt.bias": + yield f"{base}.linear_attn.dt_bias", _to_f32(t) + continue + if suffix == "ssm_conv1d.weight": + # F32, reshape to [conv_dim, 1, kernel] where conv_dim is attn_qkv output size + w = _to_f32(t) + w = w.reshape(gdn_attn_qkv_size, 1, gdn_conv_kernel) + yield f"{base}.linear_attn.conv1d.weight", w + continue + if suffix == "attn_q_norm.weight": + yield f"{base}.self_attn.q_norm.weight", _to_bf16(t) + continue + if suffix == "attn_k_norm.weight": + yield f"{base}.self_attn.k_norm.weight", _to_bf16(t) + continue + + # Quantized projections: keep packed; fuse per layer. + # Full-attention: qkv from q, k, v + if layer in full_layer_ids: + if suffix == "attn_q.weight": + qkv_buf.setdefault(layer, {})["q"] = t.packed() + elif suffix == "attn_k.weight": + qkv_buf.setdefault(layer, {})["k"] = t.packed() + elif suffix == "attn_v.weight": + qkv_buf.setdefault(layer, {})["v"] = t.packed() + elif suffix == "attn_output.weight": + yield f"{base}.self_attn.o_proj.qweight", t.packed() + else: + continue # unmapped for full-attn layers + + # Emit fused qkv once all three parts are present. + slots = qkv_buf.get(layer) + if slots is not None and "q" in slots and "k" in slots and "v" in slots: + # Determine if this is a mixed-quant group. + types = [ + quant_map.get((layer, "attn_q.weight")), + quant_map.get((layer, "attn_k.weight")), + quant_map.get((layer, "attn_v.weight")), + ] + if len(set(types)) == 1: + # Uniform quant: fuse via torch.cat along dim 0. + yield f"{base}.self_attn.qkv_proj.qweight", torch.cat( + [slots["q"], slots["k"], slots["v"]], dim=0 + ) + else: + # Mixed quant: emit GGUFMergedLinear format. + yield f"{base}.self_attn.qkv_proj.qweight_0", slots["q"] + yield f"{base}.self_attn.qkv_proj.qweight_1", slots["k"] + yield f"{base}.self_attn.qkv_proj.qweight_2", slots["v"] + del qkv_buf[layer] + + # GDN layers: in_proj from attn_qkv, attn_gate, ssm_beta, ssm_alpha + # and out_proj + else: + if suffix == "attn_qkv.weight": + in_proj_buf.setdefault(layer, {})["qkv"] = t.packed() + elif suffix == "attn_gate.weight": + in_proj_buf.setdefault(layer, {})["gate"] = t.packed() + elif suffix == "ssm_beta.weight": + in_proj_buf.setdefault(layer, {})["beta"] = t.packed() + elif suffix == "ssm_alpha.weight": + in_proj_buf.setdefault(layer, {})["alpha"] = t.packed() + elif suffix == "ssm_out.weight": + yield f"{base}.linear_attn.out_proj.qweight", t.packed() + else: + continue # unmapped for GDN layers + + # Emit fused in_proj once all four parts are present. + slots = in_proj_buf.get(layer) + if ( + slots is not None + and "qkv" in slots + and "gate" in slots + and "beta" in slots + and "alpha" in slots + ): + # Determine if this is a mixed-quant group. + types = [ + quant_map.get((layer, "attn_qkv.weight")), + quant_map.get((layer, "attn_gate.weight")), + quant_map.get((layer, "ssm_beta.weight")), + quant_map.get((layer, "ssm_alpha.weight")), + ] + if len(set(types)) == 1: + # Uniform quant: fuse via torch.cat along dim 0. + yield f"{base}.linear_attn.in_proj.qweight", torch.cat( + [ + slots["qkv"], + slots["gate"], + slots["beta"], + slots["alpha"], + ], + dim=0, + ) + else: + # Mixed quant: emit GGUFMergedLinear format. + yield f"{base}.linear_attn.in_proj.qweight_0", slots["qkv"] + yield f"{base}.linear_attn.in_proj.qweight_1", slots["gate"] + yield f"{base}.linear_attn.in_proj.qweight_2", slots["beta"] + yield f"{base}.linear_attn.in_proj.qweight_3", slots["alpha"] + del in_proj_buf[layer] + + # Shared expert gate_up: every layer gets this + if suffix == "ffn_gate_shexp.weight": + gate_up_buf.setdefault(layer, {})["gate"] = t.packed() + elif suffix == "ffn_up_shexp.weight": + gate_up_buf.setdefault(layer, {})["up"] = t.packed() + elif suffix == "ffn_down_shexp.weight": + yield f"{base}.mlp.shared_expert.down_proj.qweight", t.packed() + + # Emit fused gate_up once both parts are present. + gu = gate_up_buf.get(layer) + if gu is not None and "gate" in gu and "up" in gu: + # Determine if this is a mixed-quant group (unlikely but check). + types = [ + quant_map.get((layer, "ffn_gate_shexp.weight")), + quant_map.get((layer, "ffn_up_shexp.weight")), + ] + if len(set(types)) == 1: + # Uniform quant: fuse via torch.cat along dim 0. + yield f"{base}.mlp.shared_expert.gate_up_proj.qweight", torch.cat( + [gu["gate"], gu["up"]], dim=0 + ) + else: + # Mixed quant: emit GGUFMergedLinear format. + yield f"{base}.mlp.shared_expert.gate_up_proj.qweight_0", gu["gate"] + yield f"{base}.mlp.shared_expert.gate_up_proj.qweight_1", gu["up"] + del gate_up_buf[layer] + + # Verify no fusion buffers are incomplete. + assert not qkv_buf, f"incomplete full-attn qkv groups: {sorted(qkv_buf)}" + assert not in_proj_buf, f"incomplete GDN in_proj groups: {sorted(in_proj_buf)}" + assert not gate_up_buf, f"incomplete shared_expert gate_up groups: {sorted(gate_up_buf)}" + + +def convert_qwen35_to_gguf(model, config: ModelConfig, *, model_path: str) -> None: + """In place: replace qwen35moe's dense projections + embedding with native GGUF ops. + + Quantized in the checkpoint -> swapped: attention qkv/o (mixed-quant), linear_attn + in/out, shared-MLP gate_up/down, and the token embedding (IQ3_S, also the lm_head if + tied). Left as dense bf16 (F32 in the GGUF): all RMSNorms, the router gate + (ffn_gate_inp), the per-layer shared_expert_gate, conv1d.weight, A_log, dt_bias, + and the routed experts (served from the offload cache). + + The per-layer quant types are read from the GGUF file, not hardcoded, to support + different quant levels (IQ3_M, IQ3_XXS, IQ2_*, etc.) which may mix differently. + """ + from freetoken.layers.gguf import GGUFEmbedding, GGUFLinear, gguf_merged_or_plain + + # Scan quant types to drive layer swaps. + quant_map = _scan_quant_types(model_path) + + # Determine full-attention vs GDN layers. + full_layer_ids = { + lid + for lid in range(config.num_layers) + if isinstance(config.attention_group_for_layer(lid), FullAttentionGroupConfig) + } + + def swap_linear(owner, attr, quant_type=GGML_Q4_K): + """Replace a dense Linear with GGUFLinear.""" + lin = getattr(owner, attr) + out_features, in_features = lin.weight.shape + setattr( + owner, + attr, + GGUFLinear(in_features, out_features, quant_type, has_bias=lin.bias is not None), + ) + + inner = model.model + embed = GGUFEmbedding( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + quant_type=GGML_IQ3_S, + ) + inner.embed_tokens = embed + + for layer_idx, layer in enumerate(inner.layers.op_list): + if layer_idx in full_layer_ids: + # Full-attention layer: qkv_proj (mixed quant), o_proj (Q4_K) + types = [ + quant_map.get((layer_idx, "attn_q.weight"), GGML_IQ3_S), + quant_map.get((layer_idx, "attn_k.weight"), GGML_IQ3_S), + quant_map.get((layer_idx, "attn_v.weight"), GGML_Q4_K), + ] + output_sizes = [8192, 512, 512] + layer.self_attn.qkv_proj = gguf_merged_or_plain( + config.hidden_size, output_sizes, types, has_bias=False + ) + swap_linear(layer.self_attn, "o_proj", GGML_Q4_K) + else: + # GDN layer: in_proj (mixed quant), out_proj (IQ3_S) + types = [ + quant_map.get((layer_idx, "attn_qkv.weight"), GGML_Q4_K), + quant_map.get((layer_idx, "attn_gate.weight"), GGML_IQ3_S), + quant_map.get((layer_idx, "ssm_beta.weight"), GGML_IQ3_S), + quant_map.get((layer_idx, "ssm_alpha.weight"), GGML_IQ3_S), + ] + output_sizes = [8192, 4096, 32, 32] + layer.linear_attn.in_proj = gguf_merged_or_plain( + config.hidden_size, output_sizes, types, has_bias=False + ) + swap_linear(layer.linear_attn, "out_proj", GGML_IQ3_S) + + # Shared expert: gate_up_proj (IQ3_S, uniform), down_proj (Q4_K or IQ3_S). + gate_up_type = quant_map.get((layer_idx, "ffn_gate_shexp.weight"), GGML_IQ3_S) + down_type = quant_map.get((layer_idx, "ffn_down_shexp.weight"), GGML_Q4_K) + + # gate_up is typically uniform IQ3_S, but use gguf_merged_or_plain for safety. + up_type = quant_map.get((layer_idx, "ffn_up_shexp.weight"), GGML_IQ3_S) + layer.mlp.shared_expert.gate_up_proj = gguf_merged_or_plain( + config.hidden_size, + [config.shared_expert_intermediate_size, config.shared_expert_intermediate_size], + [gate_up_type, up_type], + has_bias=False, + ) + swap_linear(layer.mlp.shared_expert, "down_proj", down_type) + + # lm_head: use output.weight quant type (Q6_K in Ornith). + if config.tie_word_embeddings: + # Tied head: reference the embedding's qweight. + from freetoken.models.gemma4.gguf import GGUFTiedLMHead + + model.lm_head = GGUFTiedLMHead(embed, GGML_Q6_K) + else: + # Untied head: create a separate GGUFLinear (rare for qwen35moe). + swap_linear(model, "lm_head", GGML_Q6_K) + + __all__ = [ "parse_gguf_config", "gguf_name_to_freetoken", + "iter_gguf_weights", + "convert_qwen35_to_gguf", "_FUSE", "_EXPERT_SUFFIXES", ] diff --git a/python/freetoken/models/qwen3_5_moe/gguf_experts.py b/python/freetoken/models/qwen3_5_moe/gguf_experts.py new file mode 100644 index 00000000..4d2bedad --- /dev/null +++ b/python/freetoken/models/qwen3_5_moe/gguf_experts.py @@ -0,0 +1,280 @@ +"""Routed-expert host bank sources for the qwen35moe GGUF checkpoint. + +This module loads the per-expert weight tensors that are stored as GGUF stacks +and allocates them into host banks for the offload cache. The layout is a 3D +expert stack: [num_experts, out_features, in_features] in torch order. + +CRITICAL CORRECTNESS NOTE: The MoE kernel (kernel/csrc/gguf/moe_vec.cuh) +computes addressing as `blocks_per_row = ncols / qk` and +`x = vx + expert * nrows * blocks_per_row`, i.e. it assumes a FULLY PACKED +contiguous [E, nrows, blocks_per_row] layout with NO padding. So the bank +tensors must be exactly `row_bytes` wide for their own quant type — never pad +a smaller-type layer up to a larger type's stride, because the kernel would +then read every block at the wrong offset and return plausible-looking garbage. + +For qwen35moe specifically: +- ``ffn_gate_exps`` and ``ffn_up_exps`` are always IQ3_S (layers 0-39) +- ``ffn_down_exps`` is Q4_K for layers 0-4, IQ3_S for layers 5-39 + +The gate_up bank per layer is the per-expert concatenation of gate rows +then up rows along the output dimension, giving [E, 2*I, row_bytes(H, t)], +valid because gate and up share a quant type and therefore a row stride. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from freetoken.models.gguf.dequant import GGML_IQ3_S, GGML_Q4_K, GGML_NAME, row_bytes + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + + +def gguf_expert_types(model_path: str, num_layers: int) -> dict[str, list[int]]: + """Scan the GGUF tensor table and return per-layer expert quant types. + + Returns a dict with two keys: + - ``"gate_up"``: list of ``num_layers`` ggml_type enums for ``ffn_gate_exps``. + gate and up for each layer must have the same type (they are row-concatenated). + If they differ for any layer, raises a clear ValueError naming the layer and both types. + - ``"down"``: list of ``num_layers`` ggml_type enums for ``ffn_down_exps``. + + For qwen35moe: gate_up is always IQ3_S (uniformly), and down varies by layer + (Q4_K for 0-4, IQ3_S for 5-39). + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + + gate_types: list[int | None] = [None] * num_layers + up_types: list[int | None] = [None] * num_layers + down_types: list[int | None] = [None] * num_layers + + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if layer >= num_layers: + continue # skip the trailing NextN/MTP block + + if t.name.endswith("ffn_gate_exps.weight"): + gate_types[layer] = t.ggml_type + elif t.name.endswith("ffn_up_exps.weight"): + up_types[layer] = t.ggml_type + elif t.name.endswith("ffn_down_exps.weight"): + down_types[layer] = t.ggml_type + + # Validate that gate and up types agree for each layer (they must be row-concatenated). + gate_up_types: list[int] = [] + for layer in range(num_layers): + gate_t = gate_types[layer] + up_t = up_types[layer] + if gate_t is None or up_t is None: + raise ValueError( + f"missing expert tensors for layer {layer}: " + f"gate={GGML_NAME.get(gate_t, gate_t)}, up={GGML_NAME.get(up_t, up_t)}" + ) + if gate_t != up_t: + raise ValueError( + f"layer {layer}: ffn_gate_exps type {GGML_NAME.get(gate_t, gate_t)} != " + f"ffn_up_exps type {GGML_NAME.get(up_t, up_t)}; " + "cannot row-concatenate tensors with different quant types" + ) + gate_up_types.append(gate_t) + + # Validate down tensors are present. + for layer in range(num_layers): + if down_types[layer] is None: + raise ValueError(f"missing ffn_down_exps for layer {layer}") + + return { + "gate_up": gate_up_types, + "down": down_types, + } + + +def gguf_expert_specs( + config: ModelConfig, types: dict[str, list[int]] +) -> dict[str, list[tuple[tuple[int, ...], torch.dtype]]]: + """Per-layer expert bank shapes, accounting for per-layer dtype variation. + + The routed experts are stored as 3D stacks [E, out, in] in torch order. + Returns a list of (shape, dtype) tuples per expert bank: + - ``gate_up[layer]``: ``((E, 2*I, row_bytes(H, types["gate_up"][layer])), torch.uint8)`` + - ``down[layer]``: ``((E, H, row_bytes(I, types["down"][layer])), torch.uint8)`` + + where ``E=num_experts``, ``H=hidden_size``, ``I=moe_intermediate_size``. + + The row_bytes dimension varies by layer because the quant type varies, + and the MoE kernel needs to read the exact byte width for its quant type. + """ + E = config.num_experts + H = config.hidden_size + I = config.moe_intermediate_size + num_layers = config.num_layers + + gate_up_specs = [] + down_specs = [] + + for layer in range(num_layers): + gate_up_type = types["gate_up"][layer] + down_type = types["down"][layer] + + gate_up_row_bytes = row_bytes(H, gate_up_type) + down_row_bytes = row_bytes(I, down_type) + + gate_up_specs.append(((E, 2 * I, gate_up_row_bytes), torch.uint8)) + down_specs.append(((E, H, down_row_bytes), torch.uint8)) + + return { + "gate_up": gate_up_specs, + "down": down_specs, + } + + +def load_gguf_expert_sources( + model_path: str, config: ModelConfig, *, layer_sink=None +) -> dict[str, list[torch.Tensor]]: + """Per-layer host banks of the routed experts' native packed block bytes. + + Loads the three GGUF expert stacks (gate, up, down) into per-layer host banks + for the offload cache. The gate_up bank for each layer is the per-expert + concatenation of that expert's gate rows then its up rows along the output + dimension, giving [E, 2*I, row_bytes(H, t)] -- valid because gate and up + share a quant type and therefore a row stride. + + Returns a dict with two keys: + - ``"gate_up"``: list of ``num_layers`` tensors, each ``[E, 2*I, row_bytes_gate_up]`` uint8 + - ``"down"``: list of ``num_layers`` tensors, each ``[E, H, row_bytes_down]`` uint8 + + Parameters: + - ``layer_sink``: If None (serving mode), pins each completed layer via an + internally-owned PinPipeline. If given (converter mode), fires the completion + tracker into it instead -- nothing is pinned, and the sink may release banks, + so returned tensors are only valid until the sink releases them. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks + + types = gguf_expert_types(model_path, config.num_layers) + specs = gguf_expert_specs(config, types) + + L = config.num_layers + E = config.num_experts + H = config.hidden_size + I = config.moe_intermediate_size + + # Allocate the per-layer banks (lazy mmap, unpinned). + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in hb} + + # Per-layer buffers to accumulate gate and up before concatenating. + gate_buf: dict[int, torch.Tensor] = {} + up_buf: dict[int, torch.Tensor] = {} + seen_gate = set() + seen_up = set() + seen_down = set() + + def _load(sink) -> None: + # Track completion: 2 banks per layer (gate_up and down). + tracker = LayerCompletionTracker(2, hb, sink) if sink is not None else None + + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if layer >= L: + continue # skip the trailing NextN/MTP block + + if t.name.endswith("ffn_gate_exps.weight"): + # Shape from GGUF: [E, I, H] in torch order = [H, I, E] in ggml order + # t.packed() is [H*I, row_bytes(E, type)] + gate_buf[layer] = t.packed() + seen_gate.add(layer) + + elif t.name.endswith("ffn_up_exps.weight"): + # Shape from GGUF: [E, I, H] in torch order = [H, I, E] in ggml order + # t.packed() is [H*I, row_bytes(E, type)] + up_buf[layer] = t.packed() + seen_up.add(layer) + + elif t.name.endswith("ffn_down_exps.weight"): + # Shape from GGUF: [E, H, I] in torch order = [I, H, E] in ggml order + # t.packed() is [I*H, row_bytes(E, type)] + # Reshape to [E, H, row_bytes(I, type)] + down_row_bytes = specs["down"][layer][0][2] + banks["down"][layer].copy_(t.packed().reshape(E, H, down_row_bytes)) + seen_down.add(layer) + if tracker is not None: + tracker.note(layer) + + else: + continue + + # Emit gate_up bank once both gate and up are present. + if layer in gate_buf and layer in up_buf: + gate_up_row_bytes = specs["gate_up"][layer][0][2] + # Concatenate gate [H*I, row_bytes(E, type)] and up [H*I, row_bytes(E, type)] + # to get [H*2*I, row_bytes(E, type)], then reshape to [E, 2*I, row_bytes(H, type)] + combined = torch.cat([gate_buf[layer], up_buf[layer]], dim=0) + banks["gate_up"][layer].copy_(combined.reshape(E, 2 * I, gate_up_row_bytes)) + del gate_buf[layer], up_buf[layer] + if tracker is not None: + tracker.note(layer) + + # Load with or without pinning. + if layer_sink is not None: + _load(layer_sink) + elif torch.cuda.is_available(): + with PinPipeline() as pins: + _load(pins) + else: + _load(None) # CUDA-less: mmap banks stay pageable, never pinned + + # Verify all layers were loaded. + want = set(range(L)) + missing_gate = want - seen_gate + missing_up = want - seen_up + missing_down = want - seen_down + if missing_gate or missing_up or missing_down: + raise ValueError( + f"missing expert layers: gate {sorted(missing_gate)}, " + f"up {sorted(missing_up)}, down {sorted(missing_down)}" + ) + + return banks + + +def dummy_gguf_expert_sources(config: ModelConfig) -> dict[str, list[torch.Tensor]]: + """Random expert banks shaped like ``load_gguf_expert_sources`` output.""" + from freetoken.moe.host_banks import alloc_layer_banks, pin_banks + + # Use uniform IQ3_S for all layers (a simplification for the dummy). + num_layers = config.num_layers + gate_up_types = [GGML_IQ3_S] * num_layers + down_types = [GGML_IQ3_S] * num_layers + types = {"gate_up": gate_up_types, "down": down_types} + + specs = gguf_expert_specs(config, types) + L = config.num_layers + + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in hb} + + # Fill with random uint8. + for t in banks["gate_up"] + banks["down"]: + t.random_(0, 256) + + if torch.cuda.is_available(): + pin_banks(hb) # match the other dummies: pin-after-fill + + return banks + + +__all__ = [ + "gguf_expert_types", + "gguf_expert_specs", + "load_gguf_expert_sources", + "dummy_gguf_expert_sources", +] diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 0c033ca0..03ab68e6 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -107,6 +107,13 @@ class ModelSpec: parse_config="parse_gguf_config", iter_weights="iter_gguf_weights", ), + # GGUF (mixed IQ3_S/Q4_K) qwen35moe: same model classes, GGUF config + weight loaders. + "Qwen35MoeGGUFForCausalLM": ModelSpec( + "freetoken.models.qwen3_5_moe", + "Qwen3_5MoEForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + ), "GptOssForCausalLM": ModelSpec( "freetoken.models.gpt_oss", "GptOssForCausalLM", diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 8b6116ba..5b3138dc 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -49,6 +49,10 @@ class ExpertBanks: # streamed straight to its sink instead of staying materialized here) -- set by # convert.py's per-format streaming gate; ``sources`` may hold released tensors. streamed: bool = False + # For quant_format == "gguf": the (gate_up, down) ggml types the checkpoint used. + # Carried here so the engine can hand them to OffloadMoeCache, which hands them to + # the MoE kernels -- a GGUF bank's row stride is a property of the file, not the format. + gguf_expert_types: tuple[int, int] | None = field(default=None) _PARALLEL_CHUNK = 8 << 20 # default O_DIRECT chunk for the parallel reader @@ -278,10 +282,71 @@ def _dsfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, ) +def _gguf_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: + """Native GGUF routed experts of any MMVQ-capable ggml type (generalizes ``q4_0``). + + Like the q4_0 provider, the packed block bytes are streamed to the GPU and dequantized + inside the borrowed ggml MoE kernels -- no bf16 expert copy is ever materialized. The + difference is that the row stride depends on the ggml type the checkpoint picked per + bank, so the loader also reports ``(gate_up, down)`` types for the cache and kernels. + + The per-arch loader is resolved from the model package (the GGUF tensor layout is + architecture-specific), mirroring ``_model_setup_override``'s resolution. + """ + if parallel: + raise NotImplementedError( + "parallel reader not implemented for gguf: the checkpoint is a single packed " + "file (not safetensors), so the common reader does not apply" + ) + from freetoken.models.register import _load_attr, get_model_spec + + spec = get_model_spec(model_config.architectures[0]) + loader = _load_attr(spec.module, "load_gguf_expert_sources") + types_fn = _load_attr(spec.module, "gguf_expert_types") + + sink = None if dummy else layer_sink + types = types_fn(model_path, model_config.num_layers) + # One pool per bank is shared by every layer, and moe_vec.cuh addresses it without a + # padding allowance, so a bank whose type varies by layer cannot be served. Reject it + # here with the layers named rather than reading every block at the wrong offset. + resolved = {} + for name in ("gate_up", "down"): + distinct = sorted(set(types[name])) + if len(distinct) != 1: + from freetoken.models.gguf.dequant import GGML_NAME + + spread = { + GGML_NAME.get(t, t): [i for i, x in enumerate(types[name]) if x == t] + for t in distinct + } + raise NotImplementedError( + f"GGUF expert bank {name!r} mixes ggml types across layers ({spread}); the " + f"offload slot pool is one allocation with a single row stride, so this " + f"checkpoint cannot be served. Use a quant level whose expert banks are " + f"uniform (for Ornith-1.5-35B: IQ3_S or IQ3_XXS are; IQ3_M/IQ2_M/IQ2_XXS/" + f"IQ1_S split ffn_down_exps across two types)." + ) + resolved[name] = distinct[0] + + sources = loader(model_path, model_config, layer_sink=sink) + return ExpertBanks( + "gguf", + {name: sources[name] for name in _BANK_SCHEMAS["gguf"]}, + streamed=sink is not None, + gguf_expert_types=(resolved["gate_up"], resolved["down"]), + ) + + def _model_setup_override(model_config): architectures = getattr(model_config, "architectures", None) if not architectures: return None + # A model package's setup_offload_expert_banks is written for that model's *safetensors* + # quantization (qwen3_5_moe's is the NVFP4 loader). A GGUF checkpoint of the same + # architecture shares the module but not the bank layout, so the override would silently + # hijack the gguf path and fail deep in an NVFP4 reader. Route "gguf" to _gguf_banks. + if getattr(model_config, "expert_quant", "none") == "gguf": + return None from freetoken.models.register import _load_attr, get_model_spec @@ -301,6 +366,7 @@ def _model_setup_override(model_config): "nvfp4": _nvfp4_banks, "ds_fp4": _dsfp4_banks, "q4_0": _q4_0_banks, + "gguf": _gguf_banks, } @@ -395,12 +461,25 @@ def bank_bytes_estimate(model_config) -> int | None: fmt = expert_quant if expert_quant != "none" else ( getattr(model_config, "moe_weight_format", None) or "bf16" ) - per_expert = _BANK_BYTES_PER_EXPERT.get(fmt) layers = getattr(model_config, "num_moe_layers", None) experts = getattr(model_config, "num_experts", None) hidden = getattr(model_config, "hidden_size", None) inter = getattr(model_config, "moe_intermediate_size", None) - if per_expert is None or not all((layers, experts, hidden, inter)): + if not all((layers, experts, hidden, inter)): + return None + if fmt == "gguf": + # Not a fixed f(H, I) like the other formats: the row stride depends on the ggml + # type the checkpoint chose per bank, so size it from row_bytes directly. + types = getattr(model_config, "gguf_expert_types", None) + if types is None: + return None + from freetoken.models.gguf.dequant import row_bytes + + t_gate_up, t_down = types + per = 2 * inter * row_bytes(hidden, t_gate_up) + hidden * row_bytes(inter, t_down) + return layers * experts * per + per_expert = _BANK_BYTES_PER_EXPERT.get(fmt) + if per_expert is None: return None return layers * experts * per_expert(hidden, inter) diff --git a/python/freetoken/moe/fused_q4_0.py b/python/freetoken/moe/fused_q4_0.py index f0c7acde..d2041701 100644 --- a/python/freetoken/moe/fused_q4_0.py +++ b/python/freetoken/moe/fused_q4_0.py @@ -33,6 +33,7 @@ def fused_experts_gguf( topk_ids: torch.Tensor, activation: str, quant_type: int, + down_quant_type: int | None = None, ) -> torch.Tensor: """Fused GGUF MoE expert compute over any MMVQ-supported quantization type. @@ -40,15 +41,24 @@ def fused_experts_gguf( dequantization happens inside the ``ggml_moe_a8_vec`` CUDA kernel. ``quant_type`` must be in ``MOE_VEC_TYPES``, which mirrors the supported types in ``ggml_moe_a8_vec`` (gguf_kernel.cu:559). + + ``quant_type`` is the gate_up bank's type; ``down_quant_type`` defaults to it. They may + differ because gate_up and down are separate banks with separate slot pools, and + llama.cpp routinely quantizes the down projection differently from gate/up. What may + NOT differ is the type *within* one bank across layers -- that pool is one allocation. """ from freetoken.kernel.gguf import ggml_moe_a8_vec - if quant_type not in MOE_VEC_TYPES: - from freetoken.models.gguf.dequant import GGML_NAME - raise NotImplementedError( - f"fused GGUF MoE kernel does not support quant type {GGML_NAME.get(quant_type, quant_type)} " - f"(only {sorted(MOE_VEC_TYPES)} supported)" - ) + if down_quant_type is None: + down_quant_type = quant_type + for label, qt in (("gate_up", quant_type), ("down", down_quant_type)): + if qt not in MOE_VEC_TYPES: + from freetoken.models.gguf.dequant import GGML_NAME + raise NotImplementedError( + f"fused GGUF MoE kernel does not support quant type " + f"{GGML_NAME.get(qt, qt)} for the {label} bank " + f"(only {sorted(MOE_VEC_TYPES)} supported)" + ) act_fn = _ACT.get(activation) if act_fn is None: @@ -64,7 +74,7 @@ def fused_experts_gguf( gate_up = ggml_moe_a8_vec(hidden_states, gate_up_q, topk_ids, top_k, qt, n2, num_tokens) inter = act_fn(gate_up) # down: each of the num_tokens*top_k intermediate rows uses its own expert id. - out = ggml_moe_a8_vec(inter, down_q, topk_ids, 1, qt, h, num_tokens * top_k) + out = ggml_moe_a8_vec(inter, down_q, topk_ids, 1, int(down_quant_type), h, num_tokens * top_k) out = out.reshape(num_tokens, top_k, h) * topk_weights.reshape(num_tokens, top_k, 1).to( out.dtype ) diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 6ee76406..f936f579 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -45,6 +45,12 @@ # 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"), + # native GGUF experts of any MMVQ-capable ggml type (the generalization of "q4_0"): + # gate_up [L*E, 2I, row_bytes(H, t_gate_up)], down [L*E, H, row_bytes(I, t_down)], + # with the two types carried on ModelConfig.gguf_expert_types because they are a + # property of the checkpoint, not of the format. Same two banks and the same + # dequant-in-kernel grouped GEMV as q4_0 -- only the row stride is type-dependent. + "gguf": ("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) @@ -114,6 +120,12 @@ class OffloadMoeCache: # prefill). The format names its bank layout (_BANK_SCHEMAS) and which kernels # may read the banks; the cache machinery itself is layout-agnostic. quant_format: str = "bf16" + # For quant_format == "gguf": the (gate_up, down) ggml types of the two banks. Unlike + # every other format, "gguf" does not imply a row stride -- the checkpoint picks a ggml + # type per tensor -- so the MoE kernels need the types handed to them at dispatch. + # Per-bank (each bank owns its own slot pool) but NOT per-layer: one pool is shared by + # all layers and moe_vec.cuh addresses it with no padding allowance. + gguf_expert_types: tuple[int, int] | None = None # Decode mode + bank layout; per-layer CPU routing is cpu_layer_ids. "gpu": # GPU-tiled banks, all decode on GPU (stream misses over PCIe into the slot # cache, GEMM on GPU). "cpu": native (CPU-readable) banks + a CPU executor; diff --git a/tests/models/test_qwen35moe_gguf.py b/tests/models/test_qwen35moe_gguf.py new file mode 100644 index 00000000..94b3ec90 --- /dev/null +++ b/tests/models/test_qwen35moe_gguf.py @@ -0,0 +1,315 @@ +"""Unit tests for qwen35moe GGUF weight loading. + +Tests GGUFMergedLinear (mixed-quant fused projections), gguf_merged_or_plain dispatch, +the GDN in_proj geometry, and the qwen35moe name mapping (dropping MTP block). All tests +run without CUDA or kernel compilation, using monkeypatched kernel mocks that record calls +and return correctly-shaped CPU tensors. + +Reference: llama.cpp's qwen3.5 GGUF mapping (gguf-py/gguf/tensor_mapping.py) and the +Ornith-1.5-35B-A3B-GGUF checkpoint geometry (ORNITH_SPEC.md). +""" + +from __future__ import annotations + +import sys +from types import ModuleType + +import pytest +import torch + +from freetoken.layers.gguf import GGUFLinear, GGUFMergedLinear, gguf_merged_or_plain +from freetoken.models.gguf.dequant import ( + GGML_IQ3_S, + GGML_Q4_K, + BLOCK_SHAPE, + row_bytes, +) +from freetoken.models.qwen3_5_moe.gguf import gguf_name_to_freetoken + + +@pytest.fixture +def mock_kernel_module(monkeypatch): + """Replace freetoken.kernel.gguf with a mock that tracks kernel calls. + + Returns a dict tracking which kernels were called and with what arguments. + This allows tests to run without CUDA or the compiled kernel extension. + """ + call_log = { + "ggml_mul_mat_vec_a8": [], # List of calls (can be multiple) + "ggml_mul_mat_a8": [], + "ggml_dequantize": [], + } + + def make_mmvq_kernel(call_log): + """Mock ggml_mul_mat_vec_a8: GEMV kernel for small batch.""" + def kernel(qweight, x, qweight_type, out_features): + call_log["ggml_mul_mat_vec_a8"].append({ + "qweight_shape": qweight.shape, + "x_shape": x.shape, + "qweight_type": qweight_type, + "out_features": out_features, + }) + batch_size = x.shape[0] + return torch.randn(batch_size, out_features, dtype=x.dtype) + return kernel + + def make_mmq_kernel(call_log): + """Mock ggml_mul_mat_a8: MMQ kernel for large batch.""" + def kernel(qweight, x, qweight_type, out_features): + call_log["ggml_mul_mat_a8"].append({ + "qweight_shape": qweight.shape, + "x_shape": x.shape, + "qweight_type": qweight_type, + "out_features": out_features, + }) + batch_size = x.shape[0] + return torch.randn(batch_size, out_features, dtype=x.dtype) + return kernel + + def make_dequant_kernel(call_log): + """Mock ggml_dequantize: materializes weight into BF16.""" + def kernel(qweight, qweight_type, out_features, in_features, out_dtype): + call_log["ggml_dequantize"].append({ + "qweight_shape": qweight.shape, + "qweight_type": qweight_type, + "out_features": out_features, + "in_features": in_features, + "out_dtype": out_dtype, + }) + return torch.randn(out_features, in_features, dtype=out_dtype) + return kernel + + mock_module = ModuleType("freetoken.kernel.gguf") + mock_module.ggml_mul_mat_vec_a8 = make_mmvq_kernel(call_log) + mock_module.ggml_mul_mat_a8 = make_mmq_kernel(call_log) + mock_module.ggml_dequantize = make_dequant_kernel(call_log) + + monkeypatch.setitem(sys.modules, "freetoken.kernel.gguf", mock_module) + + yield call_log + + # Cleanup: ensure mock is removed so it doesn't interfere with other tests + if "freetoken.kernel.gguf" in sys.modules: + del sys.modules["freetoken.kernel.gguf"] + + +class TestMergedLinearConcatenatesOutputs: + """Test GGUFMergedLinear with mixed-quant output parts.""" + + def test_merged_linear_concatenates_outputs(self, mock_kernel_module): + """GGUFMergedLinear with [8192, 512, 512] outputs and mixed quant types + produces output of width 9216 and calls kernel once per part. + + Ornith's full-attention qkv_proj: q(IQ3_S, 8192) + k(IQ3_S, 512) + v(Q4_K, 512). + """ + in_features = 2048 + output_sizes = [8192, 512, 512] + quant_types = [GGML_IQ3_S, GGML_IQ3_S, GGML_Q4_K] + + merged = GGUFMergedLinear(in_features, output_sizes, quant_types, has_bias=False) + + # Check total output width + assert merged.out_features == sum(output_sizes) == 9216 + + # Verify each part's packed buffer has the correct row_bytes for its type + assert merged.qweight_0.shape == (8192, row_bytes(in_features, GGML_IQ3_S)) + assert merged.qweight_0.shape == (8192, 880) # IQ3_S: 2048 // 256 * 110 + + assert merged.qweight_1.shape == (512, row_bytes(in_features, GGML_IQ3_S)) + assert merged.qweight_1.shape == (512, 880) + + assert merged.qweight_2.shape == (512, row_bytes(in_features, GGML_Q4_K)) + assert merged.qweight_2.shape == (512, 1152) # Q4_K: 2048 // 256 * 144 + + # Forward pass: batch size 1 (small batch, uses MMVQ) + x = torch.randn(1, in_features, dtype=torch.bfloat16) + output = merged(x) + + # Output shape should be [1, 9216] + assert output.shape == (1, 9216) + + # Kernel should have been called 3 times (once per part), all to MMVQ + assert len(mock_kernel_module["ggml_mul_mat_vec_a8"]) == 3 + assert len(mock_kernel_module["ggml_mul_mat_a8"]) == 0 + + # Verify each call received the correct qweight and out_features + calls = mock_kernel_module["ggml_mul_mat_vec_a8"] + assert calls[0]["out_features"] == 8192 + assert calls[1]["out_features"] == 512 + assert calls[2]["out_features"] == 512 + + +class TestMergedLinearRejectsLengthMismatch: + """Test GGUFMergedLinear validation of output_sizes and quant_types lengths.""" + + def test_merged_linear_rejects_length_mismatch(self): + """GGUFMergedLinear raises ValueError when output_sizes and quant_types differ.""" + in_features = 2048 + output_sizes = [8192, 512, 512] + quant_types = [GGML_IQ3_S, GGML_Q4_K] # Only 2 types, 3 sizes + + with pytest.raises(ValueError) as excinfo: + GGUFMergedLinear(in_features, output_sizes, quant_types, has_bias=False) + + assert "length" in str(excinfo.value).lower() + + def test_merged_linear_rejects_zero_output(self): + """GGUFMergedLinear raises ValueError if any output_size is <= 0.""" + in_features = 2048 + output_sizes = [8192, 0, 512] # Zero output size + quant_types = [GGML_IQ3_S, GGML_IQ3_S, GGML_Q4_K] + + with pytest.raises(ValueError) as excinfo: + GGUFMergedLinear(in_features, output_sizes, quant_types, has_bias=False) + + assert "must be > 0" in str(excinfo.value) + + +class TestGGUFMergedOrPlainDispatch: + """Test gguf_merged_or_plain routing between GGUFLinear and GGUFMergedLinear.""" + + def test_gguf_merged_or_plain_picks_plain_when_uniform(self): + """When all quant types are identical, gguf_merged_or_plain returns GGUFLinear.""" + in_features = 2048 + output_sizes = [8192, 512, 512] + # All three parts use IQ3_S + quant_types = [GGML_IQ3_S, GGML_IQ3_S, GGML_IQ3_S] + + lin = gguf_merged_or_plain(in_features, output_sizes, quant_types, has_bias=False) + + # Should return a plain GGUFLinear, not merged + assert isinstance(lin, GGUFLinear) + assert not isinstance(lin, GGUFMergedLinear) + # Total output features should be the sum + assert lin.out_features == 9216 + # Qweight should be concatenated (single packed buffer) + assert lin.qweight.shape == (9216, row_bytes(in_features, GGML_IQ3_S)) + + def test_gguf_merged_or_plain_picks_merged_when_mixed(self): + """When quant types differ, gguf_merged_or_plain returns GGUFMergedLinear.""" + in_features = 2048 + output_sizes = [8192, 512, 512] + # Mixed: IQ3_S and Q4_K + quant_types = [GGML_IQ3_S, GGML_IQ3_S, GGML_Q4_K] + + lin = gguf_merged_or_plain(in_features, output_sizes, quant_types, has_bias=False) + + # Should return a GGUFMergedLinear + assert isinstance(lin, GGUFMergedLinear) + assert lin.out_features == 9216 + # Should have separate qweight_0, qweight_1, qweight_2 + assert hasattr(lin, "qweight_0") + assert hasattr(lin, "qweight_1") + assert hasattr(lin, "qweight_2") + + +class TestGDNGeometry: + """Test that GDN in_proj geometry matches Ornith's configuration.""" + + def test_gdn_split_matches_ornith_geometry(self): + """GDN in_proj split [8192, 4096, 32, 32] sums to 12352 and matches the arithmetic. + + From ORNITH_SPEC section 1: + - embedding_length=2048 + - head_count=16, head_count_kv=2, key_length=value_length=256 + - ssm.state_size=128, group_count=16, time_step_rank=32, inner_size=4096 + + GDN in_proj fuses four tensors: + - attn_qkv: q (2 * head_count * state_size) + k (head_count * state_size) + + v (num_v_heads * state_size) + = 2*16*128 + 16*128 + 32*128 = 8192 + - attn_gate: 4096 (= num_v_heads * state_size = 32 * 128) + - ssm_beta: 32 (= num_v_heads) + - ssm_alpha: 32 (= num_v_heads) + Total: 8192 + 4096 + 32 + 32 = 12352 + """ + # Ornith GDN parameters + head_count = 16 + state_size = 128 + num_k_heads = 16 + num_v_heads = 32 + inner_size = 4096 + + # Compute expected attn_qkv output size + attn_qkv_size = 2 * num_k_heads * state_size + num_v_heads * state_size + assert attn_qkv_size == 8192 + + # Compute expected attn_gate output size + attn_gate_size = num_v_heads * state_size + assert attn_gate_size == 4096 + + # ssm_beta and ssm_alpha outputs are both num_v_heads + ssm_beta_size = num_v_heads + ssm_alpha_size = num_v_heads + assert ssm_beta_size == 32 + assert ssm_alpha_size == 32 + + # Total in_proj output width + in_proj_split = [attn_qkv_size, attn_gate_size, ssm_beta_size, ssm_alpha_size] + in_proj_total = sum(in_proj_split) + assert in_proj_total == 12352 + + # Verify all parts are as specified + assert in_proj_split == [8192, 4096, 32, 32] + + +class TestQwenNameMapping: + """Test gguf_name_to_freetoken: the inverse of llama.cpp's tensor name mapping.""" + + def test_name_map_drops_mtp_block(self): + """gguf_name_to_freetoken drops block 40 (NextN/MTP) and nextn.* suffixes.""" + num_layers = 40 # Ornith has 40 decoder layers + 1 NextN block (41 total in file) + + # Block 40 is the NextN/MTP block and should be dropped + assert gguf_name_to_freetoken("blk.40.attn_norm.weight", num_layers) is None + assert gguf_name_to_freetoken("blk.40.ffn_gate.weight", num_layers) is None + + # Any suffix starting with "nextn." on valid layers should be dropped + assert gguf_name_to_freetoken("blk.0.nextn.something", num_layers) is None + assert gguf_name_to_freetoken("blk.5.nextn.predict.weight", num_layers) is None + + def test_name_map_real_names(self): + """gguf_name_to_freetoken maps real Ornith tensor names correctly.""" + num_layers = 40 + + # Global tensors + assert gguf_name_to_freetoken("token_embd.weight", num_layers) == "model.embed_tokens.weight" + assert gguf_name_to_freetoken("output_norm.weight", num_layers) == "model.norm.weight" + assert gguf_name_to_freetoken("output.weight", num_layers) == "lm_head.weight" + + # Layer 3 (full-attention): single-output projections + assert gguf_name_to_freetoken("blk.3.attn_q.weight", num_layers) == "model.layers.3.self_attn.q_proj.weight" + assert gguf_name_to_freetoken("blk.3.attn_k.weight", num_layers) == "model.layers.3.self_attn.k_proj.weight" + assert gguf_name_to_freetoken("blk.3.attn_v.weight", num_layers) == "model.layers.3.self_attn.v_proj.weight" + assert gguf_name_to_freetoken("blk.3.attn_output.weight", num_layers) == "model.layers.3.self_attn.o_proj.weight" + + # Layer 0 (GDN): linear attention with conv1d and SSM + assert gguf_name_to_freetoken("blk.0.ssm_conv1d.weight", num_layers) == "model.layers.0.linear_attn.conv1d.weight" + assert gguf_name_to_freetoken("blk.0.ssm_norm.weight", num_layers) == "model.layers.0.linear_attn.norm.weight" + assert gguf_name_to_freetoken("blk.0.ssm_out.weight", num_layers) == "model.layers.0.linear_attn.out_proj.weight" + assert gguf_name_to_freetoken("blk.0.ssm_a", num_layers) == "model.layers.0.linear_attn.A_log" + assert gguf_name_to_freetoken("blk.0.ssm_dt.bias", num_layers) == "model.layers.0.linear_attn.dt_bias" + + # MoE tensors: shared expert and router + assert gguf_name_to_freetoken("blk.5.ffn_gate_inp.weight", num_layers) == "model.layers.5.mlp.gate.weight" + assert gguf_name_to_freetoken("blk.5.ffn_gate_inp_shexp.weight", num_layers) == "model.layers.5.mlp.shared_expert_gate.weight" + assert gguf_name_to_freetoken("blk.5.ffn_gate_shexp.weight", num_layers) == "model.layers.5.mlp.shared_expert.gate_proj.weight" + assert gguf_name_to_freetoken("blk.5.ffn_up_shexp.weight", num_layers) == "model.layers.5.mlp.shared_expert.up_proj.weight" + assert gguf_name_to_freetoken("blk.5.ffn_down_shexp.weight", num_layers) == "model.layers.5.mlp.shared_expert.down_proj.weight" + + def test_name_map_ignores_routed_experts(self): + """gguf_name_to_freetoken ignores routed-expert stacks (handled by offload banks).""" + num_layers = 40 + + # Routed-expert stacks are skipped (offload banks read them directly) + assert gguf_name_to_freetoken("blk.0.ffn_gate_exps.weight", num_layers) is None + assert gguf_name_to_freetoken("blk.15.ffn_up_exps.weight", num_layers) is None + assert gguf_name_to_freetoken("blk.39.ffn_down_exps.weight", num_layers) is None + +__all__ = [ + "test_merged_linear_concatenates_outputs", + "test_merged_linear_rejects_length_mismatch", + "test_gguf_merged_or_plain_picks_plain_when_uniform", + "test_gdn_split_matches_ornith_geometry", + "test_name_map_drops_mtp_block", +] From fcdbd0cc2227027df8038d94b6dc4cd9f3871030 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 22:00:50 -0700 Subject: [PATCH 05/36] fix(gguf): stop _LAYER_MAP naming attributes the qwen35moe module lacks _LAYER_MAP and _FUSE were written before the module tree was known and mapped attn_q/k/v -> self_attn.q_proj/k_proj/v_proj and ffn_gate_shexp/ffn_up_shexp -> mlp.shared_expert.gate_proj/up_proj. Those attributes do not exist: Qwen3_5Attention builds one merged qkv_proj (_qkv_split [8192, 512, 512]) and _SharedExpert one merged gate_up_proj. iter_gguf_weights always handled them by fusion, so the entries were dead -- but gguf_name_to_freetoken would hand a caller an invented parameter name. Replaced both with _MERGED_PARTS, the set of suffixes that are parts of a merged projection, which gguf_name_to_freetoken now reports as None (handled elsewhere). The constant documents each merged target and its concat order, including why the GDN target is in_proj rather than in_proj_qkvz/in_proj_ba (gdn.py only splits those out on the fp8 branch, and a GGUF checkpoint is not fp8). Also: tests call .forward() explicitly -- BaseOP defines no __call__. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/qwen3_5_moe/gguf.py | 51 ++++++++++++++------- tests/models/test_qwen35moe_gguf.py | 2 +- 2 files changed, 36 insertions(+), 17 deletions(-) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index 8bdc9cb5..60a2810d 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -207,10 +207,9 @@ def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: # shared by both layer kinds "attn_norm.weight": "input_layernorm.weight", "post_attention_norm.weight": "post_attention_layernorm.weight", - # full-attention layers - "attn_q.weight": "self_attn.q_proj.weight", - "attn_k.weight": "self_attn.k_proj.weight", - "attn_v.weight": "self_attn.v_proj.weight", + # full-attention layers. attn_q/attn_k/attn_v are deliberately absent: the model has + # no q_proj/k_proj/v_proj attributes -- Qwen3_5Attention builds one merged qkv_proj + # (_qkv_split = [8192, 512, 512]) -- so they are fused by iter_gguf_weights, not renamed. "attn_output.weight": "self_attn.o_proj.weight", "attn_q_norm.weight": "self_attn.q_norm.weight", "attn_k_norm.weight": "self_attn.k_norm.weight", @@ -220,20 +219,33 @@ def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: "ssm_out.weight": "linear_attn.out_proj.weight", "ssm_a": "linear_attn.A_log", "ssm_dt.bias": "linear_attn.dt_bias", - # MoE router + shared expert + # MoE router + shared expert. ffn_gate_shexp / ffn_up_shexp are absent for the same + # reason as attn_q/k/v: _SharedExpert has a single merged gate_up_proj, so they are + # fused rather than renamed. "ffn_gate_inp.weight": "mlp.gate.weight", "ffn_gate_inp_shexp.weight": "mlp.shared_expert_gate.weight", - "ffn_gate_shexp.weight": "mlp.shared_expert.gate_proj.weight", - "ffn_up_shexp.weight": "mlp.shared_expert.up_proj.weight", "ffn_down_shexp.weight": "mlp.shared_expert.down_proj.weight", } -# Pairs llama.cpp splits that FreeToken's model code wants fused, in concat order. -# Mirrors _PT_FP8_FUSE / _PT_BF16_FUSE in weight.py. -_FUSE: dict[str, tuple[str, str]] = { - "linear_attn.in_proj_qkvz.weight": ("attn_qkv.weight", "attn_gate.weight"), - "linear_attn.in_proj_ba.weight": ("ssm_beta.weight", "ssm_alpha.weight"), -} +# Suffixes that are PARTS of a merged projection: never renamed 1:1, always combined by +# iter_gguf_weights into the merged buffer the model actually declares. Listed here so +# gguf_name_to_freetoken can report them as "handled elsewhere" (None) instead of +# inventing a parameter name that does not exist on the module. +# +# The merged targets and their concat orders: +# self_attn.qkv_proj <- attn_q, attn_k, attn_v (_qkv_split [8192, 512, 512]) +# linear_attn.in_proj <- attn_qkv, attn_gate, ssm_beta, ssm_alpha +# (_in_proj_split [conv_dim, value_dim, n_v, n_v]) +# mlp.shared_expert.gate_up_proj <- ffn_gate_shexp, ffn_up_shexp +# +# NOTE the GDN target is ``in_proj``, not ``in_proj_qkvz``/``in_proj_ba``: gdn.py only +# splits those two out on the fp8 branch (``self._fp8``), and a GGUF checkpoint sets +# attn_quant="gguf", so the single fused in_proj is what exists. +_MERGED_PARTS: frozenset[str] = frozenset({ + "attn_q.weight", "attn_k.weight", "attn_v.weight", + "attn_qkv.weight", "attn_gate.weight", "ssm_beta.weight", "ssm_alpha.weight", + "ffn_gate_shexp.weight", "ffn_up_shexp.weight", +}) # Routed-expert stacks: [num_experts, out, in] packed blocks, handled by the offload # expert-bank loader rather than yielded as ordinary parameters. @@ -253,8 +265,13 @@ def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: def gguf_name_to_freetoken(name: str, num_layers: int) -> str | None: """Map one llama.cpp tensor name to its FreeToken parameter name. - Returns ``None`` for tensors FreeToken does not consume (the NextN/MTP block, and - the routed-expert stacks, which the expert-bank loader reads directly). + Returns ``None`` for anything this function does not rename 1:1 -- the NextN/MTP + block, the routed-expert stacks (read directly by the expert-bank loader), and the + parts of a merged projection (combined by :func:`iter_gguf_weights`, which is the + only place that knows the concat order and the per-part quant types). Callers that + want full coverage accounting should treat ``None`` as "handled elsewhere", not + "unmapped": returning an invented ``q_proj``/``gate_proj`` name for a fusion part + would name an attribute the module does not have. """ if name in _GLOBAL_MAP: return _GLOBAL_MAP[name] @@ -268,6 +285,8 @@ def gguf_name_to_freetoken(name: str, num_layers: int) -> str | None: return None if suffix in _EXPERT_SUFFIXES: return None + if suffix in _MERGED_PARTS: + return None # fused by iter_gguf_weights into the merged buffer mapped = _LAYER_MAP.get(suffix) if mapped is None: return None @@ -667,6 +686,6 @@ def swap_linear(owner, attr, quant_type=GGML_Q4_K): "gguf_name_to_freetoken", "iter_gguf_weights", "convert_qwen35_to_gguf", - "_FUSE", + "_MERGED_PARTS", "_EXPERT_SUFFIXES", ] diff --git a/tests/models/test_qwen35moe_gguf.py b/tests/models/test_qwen35moe_gguf.py index 94b3ec90..ff8687cb 100644 --- a/tests/models/test_qwen35moe_gguf.py +++ b/tests/models/test_qwen35moe_gguf.py @@ -123,7 +123,7 @@ def test_merged_linear_concatenates_outputs(self, mock_kernel_module): # Forward pass: batch size 1 (small batch, uses MMVQ) x = torch.randn(1, in_features, dtype=torch.bfloat16) - output = merged(x) + output = merged.forward(x) # Output shape should be [1, 9216] assert output.shape == (1, 9216) From 1a035bca1a7027a7e4ac0e4747db934eea1cad3e Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 22:02:21 -0700 Subject: [PATCH 06/36] test(gguf): assert merged-projection parts map to None, not invented names The name-map test asserted attn_q/k/v -> self_attn.q_proj/k_proj/v_proj and ffn_gate_shexp/ffn_up_shexp -> shared_expert.gate_proj/up_proj. Those attributes do not exist on the module (one merged qkv_proj and one merged gate_up_proj), so the test was locking in names the loader could never populate. Assert None for fusion parts instead, which is the contract iter_gguf_weights actually implements. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- tests/models/test_qwen35moe_gguf.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/tests/models/test_qwen35moe_gguf.py b/tests/models/test_qwen35moe_gguf.py index ff8687cb..4c5ca9f6 100644 --- a/tests/models/test_qwen35moe_gguf.py +++ b/tests/models/test_qwen35moe_gguf.py @@ -277,11 +277,14 @@ def test_name_map_real_names(self): assert gguf_name_to_freetoken("output_norm.weight", num_layers) == "model.norm.weight" assert gguf_name_to_freetoken("output.weight", num_layers) == "lm_head.weight" - # Layer 3 (full-attention): single-output projections - assert gguf_name_to_freetoken("blk.3.attn_q.weight", num_layers) == "model.layers.3.self_attn.q_proj.weight" - assert gguf_name_to_freetoken("blk.3.attn_k.weight", num_layers) == "model.layers.3.self_attn.k_proj.weight" - assert gguf_name_to_freetoken("blk.3.attn_v.weight", num_layers) == "model.layers.3.self_attn.v_proj.weight" + # Layer 3 (full-attention). Only o_proj is a 1:1 rename. attn_q/k/v are PARTS of + # the merged qkv_proj (Qwen3_5Attention has no q_proj/k_proj/v_proj attribute), so + # the mapper reports None -- iter_gguf_weights owns their fusion because only it + # knows the concat order and the per-part quant types. Asserting a q_proj name here + # would lock in a parameter the module does not have. assert gguf_name_to_freetoken("blk.3.attn_output.weight", num_layers) == "model.layers.3.self_attn.o_proj.weight" + for part in ("attn_q.weight", "attn_k.weight", "attn_v.weight"): + assert gguf_name_to_freetoken(f"blk.3.{part}", num_layers) is None, part # Layer 0 (GDN): linear attention with conv1d and SSM assert gguf_name_to_freetoken("blk.0.ssm_conv1d.weight", num_layers) == "model.layers.0.linear_attn.conv1d.weight" @@ -293,9 +296,11 @@ def test_name_map_real_names(self): # MoE tensors: shared expert and router assert gguf_name_to_freetoken("blk.5.ffn_gate_inp.weight", num_layers) == "model.layers.5.mlp.gate.weight" assert gguf_name_to_freetoken("blk.5.ffn_gate_inp_shexp.weight", num_layers) == "model.layers.5.mlp.shared_expert_gate.weight" - assert gguf_name_to_freetoken("blk.5.ffn_gate_shexp.weight", num_layers) == "model.layers.5.mlp.shared_expert.gate_proj.weight" - assert gguf_name_to_freetoken("blk.5.ffn_up_shexp.weight", num_layers) == "model.layers.5.mlp.shared_expert.up_proj.weight" assert gguf_name_to_freetoken("blk.5.ffn_down_shexp.weight", num_layers) == "model.layers.5.mlp.shared_expert.down_proj.weight" + # gate/up shexp are parts of _SharedExpert's merged gate_up_proj -- same reasoning + # as attn_q/k/v above. + for part in ("ffn_gate_shexp.weight", "ffn_up_shexp.weight"): + assert gguf_name_to_freetoken(f"blk.5.{part}", num_layers) is None, part def test_name_map_ignores_routed_experts(self): """gguf_name_to_freetoken ignores routed-expert stacks (handled by offload banks).""" From e6c824818ab5bf8a99c25a9859e268ad75450c13 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 22:09:51 -0700 Subject: [PATCH 07/36] fix(gguf): shared-expert weights were silently dropped; conv1d/gate dtype+shape Three defects the module-tree reconciliation caught on the real checkpoint: 1. The shared-expert suffixes (ffn_gate_shexp / ffn_up_shexp / ffn_down_shexp) were handled AFTER the per-layer-kind branch, whose `else: continue` swallows any suffix it does not recognize. Every layer's shared_expert.gate_up_proj and down_proj therefore went unfilled -- 81 buffers -- while the loader reported success. Moved the handling above that branch, where both layer kinds reach it. 2. mlp.shared_expert_gate: llama.cpp stores the single-output gate as a 1-D [hidden] vector; LinearReplicated(hidden, 1) declares [1, hidden]. Now reshaped. 3. linear_attn.conv1d.weight was yielded fp32. gdn.py exempts only A_log and dt_bias from the model-dtype downcast, so the depthwise conv is bf16. Found by reconciling every name/shape/dtype iter_gguf_weights yields against the constructed module's own state_dict, which is the check that catches this class of bug before it becomes fluent garbage at generation time. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/qwen3_5_moe/gguf.py | 68 +++++++++++---------- 1 file changed, 36 insertions(+), 32 deletions(-) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index 60a2810d..46b2e841 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -358,7 +358,7 @@ def iter_gguf_weights( models). A_log and dt_bias stay float32 (gdn.py keeps recurrence-gating params in fp32). - conv1d.weight stays float32 and is reshaped to [conv_dim, 1, kernel]. + conv1d.weight is bf16 (model dtype) reshaped to [conv_dim, 1, kernel]. """ from freetoken.models.gguf.reader import iter_gguf_tensors from freetoken.utils import cached_load_hf_config @@ -444,7 +444,9 @@ def layer_of(name: str) -> int: yield f"{base}.mlp.gate.weight", _to_bf16(t) continue if suffix == "ffn_gate_inp_shexp.weight": - yield f"{base}.mlp.shared_expert_gate.weight", _to_bf16(t) + # llama.cpp stores the single-output shared-expert gate as a 1-D [hidden] + # vector; _SharedExpert's LinearReplicated(hidden, 1) declares [1, hidden]. + yield f"{base}.mlp.shared_expert_gate.weight", _to_bf16(t).reshape(1, -1) continue if suffix == "ssm_norm.weight": yield f"{base}.linear_attn.norm.weight", _to_bf16(t) @@ -456,9 +458,10 @@ def layer_of(name: str) -> int: yield f"{base}.linear_attn.dt_bias", _to_f32(t) continue if suffix == "ssm_conv1d.weight": - # F32, reshape to [conv_dim, 1, kernel] where conv_dim is attn_qkv output size - w = _to_f32(t) - w = w.reshape(gdn_attn_qkv_size, 1, gdn_conv_kernel) + # F32 in the file, but _DepthwiseConv1d allocates at the model dtype: gdn.py + # exempts only A_log / dt_bias from the downcast, so this one is bf16. Reshape + # to [conv_dim, 1, kernel] -- gdn.py's _conv_weight() does .squeeze(1). + w = _to_bf16(t).reshape(gdn_attn_qkv_size, 1, gdn_conv_kernel) yield f"{base}.linear_attn.conv1d.weight", w continue if suffix == "attn_q_norm.weight": @@ -468,6 +471,34 @@ def layer_of(name: str) -> int: yield f"{base}.self_attn.k_norm.weight", _to_bf16(t) continue + # Shared expert (present on every layer, both kinds) -- must be handled BEFORE the + # per-layer-kind branch below, whose `else: continue` swallows any suffix it does + # not recognize. + if suffix in ("ffn_gate_shexp.weight", "ffn_up_shexp.weight", "ffn_down_shexp.weight"): + if suffix == "ffn_gate_shexp.weight": + gate_up_buf.setdefault(layer, {})["gate"] = t.packed() + elif suffix == "ffn_up_shexp.weight": + gate_up_buf.setdefault(layer, {})["up"] = t.packed() + else: + yield f"{base}.mlp.shared_expert.down_proj.qweight", t.packed() + + # Emit the fused gate_up once both parts have arrived. + gu = gate_up_buf.get(layer) + if gu is not None and "gate" in gu and "up" in gu: + types = [ + quant_map.get((layer, "ffn_gate_shexp.weight")), + quant_map.get((layer, "ffn_up_shexp.weight")), + ] + if len(set(types)) == 1: + yield f"{base}.mlp.shared_expert.gate_up_proj.qweight", torch.cat( + [gu["gate"], gu["up"]], dim=0 + ) + else: + yield f"{base}.mlp.shared_expert.gate_up_proj.qweight_0", gu["gate"] + yield f"{base}.mlp.shared_expert.gate_up_proj.qweight_1", gu["up"] + del gate_up_buf[layer] + continue + # Quantized projections: keep packed; fuse per layer. # Full-attention: qkv from q, k, v if layer in full_layer_ids: @@ -554,33 +585,6 @@ def layer_of(name: str) -> int: yield f"{base}.linear_attn.in_proj.qweight_3", slots["alpha"] del in_proj_buf[layer] - # Shared expert gate_up: every layer gets this - if suffix == "ffn_gate_shexp.weight": - gate_up_buf.setdefault(layer, {})["gate"] = t.packed() - elif suffix == "ffn_up_shexp.weight": - gate_up_buf.setdefault(layer, {})["up"] = t.packed() - elif suffix == "ffn_down_shexp.weight": - yield f"{base}.mlp.shared_expert.down_proj.qweight", t.packed() - - # Emit fused gate_up once both parts are present. - gu = gate_up_buf.get(layer) - if gu is not None and "gate" in gu and "up" in gu: - # Determine if this is a mixed-quant group (unlikely but check). - types = [ - quant_map.get((layer, "ffn_gate_shexp.weight")), - quant_map.get((layer, "ffn_up_shexp.weight")), - ] - if len(set(types)) == 1: - # Uniform quant: fuse via torch.cat along dim 0. - yield f"{base}.mlp.shared_expert.gate_up_proj.qweight", torch.cat( - [gu["gate"], gu["up"]], dim=0 - ) - else: - # Mixed quant: emit GGUFMergedLinear format. - yield f"{base}.mlp.shared_expert.gate_up_proj.qweight_0", gu["gate"] - yield f"{base}.mlp.shared_expert.gate_up_proj.qweight_1", gu["up"] - del gate_up_buf[layer] - # Verify no fusion buffers are incomplete. assert not qkv_buf, f"incomplete full-attn qkv groups: {sorted(qkv_buf)}" assert not in_proj_buf, f"incomplete GDN in_proj groups: {sorted(in_proj_buf)}" From af612b9db0f347b550817d5800f07023ac7aa079 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 22:11:05 -0700 Subject: [PATCH 08/36] fix(gguf): lm_head was never loaded (output.weight dropped by the blk. filter) iter_gguf_weights handled token_embd and output_norm as globals, then bailed on `not name.startswith("blk.")` -- which silently discarded output.weight, the untied LM head. lm_head.qweight [248320, 1680] stayed at its torch.empty allocation, so the model would have generated from uninitialized memory. Ornith ships output.weight (hence tie_word_embeddings=False), so it is a real buffer that must be filled. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/qwen3_5_moe/gguf.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index 46b2e841..bdacf193 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -414,6 +414,14 @@ def layer_of(name: str) -> int: if name == "output_norm.weight": yield "model.norm.weight", _to_bf16(t) continue + if name == "output.weight": + # The untied LM head, packed (Q6_K here). Ornith ships output.weight, so + # tie_word_embeddings is False and lm_head is a real GGUFLinear that must be + # filled; when a checkpoint omits it the head aliases the embedding table and + # there is nothing to yield. + if not config.tie_word_embeddings: + yield "lm_head.qweight", t.packed() + continue if not name.startswith("blk."): continue From e6b46847f2e6562bf8833f34f49f372793fcfc2a Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 22:22:45 -0700 Subject: [PATCH 09/36] fix(gguf): read every quant type from the file; bank specs match alloc_layer_banks Both bugs surfaced only on IQ3_S -- IQ3_M happened to mask them. convert_qwen35_to_gguf hardcoded quant types for o_proj (Q4_K), out_proj (IQ3_S), the token embedding and the lm_head, and used quant_map.get(key, DEFAULT) elsewhere. Q4_K is right for IQ3_M's attn_output but IQ3_S uses IQ3_S there, so every full-attention o_proj was allocated 2304 bytes/row instead of 1760. A guessed quant type silently mis-sizes a packed buffer and the only symptom is garbage output, so the new qt() helper raises on a missing tensor instead of defaulting. _scan_quant_types now also keys the globals (token_embd/output) under layer -1 so the embedding and head are sized from the file too. gguf_expert_specs returned {name: [per-layer (shape, dtype)]}, but alloc_layer_banks takes {name: (shape, dtype)} -- one spec per bank -- so bank allocation died with "too many values to unpack". The per-layer form was wrong by construction anyway: a bank's slot pool is one allocation with one stride, so its type must be uniform. It now returns the flat form and rejects a non-uniform bank. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/qwen3_5_moe/gguf.py | 92 +++++++++++-------- .../models/qwen3_5_moe/gguf_experts.py | 59 ++++++------ 2 files changed, 82 insertions(+), 69 deletions(-) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index bdacf193..f3b40c6b 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -309,6 +309,9 @@ def _scan_quant_types(model_path: str) -> dict[tuple[int, str], int]: quant_types = {} for t in iter_gguf_tensors(model_path): if not t.name.startswith("blk."): + # Globals (token_embd.weight, output.weight, output_norm.weight) keyed under + # layer -1 so the swap can size the embedding and lm_head from the file too. + quant_types[(-1, t.name)] = t.ggml_type continue _, idx, suffix = t.name.split(".", 2) layer = int(idx) @@ -623,8 +626,25 @@ def convert_qwen35_to_gguf(model, config: ModelConfig, *, model_path: str) -> No if isinstance(config.attention_group_for_layer(lid), FullAttentionGroupConfig) } - def swap_linear(owner, attr, quant_type=GGML_Q4_K): - """Replace a dense Linear with GGUFLinear.""" + def qt(layer: int, suffix: str) -> int: + """The ggml type of one tensor, straight from the file. + + No default: a guessed type silently allocates a wrong-sized packed buffer, and the + only symptom is garbage output. (An earlier version defaulted o_proj to Q4_K, which + happened to match IQ3_M and mis-sized every full-attention o_proj on IQ3_S.) + """ + key = (layer, suffix) + if key not in quant_map: + raise ValueError( + f"GGUF {model_path}: expected tensor " + f"{suffix if layer < 0 else f'blk.{layer}.{suffix}'} is absent, so its quant " + f"type cannot be read; this checkpoint does not match the qwen35moe layout " + f"this adapter expects" + ) + return quant_map[key] + + def swap_linear(owner, attr, quant_type: int): + """Replace a dense Linear with the GGUFLinear its packed weight will land in.""" lin = getattr(owner, attr) out_features, in_features = lin.weight.shape setattr( @@ -637,60 +657,60 @@ def swap_linear(owner, attr, quant_type=GGML_Q4_K): embed = GGUFEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, - quant_type=GGML_IQ3_S, + quant_type=qt(-1, "token_embd.weight"), ) inner.embed_tokens = embed for layer_idx, layer in enumerate(inner.layers.op_list): if layer_idx in full_layer_ids: - # Full-attention layer: qkv_proj (mixed quant), o_proj (Q4_K) - types = [ - quant_map.get((layer_idx, "attn_q.weight"), GGML_IQ3_S), - quant_map.get((layer_idx, "attn_k.weight"), GGML_IQ3_S), - quant_map.get((layer_idx, "attn_v.weight"), GGML_Q4_K), - ] - output_sizes = [8192, 512, 512] + # qkv_proj: q | k | v. Mixed in every Ornith quant level (v is a K-quant while + # q/k are I-quants), so this is normally the GGUFMergedLinear path. layer.self_attn.qkv_proj = gguf_merged_or_plain( - config.hidden_size, output_sizes, types, has_bias=False + config.hidden_size, + [8192, 512, 512], + [ + qt(layer_idx, "attn_q.weight"), + qt(layer_idx, "attn_k.weight"), + qt(layer_idx, "attn_v.weight"), + ], + has_bias=False, ) - swap_linear(layer.self_attn, "o_proj", GGML_Q4_K) + swap_linear(layer.self_attn, "o_proj", qt(layer_idx, "attn_output.weight")) else: - # GDN layer: in_proj (mixed quant), out_proj (IQ3_S) - types = [ - quant_map.get((layer_idx, "attn_qkv.weight"), GGML_Q4_K), - quant_map.get((layer_idx, "attn_gate.weight"), GGML_IQ3_S), - quant_map.get((layer_idx, "ssm_beta.weight"), GGML_IQ3_S), - quant_map.get((layer_idx, "ssm_alpha.weight"), GGML_IQ3_S), - ] - output_sizes = [8192, 4096, 32, 32] + # in_proj: qkv | z | b | a, matching gdn.py's + # _in_proj_split = [conv_dim, value_dim, num_v_heads, num_v_heads]. layer.linear_attn.in_proj = gguf_merged_or_plain( - config.hidden_size, output_sizes, types, has_bias=False + config.hidden_size, + [8192, 4096, 32, 32], + [ + qt(layer_idx, "attn_qkv.weight"), + qt(layer_idx, "attn_gate.weight"), + qt(layer_idx, "ssm_beta.weight"), + qt(layer_idx, "ssm_alpha.weight"), + ], + has_bias=False, ) - swap_linear(layer.linear_attn, "out_proj", GGML_IQ3_S) + swap_linear(layer.linear_attn, "out_proj", qt(layer_idx, "ssm_out.weight")) - # Shared expert: gate_up_proj (IQ3_S, uniform), down_proj (Q4_K or IQ3_S). - gate_up_type = quant_map.get((layer_idx, "ffn_gate_shexp.weight"), GGML_IQ3_S) - down_type = quant_map.get((layer_idx, "ffn_down_shexp.weight"), GGML_Q4_K) - - # gate_up is typically uniform IQ3_S, but use gguf_merged_or_plain for safety. - up_type = quant_map.get((layer_idx, "ffn_up_shexp.weight"), GGML_IQ3_S) + # Shared expert: gate|up fuse when they share a type (they do in every quant level + # seen so far); down is independent and does vary (Q4_K on IQ3_M's first layers). + I = config.shared_expert_intermediate_size layer.mlp.shared_expert.gate_up_proj = gguf_merged_or_plain( config.hidden_size, - [config.shared_expert_intermediate_size, config.shared_expert_intermediate_size], - [gate_up_type, up_type], + [I, I], + [qt(layer_idx, "ffn_gate_shexp.weight"), qt(layer_idx, "ffn_up_shexp.weight")], has_bias=False, ) - swap_linear(layer.mlp.shared_expert, "down_proj", down_type) + swap_linear( + layer.mlp.shared_expert, "down_proj", qt(layer_idx, "ffn_down_shexp.weight") + ) - # lm_head: use output.weight quant type (Q6_K in Ornith). if config.tie_word_embeddings: - # Tied head: reference the embedding's qweight. from freetoken.models.gemma4.gguf import GGUFTiedLMHead - model.lm_head = GGUFTiedLMHead(embed, GGML_Q6_K) + model.lm_head = GGUFTiedLMHead(embed, qt(-1, "token_embd.weight")) else: - # Untied head: create a separate GGUFLinear (rare for qwen35moe). - swap_linear(model, "lm_head", GGML_Q6_K) + swap_linear(model, "lm_head", qt(-1, "output.weight")) __all__ = [ diff --git a/python/freetoken/models/qwen3_5_moe/gguf_experts.py b/python/freetoken/models/qwen3_5_moe/gguf_experts.py index 4d2bedad..623f14c8 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf_experts.py +++ b/python/freetoken/models/qwen3_5_moe/gguf_experts.py @@ -96,41 +96,34 @@ def gguf_expert_types(model_path: str, num_layers: int) -> dict[str, list[int]]: def gguf_expert_specs( config: ModelConfig, types: dict[str, list[int]] -) -> dict[str, list[tuple[tuple[int, ...], torch.dtype]]]: - """Per-layer expert bank shapes, accounting for per-layer dtype variation. +) -> dict[str, tuple[tuple[int, ...], torch.dtype]]: + """Expert bank shapes as ``{name: (shape, dtype)}`` -- ``alloc_layer_banks``' contract. - The routed experts are stored as 3D stacks [E, out, in] in torch order. - Returns a list of (shape, dtype) tuples per expert bank: - - ``gate_up[layer]``: ``((E, 2*I, row_bytes(H, types["gate_up"][layer])), torch.uint8)`` - - ``down[layer]``: ``((E, H, row_bytes(I, types["down"][layer])), torch.uint8)`` + The routed experts are 3D stacks in torch order:: - where ``E=num_experts``, ``H=hidden_size``, ``I=moe_intermediate_size``. + gate_up (E, 2*I, row_bytes(H, t_gate_up)) uint8, packed blocks + down (E, H, row_bytes(I, t_down)) uint8, packed blocks - The row_bytes dimension varies by layer because the quant type varies, - and the MoE kernel needs to read the exact byte width for its quant type. + One spec per bank, not per layer: every layer of a bank MUST share a ggml type. The + GPU slot pool is a single allocation shared by all layers and ``moe_vec.cuh`` indexes + it as ``expert * nrows * (ncols / qk)`` with no padding allowance, so two strides in + one pool would read every block at the wrong offset. A non-uniform bank is rejected + here rather than mis-decoded; ``expert_banks._gguf_banks`` raises the user-facing + error naming the offending layers. """ - E = config.num_experts - H = config.hidden_size - I = config.moe_intermediate_size - num_layers = config.num_layers - - gate_up_specs = [] - down_specs = [] - - for layer in range(num_layers): - gate_up_type = types["gate_up"][layer] - down_type = types["down"][layer] - - gate_up_row_bytes = row_bytes(H, gate_up_type) - down_row_bytes = row_bytes(I, down_type) - - gate_up_specs.append(((E, 2 * I, gate_up_row_bytes), torch.uint8)) - down_specs.append(((E, H, down_row_bytes), torch.uint8)) - - return { - "gate_up": gate_up_specs, - "down": down_specs, - } + E, H, I = config.num_experts, config.hidden_size, config.moe_intermediate_size + out = {} + for name, elems in (("gate_up", H), ("down", I)): + distinct = sorted(set(types[name])) + if len(distinct) != 1: + raise ValueError( + f"expert bank {name!r} mixes ggml types across layers ({distinct}); a bank " + f"must be uniform because its slot pool is one allocation with one stride" + ) + rb = row_bytes(elems, distinct[0]) + shape = (E, 2 * I, rb) if name == "gate_up" else (E, H, rb) + out[name] = (shape, torch.uint8) + return out def load_gguf_expert_sources( @@ -203,7 +196,7 @@ def _load(sink) -> None: # Shape from GGUF: [E, H, I] in torch order = [I, H, E] in ggml order # t.packed() is [I*H, row_bytes(E, type)] # Reshape to [E, H, row_bytes(I, type)] - down_row_bytes = specs["down"][layer][0][2] + down_row_bytes = specs["down"][0][2] banks["down"][layer].copy_(t.packed().reshape(E, H, down_row_bytes)) seen_down.add(layer) if tracker is not None: @@ -214,7 +207,7 @@ def _load(sink) -> None: # Emit gate_up bank once both gate and up are present. if layer in gate_buf and layer in up_buf: - gate_up_row_bytes = specs["gate_up"][layer][0][2] + gate_up_row_bytes = specs["gate_up"][0][2] # Concatenate gate [H*I, row_bytes(E, type)] and up [H*I, row_bytes(E, type)] # to get [H*2*I, row_bytes(E, type)], then reshape to [E, 2*I, row_bytes(H, type)] combined = torch.cat([gate_buf[layer], up_buf[layer]], dim=0) From 0231506ec46928d50d19d700169991644dd875b3 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 22:27:52 -0700 Subject: [PATCH 10/36] fix(gguf): tokenizer could not be built for qwen35moe (KeyError in transformers) transformers keys its GGUF tokenizer converters by its own model_type, so passing the GGUF arch string straight through raised KeyError('qwen35moe') and killed the detokenizer worker during load. qwen35moe is a GPT2-style BPE with merges (tokenizer.ggml.model=gpt2, pre=qwen35), which the existing qwen2 converter handles. Two adjacent gemma4-isms were also wrong for any other vocab and are now per-arch: - eos preferred the literal , and gguf_eos_token_ids hardcoded /. Neither exists in Ornith's vocab, so generation would have used the wrong eos and registered no chat stop id -- every request running to max_tokens. _STOP_TOKENS maps arch -> stop names in preference order; gemma4 keeps (, ) so its behaviour is unchanged, qwen35moe gets (<|im_end|>, <|endoftext|>), matching the file's eos_token_id 248046. - unk defaulted to the literal . Ornith has no unknown_token_id and no token, and handing PreTrainedTokenizerFast a token absent from the vocab appends it. tok_for now drops a default that is not in the vocab. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/gguf/tokenizer.py | 39 +++++++++++++++++------ 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 6d5481c1..5ff08188 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -12,8 +12,22 @@ from .reader import gguf_architecture, load_gguf_metadata -# GGUF architecture -> transformers GGUF tokenizer-converter key. -_TOKENIZER_ARCH = {"gemma4": "gemma4_text"} +# GGUF architecture -> transformers GGUF tokenizer-converter key. transformers keys its +# converters by *its own* model_type, not by the GGUF arch string, so an arch it has never +# heard of raises KeyError inside convert_gguf_tokenizer. qwen35moe is a GPT2-style BPE +# with merges (tokenizer.ggml.model == "gpt2", pre == "qwen35"), which the qwen2 converter +# handles; there is no qwen3.5-specific converter and it would be the same BPE anyway. +_TOKENIZER_ARCH = {"gemma4": "gemma4_text", "qwen35moe": "qwen2"} + +# Per-arch chat/stop tokens, in preference order: the first one present in the vocab +# becomes eos (so chat generation halts on the turn end rather than the formal document +# eos), and every one present is a stop id. Keyed by GGUF arch because these names are +# vocab-specific -- gemma4 ends a turn with , Qwen with <|im_end|>. An arch absent +# here falls back to tokenizer.ggml.eos_token_id alone. +_STOP_TOKENS: dict[str, tuple[str, ...]] = { + "gemma4": ("", ""), + "qwen35moe": ("<|im_end|>", "<|endoftext|>"), +} def load_gguf_tokenizer(model_path: str): @@ -32,13 +46,19 @@ def load_gguf_tokenizer(model_path: str): tokens = tok_dict["tokens"] - def tok_for(id_key: str, default: str) -> str: + def tok_for(id_key: str, default: str | None) -> str | None: + """The token named by ``tokenizer.ggml.``, else ``default`` if it is in the + vocab. A default that is not in this vocab is dropped: handing + PreTrainedTokenizerFast an unknown special token would append it to the vocab and + shift nothing but confuse decoding (Qwen has no at all).""" tid = meta.get(f"tokenizer.ggml.{id_key}") - return tokens[int(tid)] if tid is not None and int(tid) < len(tokens) else default + if tid is not None and int(tid) < len(tokens): + return tokens[int(tid)] + return default if default is not None and default in tokens else None - # gemma4 chat turns end with ; prefer it as eos so chat generation halts - # (the formal is also a stop id, see gguf_eos_token_ids). - turn_end = "" if "" in tokens else None + # Prefer the chat turn end as eos so chat generation halts there; the formal document + # eos stays a stop id (see gguf_eos_token_ids). + turn_end = next((t for t in _STOP_TOKENS.get(arch, ()) if t in tokens), None) tokenizer = PreTrainedTokenizerFast( tokenizer_object=fast, bos_token=tok_for("bos_token_id", ""), @@ -63,8 +83,9 @@ def gguf_eos_token_ids(model_path: str, tokenizer) -> set[int]: if eid is not None: ids.add(int(eid)) # Look the stop tokens up in the vocab directly (convert_tokens_to_ids would map an - # absent name to , wrongly adding it as a stop id). - for name in ("", ""): + # absent name to , wrongly adding it as a stop id). Names are per-arch: gemma4's + # / do not exist in a Qwen vocab and vice versa. + for name in _STOP_TOKENS.get(gguf_architecture(model_path), ()): try: ids.add(tokens.index(name)) except ValueError: From 0a9e9402c6baca2426bf5e5cde7d14f7fc344500 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 22:30:19 -0700 Subject: [PATCH 11/36] fix(gguf): the GGUF op swap was never invoked, so loading died on embed_tokens.weight convert_qwen35_to_gguf existed but nothing called it, so the module tree still held dense ops and load_state_dict raised KeyError('model.embed_tokens.weight') against a loader that yields .qweight. gemma4 performs its swap at the end of the model's __init__ behind is_gguf_model(); qwen3_5_moe now does the same. Unlike gemma4 -- whose checkpoint is uniformly Q4_0/Q6_K, so its swap needs no file access -- this swap sizes every packed buffer by the ggml type that tensor actually uses, which only the file knows. parse_gguf_config therefore publishes ModelConfig.gguf_model_path and __init__ asserts it is set. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/config.py | 5 +++++ python/freetoken/models/qwen3_5_moe/gguf.py | 7 +++++++ python/freetoken/models/qwen3_5_moe/model.py | 12 ++++++++++++ 3 files changed, 24 insertions(+) diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index 09c61151..66416f73 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -257,6 +257,11 @@ class ModelConfig: # *_M mixes do: Ornith IQ3_M splits ffn_down_exps across Q4_K and IQ3_S, while its # IQ3_S/IQ3_XXS siblings are uniform and load fine). gguf_expert_types: tuple[int, int] | None = None + # Path to the .gguf this config was parsed from, for ``expert_quant == "gguf"``. The + # model's __init__ swaps its dense ops for GGUF ops and has to size each packed buffer + # by the ggml type that tensor actually uses, which is only in the file -- unlike + # gemma4's all-Q4_0 checkpoint, llama.cpp's mixed quants pick a type per tensor. + gguf_model_path: str | None = None shared_expert_intermediate_size: int = 0 use_qk_norm: bool = False # ----- DeepSeek/GLM-style MoE extensions (default keeps other models intact) ----- diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index f3b40c6b..2da79559 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -191,6 +191,7 @@ def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: attention_groups=groups, expert_quant="gguf", gguf_expert_types=_uniform_expert_types(shim.model_path, num_layers), + gguf_model_path=shim.model_path, weight_block_size=None, attn_quant="gguf", dense_quant="gguf", @@ -602,6 +603,11 @@ def layer_of(name: str) -> int: assert not gate_up_buf, f"incomplete shared_expert gate_up groups: {sorted(gate_up_buf)}" +def is_gguf_model(config: ModelConfig) -> bool: + """True when this config came from a GGUF checkpoint (native block-quant path).""" + return getattr(config, "expert_quant", "none") == "gguf" + + def convert_qwen35_to_gguf(model, config: ModelConfig, *, model_path: str) -> None: """In place: replace qwen35moe's dense projections + embedding with native GGUF ops. @@ -718,6 +724,7 @@ def swap_linear(owner, attr, quant_type: int): "gguf_name_to_freetoken", "iter_gguf_weights", "convert_qwen35_to_gguf", + "is_gguf_model", "_MERGED_PARTS", "_EXPERT_SUFFIXES", ] diff --git a/python/freetoken/models/qwen3_5_moe/model.py b/python/freetoken/models/qwen3_5_moe/model.py index eba7fd24..12b15212 100644 --- a/python/freetoken/models/qwen3_5_moe/model.py +++ b/python/freetoken/models/qwen3_5_moe/model.py @@ -109,6 +109,18 @@ def __init__(self, config: ModelConfig): ) super().__init__() + # A GGUF checkpoint carries native block-quantized weights: swap the dense + # projections + embedding for GGUF-quant ops so the packed buffers have somewhere + # to land (routed experts stay on the offload cache). Mirrors gemma4/model.py. + from .gguf import convert_qwen35_to_gguf, is_gguf_model + + if is_gguf_model(config): + assert config.gguf_model_path is not None, ( + "expert_quant=='gguf' but ModelConfig.gguf_model_path is unset; the per-tensor " + "ggml types can only be read from the file" + ) + convert_qwen35_to_gguf(self, config, model_path=config.gguf_model_path) + def forward(self) -> torch.Tensor: output = self.model.forward(get_global_ctx().batch.input_ids) return self.lm_head.forward(output) From 4e70902b5da8ee8c3fa7fceb211f29319e549ffe Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 22:33:32 -0700 Subject: [PATCH 12/36] fix(gguf): stray gguf_expert_types kwarg on resolve_moe_cache_auto() Threading the bank types through the engine, I bulk-replaced both `quant_format=banks.quant_format` sites -- but one of them is resolve_moe_cache_auto(), which only sizes the slot pool from bank bytes and has no such parameter. Only the OffloadMoeCache constructor takes it. TypeError at cache-sizing time, after the weights and expert banks had already loaded. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/engine/engine.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 6771719d..018e9622 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -491,7 +491,6 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks) -> tuple[int kv_reserve_tokens=max(config.kv_reserve_tokens, min_reserve), page_size=page_tokens, quant_format=banks.quant_format, - gguf_expert_types=banks.gguf_expert_types, ) def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: From 6d0297f1c542d4bcbbc17af12ff70102b82e9a1a Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 22:39:39 -0700 Subject: [PATCH 13/36] fix(gguf): expert gate_up fusion interleaved experts, producing fluent garbage The MoE expert banks were built as cat([gate, up], dim=0).reshape(E, 2I, rb). gate and up each arrive from the reader as [rows, row_bytes] = [E*I, row_bytes(H)], and ggml's fastest-first dims [H, I, E] make E the slowest axis, so those rows are EXPERT-MAJOR. Concatenating on dim 0 therefore lays down every expert's gate before any up: expert 0's slice became its own gate rows plus expert 1's gate rows, and its up rows sat E*I rows away. Every routed expert in all 40 layers read another expert's weights. This loaded cleanly, ran at full speed (56 tok/s decode on a 4060), and generated confident nonsense -- 'r med media media media middle med med middle med...' -- which is why the fix comes with a byte-exact identity check rather than a reasoned argument: build the bank, then assert bank[e, :I] is gate rows [e*I, (e+1)*I) and bank[e, I:] is that same expert's up rows, for experts at both ends and the middle. The old code passes for expert 0's gate and expert 255's up and fails everywhere between, which is exactly the interleaving signature. Fixed by reshaping each part to [E, I, rb] and concatenating on the row axis (dim=1), so each expert's gate and up stay together. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- .../models/qwen3_5_moe/gguf_experts.py | 21 ++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/python/freetoken/models/qwen3_5_moe/gguf_experts.py b/python/freetoken/models/qwen3_5_moe/gguf_experts.py index 623f14c8..5948b804 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf_experts.py +++ b/python/freetoken/models/qwen3_5_moe/gguf_experts.py @@ -207,11 +207,22 @@ def _load(sink) -> None: # Emit gate_up bank once both gate and up are present. if layer in gate_buf and layer in up_buf: - gate_up_row_bytes = specs["gate_up"][0][2] - # Concatenate gate [H*I, row_bytes(E, type)] and up [H*I, row_bytes(E, type)] - # to get [H*2*I, row_bytes(E, type)], then reshape to [E, 2*I, row_bytes(H, type)] - combined = torch.cat([gate_buf[layer], up_buf[layer]], dim=0) - banks["gate_up"][layer].copy_(combined.reshape(E, 2 * I, gate_up_row_bytes)) + rb = specs["gate_up"][0][2] + # gate and up each arrive as [rows, row_bytes] = [E*I, row_bytes(H)], and + # ggml's fastest-first dims [H, I, E] make E the slowest axis, so those rows + # are EXPERT-MAJOR: expert e owns rows [e*I, (e+1)*I). + # + # The bank must be [E, 2I, row_bytes(H)] with each expert's own gate rows + # followed by its own up rows. So reshape to [E, I, rb] and concatenate on + # the ROW axis (dim=1), per expert. + # + # cat(dim=0) then reshape(E, 2I, rb) -- the obvious-looking version -- is + # wrong: it lays down every expert's gate before any up, so expert 0 would + # get its gate rows plus expert 1's gate rows, and up would be E*I rows + # away. That loads and runs at full speed and emits fluent nonsense. + g = gate_buf[layer].reshape(E, I, rb) + u = up_buf[layer].reshape(E, I, rb) + banks["gate_up"][layer].copy_(torch.cat([g, u], dim=1)) del gate_buf[layer], up_buf[layer] if tracker is not None: tracker.note(layer) From aec2db3eb36dfb52808fbf3e62bcbff6806bfc91 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 22:45:11 -0700 Subject: [PATCH 14/36] fix(gguf): missing (1+w) shift on Qwen3.5 RMSNorm weights Qwen3.5 stores these norm weights as w with an effective scale of 1+w, and FreeToken's GemmaRMSNorm multiplies by the raw weight -- so weight.py bakes the +1 in at load time (_is_gemma_norm / _GEMMA_NORM_SUFFIXES). The GGUF loader yielded them raw, so every RMSNorm in all 40 layers scaled by w instead of 1+w. w is typically 0.01-0.1, i.e. the residual stream was attenuated 10-100x at every norm, which is why generation was fluent nonsense rather than an error. Shifted: model.norm, input_layernorm, post_attention_layernorm, self_attn.q_norm, self_attn.k_norm. NOT shifted: linear_attn.norm (ssm_norm) -- the GDN gated norm is a standard w*x norm and weight.py excludes it too. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/qwen3_5_moe/gguf.py | 31 +++++++++++++++++---- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index 2da79559..1db251b9 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -320,6 +320,27 @@ def _scan_quant_types(model_path: str) -> dict[tuple[int, str], int]: return quant_types +# Qwen3.5 stores these RMSNorm weights as ``w`` with an effective scale of ``1 + w``, +# while FreeToken's GemmaRMSNorm multiplies by the raw weight -- so the +1 is baked in at +# load time. This mirrors _GEMMA_NORM_SUFFIXES / _is_gemma_norm in weight.py exactly, +# including the exclusion: linear_attn.norm (GDN gated norm, ssm_norm in the GGUF) is a +# standard w*x norm and must NOT be shifted. +# +# Getting this wrong is silent and total: w is typically 0.01-0.1, so scaling by w instead +# of 1+w attenuates every residual stream by 10-100x and the model emits fluent nonsense. +_PLUS_ONE_NORM_SUFFIXES = ( + "attn_norm.weight", # -> input_layernorm + "post_attention_norm.weight", # -> post_attention_layernorm + "attn_q_norm.weight", # -> self_attn.q_norm + "attn_k_norm.weight", # -> self_attn.k_norm +) + + +def _to_bf16_plus1(t) -> torch.Tensor: + """A (1 + weight) norm: dequantize then add 1, matching weight.py's load-time shift.""" + return _to_bf16(t) + 1.0 + + def _to_bf16(t) -> torch.Tensor: """Dequantize a GgufTensor (F32/F16/Q*) to a dense bf16 tensor of its torch shape.""" flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.bfloat16) @@ -416,7 +437,7 @@ def layer_of(name: str) -> int: yield "model.embed_tokens.qweight", t.packed() # IQ3_S packed table continue if name == "output_norm.weight": - yield "model.norm.weight", _to_bf16(t) + yield "model.norm.weight", _to_bf16_plus1(t) # (1 + w), see _PLUS_ONE_NORM_SUFFIXES continue if name == "output.weight": # The untied LM head, packed (Q6_K here). Ornith ships output.weight, so @@ -447,10 +468,10 @@ def layer_of(name: str) -> int: # A_log, dt_bias -> float32 # conv1d.weight -> float32, reshaped if suffix == "attn_norm.weight": - yield f"{base}.input_layernorm.weight", _to_bf16(t) + yield f"{base}.input_layernorm.weight", _to_bf16_plus1(t) continue if suffix == "post_attention_norm.weight": - yield f"{base}.post_attention_layernorm.weight", _to_bf16(t) + yield f"{base}.post_attention_layernorm.weight", _to_bf16_plus1(t) continue if suffix == "ffn_gate_inp.weight": yield f"{base}.mlp.gate.weight", _to_bf16(t) @@ -477,10 +498,10 @@ def layer_of(name: str) -> int: yield f"{base}.linear_attn.conv1d.weight", w continue if suffix == "attn_q_norm.weight": - yield f"{base}.self_attn.q_norm.weight", _to_bf16(t) + yield f"{base}.self_attn.q_norm.weight", _to_bf16_plus1(t) continue if suffix == "attn_k_norm.weight": - yield f"{base}.self_attn.k_norm.weight", _to_bf16(t) + yield f"{base}.self_attn.k_norm.weight", _to_bf16_plus1(t) continue # Shared expert (present on every layer, both kinds) -- must be handled BEFORE the From 853507b6c85d3c7fb4e71c6af18eb239f7e83c96 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 23:35:20 -0700 Subject: [PATCH 15/36] fix(gguf): ssm_a holds A not A_log; revert the (1+w) norm shift Two value-level bugs, both invisible to structural checks -- shapes, dtypes, byte identity and per-layer activation magnitudes were all correct while the model emitted fluent nonsense. 1. ssm_a. llama.cpp writes this tensor pre-transformed: it holds A = -exp(A_log), not A_log. Every value in Ornith IQ3_S is negative, spanning [-70.11, -0.0189]. gdn.py computes g = -A_log.exp() * softplus(a + dt_bias), so feeding it A gives -exp(-10.6) ~ -2.5e-5 and the recurrent decay gate collapses to zero in all 30 GDN layers. Now yields log(-A), so -exp(A_log) reproduces the stored A exactly, and raises if the tensor is not wholly negative rather than silently producing NaN. 2. The (1+w) RMSNorm shift added in the previous commit was wrong for GGUF and is reverted. HF stores raw w and weight.py adds 1.0, but llama.cpp's converter already folds the +1 into what it writes, so the shift double-counted. Measured means on Ornith IQ3_S: attn_norm 0.920, post_attention_norm 1.135, attn_q_norm 1.326, output_norm 2.640 -- all centred near 1, whereas raw w centres near 0.0X. Both were found by reading the actual stored values against what gdn.py and GemmaRMSNorm expect. Structural verification cannot catch a tensor that is the right shape and the wrong quantity. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/qwen3_5_moe/gguf.py | 54 +++++++++++++-------- 1 file changed, 34 insertions(+), 20 deletions(-) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index 1db251b9..65d65554 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -320,23 +320,23 @@ def _scan_quant_types(model_path: str) -> dict[tuple[int, str], int]: return quant_types -# Qwen3.5 stores these RMSNorm weights as ``w`` with an effective scale of ``1 + w``, -# while FreeToken's GemmaRMSNorm multiplies by the raw weight -- so the +1 is baked in at -# load time. This mirrors _GEMMA_NORM_SUFFIXES / _is_gemma_norm in weight.py exactly, -# including the exclusion: linear_attn.norm (GDN gated norm, ssm_norm in the GGUF) is a -# standard w*x norm and must NOT be shifted. +# NOTE on the Gemma-style (1 + weight) RMSNorm convention: do NOT apply it here. # -# Getting this wrong is silent and total: w is typically 0.01-0.1, so scaling by w instead -# of 1+w attenuates every residual stream by 10-100x and the model emits fluent nonsense. -_PLUS_ONE_NORM_SUFFIXES = ( - "attn_norm.weight", # -> input_layernorm - "post_attention_norm.weight", # -> post_attention_layernorm - "attn_q_norm.weight", # -> self_attn.q_norm - "attn_k_norm.weight", # -> self_attn.k_norm -) +# HF Qwen3.5 stores these norm weights as raw ``w`` with an effective scale of ``1 + w``, +# and FreeToken's GemmaRMSNorm multiplies by the raw buffer -- so weight.py adds 1.0 at +# load time (_is_gemma_norm / _GEMMA_NORM_SUFFIXES). llama.cpp's converter ALREADY folds +# the +1 into the tensors it writes, so a GGUF checkpoint arrives pre-shifted and adding +# it again double-counts. +# +# Measured on Ornith-1.5-35B IQ3_S (mean, min, max): +# blk.3.attn_norm.weight 0.920 0.701 1.660 +# blk.3.post_attention_norm.weight 1.135 0.004 1.279 +# blk.3.attn_q_norm.weight 1.326 0.684 1.883 +# output_norm.weight 2.640 0.763 3.484 +# Raw w would centre near 0.0X; these centre near 1, i.e. already 1+w. -def _to_bf16_plus1(t) -> torch.Tensor: +def _to_bf16(t) -> torch.Tensor: """A (1 + weight) norm: dequantize then add 1, matching weight.py's load-time shift.""" return _to_bf16(t) + 1.0 @@ -437,7 +437,7 @@ def layer_of(name: str) -> int: yield "model.embed_tokens.qweight", t.packed() # IQ3_S packed table continue if name == "output_norm.weight": - yield "model.norm.weight", _to_bf16_plus1(t) # (1 + w), see _PLUS_ONE_NORM_SUFFIXES + yield "model.norm.weight", _to_bf16(t) continue if name == "output.weight": # The untied LM head, packed (Q6_K here). Ornith ships output.weight, so @@ -468,10 +468,10 @@ def layer_of(name: str) -> int: # A_log, dt_bias -> float32 # conv1d.weight -> float32, reshaped if suffix == "attn_norm.weight": - yield f"{base}.input_layernorm.weight", _to_bf16_plus1(t) + yield f"{base}.input_layernorm.weight", _to_bf16(t) continue if suffix == "post_attention_norm.weight": - yield f"{base}.post_attention_layernorm.weight", _to_bf16_plus1(t) + yield f"{base}.post_attention_layernorm.weight", _to_bf16(t) continue if suffix == "ffn_gate_inp.weight": yield f"{base}.mlp.gate.weight", _to_bf16(t) @@ -485,7 +485,21 @@ def layer_of(name: str) -> int: yield f"{base}.linear_attn.norm.weight", _to_bf16(t) continue if suffix == "ssm_a": - yield f"{base}.linear_attn.A_log", _to_f32(t) + # llama.cpp writes this tensor already transformed: it holds A = -exp(A_log), + # not A_log. Measured on Ornith IQ3_S every value is negative, spanning + # [-70.11, -0.0189]. gdn.py computes + # g = -A_log.exp() * softplus(a + dt_bias) + # so handing it A directly gives -exp(-10.6) ~ -2.5e-5 and the recurrent decay + # gate collapses to zero in all 30 GDN layers -- healthy activation magnitudes, + # incoherent text. Invert the transform so -exp(A_log) reproduces the stored A. + a = _to_f32(t) + if not bool((a < 0).all()): + raise ValueError( + f"{name}: expected llama.cpp's pre-transformed A = -exp(A_log) (all " + f"negative); got min={float(a.min())} max={float(a.max())}, which would " + f"make log(-A) NaN" + ) + yield f"{base}.linear_attn.A_log", torch.log(-a) continue if suffix == "ssm_dt.bias": yield f"{base}.linear_attn.dt_bias", _to_f32(t) @@ -498,10 +512,10 @@ def layer_of(name: str) -> int: yield f"{base}.linear_attn.conv1d.weight", w continue if suffix == "attn_q_norm.weight": - yield f"{base}.self_attn.q_norm.weight", _to_bf16_plus1(t) + yield f"{base}.self_attn.q_norm.weight", _to_bf16(t) continue if suffix == "attn_k_norm.weight": - yield f"{base}.self_attn.k_norm.weight", _to_bf16_plus1(t) + yield f"{base}.self_attn.k_norm.weight", _to_bf16(t) continue # Shared expert (present on every layer, both kinds) -- must be handled BEFORE the From 711c34b6cbbcace9b7d02a3211fcb226c1bded40 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 23:41:53 -0700 Subject: [PATCH 16/36] fix(gguf): un-tile V heads -- llama.cpp reorders them, FreeToken expects HF order When num_k_heads != num_v_heads, llama.cpp's converter reorders every tensor that indexes the V-head dimension so ggml_repeat can replace an interleaved repeat (conversion/qwen.py, _LinearAttentionVReorderBase): HF / FreeToken (grouped by K head): [G0_v0, G0_v1, G1_v0, G1_v1, ...] GGUF (tiled for ggml): [G0_v0, G1_v0, ..., G0_v1, G1_v1, ...] FreeToken's GDN pairs K head k with V heads [k*R, (k+1)*R) -- grouped -- and the safetensors loader applies no reorder because HF is already grouped. So every K/V pairing in all 30 GDN layers was wrong. Ornith: 16 K heads, 32 V heads, R=2, head_v_dim=128. Inverted on all seven tensor groups llama.cpp touches: attn_qkv (V rows only), attn_gate, ssm_beta, ssm_alpha, ssm_a, ssm_dt.bias, ssm_conv1d (V channels only) and ssm_out (columns). Row permutations are applied directly to the packed bytes -- each output row is an independent run of quant blocks, so reordering rows never splits one. ssm_out is the exception: its COLUMNS carry the V dimension, and a 128-wide head straddles the 256-element IQ3_S blocks, so it cannot be permuted while packed. That one tensor is dequantized to dense bf16 (16 MiB per GDN layer, ~503 MiB total) and convert_qwen35_to_gguf leaves linear_attn.out_proj as a dense Linear. Like the A_log transform, this is invisible to every structural check: shapes, dtypes, byte identity against the source, and per-layer activation magnitudes are all correct while the K/V pairing is scrambled. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/qwen3_5_moe/gguf.py | 108 ++++++++++++++++++-- 1 file changed, 98 insertions(+), 10 deletions(-) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index 65d65554..fa0fec0a 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -336,6 +336,50 @@ def _scan_quant_types(model_path: str) -> dict[tuple[int, str], int]: # Raw w would centre near 0.0X; these centre near 1, i.e. already 1+w. + +# -------------------------------------------------------------------------------------- +# V-head order: llama.cpp writes TILED, FreeToken (like HF) wants GROUPED +# -------------------------------------------------------------------------------------- +# When num_k_heads != num_v_heads, llama.cpp's converter reorders every tensor that indexes +# the V-head dimension (conversion/qwen.py, _LinearAttentionVReorderBase) so ggml_repeat can +# replace an interleaved repeat: +# +# HF / FreeToken (grouped by K head): [G0_v0, G0_v1, G1_v0, G1_v1, ...] +# GGUF (tiled for ggml): [G0_v0, G1_v0, ..., G0_v1, G1_v1, ...] +# +# FreeToken's GDN pairs K head k with V heads [k*R, (k+1)*R), i.e. grouped -- the safetensors +# loader applies no reorder because HF is already grouped. So a GGUF checkpoint must be +# un-tiled on load or every K/V pairing in all 30 GDN layers is wrong. That is invisible to +# shape, dtype and magnitude checks: activations stay healthy and the text is nonsense. +# +# Ornith: num_k_heads=16, num_v_heads=32 -> num_v_per_k=2, head_v_dim=128. + + +def _ungroup_v(t: torch.Tensor, dim: int, num_k_heads: int, num_v_per_k: int, head_dim: int): + """Tiled -> grouped along ``dim``: the inverse of llama.cpp's _reorder_v_heads. + + The forward transform views the axis as [K, R, D] and swaps K/R. The inverse is the same + operation with the two counts exchanged: view as [R, K, D] and swap back. + """ + shape = list(t.shape) + if dim < 0: + dim += len(shape) + view = shape[:dim] + [num_v_per_k, num_k_heads, head_dim] + shape[dim + 1:] + out = t.reshape(*view) + perm = list(range(len(view))) + perm[dim], perm[dim + 1] = perm[dim + 1], perm[dim] + return out.permute(*perm).contiguous().reshape(*shape) + + +def _ungroup_packed_rows(packed: torch.Tensor, num_k_heads: int, num_v_per_k: int, head_dim: int): + """Un-tile whole ROWS of a packed [out, row_bytes] tensor. + + Safe on quantized data: each output row is an independent run of blocks over the input + dim, so permuting rows never splits a block. (Permuting COLUMNS would -- see ssm_out.) + """ + return _ungroup_v(packed, 0, num_k_heads, num_v_per_k, head_dim) + + def _to_bf16(t) -> torch.Tensor: """A (1 + weight) norm: dequantize then add 1, matching weight.py's load-time shift.""" return _to_bf16(t) + 1.0 @@ -416,6 +460,14 @@ def iter_gguf_weights( else 8192 ) gdn_conv_kernel = gdn_group.conv_kernel_dim if gdn_group else 4 + # V-head un-tiling geometry (see _ungroup_v). Only needed when the GDN has fewer K + # heads than V heads, which is exactly when llama.cpp reorders. + _vK = gdn_group.num_key_heads if gdn_group else 0 + _vN = gdn_group.num_value_heads if gdn_group else 0 + _vD = gdn_group.value_head_dim if gdn_group else 0 + _vR = (_vN // _vK) if _vK else 1 + _untile = bool(_vK and _vN and _vK != _vN) + _qk_rows = 2 * gdn_group.num_key_heads * gdn_group.key_head_dim if gdn_group else 0 # Scan quant types once to determine which fusion groups are mixed-quant. quant_map = _scan_quant_types(model_path) @@ -493,6 +545,8 @@ def layer_of(name: str) -> int: # gate collapses to zero in all 30 GDN layers -- healthy activation magnitudes, # incoherent text. Invert the transform so -exp(A_log) reproduces the stored A. a = _to_f32(t) + if _untile: + a = _ungroup_v(a, 0, _vK, _vR, 1) if not bool((a < 0).all()): raise ValueError( f"{name}: expected llama.cpp's pre-transformed A = -exp(A_log) (all " @@ -502,13 +556,21 @@ def layer_of(name: str) -> int: yield f"{base}.linear_attn.A_log", torch.log(-a) continue if suffix == "ssm_dt.bias": - yield f"{base}.linear_attn.dt_bias", _to_f32(t) + dt = _to_f32(t) + if _untile: + dt = _ungroup_v(dt, 0, _vK, _vR, 1) + yield f"{base}.linear_attn.dt_bias", dt continue if suffix == "ssm_conv1d.weight": # F32 in the file, but _DepthwiseConv1d allocates at the model dtype: gdn.py # exempts only A_log / dt_bias from the downcast, so this one is bf16. Reshape # to [conv_dim, 1, kernel] -- gdn.py's _conv_weight() does .squeeze(1). - w = _to_bf16(t).reshape(gdn_attn_qkv_size, 1, gdn_conv_kernel) + w = _to_bf16(t).reshape(gdn_attn_qkv_size, gdn_conv_kernel) + if _untile: + # channels are [q | k | v]; only the V block is tiled + qk, v = w[:_qk_rows], w[_qk_rows:] + w = torch.cat([qk, _ungroup_v(v, 0, _vK, _vR, _vD)], dim=0) + w = w.reshape(gdn_attn_qkv_size, 1, gdn_conv_kernel) yield f"{base}.linear_attn.conv1d.weight", w continue if suffix == "attn_q_norm.weight": @@ -585,15 +647,38 @@ def layer_of(name: str) -> int: # and out_proj else: if suffix == "attn_qkv.weight": - in_proj_buf.setdefault(layer, {})["qkv"] = t.packed() + w = t.packed() + if _untile: + # rows are [q | k | v]; only the V rows are tiled. Row permutation is + # safe on packed data -- each row is its own run of blocks. + w = torch.cat( + [w[:_qk_rows], _ungroup_packed_rows(w[_qk_rows:], _vK, _vR, _vD)], dim=0 + ) + in_proj_buf.setdefault(layer, {})["qkv"] = w elif suffix == "attn_gate.weight": - in_proj_buf.setdefault(layer, {})["gate"] = t.packed() - elif suffix == "ssm_beta.weight": - in_proj_buf.setdefault(layer, {})["beta"] = t.packed() - elif suffix == "ssm_alpha.weight": - in_proj_buf.setdefault(layer, {})["alpha"] = t.packed() + w = t.packed() + if _untile: + w = _ungroup_packed_rows(w, _vK, _vR, _vD) + in_proj_buf.setdefault(layer, {})["gate"] = w + elif suffix in ("ssm_beta.weight", "ssm_alpha.weight"): + w = t.packed() + if _untile: + # one row per V head -> head_dim 1 + w = _ungroup_packed_rows(w, _vK, _vR, 1) + in_proj_buf.setdefault(layer, {})[ + "beta" if suffix.startswith("ssm_beta") else "alpha" + ] = w elif suffix == "ssm_out.weight": - yield f"{base}.linear_attn.out_proj.qweight", t.packed() + # out_proj consumes the V dimension along its COLUMNS, and llama.cpp tiled + # those columns. A column permutation cannot be done on packed data -- a + # 128-wide head straddles the 256-element quant blocks -- so this one tensor + # is dequantized to dense bf16. Cost: out*in*2 bytes per GDN layer + # (2048*4096*2 = 16 MiB, ~503 MiB over 30 layers). convert_qwen35_to_gguf + # therefore leaves linear_attn.out_proj as a dense Linear. + w = _to_bf16(t) + if _untile: + w = _ungroup_v(w, 1, _vK, _vR, _vD) + yield f"{base}.linear_attn.out_proj.weight", w else: continue # unmapped for GDN layers @@ -731,7 +816,10 @@ def swap_linear(owner, attr, quant_type: int): ], has_bias=False, ) - swap_linear(layer.linear_attn, "out_proj", qt(layer_idx, "ssm_out.weight")) + # linear_attn.out_proj is deliberately NOT swapped: its columns index the + # V-head dimension, which llama.cpp tiled, and un-tiling columns needs dense + # values (a 128-wide head straddles the quant blocks). iter_gguf_weights yields + # it as dense bf16 ".weight", so the constructed Linear must stay dense. # Shared expert: gate|up fuse when they share a type (they do in every quant level # seen so far); down is independent and does vary (Q4_K on IQ3_M's first layers). From 811a2b3511dd5203e4b627ccfba897943908b2ac Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Sun, 23 Aug 2026 23:43:19 -0700 Subject: [PATCH 17/36] fix(gguf): ssm_out needs the CUDA dequantizer (pure-torch path is Q4_0/Q6_K only) _to_bf16 goes through models/gguf/dequant.dequantize, which implements only Q4_0 and Q6_K -- it is the reference/test path. ssm_out is IQ3_S, so materializing it dense for the column un-tiling raised NotImplementedError. _dequant_any routes through ggml_dequantize, which covers all 19 types, and returns a CPU tensor so the normal load path places it. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/qwen3_5_moe/gguf.py | 27 ++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index fa0fec0a..9569926c 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -46,6 +46,7 @@ RotaryConfig, ) from freetoken.models.gguf.dequant import ( + GGML_UNQUANTIZED as GGML_UNQUANTIZED_SET, GGML_IQ3_S, GGML_NAME, GGML_Q4_K, @@ -391,6 +392,30 @@ def _to_bf16(t) -> torch.Tensor: return flat.reshape(t.shape) +def _dequant_any(t) -> torch.Tensor: + """Dequantize a GgufTensor of ANY ggml type to dense bf16, via the CUDA kernel. + + The pure-torch ``dequantize`` in models/gguf/dequant.py implements only Q4_0 and Q6_K + (it is the reference/test path). ``ggml_dequantize`` covers all 19 quant types, so use + it for the one tensor that genuinely has to be materialized dense -- ssm_out, whose + columns need un-tiling. Round-trips through the GPU; the result is a CPU tensor so the + normal load path places it. + """ + from freetoken.kernel.gguf import ggml_dequantize + + if t.ggml_type in GGML_UNQUANTIZED_SET: + return _to_bf16(t) + if not torch.cuda.is_available(): + raise RuntimeError( + f"{t.name}: needs dense dequantization of ggml type " + f"{GGML_NAME.get(t.ggml_type, t.ggml_type)}, which only the CUDA kernel " + f"implements, but no CUDA device is available" + ) + out_f, in_f = t.shape[0], t.shape[1] + packed = t.packed().reshape(out_f, row_bytes(in_f, t.ggml_type)).cuda() + return ggml_dequantize(packed, t.ggml_type, out_f, in_f, torch.bfloat16).cpu() + + def _to_f32(t) -> torch.Tensor: """Dequantize a GgufTensor (F32/F16/Q*) to a dense float32 tensor of its torch shape.""" flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.float32) @@ -675,7 +700,7 @@ def layer_of(name: str) -> int: # is dequantized to dense bf16. Cost: out*in*2 bytes per GDN layer # (2048*4096*2 = 16 MiB, ~503 MiB over 30 layers). convert_qwen35_to_gguf # therefore leaves linear_attn.out_proj as a dense Linear. - w = _to_bf16(t) + w = _dequant_any(t) if _untile: w = _ungroup_v(w, 1, _vK, _vR, _vD) yield f"{base}.linear_attn.out_proj.weight", w From 4163a935c8b55db5ac9d48de55f0f2f57d1c47ae Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Mon, 24 Aug 2026 00:42:35 -0700 Subject: [PATCH 18/36] docs(gguf): mixed-bank error now names the cause and the --pure workaround The refusal said which layers disagreed but not why it is impossible or what to do instead. It now states that the slot pool is one allocation and moe_vec.cuh indexes it with no padding allowance, that llama.cpp's *_M and *_XXS levels raise the precision of the first few layers' ffn_down_exps, and that llama-quantize --pure produces a uniform checkpoint that loads. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/moe/expert_banks.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 5b3138dc..9a985aa6 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -320,11 +320,14 @@ def _gguf_banks(model_path, model_config, device, dtype, dummy, parallel=False, for t in distinct } raise NotImplementedError( - f"GGUF expert bank {name!r} mixes ggml types across layers ({spread}); the " - f"offload slot pool is one allocation with a single row stride, so this " - f"checkpoint cannot be served. Use a quant level whose expert banks are " - f"uniform (for Ornith-1.5-35B: IQ3_S or IQ3_XXS are; IQ3_M/IQ2_M/IQ2_XXS/" - f"IQ1_S split ffn_down_exps across two types)." + f"GGUF expert bank {name!r} mixes ggml types across layers ({spread}). " + f"The offload slot pool is a single allocation shared by every layer and " + f"moe_vec.cuh addresses it as expert * nrows * (ncols / qk) with no padding " + f"allowance, so one pool cannot hold two row strides. " + f"llama.cpp's mixed quant levels (the *_M and *_XXS families) raise the " + f"precision of the first few layers' ffn_down_exps, which is what trips this. " + f"Re-quantize with `llama-quantize --pure` to get one type throughout, or pick " + f"a level that is already uniform." ) resolved[name] = distinct[0] From 439085b396b531d0c69f3d9fcb6ec5cc8f9c99a5 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Mon, 24 Aug 2026 01:03:15 -0700 Subject: [PATCH 19/36] feat(gguf): serve the dense qwen35 variant (Qwen3.8-27B) Same hybrid GDN/full-attention decoder as qwen35moe with a plain SwiGLU MLP instead of routed experts, so it reuses the package, the model classes and this adapter. - _kv() prefixes metadata keys with the checkpoint's own general.architecture rather than a hardcoded "qwen35moe", since the dense arch string is "qwen35". - parse_gguf_config reads feed_forward_length into intermediate_size when expert_count is absent, and leaves expert_quant "none" so the engine does not go looking for offload banks that do not exist. - is_gguf_model now keys on gguf_model_path instead of expert_quant, otherwise the dense path would skip the op swap entirely. - iter_gguf_weights maps ffn_gate/ffn_up/ffn_down onto mlp.gate_up_proj and mlp.down_proj (Qwen3_5DenseMLP subclasses _SharedExpert, so the same fused shape applies). Handled before the per-layer-kind branch, whose else: continue would otherwise drop them. - convert_qwen35_to_gguf derives the qkv and in_proj split widths from the config instead of Ornith's constants. Ornith-1.5 is 16 q heads / 2 kv / 32 v; Qwen3.8-27B is 24 / 4 / 48, so the hardcoded [8192, 512, 512] and [8192, 4096, 32, 32] were wrong for anything else. Dense checkpoints have no expert banks, so the uniform-bank restriction does not apply to them: Qwen3.8-27B Q4_K_M mixes Q4_K and Q6_K across layers for attn_qkv, attn_v and ffn_down, and per-tensor quant types already handle that. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/gguf/config.py | 4 + python/freetoken/models/qwen3_5_moe/gguf.py | 112 +++++++++++++++++--- python/freetoken/models/register.py | 7 ++ 3 files changed, 109 insertions(+), 14 deletions(-) diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 29d7efda..3d269eeb 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -20,6 +20,10 @@ GGUF_ARCH_TO_REGISTRY: dict[str, str] = { "gemma4": "Gemma4GGUFForCausalLM", "qwen35moe": "Qwen35MoeGGUFForCausalLM", + # Dense sibling (Qwen3.8-27B): same hybrid GDN/full-attention decoder, a plain SwiGLU + # MLP instead of routed experts. Same model classes and the same GGUF adapter; the + # config's expert_count is absent so moe_enabled comes out False. + "qwen35": "Qwen35GGUFForCausalLM", } diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index 9569926c..b1bb7bc4 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -62,10 +62,16 @@ def _kv(shim: "GgufConfigShim", key: str, default: Any = None) -> Any: - """Read ``qwen35moe.`` from the GGUF metadata.""" - val = shim.metadata.get(f"{_ARCH}.{key}", default) + """Read ``.`` from the GGUF metadata. + + The prefix is the checkpoint's own ``general.architecture``: "qwen35moe" for the MoE + variant, "qwen35" for the dense one (e.g. Qwen3.8-27B). Same geometry keys either way. + """ + val = shim.metadata.get(f"{shim.model_type}.{key}", default) if val is None and default is None: - raise ValueError(f"GGUF {shim.model_path}: missing required key {_ARCH}.{key}") + raise ValueError( + f"GGUF {shim.model_path}: missing required key {shim.model_type}.{key}" + ) return val @@ -108,10 +114,15 @@ def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: rotary_dim = int(_kv(shim, "rope.dimension_count")) max_pos = int(_kv(shim, "context_length")) + # Dense variants (qwen35, e.g. Qwen3.8-27B) carry no expert_* keys at all: every + # decoder layer gets a plain SwiGLU MLP sized by feed_forward_length instead of the + # routed block plus shared expert. num_experts = int(_kv(shim, "expert_count", 0)) experts_per_tok = int(_kv(shim, "expert_used_count", 0)) moe_inter = int(_kv(shim, "expert_feed_forward_length", 0)) shared_inter = int(_kv(shim, "expert_shared_feed_forward_length", 0)) + dense_inter = int(_kv(shim, "feed_forward_length", 0)) + moe_enabled = num_experts > 0 # GDN geometry. state_size is the per-head dim; group_count is the number of k heads # and time_step_rank the number of v heads (see module docstring for the arithmetic @@ -173,7 +184,7 @@ def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: head_dim=head_dim, hidden_size=hidden_size, vocab_size=shim.vocab_size, - intermediate_size=0, # every layer is MoE in qwen35moe + intermediate_size=0 if moe_enabled else dense_inter, hidden_act="silu", rms_norm_eps=rms_eps, tie_word_embeddings=shim.tie_word_embeddings, @@ -183,15 +194,21 @@ def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: moe_intermediate_size=moe_inter, shared_expert_intermediate_size=shared_inter, norm_topk_prob=True, - moe_enabled=num_experts > 0, + moe_enabled=moe_enabled, use_qk_norm=True, - model_type=_ARCH, - architectures=["Qwen35MoeGGUFForCausalLM"], + model_type=shim.model_type, + architectures=list(shim.architectures), vision_config=None, image_token_id=None, attention_groups=groups, - expert_quant="gguf", - gguf_expert_types=_uniform_expert_types(shim.model_path, num_layers), + # Only the MoE variant has offload expert banks; a dense checkpoint must not + # advertise expert_quant="gguf" or the engine would go looking for banks that do + # not exist. is_gguf_model() keys on gguf_model_path instead, so the op swap still + # runs for both. + expert_quant="gguf" if moe_enabled else "none", + gguf_expert_types=( + _uniform_expert_types(shim.model_path, num_layers) if moe_enabled else None + ), gguf_model_path=shim.model_path, weight_block_size=None, attn_quant="gguf", @@ -500,7 +517,8 @@ def iter_gguf_weights( # Per-layer fusion buffers: layer -> {slot: packed[out, row_bytes]}. qkv_buf: dict[int, dict[str, torch.Tensor]] = {} # full-attn qkv in_proj_buf: dict[int, dict[str, torch.Tensor]] = {} # GDN in_proj (qkv+gate+beta+alpha) - gate_up_buf: dict[int, dict[str, torch.Tensor]] = {} # shared_expert gate_up + gate_up_buf: dict[int, dict[str, torch.Tensor]] = {} # shared_expert gate_up (MoE) + dense_mlp_buf: dict[int, dict[str, torch.Tensor]] = {} # mlp gate_up (dense qwen35) def layer_of(name: str) -> int: return int(name.split(".")[1]) @@ -605,6 +623,34 @@ def layer_of(name: str) -> int: yield f"{base}.self_attn.k_norm.weight", _to_bf16(t) continue + # Dense variant (qwen35): a plain SwiGLU MLP per layer instead of the routed block. + # Qwen3_5DenseMLP subclasses _SharedExpert, so the targets are mlp.gate_up_proj + # (gate|up fused) and mlp.down_proj. Same placement rule as the shared expert below: + # these appear on both layer kinds, so they must be consumed before the layer-kind + # branch whose `else: continue` drops anything it does not recognise. + if suffix in ("ffn_gate.weight", "ffn_up.weight", "ffn_down.weight"): + if suffix == "ffn_down.weight": + yield f"{base}.mlp.down_proj.qweight", t.packed() + else: + dense_mlp_buf.setdefault(layer, {})[ + "gate" if suffix == "ffn_gate.weight" else "up" + ] = t.packed() + d = dense_mlp_buf[layer] + if "gate" in d and "up" in d: + types = [ + quant_map.get((layer, "ffn_gate.weight")), + quant_map.get((layer, "ffn_up.weight")), + ] + if len(set(types)) == 1: + yield f"{base}.mlp.gate_up_proj.qweight", torch.cat( + [d["gate"], d["up"]], dim=0 + ) + else: + yield f"{base}.mlp.gate_up_proj.qweight_0", d["gate"] + yield f"{base}.mlp.gate_up_proj.qweight_1", d["up"] + del dense_mlp_buf[layer] + continue + # Shared expert (present on every layer, both kinds) -- must be handled BEFORE the # per-layer-kind branch below, whose `else: continue` swallows any suffix it does # not recognize. @@ -746,11 +792,17 @@ def layer_of(name: str) -> int: assert not qkv_buf, f"incomplete full-attn qkv groups: {sorted(qkv_buf)}" assert not in_proj_buf, f"incomplete GDN in_proj groups: {sorted(in_proj_buf)}" assert not gate_up_buf, f"incomplete shared_expert gate_up groups: {sorted(gate_up_buf)}" + assert not dense_mlp_buf, f"incomplete dense mlp gate_up groups: {sorted(dense_mlp_buf)}" def is_gguf_model(config: ModelConfig) -> bool: - """True when this config came from a GGUF checkpoint (native block-quant path).""" - return getattr(config, "expert_quant", "none") == "gguf" + """True when this config came from a GGUF checkpoint (native block-quant path). + + Keys on gguf_model_path rather than expert_quant: the dense qwen35 variant has no + expert banks and therefore no expert_quant="gguf", but still needs its dense ops + swapped for GGUF ops. + """ + return getattr(config, "gguf_model_path", None) is not None def convert_qwen35_to_gguf(model, config: ModelConfig, *, model_path: str) -> None: @@ -777,6 +829,25 @@ def convert_qwen35_to_gguf(model, config: ModelConfig, *, model_path: str) -> No if isinstance(config.attention_group_for_layer(lid), FullAttentionGroupConfig) } + # Split widths come from the config, not constants: Ornith-1.5 is 16 q heads / 2 kv / + # 32 v heads, Qwen3.8-27B is 24 / 4 / 48. Hardcoding either breaks the other. + _qkv_split = [ + config.num_qo_heads * config.head_dim * 2, # q is gated, hence *2 + config.num_kv_heads * config.head_dim, + config.num_kv_heads * config.head_dim, + ] + _g = config.linear_attention_group() + _in_proj_split = ( + [ + 2 * _g.num_key_heads * _g.key_head_dim + _g.num_value_heads * _g.value_head_dim, + _g.num_value_heads * _g.value_head_dim, + _g.num_value_heads, + _g.num_value_heads, + ] + if _g is not None + else [] + ) + def qt(layer: int, suffix: str) -> int: """The ggml type of one tensor, straight from the file. @@ -818,7 +889,7 @@ def swap_linear(owner, attr, quant_type: int): # q/k are I-quants), so this is normally the GGUFMergedLinear path. layer.self_attn.qkv_proj = gguf_merged_or_plain( config.hidden_size, - [8192, 512, 512], + _qkv_split, [ qt(layer_idx, "attn_q.weight"), qt(layer_idx, "attn_k.weight"), @@ -832,7 +903,7 @@ def swap_linear(owner, attr, quant_type: int): # _in_proj_split = [conv_dim, value_dim, num_v_heads, num_v_heads]. layer.linear_attn.in_proj = gguf_merged_or_plain( config.hidden_size, - [8192, 4096, 32, 32], + _in_proj_split, [ qt(layer_idx, "attn_qkv.weight"), qt(layer_idx, "attn_gate.weight"), @@ -846,6 +917,19 @@ def swap_linear(owner, attr, quant_type: int): # values (a 128-wide head straddles the quant blocks). iter_gguf_weights yields # it as dense bf16 ".weight", so the constructed Linear must stay dense. + if not config.moe_enabled: + # Dense qwen35: one SwiGLU MLP per layer (Qwen3_5DenseMLP), no routed experts + # and no shared expert. + I = config.intermediate_size + layer.mlp.gate_up_proj = gguf_merged_or_plain( + config.hidden_size, + [I, I], + [qt(layer_idx, "ffn_gate.weight"), qt(layer_idx, "ffn_up.weight")], + has_bias=False, + ) + swap_linear(layer.mlp, "down_proj", qt(layer_idx, "ffn_down.weight")) + continue + # Shared expert: gate|up fuse when they share a type (they do in every quant level # seen so far); down is independent and does vary (Q4_K on IQ3_M's first layers). I = config.shared_expert_intermediate_size diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 03ab68e6..c062df0c 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -114,6 +114,13 @@ class ModelSpec: parse_config="parse_gguf_config", iter_weights="iter_gguf_weights", ), + # Dense qwen35 GGUF (Qwen3.8-27B): same package and classes, moe_enabled==False. + "Qwen35GGUFForCausalLM": ModelSpec( + "freetoken.models.qwen3_5_moe", + "Qwen3_5MoEForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + ), "GptOssForCausalLM": ModelSpec( "freetoken.models.gpt_oss", "GptOssForCausalLM", From 39a8c6e09bb51ad5c439169fedc5c2a3d5d99657 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Mon, 24 Aug 2026 01:39:42 -0700 Subject: [PATCH 20/36] docs: document the GGUF architectures and quant types in models.md models.md still said native GGUF was Gemma-4 only, which stopped being true once qwen35moe and qwen35 landed. Adds a GGUF section listing the architectures by their general.architecture string, the quant types with the prefill and decode path per family, and the two constraints a user hits when picking a file: routed-expert banks must use one ggml type across all layers (llama-quantize --pure avoids it, dense models are exempt), and GGUF paths are TP=1 with any NextN/MTP block dropped. --- docs/models.md | 36 +++++++++++++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/docs/models.md b/docs/models.md index e4850a12..21120885 100644 --- a/docs/models.md +++ b/docs/models.md @@ -1,8 +1,9 @@ # Supported models -FreeToken loads HF safetensors checkpoints directly (plus native GGUF for -Gemma-4). The checkpoints below are known-good — the prebuilt kernels are tuned -for them; other checkpoints of the same architectures work too. +FreeToken loads HF safetensors checkpoints directly, plus native GGUF for the +architectures listed under [GGUF](#gguf) below. The checkpoints below are +known-good — the prebuilt kernels are tuned for them; other checkpoints of the +same architectures work too. | Model | HF checkpoints | |---|---| @@ -17,6 +18,35 @@ for them; other checkpoints of the same architectures work too. | MiniMax-M2.5 | [nvidia/MiniMax-M2.5-NVFP4](https://huggingface.co/nvidia/MiniMax-M2.5-NVFP4) | | Muse-Glimmer | [meta-models/Muse-Glimmer-30B](https://huggingface.co/meta-models/Muse-Glimmer-30B), [RedHatAI/Muse-Glimmer-30B-NVFP4](https://huggingface.co/RedHatAI/Muse-Glimmer-30B-NVFP4) | +## GGUF + +Native GGUF, meaning the block-quantized weights are kept packed and dequantized inside +the kernels rather than expanded to bf16 at load. + +| GGUF `general.architecture` | Covers | +|---|---| +| `gemma4` | Gemma-4 | +| `qwen35moe` | Qwen3.5 / Qwen3.6 MoE (e.g. Qwen3.6-35B-A3B, Qwen3.5-122B-A10B) | +| `qwen35` | Qwen3.5 / Qwen3.6 dense (e.g. Qwen3.6-27B, Qwen3.5-9B) | + +Quant types follow what the vendored kernels in `csrc/gguf/` implement: + +- Standard and K-quants (Q4_0, Q4_1, Q5_0, Q5_1, Q8_0, Q2_K through Q6_K) use MMQ for + prefill and MMVQ for decode. +- I-quants (IQ1_S, IQ1_M, IQ2_XXS, IQ2_XS, IQ2_S, IQ3_XXS, IQ3_S, IQ4_NL, IQ4_XS) have no + MMQ kernel, so prefill dequantizes and runs a plain matmul; decode uses MMVQ. + +Two constraints worth knowing before picking a file: + +- A MoE checkpoint's routed-expert banks must use one ggml type across every layer. The GPU + slot pool is a single allocation with a single row stride, so a bank that changes type + between layers cannot be served and the load fails with the offending layers named. + llama.cpp's `_M` and `_XXS` levels raise the precision of the first few layers' + `ffn_down_exps` and hit this; `llama-quantize --pure` produces a checkpoint that loads. + Dense models have no expert banks and are unaffected. +- GGUF paths are TP=1 only, and a NextN/MTP block in the checkpoint is dropped (served + text-only, no speculative decoding). + ## MoE backends `ft serve --moe-backend {auto,fused,offload,cpu,hybrid}`: From 08f702efc69e994f72d3e281ae5dfd91d4132401 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Mon, 24 Aug 2026 01:57:38 -0700 Subject: [PATCH 21/36] fix(gguf): the no-routed-experts assert fired on dense checkpoints iter_gguf_weights asserted include_moe_experts was False unconditionally, because the MoE variant serves its routed experts from the offload cache. A dense qwen35 checkpoint has no routed experts, so the engine correctly asks for everything and the assert tripped during load. Now conditional on the checkpoint actually having experts. --- python/freetoken/models/qwen3_5_moe/gguf.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index b1bb7bc4..a7dd96db 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -474,9 +474,13 @@ def iter_gguf_weights( from freetoken.models.gguf.reader import iter_gguf_tensors from freetoken.utils import cached_load_hf_config - assert not include_moe_experts, ( - "qwen35moe GGUF stores experts as IQ3_S and only supports the offload backend; " - "experts are loaded into the offload cache via the expert-bank loader." + # Only the MoE variant keeps its routed experts out of this iterator; they come from + # the offload cache instead. A dense qwen35 checkpoint has no routed experts at all, so + # the engine legitimately asks for "everything" and the assert must not fire. + config_moe = int(_kv(cached_load_hf_config(model_path), "expert_count", 0)) > 0 + assert not (config_moe and include_moe_experts), ( + "qwen35moe GGUF keeps its routed experts in the offload cache; they are loaded by " + "the expert-bank loader, not by iter_gguf_weights." ) assert include_non_moe _require_tp1("weight loading") From bd49652f00e1145d142abb5f7da937c89b4925e1 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Mon, 24 Aug 2026 01:56:27 -0700 Subject: [PATCH 22/36] fix(gguf): tokenizer arch map was missing the dense qwen35 entry Same failure as qwen35moe had: convert_gguf_tokenizer is keyed on transformers' own model_type, so an unmapped GGUF arch string raises KeyError and kills the detokenizer worker during load. The dense variant shares the MoE one's tokenizer exactly (gpt2 BPE, pre=qwen35, 248320 tokens, eos <|im_end|>), so it maps to the same qwen2 converter and the same stop tokens. (cherry picked from commit e43bf646c899125e227ad8569398666c70c6b971) --- python/freetoken/models/gguf/tokenizer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 5ff08188..03b24616 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -17,7 +17,7 @@ # heard of raises KeyError inside convert_gguf_tokenizer. qwen35moe is a GPT2-style BPE # with merges (tokenizer.ggml.model == "gpt2", pre == "qwen35"), which the qwen2 converter # handles; there is no qwen3.5-specific converter and it would be the same BPE anyway. -_TOKENIZER_ARCH = {"gemma4": "gemma4_text", "qwen35moe": "qwen2"} +_TOKENIZER_ARCH = {"gemma4": "gemma4_text", "qwen35moe": "qwen2", "qwen35": "qwen2"} # Per-arch chat/stop tokens, in preference order: the first one present in the vocab # becomes eos (so chat generation halts on the turn end rather than the formal document @@ -27,6 +27,8 @@ _STOP_TOKENS: dict[str, tuple[str, ...]] = { "gemma4": ("", ""), "qwen35moe": ("<|im_end|>", "<|endoftext|>"), + # Dense sibling: same vocab and same chat markers as the MoE variant. + "qwen35": ("<|im_end|>", "<|endoftext|>"), } From 834ed4d92f193be68aee6696829d2a64d4bfb834 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Mon, 24 Aug 2026 02:38:20 -0700 Subject: [PATCH 23/36] feat(gguf): multi-shard GGUF reading + qwen3moe architecture Multi-shard was the availability blocker: every large model ships split (-00001-of-000NN), so quant-type support was unreachable for exactly the checkpoints it mattered for. reader.py now resolves any shard (or a directory of shards) to the ordered set, reads metadata/arch/tokenizer from shard 1 only, and aggregates tensor tables across shards. Validated against llama.cpp's convention: split.no is 0-based while filenames are 1-based, and split.tensors.count is the TOTAL across shards, not per shard. A missing or extra shard raises with the indices named rather than loading a partial model. qwen3moe support reuses the qwen35moe adapter's shape with the GDN complexity removed: plain attention on every layer, no shared expert, no MTP, and no (1+w) norm shift (qwen3_moe uses RMSNormFused, and plain Qwen3 GGUFs are not pre-shifted). Merged qkv still routes through GGUFMergedLinear because Qwen3 quants mix types across q/k/v. Also fixes the same defect qwen35moe had: Qwen3MoeForCausalLM never called convert_qwen3moe_to_gguf, so the module kept dense ops while the loader yielded .qweight. Caught by the review pass, not by construction. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_014h3QjryXQL6PXJWQdA6tvu --- python/freetoken/models/gguf/config.py | 1 + python/freetoken/models/gguf/reader.py | 309 ++++++++-- python/freetoken/models/gguf/tokenizer.py | 9 +- .../models/qwen3_5_moe/gguf_experts.py | 8 +- python/freetoken/models/qwen3_moe/__init__.py | 6 + python/freetoken/models/qwen3_moe/gguf.py | 477 +++++++++++++++ .../models/qwen3_moe/gguf_experts.py | 283 +++++++++ python/freetoken/models/qwen3_moe/model.py | 13 + python/freetoken/models/register.py | 8 + tests/models/test_gguf_shards.py | 559 ++++++++++++++++++ 10 files changed, 1627 insertions(+), 46 deletions(-) create mode 100644 python/freetoken/models/qwen3_moe/gguf.py create mode 100644 python/freetoken/models/qwen3_moe/gguf_experts.py create mode 100644 tests/models/test_gguf_shards.py diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 3d269eeb..1ab5704c 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -24,6 +24,7 @@ # MLP instead of routed experts. Same model classes and the same GGUF adapter; the # config's expert_count is absent so moe_enabled comes out False. "qwen35": "Qwen35GGUFForCausalLM", + "qwen3moe": "Qwen3MoeGGUFForCausalLM", } diff --git a/python/freetoken/models/gguf/reader.py b/python/freetoken/models/gguf/reader.py index b950d929..9b58ec9d 100644 --- a/python/freetoken/models/gguf/reader.py +++ b/python/freetoken/models/gguf/reader.py @@ -6,12 +6,42 @@ reversed), the ggml quant type, and a zero-copy ``uint8`` view of the packed block bytes laid out as ``[rows, row_bytes]`` (rows = product of all but the fastest ggml dim; row_bytes spans whole quant blocks of the fastest dim). + +Multi-shard GGUF support (llama.cpp split convention, per ground truth in spec): + +Filenames follow ``-%05d-of-%05d.gguf``, both numbers 1-based. The split layout is: + + - Shard 1 (``-00001-of-000NN``) holds the FULL KV metadata: ``general.architecture``, + all ``.*`` config keys, all ``tokenizer.*`` keys. It also carries tensor count + and a ``split.no = 0`` marker (0-based, even though filenames are 1-based). + - Shards 2..N (``-00002-of-000NN`` to ``-000NN-of-000NN``) carry exactly 3 metadata keys: + ``split.no`` (1..N-1), ``split.count`` (always NN), and ``split.tensors.count`` (the + TOTAL tensor count across all shards, not per-shard). They list no architecture keys. + - Tensor distribution: e.g. Hy3 IQ1_M (1298 total) splits as shard 1 with 694 tensors, + shard 2 with 604 tensors. ``split.tensors.count`` is always 1298. + +Examples: + + - A bare ``.gguf`` file (no shard marker) -> no change, single-file path throughout. + - ``model-00001-of-00002.gguf`` passed to any function -> caller gets shard 1 metadata + and tensors from both shards 1..2 in order. + - A directory containing ``model-00001-of-00002.gguf`` -> caller passes the dir, + is_gguf_path resolves it, downstream gets the first-shard path. + +Validation on open: if shard 1 declares ``split.count = N``, all N shards must exist +(1..N), no gaps. Also assert summed tensor count across all shards equals +``split.tensors.count``. Both keys are read from shard 1 only. + +Shard readers are cached per path (one per shard file), so opening ``-00002-of-00002`` +after ``-00001-of-00002`` will reuse the first-shard reader (no double-load of shard 1). """ from __future__ import annotations import functools +import glob import os +import re import struct from dataclasses import dataclass from typing import Any, Iterator @@ -20,11 +50,114 @@ import torch +def gguf_shards(path: str) -> list[str]: + r"""Return the ordered list of shard paths given any shard's path (or a plain .gguf). + + Matches the llama.cpp pattern ``(?P.+)-(\d{5})-of-(\d{5})\.gguf$`` on the + basename. A non-shard path returns ``[path]``. For a shard path, globs sibling shards, + sorts by index, and validates the set is complete 1..N with none missing. + + Raises a clear error naming the missing indices if any are absent (truncated downloads + are the common failure case and must not load silently). + """ + basename = os.path.basename(path) + match = re.match(r"(?P.+)-(\d{5})-of-(\d{5})\.gguf$", basename) + if not match: + # Not a shard file; return as single-file path. + return [path] + + base, shard_idx_str, total_shards_str = match.group("base"), match.group(2), match.group(3) + total_shards = int(total_shards_str) + shard_dir = os.path.dirname(path) + + # Glob all sibling shards + pattern = os.path.join(shard_dir, f"{base}-?????.of-{total_shards_str}.gguf") + found_shards = sorted(glob.glob(pattern)) + + # Parse indices and validate completeness + shard_indices = set() + shard_map = {} # index -> path + for shard_path in found_shards: + shard_basename = os.path.basename(shard_path) + shard_match = re.match(rf"{re.escape(base)}-(\d{{5}})-of-{total_shards_str}\.gguf$", shard_basename) + if shard_match: + idx = int(shard_match.group(1)) + shard_indices.add(idx) + shard_map[idx] = shard_path + + # Verify complete range 1..N + expected = set(range(1, total_shards + 1)) + if shard_indices != expected: + missing = sorted(expected - shard_indices) + raise ValueError( + f"Incomplete shard set for {base}: expected shards 1..{total_shards}, " + f"missing {missing}. (Truncated download?)" + ) + + # Return in order 1..N + return [shard_map[i] for i in range(1, total_shards + 1)] + + +def resolve_gguf_path(model_path: str) -> str | None: + """Resolve a path to the first shard (shard 1) of a GGUF file. + + Accepts: + - A single ``.gguf`` file -> returns it as-is. + - A shard file (e.g., ``-00002-of-00002.gguf``) -> returns shard 1 path. + - A directory containing exactly one shard 1 file -> returns that file path. + + Returns ``None`` if the path is none of the above. + """ + if not isinstance(model_path, str): + return None + + # Case 1: A single .gguf file (not a shard) + if os.path.isfile(model_path) and model_path.endswith(".gguf"): + basename = os.path.basename(model_path) + if not re.match(r".+-\d{5}-of-\d{5}\.gguf$", basename): + # Plain .gguf, not a shard + return model_path + + # Case 2: A shard file or a directory + if os.path.isfile(model_path) and model_path.endswith(".gguf"): + # It's a shard file; get shard 1 + shards = gguf_shards(model_path) + return shards[0] if shards else None + + if os.path.isdir(model_path): + # Look for exactly one shard-1 file in the directory + pattern = os.path.join(model_path, "*-00001-of-?????.gguf") + candidates = glob.glob(pattern) + if len(candidates) == 1: + return candidates[0] + + return None + + def is_gguf_path(model_path: str) -> bool: - """A single ``.gguf`` file (the only GGUF layout FreeToken loads directly).""" - return isinstance(model_path, str) and os.path.isfile(model_path) and model_path.endswith( - ".gguf" - ) + """A ``.gguf`` file or directory, supporting single files and multi-shard layouts. + + Accepts: + - A single ``.gguf`` file. + - Any shard of a multi-shard ``.gguf`` (e.g., shard 2 of 5). + - A directory containing exactly one shard 1 file. + + Returns ``True`` only if one of these conditions holds. + """ + if not isinstance(model_path, str): + return False + + # Case 1: A .gguf file (single or shard) + if os.path.isfile(model_path) and model_path.endswith(".gguf"): + return True + + # Case 2: A directory with a shard 1 file + if os.path.isdir(model_path): + pattern = os.path.join(model_path, "*-00001-of-?????.gguf") + candidates = glob.glob(pattern) + return len(candidates) == 1 + + return False # Canonical name of the metadata-only GGUF that ``convert_checkpoint`` drops into an FTW @@ -42,18 +175,24 @@ def is_gguf_path(model_path: str) -> bool: def gguf_config_source(model_path: str) -> str | None: """The ``.gguf`` file to source config/tokenizer/metadata from, or ``None``. - A bare ``.gguf`` file resolves to itself; an FTW dir carrying a - :data:`FTW_METADATA_GGUF` resolves to that embedded metadata file. This is the single - seam config/tokenizer dispatch uses to decide "this checkpoint is GGUF-config-sourced" - -- a real file and a converted-FTW dir both land on a genuine ``.gguf`` path the reader - can parse, so no downstream code learns about the FTW wrapper. + A bare ``.gguf`` file or any shard resolves to the first shard; an FTW dir carrying + a :data:`FTW_METADATA_GGUF` resolves to that embedded metadata file. A directory + containing shards resolves to shard 1. This is the single seam config/tokenizer + dispatch uses to decide "this checkpoint is GGUF-config-sourced" -- a real file, a + shard file, a shard directory, and a converted-FTW dir all land on a genuine ``.gguf`` + path the reader can parse, so no downstream code learns about the layout. """ - if is_gguf_path(model_path): - return model_path + # Case 1: Check for FTW metadata file first (highest priority) if isinstance(model_path, str) and os.path.isdir(model_path): cand = os.path.join(model_path, FTW_METADATA_GGUF) if os.path.isfile(cand): return cand + + # Case 2: Try to resolve to a GGUF (single, shard, or directory) + resolved = resolve_gguf_path(model_path) + if resolved is not None: + return resolved + return None @@ -121,61 +260,147 @@ def _field_value(reader, name: str) -> Any: @functools.cache def _reader(model_path: str): + """Get or create a GGUFReader for the given path, with shard validation. + + For single-shard files, this is a pass-through. For shard 1 of a multi-shard set, + this validates that: + 1. All shards 1..N are present and complete (no missing indices). + 2. The summed tensor count across all shards matches split.tensors.count (if present). + """ import gguf - return gguf.GGUFReader(model_path) + reader = gguf.GGUFReader(model_path) + + # Check if this is shard 1 of a multi-shard set + split_count = _field_value(reader, "split.count") + split_no = _field_value(reader, "split.no") + + if split_count is not None and split_no == 0: + # This is shard 1 of a multi-shard set; validate completeness + try: + shards = gguf_shards(model_path) + if len(shards) != split_count: + raise ValueError( + f"GGUF shard validation: {model_path} declares split.count={split_count}, " + f"but found {len(shards)} shards" + ) + + # Validate tensor count sum if split.tensors.count is declared + split_tensors_count = _field_value(reader, "split.tensors.count") + if split_tensors_count is not None: + total_tensor_count = 0 + for shard_path in shards: + shard_reader = gguf.GGUFReader(shard_path) + total_tensor_count += len(shard_reader.tensors) + + if total_tensor_count != split_tensors_count: + raise ValueError( + f"GGUF shard validation: {model_path} declares " + f"split.tensors.count={split_tensors_count}, but summed tensor count " + f"across all shards is {total_tensor_count}" + ) + except ValueError: + raise + + return reader @functools.cache def load_gguf_metadata(model_path: str) -> dict[str, Any]: - """All GGUF KV metadata as ``{field_name: python_value}`` (arrays -> lists).""" - reader = _reader(model_path) + """All GGUF KV metadata as ``{field_name: python_value}`` (arrays -> lists). + + Metadata is read from shard 1 only. If the caller passes any other shard, this + function resolves to shard 1 first. + """ + shard1_path = resolve_gguf_path(model_path) + if shard1_path is None: + raise ValueError(f"Cannot resolve GGUF path: {model_path}") + + reader = _reader(shard1_path) return {name: field.contents() for name, field in reader.fields.items()} def gguf_architecture(model_path: str) -> str: - arch = _field_value(_reader(model_path), "general.architecture") + """The model architecture string (e.g., "qwen3moe", "qwen35moe"). + + Architecture is read from shard 1 only. If the caller passes any other shard, + this function resolves to shard 1 first. + """ + shard1_path = resolve_gguf_path(model_path) + if shard1_path is None: + raise ValueError(f"Cannot resolve GGUF path: {model_path}") + + arch = _field_value(_reader(shard1_path), "general.architecture") if arch is None: - raise ValueError(f"GGUF file {model_path} has no general.architecture") + raise ValueError(f"GGUF file {shard1_path} has no general.architecture") return str(arch) def iter_gguf_tensors(model_path: str) -> Iterator[GgufTensor]: - """Yield every tensor with its torch shape, ggml type, and packed block bytes.""" + """Yield every tensor with its torch shape, ggml type, and packed block bytes. + + For multi-shard files, yields tensors from shard 1, then shard 2, ..., in order. + Single-shard files take exactly the same code path (gguf_shards returns [path]). + """ import gguf - reader = _reader(model_path) - for t in reader.tensors: - ne = [int(s) for s in t.shape] # ggml order, fastest dim first - torch_shape = tuple(reversed(ne)) - block, type_size = gguf.GGML_QUANT_SIZES[t.tensor_type] - n_fast = ne[0] - if n_fast % block != 0: - raise ValueError( - f"{t.name}: fastest dim {n_fast} not a multiple of block {block} " - f"for {t.tensor_type.name}" + shard1_path = resolve_gguf_path(model_path) + if shard1_path is None: + raise ValueError(f"Cannot resolve GGUF path: {model_path}") + + # Get all shard paths in order + shards = gguf_shards(shard1_path) + + # Iterate over each shard and yield tensors + for shard_path in shards: + reader = _reader(shard_path) + for t in reader.tensors: + ne = [int(s) for s in t.shape] # ggml order, fastest dim first + torch_shape = tuple(reversed(ne)) + block, type_size = gguf.GGML_QUANT_SIZES[t.tensor_type] + n_fast = ne[0] + if n_fast % block != 0: + raise ValueError( + f"{t.name}: fastest dim {n_fast} not a multiple of block {block} " + f"for {t.tensor_type.name}" + ) + row_bytes = n_fast // block * type_size + rows = int(np.prod(ne[1:])) if len(ne) > 1 else 1 + # gguf-py returns quantized tensors as raw uint8 but F32/F16 as typed arrays; + # normalize everything to a flat byte view before shaping into [rows, row_bytes]. + flat = np.ascontiguousarray(t.data).reshape(-1).view(np.uint8) + raw = flat.reshape(rows, row_bytes) + yield GgufTensor( + name=t.name, + shape=torch_shape, + ggml_type=int(t.tensor_type), + rows=rows, + row_bytes=row_bytes, + _raw=raw, ) - row_bytes = n_fast // block * type_size - rows = int(np.prod(ne[1:])) if len(ne) > 1 else 1 - # gguf-py returns quantized tensors as raw uint8 but F32/F16 as typed arrays; - # normalize everything to a flat byte view before shaping into [rows, row_bytes]. - flat = np.ascontiguousarray(t.data).reshape(-1).view(np.uint8) - raw = flat.reshape(rows, row_bytes) - yield GgufTensor( - name=t.name, - shape=torch_shape, - ggml_type=int(t.tensor_type), - rows=rows, - row_bytes=row_bytes, - _raw=raw, - ) def gguf_tensor_names(model_path: str) -> set[str]: - return {t.name for t in _reader(model_path).tensors} + """The union of tensor names across all shards. + + For multi-shard files, returns the union of tensor names from shard 1, shard 2, etc. + Single-shard files take exactly the same code path. + """ + shard1_path = resolve_gguf_path(model_path) + if shard1_path is None: + raise ValueError(f"Cannot resolve GGUF path: {model_path}") + + shards = gguf_shards(shard1_path) + names = set() + for shard_path in shards: + reader = _reader(shard_path) + names.update(t.name for t in reader.tensors) + return names __all__ = [ + "gguf_shards", + "resolve_gguf_path", "is_gguf_path", "FTW_METADATA_GGUF", "OUTPUT_WEIGHT_PRESENT_KV", diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 03b24616..ddc4433b 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -17,7 +17,13 @@ # heard of raises KeyError inside convert_gguf_tokenizer. qwen35moe is a GPT2-style BPE # with merges (tokenizer.ggml.model == "gpt2", pre == "qwen35"), which the qwen2 converter # handles; there is no qwen3.5-specific converter and it would be the same BPE anyway. -_TOKENIZER_ARCH = {"gemma4": "gemma4_text", "qwen35moe": "qwen2", "qwen35": "qwen2"} +# qwen3moe is also a GPT2-style BPE (tokenizer.ggml.model == "gpt2", pre == "qwen3"). +_TOKENIZER_ARCH = { + "gemma4": "gemma4_text", + "qwen35moe": "qwen2", + "qwen35": "qwen2", + "qwen3moe": "qwen2", +} # Per-arch chat/stop tokens, in preference order: the first one present in the vocab # becomes eos (so chat generation halts on the turn end rather than the formal document @@ -29,6 +35,7 @@ "qwen35moe": ("<|im_end|>", "<|endoftext|>"), # Dense sibling: same vocab and same chat markers as the MoE variant. "qwen35": ("<|im_end|>", "<|endoftext|>"), + "qwen3moe": ("<|im_end|>", "<|endoftext|>"), } diff --git a/python/freetoken/models/qwen3_5_moe/gguf_experts.py b/python/freetoken/models/qwen3_5_moe/gguf_experts.py index 5948b804..1c2872af 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf_experts.py +++ b/python/freetoken/models/qwen3_5_moe/gguf_experts.py @@ -193,9 +193,11 @@ def _load(sink) -> None: seen_up.add(layer) elif t.name.endswith("ffn_down_exps.weight"): - # Shape from GGUF: [E, H, I] in torch order = [I, H, E] in ggml order - # t.packed() is [I*H, row_bytes(E, type)] - # Reshape to [E, H, row_bytes(I, type)] + # torch shape [E, H, I] = ggml dims [I, H, E] with I fastest, so the + # reader hands back [rows, row_bytes] = [E*H, row_bytes(I)] with rows in + # expert-major order. Reshaping to [E, H, row_bytes(I)] is therefore a + # plain view, no data movement. (The row_bytes is over I, the fastest + # dim, not over E.) down_row_bytes = specs["down"][0][2] banks["down"][layer].copy_(t.packed().reshape(E, H, down_row_bytes)) seen_down.add(layer) diff --git a/python/freetoken/models/qwen3_moe/__init__.py b/python/freetoken/models/qwen3_moe/__init__.py index 59c74fb0..af3d4f77 100644 --- a/python/freetoken/models/qwen3_moe/__init__.py +++ b/python/freetoken/models/qwen3_moe/__init__.py @@ -1,10 +1,16 @@ from .config import parse_config from .model import Qwen3MoeForCausalLM from .weight import iter_weights, iter_weights_parallel +from .gguf import parse_gguf_config, iter_gguf_weights +from .gguf_experts import gguf_expert_types, load_gguf_expert_sources __all__ = [ "Qwen3MoeForCausalLM", "parse_config", "iter_weights", "iter_weights_parallel", + "parse_gguf_config", + "iter_gguf_weights", + "gguf_expert_types", + "load_gguf_expert_sources", ] diff --git a/python/freetoken/models/qwen3_moe/gguf.py b/python/freetoken/models/qwen3_moe/gguf.py new file mode 100644 index 00000000..a3640d99 --- /dev/null +++ b/python/freetoken/models/qwen3_moe/gguf.py @@ -0,0 +1,477 @@ +"""Qwen3-MoE GGUF adapter: build the FreeToken ``ModelConfig`` and stream weights +from a llama.cpp ``qwen3moe`` checkpoint. + +Qwen3-MoE is simpler than Qwen3.5-MoE: all layers use standard full attention (no GDN), +there is no shared expert, and norms are plain (no Gemma-style (1+w) shift). The tensors +map directly from llama.cpp's gguf-py/gguf/tensor_mapping.py without complex fusion layers. + +Verified against unsloth/Qwen3-235B-A22B-GGUF (Q4_K_M) and unsloth/Qwen3-30B-A3B-GGUF. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Iterator + +import torch + +from freetoken.models.config import ( + FullAttentionGroupConfig, + ModelConfig, + RotaryConfig, +) +from freetoken.models.gguf.dequant import ( + GGML_UNQUANTIZED as GGML_UNQUANTIZED_SET, + GGML_NAME, + dequantize, + row_bytes, +) + +if TYPE_CHECKING: + from freetoken.models.gguf.config import GgufConfigShim + + +def _kv(shim: "GgufConfigShim", key: str, default: Any = None) -> Any: + """Read ``.`` from the GGUF metadata. + + The prefix is the checkpoint's own ``general.architecture``: "qwen3moe" for the MoE + variant. Metadata keys follow the llama.cpp convention with this prefix. + """ + val = shim.metadata.get(f"{shim.model_type}.{key}", default) + if val is None and default is None: + raise ValueError( + f"GGUF {shim.model_path}: missing required key {shim.model_type}.{key}" + ) + return val + + +def _uniform_expert_types(model_path: str, num_layers: int) -> tuple[int, int] | None: + """``(gate_up, down)`` ggml types of the routed-expert banks, or None if not uniform. + + The offload slot pool is one allocation per bank shared by every layer, and + ``moe_vec.cuh`` addresses it as ``expert * nrows * (ncols / qk)`` with no padding + allowance -- so a bank whose type varies by layer cannot be served. We return None + rather than raising here because ``parse_gguf_config`` also runs for metadata-only + inspection; ``expert_banks._gguf_banks`` is where the load actually fails, with the + offending layers named. (llama.cpp's *_M mixes hit this.) + """ + from .gguf_experts import gguf_expert_types + + try: + types = gguf_expert_types(model_path, num_layers) + except Exception: + return None + gate_up, down = set(types["gate_up"]), set(types["down"]) + if len(gate_up) != 1 or len(down) != 1: + return None + return (next(iter(gate_up)), next(iter(down))) + + +def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + """Build ModelConfig from qwen3moe GGUF metadata. + + All layers are standard full attention. Attention groups contain a single + FullAttentionGroupConfig spanning all layers. + """ + num_layers = int(_kv(shim, "block_count")) + + hidden_size = int(_kv(shim, "embedding_length")) + num_qo_heads = int(_kv(shim, "attention.head_count")) + num_kv_heads = int(_kv(shim, "attention.head_count_kv")) + head_dim = int(_kv(shim, "attention.key_length")) + rms_eps = float(_kv(shim, "attention.layer_norm_rms_epsilon")) + rope_base = float(_kv(shim, "rope.freq_base")) + max_pos = int(_kv(shim, "context_length")) + + # Routed expert configuration. + num_experts = int(_kv(shim, "expert_count")) + experts_per_tok = int(_kv(shim, "expert_used_count")) + moe_inter = int(_kv(shim, "expert_feed_forward_length")) + dense_inter = int(_kv(shim, "feed_forward_length")) + moe_enabled = num_experts > 0 + + # Rotary embedding configuration: use rope.dimension_count if present, else head_dim. + rotary_dim = int(_kv(shim, "rope.dimension_count", head_dim)) + + rotary = RotaryConfig( + head_dim=head_dim, + rotary_dim=rotary_dim, + max_position=max_pos, + base=rope_base, + scaling=None, + ) + + # All layers are full attention: single FullAttentionGroupConfig. + groups = ( + FullAttentionGroupConfig( + name="full", + layer_ids=tuple(range(num_layers)), + num_kv_heads=num_kv_heads, + head_dim=head_dim, + rotary_config=rotary, + ), + ) + + return ModelConfig( + num_layers=num_layers, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + hidden_size=hidden_size, + vocab_size=shim.vocab_size, + intermediate_size=0 if moe_enabled else dense_inter, + hidden_act="silu", + rms_norm_eps=rms_eps, + tie_word_embeddings=shim.tie_word_embeddings, + rotary_config=rotary, + num_experts=num_experts, + num_experts_per_tok=experts_per_tok, + moe_intermediate_size=moe_inter, + norm_topk_prob=True, + moe_enabled=moe_enabled, + use_qk_norm=True, + model_type=shim.model_type, + architectures=list(shim.architectures), + vision_config=None, + image_token_id=None, + attention_groups=groups, + expert_quant="gguf" if moe_enabled else "none", + gguf_expert_types=( + _uniform_expert_types(shim.model_path, num_layers) if moe_enabled else None + ), + gguf_model_path=shim.model_path, + weight_block_size=None, + attn_quant="gguf", + dense_quant="gguf", + lm_head_quant="gguf", + ) + + +# -------------------------------------------------------------------------------------- +# Tensor-name mapping (inverse of llama.cpp gguf-py/gguf/tensor_mapping.py for qwen3moe) +# -------------------------------------------------------------------------------------- + +# Per-layer 1:1 renames that need no reshaping or fusing. +_LAYER_MAP: dict[str, str] = { + "attn_norm.weight": "input_layernorm.weight", + "ffn_norm.weight": "post_attention_layernorm.weight", + "attn_output.weight": "self_attn.o_proj.weight", + "attn_q_norm.weight": "self_attn.q_norm.weight", + "attn_k_norm.weight": "self_attn.k_norm.weight", + "ffn_gate_inp.weight": "mlp.gate.weight", +} + +# Suffixes that are PARTS of a merged projection: never renamed 1:1, always combined by +# iter_gguf_weights into the merged buffer the model actually declares. +_MERGED_PARTS: frozenset[str] = frozenset({ + "attn_q.weight", "attn_k.weight", "attn_v.weight", +}) + +# Routed-expert stacks: [num_experts, out, in] packed blocks, handled by the offload +# expert-bank loader rather than yielded as ordinary parameters. +_EXPERT_SUFFIXES = ( + "ffn_gate_exps.weight", + "ffn_up_exps.weight", + "ffn_down_exps.weight", +) + +_GLOBAL_MAP: dict[str, str] = { + "token_embd.weight": "model.embed_tokens.weight", + "output_norm.weight": "model.norm.weight", + "output.weight": "lm_head.weight", +} + + +def gguf_name_to_freetoken(name: str, num_layers: int) -> str | None: + """Map one llama.cpp tensor name to its FreeToken parameter name. + + Returns ``None`` for the routed-expert stacks (read directly by the expert-bank + loader) and the parts of a merged projection (combined by :func:`iter_gguf_weights`). + """ + if name in _GLOBAL_MAP: + return _GLOBAL_MAP[name] + if not name.startswith("blk."): + return None + _, idx, suffix = name.split(".", 2) + layer = int(idx) + if layer >= num_layers: + return None # out-of-bounds (should not happen in qwen3moe) + if suffix in _EXPERT_SUFFIXES: + return None + if suffix in _MERGED_PARTS: + return None # fused by iter_gguf_weights into the merged buffer + mapped = _LAYER_MAP.get(suffix) + if mapped is None: + return None + return f"model.layers.{layer}.{mapped}" + + +# -------------------------------------------------------------------------------------- +# Weight loading: GGUF tensor names -> FreeToken qwen3moe module params. +# -------------------------------------------------------------------------------------- + + +def _scan_quant_types(model_path: str) -> dict[tuple[int, str], int]: + """Scan GGUF tensor table once and return {(layer, suffix): ggml_type}. + + This allows us to detect which groups are mixed-quant without hardcoding. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + + quant_types = {} + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + # Globals (token_embd.weight, output.weight, output_norm.weight) keyed under + # layer -1 so the swap can size the embedding and lm_head from the file too. + quant_types[(-1, t.name)] = t.ggml_type + continue + _, idx, suffix = t.name.split(".", 2) + layer = int(idx) + quant_types[(layer, suffix)] = t.ggml_type + return quant_types + + +def _to_bf16(t) -> torch.Tensor: + """Dequantize a GgufTensor (F32/F16/Q*) to a dense bf16 tensor of its torch shape. + + Unlike Qwen3.5-MoE's Gemma-style norms which apply a (1+w) shift at load time, qwen3moe + norms are plain. Dequantize as-is without adding 1.0. + """ + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.bfloat16) + return flat.reshape(t.shape) + + +def _to_f32(t) -> torch.Tensor: + """Dequantize a GgufTensor (F32/F16/Q*) to a dense float32 tensor of its torch shape.""" + flat = dequantize(t.packed().reshape(-1), t.ggml_type, torch.float32) + return flat.reshape(t.shape) + + +def _require_tp1(what: str) -> None: + """GGUF quant layers / expert banks are not sharded; reject TP>1 with a clear error.""" + from freetoken.distributed import get_tp_info + + if get_tp_info().size > 1: + raise NotImplementedError( + f"qwen3moe GGUF {what} currently supports TP=1 only " + "(GGUF quant layers and expert banks are not tensor-parallel sharded)." + ) + + +def iter_gguf_weights( + model_path: str, + device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """Yield (param_name, tensor) for every non-expert qwen3moe parameter. + + Quantized projections (attention qkv/o) stay in their native packed block layout and + are yielded as ``.qweight`` or ``.qweight_`` for mixed-quant groups; norms and the + router gate dequantize to bf16. q/k/v are fused by concatenating packed rows. + + Routed experts are served from the offload cache (asserts the offload contract). + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + from freetoken.utils import cached_load_hf_config + + # Only the MoE variant keeps its routed experts out of this iterator; they come from + # the offload cache instead. + config_moe = int(_kv(cached_load_hf_config(model_path), "expert_count", 0)) > 0 + assert not (config_moe and include_moe_experts), ( + "qwen3moe GGUF keeps its routed experts in the offload cache; they are loaded by " + "the expert-bank loader, not by iter_gguf_weights." + ) + assert include_non_moe + _require_tp1("weight loading") + + # Parse config to get the number of layers and other geometry. + config = parse_gguf_config(cached_load_hf_config(model_path)) + + # Scan quant types once to determine which fusion groups are mixed-quant. + quant_map = _scan_quant_types(model_path) + + # Per-layer fusion buffer for qkv: layer -> {slot: packed[out, row_bytes]}. + qkv_buf: dict[int, dict[str, torch.Tensor]] = {} + + def layer_of(name: str) -> int: + return int(name.split(".")[1]) + + for t in iter_gguf_tensors(model_path): + name = t.name + layer = layer_of(name) if name.startswith("blk.") else None + + # Global tensors + if name == "token_embd.weight": + yield "model.embed_tokens.qweight", t.packed() + continue + if name == "output_norm.weight": + yield "model.norm.weight", _to_bf16(t) + continue + if name == "output.weight": + if not config.tie_word_embeddings: + yield "lm_head.qweight", t.packed() + continue + if not name.startswith("blk."): + continue + + # Skip out-of-bounds layers (should not happen in qwen3moe). + if layer >= config.num_layers: + continue + + # Skip routed-expert stacks (offload banks). + if any(name.endswith(sfx) for sfx in _EXPERT_SUFFIXES): + continue + + suffix = name.split(".", 2)[2] # after "blk.N." + base = f"model.layers.{layer}" + + # Scalar/norm tensors: dequant to bf16. + if suffix == "attn_norm.weight": + yield f"{base}.input_layernorm.weight", _to_bf16(t) + continue + if suffix == "ffn_norm.weight": + yield f"{base}.post_attention_layernorm.weight", _to_bf16(t) + continue + if suffix == "ffn_gate_inp.weight": + yield f"{base}.mlp.gate.weight", _to_bf16(t) + continue + if suffix == "attn_q_norm.weight": + yield f"{base}.self_attn.q_norm.weight", _to_bf16(t) + continue + if suffix == "attn_k_norm.weight": + yield f"{base}.self_attn.k_norm.weight", _to_bf16(t) + continue + + # Quantized projections: keep packed; fuse per layer. + # All layers are full-attention: fuse q, k, v into qkv_proj. + if suffix == "attn_q.weight": + qkv_buf.setdefault(layer, {})["q"] = t.packed() + elif suffix == "attn_k.weight": + qkv_buf.setdefault(layer, {})["k"] = t.packed() + elif suffix == "attn_v.weight": + qkv_buf.setdefault(layer, {})["v"] = t.packed() + elif suffix == "attn_output.weight": + yield f"{base}.self_attn.o_proj.qweight", t.packed() + else: + continue # unmapped suffix + + # Emit fused qkv once all three parts are present. + slots = qkv_buf.get(layer) + if slots is not None and "q" in slots and "k" in slots and "v" in slots: + # Determine if this is a mixed-quant group. + types = [ + quant_map.get((layer, "attn_q.weight")), + quant_map.get((layer, "attn_k.weight")), + quant_map.get((layer, "attn_v.weight")), + ] + if len(set(types)) == 1: + # Uniform quant: fuse via torch.cat along dim 0. + yield f"{base}.self_attn.qkv_proj.qweight", torch.cat( + [slots["q"], slots["k"], slots["v"]], dim=0 + ) + else: + # Mixed quant: emit GGUFMergedLinear format. + yield f"{base}.self_attn.qkv_proj.qweight_0", slots["q"] + yield f"{base}.self_attn.qkv_proj.qweight_1", slots["k"] + yield f"{base}.self_attn.qkv_proj.qweight_2", slots["v"] + del qkv_buf[layer] + + # Verify no fusion buffers are incomplete. + assert not qkv_buf, f"incomplete qkv groups: {sorted(qkv_buf)}" + + +def is_gguf_model(config: ModelConfig) -> bool: + """True when this config came from a GGUF checkpoint (native block-quant path). + + Keys on gguf_model_path rather than expert_quant: ensures the op swap runs. + """ + return getattr(config, "gguf_model_path", None) is not None + + +def convert_qwen3moe_to_gguf(model, config: ModelConfig, *, model_path: str) -> None: + """In place: replace qwen3moe's dense projections + embedding with native GGUF ops. + + Quantized in the checkpoint -> swapped: attention qkv/o (mixed-quant). Left as dense + bf16 (F32 in the GGUF): all RMSNorms, the router gate, and the routed experts + (served from the offload cache). + + The per-layer quant types are read from the GGUF file, not hardcoded, to support + different quant levels. + """ + from freetoken.layers.gguf import GGUFEmbedding, GGUFLinear, gguf_merged_or_plain + + # Scan quant types to drive layer swaps. + quant_map = _scan_quant_types(model_path) + + # Split widths for qkv come from the config. + _qkv_split = [ + config.num_qo_heads * config.head_dim, + config.num_kv_heads * config.head_dim, + config.num_kv_heads * config.head_dim, + ] + + def qt(layer: int, suffix: str) -> int: + """The ggml type of one tensor, straight from the file. + + No default: a guessed type silently allocates a wrong-sized packed buffer. + """ + key = (layer, suffix) + if key not in quant_map: + raise ValueError( + f"GGUF {model_path}: expected tensor " + f"{suffix if layer < 0 else f'blk.{layer}.{suffix}'} is absent, so its quant " + f"type cannot be read; this checkpoint does not match the qwen3moe layout " + f"this adapter expects" + ) + return quant_map[key] + + def swap_linear(owner, attr, quant_type: int): + """Replace a dense Linear with the GGUFLinear its packed weight will land in.""" + lin = getattr(owner, attr) + out_features, in_features = lin.weight.shape + setattr( + owner, + attr, + GGUFLinear(in_features, out_features, quant_type, has_bias=lin.bias is not None), + ) + + inner = model.model + embed = GGUFEmbedding( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + quant_type=qt(-1, "token_embd.weight"), + ) + inner.embed_tokens = embed + + for layer_idx, layer in enumerate(inner.layers.op_list): + # qkv_proj: q | k | v. + layer.self_attn.qkv_proj = gguf_merged_or_plain( + config.hidden_size, + _qkv_split, + [ + qt(layer_idx, "attn_q.weight"), + qt(layer_idx, "attn_k.weight"), + qt(layer_idx, "attn_v.weight"), + ], + has_bias=False, + ) + swap_linear(layer.self_attn, "o_proj", qt(layer_idx, "attn_output.weight")) + + if config.tie_word_embeddings: + from freetoken.models.gemma4.gguf import GGUFTiedLMHead + + model.lm_head = GGUFTiedLMHead(embed, qt(-1, "token_embd.weight")) + else: + swap_linear(model, "lm_head", qt(-1, "output.weight")) + + +__all__ = [ + "parse_gguf_config", + "gguf_name_to_freetoken", + "iter_gguf_weights", + "convert_qwen3moe_to_gguf", + "is_gguf_model", + "_MERGED_PARTS", + "_EXPERT_SUFFIXES", +] diff --git a/python/freetoken/models/qwen3_moe/gguf_experts.py b/python/freetoken/models/qwen3_moe/gguf_experts.py new file mode 100644 index 00000000..675cc5d2 --- /dev/null +++ b/python/freetoken/models/qwen3_moe/gguf_experts.py @@ -0,0 +1,283 @@ +"""Routed-expert host bank sources for the qwen3moe GGUF checkpoint. + +This module loads the per-expert weight tensors that are stored as GGUF stacks +and allocates them into host banks for the offload cache. The layout is a 3D +expert stack: [num_experts, out_features, in_features] in torch order. + +CRITICAL CORRECTNESS NOTE: The MoE kernel (kernel/csrc/gguf/moe_vec.cuh) +computes addressing as `blocks_per_row = ncols / qk` and +`x = vx + expert * nrows * blocks_per_row`, i.e. it assumes a FULLY PACKED +contiguous [E, nrows, blocks_per_row] layout with NO padding. So the bank +tensors must be exactly `row_bytes` wide for their own quant type — never pad +a smaller-type layer up to a larger type's stride, because the kernel would +then read every block at the wrong offset and return plausible-looking garbage. + +For qwen3moe: +- ``ffn_gate_exps`` and ``ffn_up_exps`` must share a quant type per layer +- ``ffn_down_exps`` can vary independently per layer +- The gate_up bank per layer is the per-expert concatenation of gate rows + then up rows along the output dimension, giving [E, 2*I, row_bytes(H, t)], + valid because gate and up share a quant type and therefore a row stride. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from freetoken.models.gguf.dequant import GGML_NAME, row_bytes + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + + +def gguf_expert_types(model_path: str, num_layers: int) -> dict[str, list[int]]: + """Scan the GGUF tensor table and return per-layer expert quant types. + + Returns a dict with two keys: + - ``"gate_up"``: list of ``num_layers`` ggml_type enums for ``ffn_gate_exps``. + gate and up for each layer must have the same type (they are row-concatenated). + If they differ for any layer, raises a clear ValueError naming the layer and both types. + - ``"down"``: list of ``num_layers`` ggml_type enums for ``ffn_down_exps``. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + + gate_types: list[int | None] = [None] * num_layers + up_types: list[int | None] = [None] * num_layers + down_types: list[int | None] = [None] * num_layers + + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if layer >= num_layers: + continue # skip out-of-bounds layers + + if t.name.endswith("ffn_gate_exps.weight"): + gate_types[layer] = t.ggml_type + elif t.name.endswith("ffn_up_exps.weight"): + up_types[layer] = t.ggml_type + elif t.name.endswith("ffn_down_exps.weight"): + down_types[layer] = t.ggml_type + + # Validate that gate and up types agree for each layer (they must be row-concatenated). + gate_up_types: list[int] = [] + for layer in range(num_layers): + gate_t = gate_types[layer] + up_t = up_types[layer] + if gate_t is None or up_t is None: + raise ValueError( + f"missing expert tensors for layer {layer}: " + f"gate={GGML_NAME.get(gate_t, gate_t)}, up={GGML_NAME.get(up_t, up_t)}" + ) + if gate_t != up_t: + raise ValueError( + f"layer {layer}: ffn_gate_exps type {GGML_NAME.get(gate_t, gate_t)} != " + f"ffn_up_exps type {GGML_NAME.get(up_t, up_t)}; " + "cannot row-concatenate tensors with different quant types" + ) + gate_up_types.append(gate_t) + + # Validate down tensors are present. + for layer in range(num_layers): + if down_types[layer] is None: + raise ValueError(f"missing ffn_down_exps for layer {layer}") + + return { + "gate_up": gate_up_types, + "down": down_types, + } + + +def gguf_expert_specs( + config: ModelConfig, types: dict[str, list[int]] +) -> dict[str, tuple[tuple[int, ...], torch.dtype]]: + """Expert bank shapes as ``{name: (shape, dtype)}`` -- ``alloc_layer_banks``' contract. + + The routed experts are 3D stacks in torch order:: + + gate_up (E, 2*I, row_bytes(H, t_gate_up)) uint8, packed blocks + down (E, H, row_bytes(I, t_down)) uint8, packed blocks + + One spec per bank, not per layer: every layer of a bank MUST share a ggml type. The + GPU slot pool is a single allocation shared by all layers and ``moe_vec.cuh`` indexes + it as ``expert * nrows * (ncols / qk)`` with no padding allowance, so two strides in + one pool would read every block at the wrong offset. A non-uniform bank is rejected + here rather than mis-decoded; ``expert_banks._gguf_banks`` raises the user-facing + error naming the offending layers. + """ + E, H, I = config.num_experts, config.hidden_size, config.moe_intermediate_size + out = {} + for name, elems in (("gate_up", H), ("down", I)): + distinct = sorted(set(types[name])) + if len(distinct) != 1: + raise ValueError( + f"expert bank {name!r} mixes ggml types across layers ({distinct}); a bank " + f"must be uniform because its slot pool is one allocation with one stride" + ) + rb = row_bytes(elems, distinct[0]) + shape = (E, 2 * I, rb) if name == "gate_up" else (E, H, rb) + out[name] = (shape, torch.uint8) + return out + + +def load_gguf_expert_sources( + model_path: str, config: ModelConfig, *, layer_sink=None +) -> dict[str, list[torch.Tensor]]: + """Per-layer host banks of the routed experts' native packed block bytes. + + Loads the three GGUF expert stacks (gate, up, down) into per-layer host banks + for the offload cache. The gate_up bank for each layer is the per-expert + concatenation of that expert's gate rows then its up rows along the output + dimension, giving [E, 2*I, row_bytes(H, t)] -- valid because gate and up + share a quant type and therefore a row stride. + + Returns a dict with two keys: + - ``"gate_up"``: list of ``num_layers`` tensors, each ``[E, 2*I, row_bytes_gate_up]`` uint8 + - ``"down"``: list of ``num_layers`` tensors, each ``[E, H, row_bytes_down]`` uint8 + + Parameters: + - ``layer_sink``: If None (serving mode), pins each completed layer via an + internally-owned PinPipeline. If given (converter mode), fires the completion + tracker into it instead -- nothing is pinned, and the sink may release banks, + so returned tensors are only valid until the sink releases them. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks + + types = gguf_expert_types(model_path, config.num_layers) + specs = gguf_expert_specs(config, types) + + L = config.num_layers + E = config.num_experts + H = config.hidden_size + I = config.moe_intermediate_size + + # Allocate the per-layer banks (lazy mmap, unpinned). + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in hb} + + # Per-layer buffers to accumulate gate and up before concatenating. + gate_buf: dict[int, torch.Tensor] = {} + up_buf: dict[int, torch.Tensor] = {} + seen_gate = set() + seen_up = set() + seen_down = set() + + def _load(sink) -> None: + # Track completion: 2 banks per layer (gate_up and down). + tracker = LayerCompletionTracker(2, hb, sink) if sink is not None else None + + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if layer >= L: + continue # skip out-of-bounds layers + + if t.name.endswith("ffn_gate_exps.weight"): + # Shape from GGUF: [E, I, H] in torch order = [H, I, E] in ggml order + # t.packed() is [H*I, row_bytes(E, type)] + gate_buf[layer] = t.packed() + seen_gate.add(layer) + + elif t.name.endswith("ffn_up_exps.weight"): + # Shape from GGUF: [E, I, H] in torch order = [H, I, E] in ggml order + # t.packed() is [H*I, row_bytes(E, type)] + up_buf[layer] = t.packed() + seen_up.add(layer) + + elif t.name.endswith("ffn_down_exps.weight"): + # torch shape [E, H, I] = ggml dims [I, H, E] with I fastest, so the + # reader hands back [rows, row_bytes] = [E*H, row_bytes(I)] with rows in + # expert-major order. Reshaping to [E, H, row_bytes(I)] is therefore a + # plain view, no data movement. (The row_bytes is over I, the fastest + # dim, not over E.) + down_row_bytes = specs["down"][0][2] + banks["down"][layer].copy_(t.packed().reshape(E, H, down_row_bytes)) + seen_down.add(layer) + if tracker is not None: + tracker.note(layer) + + else: + continue + + # Emit gate_up bank once both gate and up are present. + if layer in gate_buf and layer in up_buf: + rb = specs["gate_up"][0][2] + # gate and up each arrive as [rows, row_bytes] = [E*I, row_bytes(H)], and + # ggml's fastest-first dims [H, I, E] make E the slowest axis, so those rows + # are EXPERT-MAJOR: expert e owns rows [e*I, (e+1)*I). + # + # The bank must be [E, 2I, row_bytes(H)] with each expert's own gate rows + # followed by its own up rows. So reshape to [E, I, rb] and concatenate on + # the ROW axis (dim=1), per expert. + # + # cat(dim=0) then reshape(E, 2I, rb) -- the obvious-looking version -- is + # wrong: it lays down every expert's gate before any up, so expert 0 would + # get its gate rows plus expert 1's gate rows, and up would be E*I rows + # away. That loads and runs at full speed and emits fluent nonsense. + g = gate_buf[layer].reshape(E, I, rb) + u = up_buf[layer].reshape(E, I, rb) + banks["gate_up"][layer].copy_(torch.cat([g, u], dim=1)) + del gate_buf[layer], up_buf[layer] + if tracker is not None: + tracker.note(layer) + + # Load with or without pinning. + if layer_sink is not None: + _load(layer_sink) + elif torch.cuda.is_available(): + with PinPipeline() as pins: + _load(pins) + else: + _load(None) # CUDA-less: mmap banks stay pageable, never pinned + + # Verify all layers were loaded. + want = set(range(L)) + missing_gate = want - seen_gate + missing_up = want - seen_up + missing_down = want - seen_down + if missing_gate or missing_up or missing_down: + raise ValueError( + f"missing expert layers: gate {sorted(missing_gate)}, " + f"up {sorted(missing_up)}, down {sorted(missing_down)}" + ) + + return banks + + +def dummy_gguf_expert_sources(config: ModelConfig) -> dict[str, list[torch.Tensor]]: + """Random expert banks shaped like ``load_gguf_expert_sources`` output.""" + from freetoken.models.gguf.dequant import GGML_IQ3_S + from freetoken.moe.host_banks import alloc_layer_banks, pin_banks + + # Use uniform IQ3_S for all layers (a simplification for the dummy). + num_layers = config.num_layers + gate_up_types = [GGML_IQ3_S] * num_layers + down_types = [GGML_IQ3_S] * num_layers + types = {"gate_up": gate_up_types, "down": down_types} + + specs = gguf_expert_specs(config, types) + L = config.num_layers + + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in hb} + + # Fill with random uint8. + for t in banks["gate_up"] + banks["down"]: + t.random_(0, 256) + + if torch.cuda.is_available(): + pin_banks(hb) # match the other dummies: pin-after-fill + + return banks + + +__all__ = [ + "gguf_expert_types", + "gguf_expert_specs", + "load_gguf_expert_sources", + "dummy_gguf_expert_sources", +] diff --git a/python/freetoken/models/qwen3_moe/model.py b/python/freetoken/models/qwen3_moe/model.py index e37437d2..4fa79139 100644 --- a/python/freetoken/models/qwen3_moe/model.py +++ b/python/freetoken/models/qwen3_moe/model.py @@ -75,6 +75,19 @@ def __init__(self, config: ModelConfig): ) super().__init__() + # A GGUF checkpoint carries native block-quantized weights: swap the dense + # projections + embedding for GGUF-quant ops so the packed buffers have + # somewhere to land (routed experts stay on the offload cache). Mirrors + # gemma4/model.py and qwen3_5_moe/model.py. + from .gguf import convert_qwen3moe_to_gguf, is_gguf_model + + if is_gguf_model(config): + assert config.gguf_model_path is not None, ( + "expert_quant=='gguf' but ModelConfig.gguf_model_path is unset; the " + "per-tensor ggml types can only be read from the file" + ) + convert_qwen3moe_to_gguf(self, config, model_path=config.gguf_model_path) + def forward(self) -> torch.Tensor: output = self.model.forward(get_global_ctx().batch.input_ids) logits = self.lm_head.forward(output) diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index c062df0c..48bd93e7 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -121,6 +121,14 @@ class ModelSpec: parse_config="parse_gguf_config", iter_weights="iter_gguf_weights", ), + # GGUF qwen3moe: simpler than qwen3.5moe (full attention, no GDN/shared expert, plain + # norms), same model class as qwen3_moe HF, GGUF config + weight loaders. + "Qwen3MoeGGUFForCausalLM": ModelSpec( + "freetoken.models.qwen3_moe", + "Qwen3MoeForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + ), "GptOssForCausalLM": ModelSpec( "freetoken.models.gpt_oss", "GptOssForCausalLM", diff --git a/tests/models/test_gguf_shards.py b/tests/models/test_gguf_shards.py new file mode 100644 index 00000000..513cabbf --- /dev/null +++ b/tests/models/test_gguf_shards.py @@ -0,0 +1,559 @@ +"""Tests for GGUF multi-shard discovery and aggregation, plus qwen3moe mapping. + +This module verifies that multi-shard GGUF files (following llama.cpp's split +convention) are correctly discovered, validated, and aggregated by +freetoken.models.gguf.reader, and that the qwen3moe tensor-name mapping is correct. + +Test fixtures build synthetic GGUF files on disk using gguf.GGUFWriter where +possible, or raw GGUF bytes where needed. The tensor data is tiny (a few F32 +values) to keep I/O fast; the focus is on metadata, discovery, and naming, not +quantization or size. + +llama.cpp's shard layout (ground truth from measured checksums): + - Shard 1: full metadata (general.architecture, .*, tokenizer.*) + split.no=0 + - Shards 2..N: exactly 3 keys (split.no, split.count, split.tensors.count) + - split.tensors.count is the TOTAL across all shards, not per-shard + - Filenames are 1-based; split.no is 0-based +""" + +from __future__ import annotations + +import os +import struct +import tempfile +from pathlib import Path + +import pytest + + +class TestSingleFileUnchanged: + """A plain one-file .gguf still reports correct arch, metadata, tensors. + + This is the regression guard: single-file behavior must be untouched by + multi-shard support. + """ + + def test_single_file_unchanged(self, tmp_path): + """Verify a single .gguf file works unchanged.""" + from freetoken.models.gguf.reader import ( + gguf_architecture, + is_gguf_path, + load_gguf_metadata, + gguf_tensor_names, + iter_gguf_tensors, + ) + + # Write a minimal GGUF file (no shards) + gguf_path = self._write_minimal_gguf(tmp_path, "single.gguf") + + # Check single-file detection + assert is_gguf_path(str(gguf_path)) + + # Check architecture is readable + arch = gguf_architecture(str(gguf_path)) + assert arch == "test_model" + + # Check metadata + metadata = load_gguf_metadata(str(gguf_path)) + assert metadata.get("general.architecture") == "test_model" + assert metadata.get("test_model.hidden_size") == 256 + + # Check tensor names + names = gguf_tensor_names(str(gguf_path)) + assert "model.embed.weight" in names + assert len(names) == 1 + + # Check tensor iteration + tensors = list(iter_gguf_tensors(str(gguf_path))) + assert len(tensors) == 1 + assert tensors[0].name == "model.embed.weight" + assert tensors[0].shape == (256, 64) + + def _write_minimal_gguf(self, tmp_path: Path, filename: str) -> Path: + """Write a minimal single-file GGUF with one tensor.""" + try: + import gguf + except ImportError: + pytest.skip("gguf package not available") + + gguf_path = tmp_path / filename + writer = gguf.GGUFWriter(str(gguf_path)) + + # Add metadata + writer.add_string("general.architecture", "test_model") + writer.add_uint32("test_model.hidden_size", 256) + writer.add_uint32("test_model.num_layers", 2) + + # Add one tensor (embedding) + import numpy as np + + emb_data = np.random.randn(256, 64).astype(np.float32) + writer.add_tensor("model.embed.weight", emb_data) + + writer.write_header_and_data(str(gguf_path)) + return gguf_path + + +class TestShardDiscovery: + """Given a 3-shard set, gguf_shards() returns all three in index order.""" + + def test_shard_discovery_orders_and_completes(self, tmp_path): + """Discover all 3 shards when given any shard or the directory.""" + from freetoken.models.gguf.reader import ( + gguf_shards, + is_gguf_path, + gguf_tensor_names, + ) + + # Write 3 shards + shard_paths = self._write_3_shards(tmp_path) + + # Test 1: discover from shard 1 + result = gguf_shards(str(shard_paths[0])) + assert len(result) == 3 + assert [Path(p).name for p in result] == [ + "model-00001-of-00003.gguf", + "model-00002-of-00003.gguf", + "model-00003-of-00003.gguf", + ] + + # Test 2: discover from shard 2 + result = gguf_shards(str(shard_paths[1])) + assert len(result) == 3 + assert result[0] == str(shard_paths[0]) # first shard returned + + # Test 3: discover from shard 3 + result = gguf_shards(str(shard_paths[2])) + assert len(result) == 3 + + # Test 4: directory resolution + assert is_gguf_path(str(tmp_path)) + + # Test 5: tensors are aggregated across shards + names = gguf_tensor_names(str(shard_paths[0])) + assert names == { + "token_embd.weight", # shard 1 + "blk.0.attn_norm.weight", # shard 2 + "blk.1.attn_norm.weight", # shard 3 + } + + def _write_3_shards(self, tmp_path: Path) -> list[Path]: + """Write a 3-shard GGUF set with metadata and tensors distributed.""" + try: + import gguf + except ImportError: + pytest.skip("gguf package not available") + + import numpy as np + + shard_paths = [] + + # Shard 1: full metadata + 1 tensor + shard1 = tmp_path / "model-00001-of-00003.gguf" + writer = gguf.GGUFWriter(str(shard1)) + writer.add_string("general.architecture", "qwen3moe") + writer.add_uint32("qwen3moe.block_count", 2) + writer.add_uint32("qwen3moe.embedding_length", 128) + writer.add_uint32("qwen3moe.attention.head_count", 4) + writer.add_uint32("qwen3moe.attention.head_count_kv", 2) + writer.add_uint32("qwen3moe.attention.key_length", 32) + writer.add_float32("qwen3moe.attention.layer_norm_rms_epsilon", 1e-6) + writer.add_float32("qwen3moe.rope.freq_base", 10000.0) + writer.add_uint32("qwen3moe.context_length", 4096) + writer.add_uint32("qwen3moe.expert_count", 8) + writer.add_uint32("qwen3moe.expert_used_count", 2) + writer.add_uint32("qwen3moe.expert_feed_forward_length", 512) + writer.add_uint32("qwen3moe.feed_forward_length", 512) + # Split metadata for shard 1 + writer.add_uint32("split.no", 0) + writer.add_uint32("split.count", 3) + writer.add_uint32("split.tensors.count", 3) + # One tensor in shard 1 + emb = np.random.randn(1024, 128).astype(np.float32) + writer.add_tensor("token_embd.weight", emb) + writer.write_header_and_data(str(shard1)) + shard_paths.append(shard1) + + # Shard 2: minimal metadata + 1 tensor + shard2 = tmp_path / "model-00002-of-00003.gguf" + writer = gguf.GGUFWriter(str(shard2)) + writer.add_uint32("split.no", 1) + writer.add_uint32("split.count", 3) + writer.add_uint32("split.tensors.count", 3) + # One tensor in shard 2 + norm = np.random.randn(128).astype(np.float32) + writer.add_tensor("blk.0.attn_norm.weight", norm) + writer.write_header_and_data(str(shard2)) + shard_paths.append(shard2) + + # Shard 3: minimal metadata + 1 tensor + shard3 = tmp_path / "model-00003-of-00003.gguf" + writer = gguf.GGUFWriter(str(shard3)) + writer.add_uint32("split.no", 2) + writer.add_uint32("split.count", 3) + writer.add_uint32("split.tensors.count", 3) + # One tensor in shard 3 + norm = np.random.randn(128).astype(np.float32) + writer.add_tensor("blk.1.attn_norm.weight", norm) + writer.write_header_and_data(str(shard3)) + shard_paths.append(shard3) + + return shard_paths + + +class TestMissingShard: + """Missing shard must raise an error naming the missing index.""" + + def test_missing_shard_raises(self, tmp_path): + """Delete the middle shard; opening must raise with missing index named.""" + from freetoken.models.gguf.reader import gguf_shards + + # Write 3 shards + shard1 = self._write_minimal_shards(tmp_path, count=3)[0] + + # Delete shard 2 + shard2 = tmp_path / "model-00002-of-00003.gguf" + shard2.unlink() + + # Attempt to discover shards from shard 1 should raise + with pytest.raises(ValueError) as exc_info: + gguf_shards(str(shard1)) + + error_msg = str(exc_info.value) + assert "missing" in error_msg.lower() + assert "2" in error_msg + + def _write_minimal_shards( + self, tmp_path: Path, count: int + ) -> list[Path]: + """Write a minimal set of shards with just metadata.""" + try: + import gguf + except ImportError: + pytest.skip("gguf package not available") + + import numpy as np + + shard_paths = [] + for i in range(count): + shard_num = i + 1 # 1-based + shard_file = tmp_path / f"model-{shard_num:05d}-of-{count:05d}.gguf" + writer = gguf.GGUFWriter(str(shard_file)) + + if i == 0: + # Shard 1: full metadata + writer.add_string("general.architecture", "qwen3moe") + writer.add_uint32("qwen3moe.block_count", 1) + writer.add_uint32("qwen3moe.embedding_length", 128) + writer.add_uint32("qwen3moe.attention.head_count", 4) + writer.add_uint32("qwen3moe.attention.head_count_kv", 2) + writer.add_uint32("qwen3moe.attention.key_length", 32) + writer.add_float32("qwen3moe.attention.layer_norm_rms_epsilon", 1e-6) + writer.add_float32("qwen3moe.rope.freq_base", 10000.0) + writer.add_uint32("qwen3moe.context_length", 4096) + writer.add_uint32("qwen3moe.expert_count", 8) + writer.add_uint32("qwen3moe.expert_used_count", 2) + writer.add_uint32("qwen3moe.expert_feed_forward_length", 512) + writer.add_uint32("qwen3moe.feed_forward_length", 512) + + # Split metadata (all shards) + writer.add_uint32("split.no", i) + writer.add_uint32("split.count", count) + writer.add_uint32("split.tensors.count", count) + + # Add a dummy tensor + tensor_data = np.random.randn(32).astype(np.float32) + writer.add_tensor(f"blk.0.data_{i}.weight", tensor_data) + writer.write_header_and_data(str(shard_file)) + shard_paths.append(shard_file) + + return shard_paths + + +class TestMetadataFromShardOne: + """Shard 1 carries general.architecture and arch keys; others don't.""" + + def test_metadata_comes_from_shard_one(self, tmp_path): + """Reading arch/metadata from any shard returns shard 1's values.""" + from freetoken.models.gguf.reader import ( + gguf_architecture, + load_gguf_metadata, + ) + + shard_paths = self._write_3_shards_with_metadata(tmp_path) + + # Test from shard 1 + arch1 = gguf_architecture(str(shard_paths[0])) + meta1 = load_gguf_metadata(str(shard_paths[0])) + assert arch1 == "qwen3moe" + assert meta1.get("qwen3moe.block_count") == 2 + + # Test from shard 2: should get shard 1's arch + arch2 = gguf_architecture(str(shard_paths[1])) + meta2 = load_gguf_metadata(str(shard_paths[1])) + assert arch2 == "qwen3moe" + assert meta2.get("qwen3moe.block_count") == 2 + + # Test from shard 3: should get shard 1's arch + arch3 = gguf_architecture(str(shard_paths[2])) + meta3 = load_gguf_metadata(str(shard_paths[2])) + assert arch3 == "qwen3moe" + assert meta3.get("qwen3moe.block_count") == 2 + + def _write_3_shards_with_metadata(self, tmp_path: Path) -> list[Path]: + """Write 3 shards where only shard 1 has arch metadata.""" + try: + import gguf + except ImportError: + pytest.skip("gguf package not available") + + import numpy as np + + shard_paths = [] + + # Shard 1: full metadata + shard1 = tmp_path / "model-00001-of-00003.gguf" + writer = gguf.GGUFWriter(str(shard1)) + writer.add_string("general.architecture", "qwen3moe") + writer.add_uint32("qwen3moe.block_count", 2) + writer.add_uint32("qwen3moe.embedding_length", 128) + writer.add_uint32("qwen3moe.attention.head_count", 4) + writer.add_uint32("qwen3moe.attention.head_count_kv", 2) + writer.add_uint32("qwen3moe.attention.key_length", 32) + writer.add_float32("qwen3moe.attention.layer_norm_rms_epsilon", 1e-6) + writer.add_float32("qwen3moe.rope.freq_base", 10000.0) + writer.add_uint32("qwen3moe.context_length", 4096) + writer.add_uint32("qwen3moe.expert_count", 8) + writer.add_uint32("qwen3moe.expert_used_count", 2) + writer.add_uint32("qwen3moe.expert_feed_forward_length", 512) + writer.add_uint32("qwen3moe.feed_forward_length", 512) + writer.add_uint32("split.no", 0) + writer.add_uint32("split.count", 3) + writer.add_uint32("split.tensors.count", 2) + # One tensor + data = np.random.randn(128).astype(np.float32) + writer.add_tensor("blk.0.attn_norm.weight", data) + writer.write_header_and_data(str(shard1)) + shard_paths.append(shard1) + + # Shard 2: only split metadata (no arch keys) + shard2 = tmp_path / "model-00002-of-00003.gguf" + writer = gguf.GGUFWriter(str(shard2)) + writer.add_uint32("split.no", 1) + writer.add_uint32("split.count", 3) + writer.add_uint32("split.tensors.count", 2) + # One tensor + data = np.random.randn(128).astype(np.float32) + writer.add_tensor("blk.1.attn_norm.weight", data) + writer.write_header_and_data(str(shard2)) + shard_paths.append(shard2) + + # Shard 3: only split metadata + shard3 = tmp_path / "model-00003-of-00003.gguf" + writer = gguf.GGUFWriter(str(shard3)) + writer.add_uint32("split.no", 2) + writer.add_uint32("split.count", 3) + writer.add_uint32("split.tensors.count", 2) + # Minimal tensor to pass validation + data = np.random.randn(1).astype(np.float32) + writer.add_tensor("placeholder", data) + writer.write_header_and_data(str(shard3)) + shard_paths.append(shard3) + + return shard_paths + + +class TestTensorAggregation: + """iter_gguf_tensors yields union across all shards in shard order.""" + + def test_tensors_aggregate_across_shards(self, tmp_path): + """Tensors from all shards are yielded in shard order, total matches count.""" + from freetoken.models.gguf.reader import ( + iter_gguf_tensors, + gguf_tensor_names, + ) + + shard_paths = self._write_3_shards_tensors(tmp_path) + + # Check tensor iteration from shard 1 + tensors = list(iter_gguf_tensors(str(shard_paths[0]))) + assert len(tensors) == 3 + assert tensors[0].name == "blk.0.w1" + assert tensors[1].name == "blk.0.w2" + assert tensors[2].name == "blk.1.w1" + + # Check union of names + names = gguf_tensor_names(str(shard_paths[0])) + assert names == {"blk.0.w1", "blk.0.w2", "blk.1.w1"} + assert len(names) == 3 + + # Check from shard 2: should still get all 3 tensors + tensors_from_s2 = list(iter_gguf_tensors(str(shard_paths[1]))) + assert len(tensors_from_s2) == 3 + + def _write_3_shards_tensors(self, tmp_path: Path) -> list[Path]: + """Write 3 shards with tensors distributed across them.""" + try: + import gguf + except ImportError: + pytest.skip("gguf package not available") + + import numpy as np + + shard_paths = [] + + # Shard 1: 2 tensors + shard1 = tmp_path / "model-00001-of-00003.gguf" + writer = gguf.GGUFWriter(str(shard1)) + writer.add_string("general.architecture", "qwen3moe") + writer.add_uint32("qwen3moe.block_count", 2) + writer.add_uint32("qwen3moe.embedding_length", 128) + writer.add_uint32("qwen3moe.attention.head_count", 4) + writer.add_uint32("qwen3moe.attention.head_count_kv", 2) + writer.add_uint32("qwen3moe.attention.key_length", 32) + writer.add_float32("qwen3moe.attention.layer_norm_rms_epsilon", 1e-6) + writer.add_float32("qwen3moe.rope.freq_base", 10000.0) + writer.add_uint32("qwen3moe.context_length", 4096) + writer.add_uint32("qwen3moe.expert_count", 8) + writer.add_uint32("qwen3moe.expert_used_count", 2) + writer.add_uint32("qwen3moe.expert_feed_forward_length", 512) + writer.add_uint32("qwen3moe.feed_forward_length", 512) + writer.add_uint32("split.no", 0) + writer.add_uint32("split.count", 3) + writer.add_uint32("split.tensors.count", 3) + writer.add_tensor("blk.0.w1", np.ones((8, 4), dtype=np.float32)) + writer.add_tensor("blk.0.w2", np.ones((4, 8), dtype=np.float32)) + writer.write_header_and_data(str(shard1)) + shard_paths.append(shard1) + + # Shard 2: 1 tensor + shard2 = tmp_path / "model-00002-of-00003.gguf" + writer = gguf.GGUFWriter(str(shard2)) + writer.add_uint32("split.no", 1) + writer.add_uint32("split.count", 3) + writer.add_uint32("split.tensors.count", 3) + writer.add_tensor("blk.1.w1", np.ones((8, 4), dtype=np.float32)) + writer.write_header_and_data(str(shard2)) + shard_paths.append(shard2) + + # Shard 3: 0 tensors + shard3 = tmp_path / "model-00003-of-00003.gguf" + writer = gguf.GGUFWriter(str(shard3)) + writer.add_uint32("split.no", 2) + writer.add_uint32("split.count", 3) + writer.add_uint32("split.tensors.count", 3) + writer.write_header_and_data(str(shard3)) + shard_paths.append(shard3) + + return shard_paths + + +class TestDeclaredCountMismatch: + """Shard 1 says split.count=N but only M Path: + """Write 2 shards but declare split.count=3.""" + try: + import gguf + except ImportError: + pytest.skip("gguf package not available") + + import numpy as np + + # Shard 1: declares count=3 but we'll only write 2 + shard1 = tmp_path / "model-00001-of-00003.gguf" + writer = gguf.GGUFWriter(str(shard1)) + writer.add_string("general.architecture", "qwen3moe") + writer.add_uint32("qwen3moe.block_count", 1) + writer.add_uint32("qwen3moe.embedding_length", 128) + writer.add_uint32("qwen3moe.attention.head_count", 4) + writer.add_uint32("qwen3moe.attention.head_count_kv", 2) + writer.add_uint32("qwen3moe.attention.key_length", 32) + writer.add_float32("qwen3moe.attention.layer_norm_rms_epsilon", 1e-6) + writer.add_float32("qwen3moe.rope.freq_base", 10000.0) + writer.add_uint32("qwen3moe.context_length", 4096) + writer.add_uint32("qwen3moe.expert_count", 8) + writer.add_uint32("qwen3moe.expert_used_count", 2) + writer.add_uint32("qwen3moe.expert_feed_forward_length", 512) + writer.add_uint32("qwen3moe.feed_forward_length", 512) + writer.add_uint32("split.no", 0) + writer.add_uint32("split.count", 3) # DECLARED 3 + writer.add_uint32("split.tensors.count", 1) + writer.add_tensor("data", np.ones((1,), dtype=np.float32)) + writer.write_header_and_data(str(shard1)) + + # Shard 2: exists + shard2 = tmp_path / "model-00002-of-00003.gguf" + writer = gguf.GGUFWriter(str(shard2)) + writer.add_uint32("split.no", 1) + writer.add_uint32("split.count", 3) + writer.add_uint32("split.tensors.count", 1) + writer.add_tensor("data", np.ones((1,), dtype=np.float32)) + writer.write_header_and_data(str(shard2)) + + # Shard 3: does NOT exist (this is the error case) + + return shard1 + + +class TestQwen3MoeMappingFFNNorm: + """Test qwen3moe's tensor-name mapping: ffn_norm -> post_attention_layernorm. + + This is a pure-function test with no files or fixtures: it just verifies + that the name mapping from llama.cpp's GGUF naming to FreeToken's module + names is correct. + """ + + def test_qwen3moe_maps_ffn_norm_to_post_attention(self): + """Verify blk.N.ffn_norm.weight maps to post_attention_layernorm.""" + from freetoken.models.qwen3_moe.gguf import gguf_name_to_freetoken + + # The critical mapping: ffn_norm -> post_attention_layernorm + mapped = gguf_name_to_freetoken("blk.0.ffn_norm.weight", num_layers=2) + assert mapped == "model.layers.0.post_attention_layernorm.weight" + + # Verify on multiple layers + mapped = gguf_name_to_freetoken("blk.1.ffn_norm.weight", num_layers=2) + assert mapped == "model.layers.1.post_attention_layernorm.weight" + + def test_qwen3moe_merged_projections_handled(self): + """Verify attn_q/k/v are reported as merged-projection parts (None). + + qwen3moe's qkv_proj is a merged projection, so the individual + attn_q/attn_k/attn_v parts return None (handled by iter_gguf_weights). + """ + from freetoken.models.qwen3_moe.gguf import gguf_name_to_freetoken + + # These are parts of the merged qkv_proj, so they return None + assert gguf_name_to_freetoken("blk.0.attn_q.weight", num_layers=2) is None + assert gguf_name_to_freetoken("blk.0.attn_k.weight", num_layers=2) is None + assert gguf_name_to_freetoken("blk.0.attn_v.weight", num_layers=2) is None + + def test_qwen3moe_expert_suffixes_handled(self): + """Verify expert stacks (ffn_*_exps) return None (handled by expert-bank loader).""" + from freetoken.models.qwen3_moe.gguf import gguf_name_to_freetoken + + # Expert stacks are handled by the offload expert-bank loader + assert ( + gguf_name_to_freetoken("blk.0.ffn_gate_exps.weight", num_layers=2) is None + ) + assert ( + gguf_name_to_freetoken("blk.0.ffn_up_exps.weight", num_layers=2) is None + ) + assert ( + gguf_name_to_freetoken("blk.0.ffn_down_exps.weight", num_layers=2) is None + ) From ac12e7227cab9dc966c351486dc3eb8ef5c6d24c Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Mon, 24 Aug 2026 02:39:04 -0700 Subject: [PATCH 24/36] fix(gguf): shard glob typo made every shard set look incomplete The sibling-shard glob was f"{base}-?????.of-{total}.gguf" with a literal dot where the -of- separator belongs. ????? consumes the index, then the dot has to match the hyphen of -of-, so the pattern matched nothing and a complete set reported every index missing. Verified against a real 3-shard split of Ornith IQ3_S produced by llama-gguf-split. Also resolves a directory of shards to its first shard, since a user is more likely to point at the download folder than at a specific shard, and raises if the folder holds more than one split model. --- python/freetoken/models/gguf/reader.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/python/freetoken/models/gguf/reader.py b/python/freetoken/models/gguf/reader.py index 9b58ec9d..449871fb 100644 --- a/python/freetoken/models/gguf/reader.py +++ b/python/freetoken/models/gguf/reader.py @@ -60,6 +60,19 @@ def gguf_shards(path: str) -> list[str]: Raises a clear error naming the missing indices if any are absent (truncated downloads are the common failure case and must not load silently). """ + # A directory: find the first shard inside it and continue from there. Users routinely + # pass the folder a split model was downloaded into rather than a specific shard. + if os.path.isdir(path): + first = sorted(glob.glob(os.path.join(path, "*-00001-of-?????.gguf"))) + if len(first) > 1: + raise ValueError( + f"{path}: contains {len(first)} different split models " + f"({[os.path.basename(f) for f in first]}); point at one shard instead" + ) + if not first: + return [path] + path = first[0] + basename = os.path.basename(path) match = re.match(r"(?P.+)-(\d{5})-of-(\d{5})\.gguf$", basename) if not match: @@ -71,7 +84,7 @@ def gguf_shards(path: str) -> list[str]: shard_dir = os.path.dirname(path) # Glob all sibling shards - pattern = os.path.join(shard_dir, f"{base}-?????.of-{total_shards_str}.gguf") + pattern = os.path.join(shard_dir, f"{base}-?????-of-{total_shards_str}.gguf") found_shards = sorted(glob.glob(pattern)) # Parse indices and validate completeness From 89dff4bc4ed60041150d4614098d8990a06eaff3 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Mon, 24 Aug 2026 05:45:12 -0700 Subject: [PATCH 25/36] test(gguf): rewrite the shard tests to emit GGUF bytes directly The generated tests called gguf.GGUFWriter(path) without its required arch positional, so six of seven failed on a TypeError before testing anything. Hand-writing the format keeps these independent of that writer's API and is only about forty lines. Covers the three properties of llama.cpp's split convention that are easy to get wrong and invisible when you do: split.no is 0-based while filenames are 1-based, split.tensors.count is the total across shards rather than per shard, and only shard 1 carries the metadata. There is also a fixture guard asserting shard 2 genuinely lacks general.architecture, since otherwise the shard-1-resolution test would pass vacuously. --- tests/models/test_gguf_shards.py | 683 ++++++++----------------------- 1 file changed, 162 insertions(+), 521 deletions(-) diff --git a/tests/models/test_gguf_shards.py b/tests/models/test_gguf_shards.py index 513cabbf..1bef402f 100644 --- a/tests/models/test_gguf_shards.py +++ b/tests/models/test_gguf_shards.py @@ -1,559 +1,200 @@ -"""Tests for GGUF multi-shard discovery and aggregation, plus qwen3moe mapping. - -This module verifies that multi-shard GGUF files (following llama.cpp's split -convention) are correctly discovered, validated, and aggregated by -freetoken.models.gguf.reader, and that the qwen3moe tensor-name mapping is correct. - -Test fixtures build synthetic GGUF files on disk using gguf.GGUFWriter where -possible, or raw GGUF bytes where needed. The tensor data is tiny (a few F32 -values) to keep I/O fast; the focus is on metadata, discovery, and naming, not -quantization or size. - -llama.cpp's shard layout (ground truth from measured checksums): - - Shard 1: full metadata (general.architecture, .*, tokenizer.*) + split.no=0 - - Shards 2..N: exactly 3 keys (split.no, split.count, split.tensors.count) - - split.tensors.count is the TOTAL across all shards, not per-shard - - Filenames are 1-based; split.no is 0-based +"""Multi-shard GGUF reading: discovery, shard-1 metadata, tensor aggregation. + +Large GGUF checkpoints ship split (``-00001-of-000NN``), and llama.cpp's convention has +three properties that are easy to get wrong and nearly invisible when you do: + +* ``split.no`` is 0-BASED while the filenames are 1-BASED. Shard ``-00002-of-00003`` + carries ``split.no = 1``. +* ``split.tensors.count`` is the TOTAL across every shard, not this shard's count. +* Only shard 1 carries the real metadata. Later shards hold exactly three ``split.*`` keys + and no ``general.architecture`` at all, so anything that reads arch or tokenizer from an + arbitrary shard gets nothing. + +The fixtures below write GGUF bytes directly rather than going through ``gguf.GGUFWriter``: +the format is small, and hand-writing it keeps these tests independent of that writer's +API (which takes a required ``arch`` positional and has moved around between releases). +Layout per the spec: magic, uint32 version, uint64 tensor_count, uint64 kv_count, the KV +pairs, the tensor infos, then padding to ``general.alignment`` and the tensor data. """ from __future__ import annotations -import os import struct -import tempfile from pathlib import Path import pytest +from freetoken.models.gguf.reader import ( + gguf_architecture, + gguf_shards, + gguf_tensor_names, + is_gguf_path, + iter_gguf_tensors, + load_gguf_metadata, +) -class TestSingleFileUnchanged: - """A plain one-file .gguf still reports correct arch, metadata, tensors. - - This is the regression guard: single-file behavior must be untouched by - multi-shard support. - """ +# GGUF value type tags +_UINT32, _UINT64, _STRING = 4, 10, 8 +_F32_TENSOR_TYPE = 0 +_ALIGN = 32 - def test_single_file_unchanged(self, tmp_path): - """Verify a single .gguf file works unchanged.""" - from freetoken.models.gguf.reader import ( - gguf_architecture, - is_gguf_path, - load_gguf_metadata, - gguf_tensor_names, - iter_gguf_tensors, - ) - # Write a minimal GGUF file (no shards) - gguf_path = self._write_minimal_gguf(tmp_path, "single.gguf") +def _u32(v: int) -> bytes: + return struct.pack(" bytes: + return struct.pack(" bytes: + raw = s.encode("utf-8") + return _u64(len(raw)) + raw - # Check tensor iteration - tensors = list(iter_gguf_tensors(str(gguf_path))) - assert len(tensors) == 1 - assert tensors[0].name == "model.embed.weight" - assert tensors[0].shape == (256, 64) - def _write_minimal_gguf(self, tmp_path: Path, filename: str) -> Path: - """Write a minimal single-file GGUF with one tensor.""" - try: - import gguf - except ImportError: - pytest.skip("gguf package not available") +def _kv(key: str, tag: int, value) -> bytes: + out = _string(key) + _u32(tag) + if tag == _STRING: + return out + _string(value) + if tag == _UINT32: + return out + _u32(value) + if tag == _UINT64: + return out + _u64(value) + raise AssertionError(f"unhandled tag {tag}") - gguf_path = tmp_path / filename - writer = gguf.GGUFWriter(str(gguf_path)) - # Add metadata - writer.add_string("general.architecture", "test_model") - writer.add_uint32("test_model.hidden_size", 256) - writer.add_uint32("test_model.num_layers", 2) +def _write_gguf(path: Path, kvs: list[bytes], tensors: list[tuple[str, int]]) -> None: + """Write a GGUF with ``tensors`` as [(name, n_elements)], each F32. - # Add one tensor (embedding) - import numpy as np + Tensor data is written contiguously after the aligned header; the values themselves are + irrelevant here since these tests only exercise discovery, metadata and the tensor + table. + """ + head = b"GGUF" + _u32(3) + _u64(len(tensors)) + _u64(len(kvs)) + head += b"".join(kvs) + offset = 0 + infos = b"" + for name, n in tensors: + infos += _string(name) + _u32(1) + _u64(n) + _u32(_F32_TENSOR_TYPE) + _u64(offset) + nbytes = n * 4 + offset += (nbytes + _ALIGN - 1) // _ALIGN * _ALIGN + body = head + infos + pad = (-len(body)) % _ALIGN + body += b"\0" * pad + body += b"\0" * offset + path.write_bytes(body) + + +def _full_kvs(arch: str = "qwen3moe", *, extra: list[bytes] | None = None) -> list[bytes]: + """Shard 1's KV block: the real metadata.""" + kvs = [ + _kv("general.architecture", _STRING, arch), + _kv("general.alignment", _UINT32, _ALIGN), + _kv(f"{arch}.block_count", _UINT32, 4), + _kv(f"{arch}.embedding_length", _UINT32, 128), + ] + return kvs + (extra or []) + + +def _split_kvs(no: int, count: int, total_tensors: int) -> list[bytes]: + """A non-first shard's KV block: exactly the three split keys, no architecture.""" + return [ + _kv("split.no", _UINT32, no), + _kv("split.count", _UINT32, count), + _kv("split.tensors.count", _UINT32, total_tensors), + ] + + +def _make_split(tmp_path: Path, base: str, per_shard: list[list[str]], *, + declared_count: int | None = None) -> list[Path]: + """Write a split set; returns the shard paths in order.""" + n = len(per_shard) + declared = declared_count if declared_count is not None else n + total = sum(len(names) for names in per_shard) + paths = [] + for i, names in enumerate(per_shard): + p = tmp_path / f"{base}-{i + 1:05d}-of-{n:05d}.gguf" + kvs = (_full_kvs() + _split_kvs(0, declared, total)) if i == 0 \ + else _split_kvs(i, declared, total) + _write_gguf(p, kvs, [(nm, 8) for nm in names]) + paths.append(p) + return paths - emb_data = np.random.randn(256, 64).astype(np.float32) - writer.add_tensor("model.embed.weight", emb_data) - writer.write_header_and_data(str(gguf_path)) - return gguf_path +class TestSingleFileUnchanged: + def test_single_file_unchanged(self, tmp_path: Path): + """A plain one-file GGUF must behave exactly as before the shard work.""" + p = tmp_path / "single.gguf" + _write_gguf(p, _full_kvs(), [("token_embd.weight", 8), ("output.weight", 8)]) + assert is_gguf_path(str(p)) + assert gguf_shards(str(p)) == [str(p)] + assert gguf_architecture(str(p)) == "qwen3moe" + assert load_gguf_metadata(str(p))["qwen3moe.block_count"] == 4 + assert gguf_tensor_names(str(p)) == {"token_embd.weight", "output.weight"} + assert len(list(iter_gguf_tensors(str(p)))) == 2 class TestShardDiscovery: - """Given a 3-shard set, gguf_shards() returns all three in index order.""" - - def test_shard_discovery_orders_and_completes(self, tmp_path): - """Discover all 3 shards when given any shard or the directory.""" - from freetoken.models.gguf.reader import ( - gguf_shards, - is_gguf_path, - gguf_tensor_names, - ) - - # Write 3 shards - shard_paths = self._write_3_shards(tmp_path) - - # Test 1: discover from shard 1 - result = gguf_shards(str(shard_paths[0])) - assert len(result) == 3 - assert [Path(p).name for p in result] == [ - "model-00001-of-00003.gguf", - "model-00002-of-00003.gguf", - "model-00003-of-00003.gguf", - ] - - # Test 2: discover from shard 2 - result = gguf_shards(str(shard_paths[1])) - assert len(result) == 3 - assert result[0] == str(shard_paths[0]) # first shard returned - - # Test 3: discover from shard 3 - result = gguf_shards(str(shard_paths[2])) - assert len(result) == 3 - - # Test 4: directory resolution - assert is_gguf_path(str(tmp_path)) - - # Test 5: tensors are aggregated across shards - names = gguf_tensor_names(str(shard_paths[0])) - assert names == { - "token_embd.weight", # shard 1 - "blk.0.attn_norm.weight", # shard 2 - "blk.1.attn_norm.weight", # shard 3 - } - - def _write_3_shards(self, tmp_path: Path) -> list[Path]: - """Write a 3-shard GGUF set with metadata and tensors distributed.""" - try: - import gguf - except ImportError: - pytest.skip("gguf package not available") - - import numpy as np - - shard_paths = [] - - # Shard 1: full metadata + 1 tensor - shard1 = tmp_path / "model-00001-of-00003.gguf" - writer = gguf.GGUFWriter(str(shard1)) - writer.add_string("general.architecture", "qwen3moe") - writer.add_uint32("qwen3moe.block_count", 2) - writer.add_uint32("qwen3moe.embedding_length", 128) - writer.add_uint32("qwen3moe.attention.head_count", 4) - writer.add_uint32("qwen3moe.attention.head_count_kv", 2) - writer.add_uint32("qwen3moe.attention.key_length", 32) - writer.add_float32("qwen3moe.attention.layer_norm_rms_epsilon", 1e-6) - writer.add_float32("qwen3moe.rope.freq_base", 10000.0) - writer.add_uint32("qwen3moe.context_length", 4096) - writer.add_uint32("qwen3moe.expert_count", 8) - writer.add_uint32("qwen3moe.expert_used_count", 2) - writer.add_uint32("qwen3moe.expert_feed_forward_length", 512) - writer.add_uint32("qwen3moe.feed_forward_length", 512) - # Split metadata for shard 1 - writer.add_uint32("split.no", 0) - writer.add_uint32("split.count", 3) - writer.add_uint32("split.tensors.count", 3) - # One tensor in shard 1 - emb = np.random.randn(1024, 128).astype(np.float32) - writer.add_tensor("token_embd.weight", emb) - writer.write_header_and_data(str(shard1)) - shard_paths.append(shard1) - - # Shard 2: minimal metadata + 1 tensor - shard2 = tmp_path / "model-00002-of-00003.gguf" - writer = gguf.GGUFWriter(str(shard2)) - writer.add_uint32("split.no", 1) - writer.add_uint32("split.count", 3) - writer.add_uint32("split.tensors.count", 3) - # One tensor in shard 2 - norm = np.random.randn(128).astype(np.float32) - writer.add_tensor("blk.0.attn_norm.weight", norm) - writer.write_header_and_data(str(shard2)) - shard_paths.append(shard2) - - # Shard 3: minimal metadata + 1 tensor - shard3 = tmp_path / "model-00003-of-00003.gguf" - writer = gguf.GGUFWriter(str(shard3)) - writer.add_uint32("split.no", 2) - writer.add_uint32("split.count", 3) - writer.add_uint32("split.tensors.count", 3) - # One tensor in shard 3 - norm = np.random.randn(128).astype(np.float32) - writer.add_tensor("blk.1.attn_norm.weight", norm) - writer.write_header_and_data(str(shard3)) - shard_paths.append(shard3) - - return shard_paths + def test_discovery_from_any_shard_or_directory(self, tmp_path: Path): + paths = _make_split(tmp_path, "m", [["a.weight"], ["b.weight"], ["c.weight"]]) + want = [str(p) for p in paths] + for handed in (paths[0], paths[1], paths[2]): + assert gguf_shards(str(handed)) == want, f"from {handed.name}" + # a user pointing at the folder must work too + assert gguf_shards(str(tmp_path)) == want + + def test_every_shard_is_a_gguf_path(self, tmp_path: Path): + paths = _make_split(tmp_path, "m", [["a.weight"], ["b.weight"]]) + for p in paths: + assert is_gguf_path(str(p)) class TestMissingShard: - """Missing shard must raise an error naming the missing index.""" - - def test_missing_shard_raises(self, tmp_path): - """Delete the middle shard; opening must raise with missing index named.""" - from freetoken.models.gguf.reader import gguf_shards - - # Write 3 shards - shard1 = self._write_minimal_shards(tmp_path, count=3)[0] - - # Delete shard 2 - shard2 = tmp_path / "model-00002-of-00003.gguf" - shard2.unlink() - - # Attempt to discover shards from shard 1 should raise - with pytest.raises(ValueError) as exc_info: - gguf_shards(str(shard1)) - - error_msg = str(exc_info.value) - assert "missing" in error_msg.lower() - assert "2" in error_msg - - def _write_minimal_shards( - self, tmp_path: Path, count: int - ) -> list[Path]: - """Write a minimal set of shards with just metadata.""" - try: - import gguf - except ImportError: - pytest.skip("gguf package not available") - - import numpy as np - - shard_paths = [] - for i in range(count): - shard_num = i + 1 # 1-based - shard_file = tmp_path / f"model-{shard_num:05d}-of-{count:05d}.gguf" - writer = gguf.GGUFWriter(str(shard_file)) - - if i == 0: - # Shard 1: full metadata - writer.add_string("general.architecture", "qwen3moe") - writer.add_uint32("qwen3moe.block_count", 1) - writer.add_uint32("qwen3moe.embedding_length", 128) - writer.add_uint32("qwen3moe.attention.head_count", 4) - writer.add_uint32("qwen3moe.attention.head_count_kv", 2) - writer.add_uint32("qwen3moe.attention.key_length", 32) - writer.add_float32("qwen3moe.attention.layer_norm_rms_epsilon", 1e-6) - writer.add_float32("qwen3moe.rope.freq_base", 10000.0) - writer.add_uint32("qwen3moe.context_length", 4096) - writer.add_uint32("qwen3moe.expert_count", 8) - writer.add_uint32("qwen3moe.expert_used_count", 2) - writer.add_uint32("qwen3moe.expert_feed_forward_length", 512) - writer.add_uint32("qwen3moe.feed_forward_length", 512) - - # Split metadata (all shards) - writer.add_uint32("split.no", i) - writer.add_uint32("split.count", count) - writer.add_uint32("split.tensors.count", count) - - # Add a dummy tensor - tensor_data = np.random.randn(32).astype(np.float32) - writer.add_tensor(f"blk.0.data_{i}.weight", tensor_data) - writer.write_header_and_data(str(shard_file)) - shard_paths.append(shard_file) - - return shard_paths + def test_missing_shard_raises_naming_the_index(self, tmp_path: Path): + """A truncated download must fail loudly, never load as a partial model.""" + paths = _make_split(tmp_path, "m", [["a.weight"], ["b.weight"], ["c.weight"]]) + paths[1].unlink() # drop the middle shard + with pytest.raises(Exception) as e: + gguf_shards(str(paths[0])) + assert "2" in str(e.value), f"error should name the missing index: {e.value}" class TestMetadataFromShardOne: - """Shard 1 carries general.architecture and arch keys; others don't.""" - - def test_metadata_comes_from_shard_one(self, tmp_path): - """Reading arch/metadata from any shard returns shard 1's values.""" - from freetoken.models.gguf.reader import ( - gguf_architecture, - load_gguf_metadata, - ) - - shard_paths = self._write_3_shards_with_metadata(tmp_path) - - # Test from shard 1 - arch1 = gguf_architecture(str(shard_paths[0])) - meta1 = load_gguf_metadata(str(shard_paths[0])) - assert arch1 == "qwen3moe" - assert meta1.get("qwen3moe.block_count") == 2 - - # Test from shard 2: should get shard 1's arch - arch2 = gguf_architecture(str(shard_paths[1])) - meta2 = load_gguf_metadata(str(shard_paths[1])) - assert arch2 == "qwen3moe" - assert meta2.get("qwen3moe.block_count") == 2 - - # Test from shard 3: should get shard 1's arch - arch3 = gguf_architecture(str(shard_paths[2])) - meta3 = load_gguf_metadata(str(shard_paths[2])) - assert arch3 == "qwen3moe" - assert meta3.get("qwen3moe.block_count") == 2 - - def _write_3_shards_with_metadata(self, tmp_path: Path) -> list[Path]: - """Write 3 shards where only shard 1 has arch metadata.""" - try: - import gguf - except ImportError: - pytest.skip("gguf package not available") - - import numpy as np - - shard_paths = [] - - # Shard 1: full metadata - shard1 = tmp_path / "model-00001-of-00003.gguf" - writer = gguf.GGUFWriter(str(shard1)) - writer.add_string("general.architecture", "qwen3moe") - writer.add_uint32("qwen3moe.block_count", 2) - writer.add_uint32("qwen3moe.embedding_length", 128) - writer.add_uint32("qwen3moe.attention.head_count", 4) - writer.add_uint32("qwen3moe.attention.head_count_kv", 2) - writer.add_uint32("qwen3moe.attention.key_length", 32) - writer.add_float32("qwen3moe.attention.layer_norm_rms_epsilon", 1e-6) - writer.add_float32("qwen3moe.rope.freq_base", 10000.0) - writer.add_uint32("qwen3moe.context_length", 4096) - writer.add_uint32("qwen3moe.expert_count", 8) - writer.add_uint32("qwen3moe.expert_used_count", 2) - writer.add_uint32("qwen3moe.expert_feed_forward_length", 512) - writer.add_uint32("qwen3moe.feed_forward_length", 512) - writer.add_uint32("split.no", 0) - writer.add_uint32("split.count", 3) - writer.add_uint32("split.tensors.count", 2) - # One tensor - data = np.random.randn(128).astype(np.float32) - writer.add_tensor("blk.0.attn_norm.weight", data) - writer.write_header_and_data(str(shard1)) - shard_paths.append(shard1) - - # Shard 2: only split metadata (no arch keys) - shard2 = tmp_path / "model-00002-of-00003.gguf" - writer = gguf.GGUFWriter(str(shard2)) - writer.add_uint32("split.no", 1) - writer.add_uint32("split.count", 3) - writer.add_uint32("split.tensors.count", 2) - # One tensor - data = np.random.randn(128).astype(np.float32) - writer.add_tensor("blk.1.attn_norm.weight", data) - writer.write_header_and_data(str(shard2)) - shard_paths.append(shard2) - - # Shard 3: only split metadata - shard3 = tmp_path / "model-00003-of-00003.gguf" - writer = gguf.GGUFWriter(str(shard3)) - writer.add_uint32("split.no", 2) - writer.add_uint32("split.count", 3) - writer.add_uint32("split.tensors.count", 2) - # Minimal tensor to pass validation - data = np.random.randn(1).astype(np.float32) - writer.add_tensor("placeholder", data) - writer.write_header_and_data(str(shard3)) - shard_paths.append(shard3) - - return shard_paths + def test_metadata_resolves_to_shard_one(self, tmp_path: Path): + """Later shards carry no architecture, so reads must resolve back to shard 1.""" + paths = _make_split(tmp_path, "m", [["a.weight"], ["b.weight"], ["c.weight"]]) + for p in paths: + assert gguf_architecture(str(p)) == "qwen3moe", f"from {p.name}" + assert load_gguf_metadata(str(p))["qwen3moe.block_count"] == 4, f"from {p.name}" + + def test_later_shards_really_lack_arch(self, tmp_path: Path): + """Guards the fixture itself: if shard 2 carried arch, the test above proves nothing.""" + paths = _make_split(tmp_path, "m", [["a.weight"], ["b.weight"]]) + import gguf as gguf_pkg + + r = gguf_pkg.GGUFReader(str(paths[1])) + assert "general.architecture" not in r.fields + assert "split.no" in r.fields class TestTensorAggregation: - """iter_gguf_tensors yields union across all shards in shard order.""" - - def test_tensors_aggregate_across_shards(self, tmp_path): - """Tensors from all shards are yielded in shard order, total matches count.""" - from freetoken.models.gguf.reader import ( - iter_gguf_tensors, - gguf_tensor_names, - ) - - shard_paths = self._write_3_shards_tensors(tmp_path) - - # Check tensor iteration from shard 1 - tensors = list(iter_gguf_tensors(str(shard_paths[0]))) - assert len(tensors) == 3 - assert tensors[0].name == "blk.0.w1" - assert tensors[1].name == "blk.0.w2" - assert tensors[2].name == "blk.1.w1" - - # Check union of names - names = gguf_tensor_names(str(shard_paths[0])) - assert names == {"blk.0.w1", "blk.0.w2", "blk.1.w1"} - assert len(names) == 3 - - # Check from shard 2: should still get all 3 tensors - tensors_from_s2 = list(iter_gguf_tensors(str(shard_paths[1]))) - assert len(tensors_from_s2) == 3 - - def _write_3_shards_tensors(self, tmp_path: Path) -> list[Path]: - """Write 3 shards with tensors distributed across them.""" - try: - import gguf - except ImportError: - pytest.skip("gguf package not available") - - import numpy as np - - shard_paths = [] - - # Shard 1: 2 tensors - shard1 = tmp_path / "model-00001-of-00003.gguf" - writer = gguf.GGUFWriter(str(shard1)) - writer.add_string("general.architecture", "qwen3moe") - writer.add_uint32("qwen3moe.block_count", 2) - writer.add_uint32("qwen3moe.embedding_length", 128) - writer.add_uint32("qwen3moe.attention.head_count", 4) - writer.add_uint32("qwen3moe.attention.head_count_kv", 2) - writer.add_uint32("qwen3moe.attention.key_length", 32) - writer.add_float32("qwen3moe.attention.layer_norm_rms_epsilon", 1e-6) - writer.add_float32("qwen3moe.rope.freq_base", 10000.0) - writer.add_uint32("qwen3moe.context_length", 4096) - writer.add_uint32("qwen3moe.expert_count", 8) - writer.add_uint32("qwen3moe.expert_used_count", 2) - writer.add_uint32("qwen3moe.expert_feed_forward_length", 512) - writer.add_uint32("qwen3moe.feed_forward_length", 512) - writer.add_uint32("split.no", 0) - writer.add_uint32("split.count", 3) - writer.add_uint32("split.tensors.count", 3) - writer.add_tensor("blk.0.w1", np.ones((8, 4), dtype=np.float32)) - writer.add_tensor("blk.0.w2", np.ones((4, 8), dtype=np.float32)) - writer.write_header_and_data(str(shard1)) - shard_paths.append(shard1) - - # Shard 2: 1 tensor - shard2 = tmp_path / "model-00002-of-00003.gguf" - writer = gguf.GGUFWriter(str(shard2)) - writer.add_uint32("split.no", 1) - writer.add_uint32("split.count", 3) - writer.add_uint32("split.tensors.count", 3) - writer.add_tensor("blk.1.w1", np.ones((8, 4), dtype=np.float32)) - writer.write_header_and_data(str(shard2)) - shard_paths.append(shard2) - - # Shard 3: 0 tensors - shard3 = tmp_path / "model-00003-of-00003.gguf" - writer = gguf.GGUFWriter(str(shard3)) - writer.add_uint32("split.no", 2) - writer.add_uint32("split.count", 3) - writer.add_uint32("split.tensors.count", 3) - writer.write_header_and_data(str(shard3)) - shard_paths.append(shard3) - - return shard_paths + def test_tensors_aggregate_across_shards(self, tmp_path: Path): + per = [["a.weight", "b.weight"], ["c.weight"], ["d.weight", "e.weight"]] + paths = _make_split(tmp_path, "m", per) + expected = {n for names in per for n in names} + for p in paths: + assert gguf_tensor_names(str(p)) == expected, f"from {p.name}" + names = [t.name for t in iter_gguf_tensors(str(paths[0]))] + assert names == [n for names_ in per for n in names_], "shard order must be preserved" + assert len(names) == load_gguf_metadata(str(paths[0]))["split.tensors.count"] class TestDeclaredCountMismatch: - """Shard 1 says split.count=N but only M Path: - """Write 2 shards but declare split.count=3.""" - try: - import gguf - except ImportError: - pytest.skip("gguf package not available") - - import numpy as np - - # Shard 1: declares count=3 but we'll only write 2 - shard1 = tmp_path / "model-00001-of-00003.gguf" - writer = gguf.GGUFWriter(str(shard1)) - writer.add_string("general.architecture", "qwen3moe") - writer.add_uint32("qwen3moe.block_count", 1) - writer.add_uint32("qwen3moe.embedding_length", 128) - writer.add_uint32("qwen3moe.attention.head_count", 4) - writer.add_uint32("qwen3moe.attention.head_count_kv", 2) - writer.add_uint32("qwen3moe.attention.key_length", 32) - writer.add_float32("qwen3moe.attention.layer_norm_rms_epsilon", 1e-6) - writer.add_float32("qwen3moe.rope.freq_base", 10000.0) - writer.add_uint32("qwen3moe.context_length", 4096) - writer.add_uint32("qwen3moe.expert_count", 8) - writer.add_uint32("qwen3moe.expert_used_count", 2) - writer.add_uint32("qwen3moe.expert_feed_forward_length", 512) - writer.add_uint32("qwen3moe.feed_forward_length", 512) - writer.add_uint32("split.no", 0) - writer.add_uint32("split.count", 3) # DECLARED 3 - writer.add_uint32("split.tensors.count", 1) - writer.add_tensor("data", np.ones((1,), dtype=np.float32)) - writer.write_header_and_data(str(shard1)) - - # Shard 2: exists - shard2 = tmp_path / "model-00002-of-00003.gguf" - writer = gguf.GGUFWriter(str(shard2)) - writer.add_uint32("split.no", 1) - writer.add_uint32("split.count", 3) - writer.add_uint32("split.tensors.count", 1) - writer.add_tensor("data", np.ones((1,), dtype=np.float32)) - writer.write_header_and_data(str(shard2)) - - # Shard 3: does NOT exist (this is the error case) - - return shard1 - - -class TestQwen3MoeMappingFFNNorm: - """Test qwen3moe's tensor-name mapping: ffn_norm -> post_attention_layernorm. - - This is a pure-function test with no files or fixtures: it just verifies - that the name mapping from llama.cpp's GGUF naming to FreeToken's module - names is correct. - """ - - def test_qwen3moe_maps_ffn_norm_to_post_attention(self): - """Verify blk.N.ffn_norm.weight maps to post_attention_layernorm.""" - from freetoken.models.qwen3_moe.gguf import gguf_name_to_freetoken - - # The critical mapping: ffn_norm -> post_attention_layernorm - mapped = gguf_name_to_freetoken("blk.0.ffn_norm.weight", num_layers=2) - assert mapped == "model.layers.0.post_attention_layernorm.weight" - - # Verify on multiple layers - mapped = gguf_name_to_freetoken("blk.1.ffn_norm.weight", num_layers=2) - assert mapped == "model.layers.1.post_attention_layernorm.weight" - - def test_qwen3moe_merged_projections_handled(self): - """Verify attn_q/k/v are reported as merged-projection parts (None). - - qwen3moe's qkv_proj is a merged projection, so the individual - attn_q/attn_k/attn_v parts return None (handled by iter_gguf_weights). - """ - from freetoken.models.qwen3_moe.gguf import gguf_name_to_freetoken - - # These are parts of the merged qkv_proj, so they return None - assert gguf_name_to_freetoken("blk.0.attn_q.weight", num_layers=2) is None - assert gguf_name_to_freetoken("blk.0.attn_k.weight", num_layers=2) is None - assert gguf_name_to_freetoken("blk.0.attn_v.weight", num_layers=2) is None - - def test_qwen3moe_expert_suffixes_handled(self): - """Verify expert stacks (ffn_*_exps) return None (handled by expert-bank loader).""" - from freetoken.models.qwen3_moe.gguf import gguf_name_to_freetoken - - # Expert stacks are handled by the offload expert-bank loader - assert ( - gguf_name_to_freetoken("blk.0.ffn_gate_exps.weight", num_layers=2) is None - ) - assert ( - gguf_name_to_freetoken("blk.0.ffn_up_exps.weight", num_layers=2) is None - ) - assert ( - gguf_name_to_freetoken("blk.0.ffn_down_exps.weight", num_layers=2) is None - ) + def test_declared_count_mismatch_raises(self, tmp_path: Path): + """shard 1 says 3 shards but only 2 exist on disk.""" + tmp = tmp_path / "sub" + tmp.mkdir() + _make_split(tmp, "m", [["a.weight"], ["b.weight"]], declared_count=3) + first = tmp / "m-00001-of-00002.gguf" + with pytest.raises(Exception): + list(iter_gguf_tensors(str(first))) From 0f26f7ecb17eb410c349e4b4e65e9c1adac7d451 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Mon, 24 Aug 2026 05:47:42 -0700 Subject: [PATCH 26/36] docs: multi-shard and qwen3moe, with the measured results Adds the qwen3moe row and the split-checkpoint paragraph to models.md, and updates the fork README's verified table with the sharded Ornith run and Qwen3-30B-A3B. Corrects the Ornith IQ3_XXS row to 6/6: the earlier 5/6 was a truncated preview in my harness, not a wrong answer. --- docs/models.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/models.md b/docs/models.md index 21120885..58867abb 100644 --- a/docs/models.md +++ b/docs/models.md @@ -26,9 +26,15 @@ the kernels rather than expanded to bf16 at load. | GGUF `general.architecture` | Covers | |---|---| | `gemma4` | Gemma-4 | +| `qwen3moe` | Qwen3 MoE (e.g. Qwen3-235B-A22B, Qwen3-30B-A3B) | | `qwen35moe` | Qwen3.5 / Qwen3.6 MoE (e.g. Qwen3.6-35B-A3B, Qwen3.5-122B-A10B) | | `qwen35` | Qwen3.5 / Qwen3.6 dense (e.g. Qwen3.6-27B, Qwen3.5-9B) | +Split checkpoints load: point `--model` at any shard of a `-00001-of-000NN` set, or at the +directory holding them. Metadata, config and tokenizer are read from shard 1 (later shards +carry only `split.*` keys), and the tensor tables are aggregated across the set. A missing +shard raises with the index named rather than loading a partial model. + Quant types follow what the vendored kernels in `csrc/gguf/` implement: - Standard and K-quants (Q4_0, Q4_1, Q5_0, Q5_1, Q8_0, Q2_K through Q6_K) use MMQ for From 6f6c8640145eeca9df013e383ff51bf6bbff22f9 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Mon, 24 Aug 2026 12:34:59 -0700 Subject: [PATCH 27/36] feat(cpu-moe): Q4_K and Q6_K expert kernels for the cpu/hybrid backends Q4_0 was the only GGUF type with a CPU dot kernel, so --moe-backend cpu and hybrid refused every K-quant checkpoint and pushed those users onto offload. That is backwards: the people who need the hybrid backend most are the ones short of VRAM, and they are running exactly these quants. Adds scalar W4A16 dots for Q4_K and Q6_K, mirroring dequantize_block_q4_K and dequantize_block_q6_K in gguf/dequantize.cuh element for element. Correctness first; an AVX2/VNNI version is a separate job. Both skip the int8 activation path Q4_0 uses, since the super-block scale structure does not map onto it. Also wires quant_format == "gguf" through to a concrete format. The GGUF offload path tags the cache "gguf" and carries the ggml types separately, so the executor was matching on format strings that nothing produces and would have refused a real checkpoint even for Q4_0. _resolve_gguf_format maps the bank types onto one weight format and refuses with a specific reason when it cannot: mixed banks (Q4_K_M puts ffn_down_exps in Q6_K) get a different message from types with no kernel at all, because the fix differs. Tests check both kernels against the vendored CUDA dequant plus the production bf16 GPU decode on byte-identical banks, at three batch sizes. Blocks are synthesized rather than carved from a checkpoint so this runs without a multi-GB download. Worth recording what these tests caught, since neither bug is visible structurally and both give you a model that loads and runs at full speed: Q4_K quants are unsigned 0..15 with a per-sub-block min, not Q4_0's q-8, and a byte's high nibble is element l+32 rather than l+1. Q6_K's qh shift is per group, and an extra term in the lane index corrupted the upper half of every group while still correlating 0.17, which reads like a tolerance problem rather than a wrong kernel. --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 176 ++++++++++++++- python/freetoken/moe/cpu_executor.py | 126 +++++++++-- tests/moe/test_cpu_moe_kquant.py | 200 ++++++++++++++++++ 3 files changed, 486 insertions(+), 16 deletions(-) create mode 100644 tests/moe/test_cpu_moe_kquant.py diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 880e8637..68e6324c 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1212,7 +1212,151 @@ q4dot_fn select_q4dot() { return q4_0_dot_i8_scalar; } -enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4 }; +// ----------------------------- Q4_K and Q6_K (W4A16) ---------------------------- +// Correctness-first scalar K-quant expert dot kernels: dequantize blocks and accumulate +// against bf16 activation rows. AVX2/VNNI optimizations are a deliberate follow-up. +// Reference: llama.cpp ggml-quants.c, ggml-cuda/convert.cu, models/gguf/dequant.py. + +// Helper: extract 6-bit scale and min from Q4_K's packed scales array. +// Q4_K packs scales and mins using 6 bits each, packed into 12 bytes for 256 elements. +inline void get_scale_min_k4(int j, const uint8_t* q, uint8_t& scale, uint8_t& minv) { + if (j < 4) { + scale = q[j] & 63; + minv = q[j + 4] & 63; + } else { + scale = (q[j + 4] & 0x0F) | ((q[j - 4] >> 6) << 4); + minv = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); + } +} + +float q4_k_dot_f32_scalar(const uint8_t* w, const bf16_t* x, int K) { + // Q4_K, 256-element super-blocks of 144 bytes: half2 dm | scales[12] | qs[128]. + // dm.x scales the quants, dm.y scales the per-sub-block minimum that is SUBTRACTED. + // + // Two things here are easy to get wrong and both produce fluent-looking garbage rather + // than an obvious failure, so this mirrors dequantize_block_q4_K in gguf/dequantize.cuh + // element for element: + // * the quants are UNSIGNED 0..15. There is no -8 bias; that is Q4_0's encoding. Q4_K + // centres the range with the per-sub-block min instead. + // * within a 64-element group the low nibble of byte l is element l and the HIGH nibble + // is element l+32, not l+1. The two nibbles of a byte are 32 apart, and they carry + // different scale/min pairs (is+0 vs is+1). + float acc = 0.0f; + const int nb = K / 256; + + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 144; + + uint16_t dh_scale, dh_min; + std::memcpy(&dh_scale, blk, sizeof(uint16_t)); + std::memcpy(&dh_min, blk + 2, sizeof(uint16_t)); + const float dall = fp16_to_f32(dh_scale); + const float dmin = fp16_to_f32(dh_min); + + const uint8_t* scales = blk + 4; // 12 bytes of packed 6-bit scales + mins + const uint8_t* qs = blk + 16; // 128 bytes of 4-bit quants + const bf16_t* xb = x + (size_t)256 * b; + + for (int il = 0; il < 4; ++il) { // four 64-element groups + const int is = 2 * il; + + uint8_t sc, m; + get_scale_min_k4(is + 0, scales, sc, m); + const float d1 = dall * sc; + const float m1 = dmin * m; + get_scale_min_k4(is + 1, scales, sc, m); + const float d2 = dall * sc; + const float m2 = dmin * m; + + for (int ir = 0; ir < 8; ++ir) { + const uint8_t* q = qs + 32 * il + 4 * ir; + const bf16_t* y = xb + 64 * il + 4 * ir; + for (int l = 0; l < 4; ++l) { + acc += (d1 * (float)(q[l] & 0xF) - m1) * bf16_to_f32(y[l]); + acc += (d2 * (float)(q[l] >> 4) - m2) * bf16_to_f32(y[l + 32]); + } + } + } + } + + return acc; +} + +float q6_k_dot_f32_scalar(const uint8_t* w, const bf16_t* x, int K) { + // Q6_K: 256-element blocks, 210 bytes each. + // Block layout: ql[128] | qh[64] | scales[16] | d (fp16) + // ql: lower 4 bits of 6-bit quant values (128 bytes) + // qh: upper 2 bits of 6-bit quant values, packed (64 bytes, 2 bits per element) + // scales: 16 int8 sub-scales (8 per 128-element half) + // d: fp16 block scale + // Dequant: q in [-32,31], w = d * scales[is] * (q - 32) + // Reference: ggml-cuda/convert.cu::dequantize_block_q6_K and models/gguf/dequant.py::dequant_q6_k + + float acc = 0.0f; + const int nb = K / 256; // number of Q6_K blocks + + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 210; + + const uint8_t* ql = blk; // 128 bytes (lower 4 bits) + const uint8_t* qh = blk + 128; // 64 bytes (upper 2 bits) + const int8_t* scales = (const int8_t*)(blk + 192); // 16 int8 sub-scales + + uint16_t dh; + std::memcpy(&dh, blk + 208, sizeof(uint16_t)); + const float d = fp16_to_f32(dh); + + const int elem_base = 256 * b; + + // Process in two 128-element halves + for (int h = 0; h < 2; ++h) { + const uint8_t* ql_h = ql + 64 * h; // 64 bytes for this half + const uint8_t* qh_h = qh + 32 * h; // 32 bytes for this half + const int8_t* sc_h = scales + 8 * h; // 8 scales for this half + const int h_base = elem_base + 128 * h; + + // Process 4 groups of 32 elements, each group uses 2 of the 8 scales + for (int g = 0; g < 4; ++g) { + const int8_t* sc_g = sc_h + 2 * g; // 2 scales for this group + const int qh_bits_base = 2 * g; // Starting bit position in qh bytes + + // Group 1: ql_h[0:32] low nibbles, qh high bits [0:2] + // Group 2: ql_h[32:64] low nibbles, qh high bits [2:4] + // Group 3: ql_h[0:32] high nibbles, qh high bits [4:6] + // Group 4: ql_h[32:64] high nibbles, qh high bits [6:8] + const bool use_hi_nibble = (g >= 2); + const int ql_offset = (g % 2 == 1) ? 32 : 0; + + for (int l = 0; l < 32; ++l) { + const uint8_t ql_val = ql_h[ql_offset + l]; + const uint8_t qh_val = qh_h[l]; + + // Extract the 6-bit quant value: 4 bits from ql, 2 bits from qh + const int q_lo = use_hi_nibble ? (ql_val >> 4) : (ql_val & 0x0F); + // Shift is per-GROUP only. An extra term in l here silently corrupts the + // upper half of every group: see dequantize_block_q6_K, which reads + // (qh >> 2*g) & 3 for all 32 lanes of the group. + const int q_hi = (qh_val >> qh_bits_base) & 3; + const int q = (q_lo | (q_hi << 4)) - 32; + + // Select scale: use sc_g[0] for first 16 elements, sc_g[1] for next 16 + const int sc_idx = l / 16; + const int scale = sc_g[sc_idx]; + + // Calculate element index + const int elem_offset = 32 * g + l; + const int elem_idx = h_base + elem_offset; + + acc += d * scale * q * bf16_to_f32(x[elem_idx]); + } + } + } + } + + return acc; +} + +enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4, WF_Q4_K = 5, WF_Q6_K = 6 }; // Each ctor pointer arg is the address of a CPU int64 array of length // num_layers (one base address per layer, built by cpu_executor.py's @@ -1259,7 +1403,7 @@ struct CpuMoeExecutor { // it to a captured GPU elementwise kernel removes it while keeping the official // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py). bool input_prequant = false; - // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I). + // K-quant packed-row byte strides (Q4_0: H/32*18, Q4_K: H/256*144, Q6_K: H/256*210 for gate_up over K=H). int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0; float e2m1_lut[16]; float e4m3_lut[256]; @@ -1381,6 +1525,18 @@ struct CpuMoeExecutor { q4_gu_row_bytes = (H / 32) * 18; // K = H (gate_up rows) q4_dn_row_bytes = (I / 32) * 18; // K = I (down rows) } + if (weight_format == WF_Q4_K) { + if (H % 256 != 0 || I % 256 != 0) + throw std::runtime_error("Q4_K CPU MoE requires H and I to be multiples of 256"); + q4_gu_row_bytes = (H / 256) * 144; // K = H (gate_up rows) + q4_dn_row_bytes = (I / 256) * 144; // K = I (down rows) + } + if (weight_format == WF_Q6_K) { + if (H % 256 != 0 || I % 256 != 0) + throw std::runtime_error("Q6_K CPU MoE requires H and I to be multiples of 256"); + q4_gu_row_bytes = (H / 256) * 210; // K = H (gate_up rows) + q4_dn_row_bytes = (I / 256) * 210; // K = I (down rows) + } isa = c.name; // nvfp4 (AVX-VNNI only): W4A8 int8 decode when the CPU supports it. q4_0 is always // W4A8 (activations pre-quantized to Q8_0); select_q4dot picks VPDPBUSD / VPMADDUBSW @@ -1486,6 +1642,14 @@ struct CpuMoeExecutor { gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; return q4dot(w, xi8, xas, H); // W4A8: int8 activations (Q8_0), scale in xas } + if (fmt == WF_Q4_K) { + const uint8_t* w = gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; + return q4_k_dot_f32_scalar(w, x, H); // W4A16: bf16 activations, K-quant dequant + } + if (fmt == WF_Q6_K) { + const uint8_t* w = gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; + return q6_k_dot_f32_scalar(w, x, H); // W4A16: bf16 activations, K-quant dequant + } const size_t r = (size_t)e * (2 * I) + row; if (use_vnni) return nvi8dot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), @@ -1508,6 +1672,14 @@ struct CpuMoeExecutor { const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes; return q4dot(w, gi8, gas, I); // W4A8: int8 activations (Q8_0), scale in gas } + if (fmt == WF_Q4_K) { + const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes; + return q4_k_dot_f32_scalar(w, g, I); // W4A16: bf16 activations, K-quant dequant + } + if (fmt == WF_Q6_K) { + const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes; + return q6_k_dot_f32_scalar(w, g, I); // W4A16: bf16 activations, K-quant dequant + } const size_t r = (size_t)e * H + row; if (use_vnni) return nvi8dot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16), diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index 6fe1efc0..0752b19d 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -67,10 +67,56 @@ "swigluoai": 3, } -# Weight-format ids must match WFmt in csrc/cpu_moe/cpu_moe_ext.cpp:1215. -# Q4_0 is the only GGUF format with AVX/VNNI kernels; adding more requires new -# intrinsics work (out of scope). The CPU path for other GGUF formats is not implemented. -_WFMT_IDS = {"bf16": 0, "nvfp4": 1, "mxfp4_triton": 2, "ds_fp4": 3, "q4_0": 4} +# Weight-format ids must match WFmt in csrc/cpu_moe/cpu_moe_ext.cpp. +_WFMT_IDS = {"bf16": 0, "nvfp4": 1, "mxfp4_triton": 2, "ds_fp4": 3, "q4_0": 4, "q4_k": 5, "q6_k": 6} + +# (elements per block, bytes per block) for the K-quant expert banks the CPU GEMV reads in +# place. Must match ggml-common.h and the q4_gu_row_bytes arithmetic in cpu_moe_ext.cpp: +# block_q4_K is 144 bytes and block_q6_K is 210 bytes, both over QK_K = 256 elements. +_GGUF_KQUANT_BLOCK = {"q4_k": (256, 144), "q6_k": (256, 210)} + +# quant_format == "gguf" names a container, not a layout: the checkpoint picks a ggml type +# per tensor, so the concrete CPU format has to be recovered from the bank types. Only +# types with a CPU dot kernel appear here; everything else has to stay on --moe-backend +# offload, where the GPU dequantizes. +_GGML_TO_CPU_FMT = {2: "q4_0", 12: "q4_k", 14: "q6_k"} + + +def _resolve_gguf_format(cache) -> str: + """Map a GGUF checkpoint's expert bank types onto one CPU weight format. + + The C++ executor takes a single ``weight_format`` for both banks, so a checkpoint whose + gate_up and down banks use different ggml types cannot run here even when both types + are individually supported. That combination is common (Q4_K_M bumps ffn_down_exps to + Q6_K), so it gets its own message rather than being lumped in with unsupported types. + """ + types = getattr(cache, "gguf_expert_types", None) + if not types: + raise NotImplementedError( + "--moe-backend cpu/hybrid needs the GGUF expert bank types, but this cache " + "did not record them; use --moe-backend offload." + ) + gate_up, down = int(types[0]), int(types[1]) + + def name(t: int) -> str: + from freetoken.models.gguf.dequant import GGML_NAME + + return GGML_NAME.get(t, f"type {t}") + + if gate_up != down: + raise NotImplementedError( + f"--moe-backend cpu/hybrid runs one weight format for both expert banks, but " + f"this checkpoint stores gate_up as {name(gate_up)} and down as {name(down)}. " + f"Mixed-type banks (Q4_K_M and the _M/_XXS mixes do this) need " + f"--moe-backend offload; a --pure requantization would also make it uniform." + ) + if gate_up not in _GGML_TO_CPU_FMT: + raise NotImplementedError( + f"--moe-backend cpu/hybrid has no CPU kernel for {name(gate_up)} experts " + f"(supported: {', '.join(sorted(set(_GGML_TO_CPU_FMT.values())))}); use " + f"--moe-backend offload, which dequantizes on the GPU and covers every type." + ) + return _GGML_TO_CPU_FMT[gate_up] def compiled_extension_supports(activation: str) -> bool: @@ -162,10 +208,15 @@ def __init__( from freetoken.kernel import _cpu_moe fmt = cache.quant_format + # "gguf" is a container tag; resolve it to the concrete per-type CPU format first so + # everything downstream (the _WFMT_IDS gate, _resolve_banks, the C++ weight_format) + # sees one layout name. + if fmt == "gguf": + fmt = _resolve_gguf_format(cache) if fmt not in _WFMT_IDS: raise NotImplementedError( f"--moe-backend cpu/hybrid computes experts on the CPU and supports " - f"{sorted(_WFMT_IDS)} formats, but this checkpoint's experts are " + f"{sorted(_WFMT_IDS.keys())} formats, but this checkpoint's experts are " f"{fmt!r}; use --moe-backend offload (GPU-side dequant) instead." ) if activation not in _ACT_IDS: @@ -369,21 +420,15 @@ def _resolve_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: if fmt == "q4_0": return self._resolve_q4_0_banks(banks) + if fmt in _GGUF_KQUANT_BLOCK: + return self._resolve_kquant_banks(banks, fmt) + if fmt == "mxfp4_triton": return self._resolve_mxfp4_banks(banks) if fmt == "ds_fp4": return self._resolve_dsfp4_banks(banks) - # Detect unsupported GGUF formats (any ggml type name that isn't q4_0). - # GGUF format strings follow the pattern of ggml type names (q*, iq*). - if fmt.startswith(("q", "iq")) and fmt != "q4_0": - raise NotImplementedError( - f"the CPU/hybrid MoE backend does not support GGUF format {fmt!r}; " - f"it has AVX/VNNI kernels for Q4_0 only. Use --moe-backend fused for GPU " - f"dequantization, or --moe-backend offload to stream experts to the GPU." - ) - # nvfp4: packed e2m1 (2/byte) + fp8-e4m3 per-16 block scales + fp16 row globals. gup, gus, gug = banks["gate_up_packed"], banks["gate_up_scale"], banks["gate_up_global"] dnp, dns, dng = banks["down_packed"], banks["down_scale"], banks["down_global"] @@ -437,6 +482,59 @@ def _resolve_q4_0_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: ) return ptrs, (H, I) + def _resolve_kquant_banks(self, banks: dict, fmt: str) -> tuple[dict, tuple[int, int]]: + """Native GGUF K-quant expert banks (Q4_K, Q6_K), same schema as Q4_0 but with a + 256-element block instead of 32. + + These share Q4_0's contract: the banks handed here are byte-identical to the ones + the GPU offload path streams, and the C++ GEMV dequantizes a block inside the + K-loop rather than materialising the row. The only per-format quantities are the + block geometry, so the checks below are Q4_0's with (32, 18) parameterised out. + + Unlike Q4_0 these run W4A16 (see ``use_q4a8`` in cpu_moe_ext.cpp): the K-quant + scalar dots read the bf16 activation directly, since the super-block scale + structure does not map onto the int8 activation path. + """ + qk, blk = _GGUF_KQUANT_BLOCK[fmt] + gate_up, down = banks["gate_up"], banks["down"] + if gate_up[0].dtype != torch.uint8 or down[0].dtype != torch.uint8: + raise TypeError( + f"{fmt} expert banks must be raw packed bytes (uint8), got " + f"gate_up={gate_up[0].dtype} down={down[0].dtype}" + ) + I = int(gate_up[0].shape[1] // 2) + H = int(down[0].shape[1]) + if gate_up[0].shape[1] != 2 * I: + raise ValueError(f"gate_up must be a fused [S, 2I, ...] bank, got {gate_up[0].shape}") + # A partial block has no representation in the format, so a non-multiple here means + # the bank was built wrong; the C++ row arithmetic would silently truncate it. + if H % qk or I % qk: + raise ValueError( + f"{fmt} needs H and I to be multiples of {qk} (block size), got H={H} I={I}" + ) + want_gu, want_dn = (H // qk) * blk, (I // qk) * blk + if int(gate_up[0].shape[2]) != want_gu: + raise ValueError( + f"{fmt} gate_up row is {int(gate_up[0].shape[2])} bytes, expected {want_gu} " + f"for K={H}" + ) + if int(down[0].shape[2]) != want_dn: + raise ValueError( + f"{fmt} down row is {int(down[0].shape[2])} bytes, expected {want_dn} " + f"for K={I}" + ) + ptrs = dict( + gate_up_ptr=self._make_table(gate_up).data_ptr(), + down_ptr=self._make_table(down).data_ptr(), + gate_up_scale_ptr=0, + gate_up_global_ptr=0, + down_scale_ptr=0, + down_global_ptr=0, + gate_up_bias_ptr=0, + down_bias_ptr=0, + ) + return ptrs, (H, I) + def _resolve_mxfp4_banks(self, banks: dict) -> tuple[dict, tuple[int, int]]: """gpt-oss mxfp4 ``mxfp4_triton`` schema: transposed split-K blocks/scales (N innermost) + per-output-row biases. The C++ kernel streams K and diff --git a/tests/moe/test_cpu_moe_kquant.py b/tests/moe/test_cpu_moe_kquant.py new file mode 100644 index 00000000..f4134291 --- /dev/null +++ b/tests/moe/test_cpu_moe_kquant.py @@ -0,0 +1,200 @@ +"""CPU MoE executor -- native GGUF K-quant experts (Q4_K, Q6_K). + +Companion to test_cpu_moe_q4_0.py. Same contract: the CPU GEMV reads the *same* packed +banks the GPU offload path streams and dequantizes a block inside the K-loop, so it is +checked against the canonical CUDA dequant + the production bf16 GPU decode on +byte-identical banks. Both sides are W4A16, so the only spread is weight bf16-rounding and +reduction order. + +Blocks are synthesized rather than carved out of a real checkpoint: every byte pattern is +a legal Q4_K/Q6_K block once the fp16 scale fields hold finite values, and a self-contained +fixture keeps this runnable without a multi-GB download. + +Three bugs these tests exist to catch, all of which produce a model that loads, runs at +full speed and emits fluent nonsense: + +* Q4_K quants are UNSIGNED 0..15 offset by a per-sub-block min. Reusing Q4_0's ``q - 8`` + bias reads as a plausible kernel and destroys the output (cosine ~0.008). +* Within a Q4_K 64-element group, a byte's high nibble is element l+32, not l+1, and it + carries a different scale/min pair. +* The Q6_K ``qh`` shift is per-group only. Adding a term in the lane index corrupts + exactly the upper half of every group, which still correlates ~0.17 and so can be + mistaken for a tolerance problem. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +QK_K = 256 +_Q4_K_BYTES, _Q6_K_BYTES = 144, 210 +GGML_Q4_0, GGML_Q4_K, GGML_Q6_K = 2, 12, 14 + + +def _fp16_bytes(vals: torch.Tensor) -> torch.Tensor: + """[...] float -> [..., 2] uint8 little-endian fp16.""" + return vals.to(torch.float16).view(torch.uint8).reshape(*vals.shape, 2) + + +def _make_q4_k_rows(S: int, OUT: int, K: int, gen) -> torch.Tensor: + """Random valid Q4_K rows: [S, OUT, K//256*144]. + + Layout per block (ggml-common.h): half2 dm | uint8 scales[12] | uint8 qs[128]. The + scale and quant bytes are unconstrained, so only dm has to be finite and sanely scaled. + """ + nb = K // QK_K + d = 0.02 + 0.03 * torch.rand(S, OUT, nb, generator=gen) + dmin = 0.01 + 0.02 * torch.rand(S, OUT, nb, generator=gen) + dm = torch.cat([_fp16_bytes(d), _fp16_bytes(dmin)], dim=-1) # [S, OUT, nb, 4] + rest = torch.randint(0, 256, (S, OUT, nb, 140), dtype=torch.uint8, generator=gen) + return torch.cat([dm, rest], dim=-1).reshape(S, OUT, nb * _Q4_K_BYTES).contiguous() + + +def _make_q6_k_rows(S: int, OUT: int, K: int, gen) -> torch.Tensor: + """Random valid Q6_K rows: [S, OUT, K//256*210]. + + Layout per block: uint8 ql[128] | uint8 qh[64] | int8 scales[16] | half d. + """ + nb = K // QK_K + body = torch.randint(0, 256, (S, OUT, nb, 208), dtype=torch.uint8, generator=gen) + d = 0.01 + 0.02 * torch.rand(S, OUT, nb, generator=gen) + return torch.cat([body, _fp16_bytes(d)], dim=-1).reshape( + S, OUT, nb * _Q6_K_BYTES).contiguous() + + +_MAKERS = {"q4_k": (_make_q4_k_rows, GGML_Q4_K), "q6_k": (_make_q6_k_rows, GGML_Q6_K)} + + +def _make_cache(fmt: str, L: int, E: int, H: int, I: int, seed: int = 0, + *, as_gguf: bool = False): + """Pinned host banks in the native K-quant schema, as the offload path builds them.""" + from freetoken.kernel.pinned import alloc_pinned_tensor + + make, ggml_type = _MAKERS[fmt] + gen = torch.Generator().manual_seed(seed) + S = L * E + + def rows(OUT, K): + packed = make(S, OUT, K, gen) + pinned = alloc_pinned_tensor(*packed.shape, dtype=torch.uint8) + pinned.copy_(packed) + return pinned + + return SimpleNamespace( + # A real GGUF checkpoint tags the cache "gguf" and carries the ggml types + # separately; the executor resolves that to the concrete format. + quant_format="gguf" if as_gguf else fmt, + gguf_expert_types=(ggml_type, ggml_type) if as_gguf else None, + bank_sources={"gate_up": list(rows(2 * I, H).split(E)), + "down": list(rows(H, I).split(E))}, + num_layers=L, + num_experts=E, + decode_target="cpu", + cpu_executor=None, + ) + + +def _dequant_bank(packed: torch.Tensor, ggml_type: int, K: int, dev) -> torch.Tensor: + """[S, OUT, row_bytes] packed -> [S, OUT, K] bf16 via the vendored CUDA dequant.""" + from freetoken.kernel.gguf import ggml_dequantize + + S, OUT, row_bytes = packed.shape + flat = ggml_dequantize( + packed.reshape(-1, row_bytes).to(dev).contiguous(), ggml_type, S * OUT, K, + torch.bfloat16, + ) + return flat.reshape(S, OUT, K) + + +@pytest.mark.parametrize("fmt", ["q4_k", "q6_k"]) +@pytest.mark.parametrize("bs", [1, 3, 8]) +def test_cpu_decode_kquant_matches_dequant_then_gpu(fmt, bs): + """CPU inline-dequant K-quant GEMV vs. the CUDA dequant + bf16 GPU decode.""" + from freetoken.moe.cpu_executor import CpuMoeExecutor + from freetoken.moe.fused import fused_experts_decode_impl + + L, E, H, I, top_k, layer = 2, 8, 512, 256, 4, 1 + dev = torch.device("cuda") + cache = _make_cache(fmt, L, E, H, I, seed=100 + bs) + ex = CpuMoeExecutor(cache, top_k=top_k, activation="silu", + apply_router_weight_on_input=False, num_threads=0, + max_tokens=bs, device=dev) + + torch.manual_seed(400 + bs) + hidden = torch.randn(bs, H, device=dev, dtype=torch.bfloat16) * 0.5 + ids = torch.stack([torch.randperm(E, device=dev)[:top_k] + for _ in range(bs)]).to(torch.int32) + w = torch.rand(bs, top_k, device=dev, dtype=torch.float32) + + cpu_out = ex.decode(layer, hidden, w, ids).float() + torch.cuda.synchronize() + + ggml_type = _MAKERS[fmt][1] + gu = _dequant_bank(cache.bank_sources["gate_up"][layer], ggml_type, H, dev) + dn = _dequant_bank(cache.bank_sources["down"][layer], ggml_type, I, dev) + gpu_out = fused_experts_decode_impl(hidden, gu, dn, w, ids.clone(), "silu", False).float() + + cos = torch.nn.functional.cosine_similarity( + cpu_out.flatten(), gpu_out.flatten(), dim=0).item() + rel = ((cpu_out - gpu_out).abs().max() / (gpu_out.abs().max() + 1e-6)).item() + assert cos > 0.999, f"{fmt} bs={bs}: cosine {cos} (rel {rel})" + assert rel < 5e-2, f"{fmt} bs={bs}: rel {rel} (cosine {cos})" + + +@pytest.mark.parametrize("fmt", ["q4_k", "q6_k"]) +def test_gguf_cache_resolves_to_the_cpu_kernel(fmt): + """A checkpoint tagged quant_format 'gguf' must reach the same kernel. + + The bank types live on the cache rather than in the format tag, so without the bridge + the executor refuses every real GGUF checkpoint while accepting the literal 'q4_k' + string that nothing actually produces. + """ + from freetoken.moe.cpu_executor import CpuMoeExecutor + + L, E, H, I = 2, 8, 512, 256 + cache = _make_cache(fmt, L, E, H, I, seed=7, as_gguf=True) + ex = CpuMoeExecutor(cache, top_k=4, activation="silu", + apply_router_weight_on_input=False, num_threads=0, + max_tokens=2, device=torch.device("cuda")) + assert ex.quant_format == fmt + + +class TestGgufFormatResolution: + """The bridge's refusals. Each must name what is wrong and what to use instead.""" + + def test_mixed_banks_refused(self): + """Q4_K_M stores gate_up Q4_K and down Q6_K; one weight_format cannot serve both.""" + from freetoken.moe.cpu_executor import _resolve_gguf_format + + c = SimpleNamespace(quant_format="gguf", + gguf_expert_types=(GGML_Q4_K, GGML_Q6_K)) + with pytest.raises(NotImplementedError, match="(?i)mixed-type"): + _resolve_gguf_format(c) + + def test_uniform_but_unsupported_type_refused(self): + """IQ3_S banks are uniform but have no CPU dot kernel; offload must be named.""" + from freetoken.moe.cpu_executor import _resolve_gguf_format + + c = SimpleNamespace(quant_format="gguf", gguf_expert_types=(21, 21)) + with pytest.raises(NotImplementedError, match="(?i)no cpu kernel"): + _resolve_gguf_format(c) + + @pytest.mark.parametrize("t,want", [(GGML_Q4_0, "q4_0"), (GGML_Q4_K, "q4_k"), + (GGML_Q6_K, "q6_k")]) + def test_uniform_supported_types_resolve(self, t, want): + from freetoken.moe.cpu_executor import _resolve_gguf_format + + c = SimpleNamespace(quant_format="gguf", gguf_expert_types=(t, t)) + assert _resolve_gguf_format(c) == want + + def test_missing_types_refused(self): + from freetoken.moe.cpu_executor import _resolve_gguf_format + + c = SimpleNamespace(quant_format="gguf", gguf_expert_types=None) + with pytest.raises(NotImplementedError): + _resolve_gguf_format(c) From b2f84751826cb380156fad4fd36e613bfb454625 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Mon, 24 Aug 2026 21:35:43 -0700 Subject: [PATCH 28/36] fix(gguf): untied GGUF lm_head computed logits for every prefill position A GGUF checkpoint with untied embeddings had its lm_head swapped for a plain GGUFLinear, which has no reason to know it is the head and so projected the whole prefill hidden state. ParallelLMHead and gemma4's GGUFTiedLMHead both already drop everything but the last position of each sequence, so only the untied GGUF path was affected, and only since the qwen35moe/qwen3moe adapters landed. Logits are [tokens, vocab] and on a large-vocabulary model that is the biggest tensor in the model. Ornith-1.5's vocab is 248,320, which in bf16 is 486 KiB of logits per token, so a 1,800-token prompt asked for a single 894 MB allocation. Every row but one was then thrown away by the sampler. Measured on an RTX 4060 Laptop (8 GB), Ornith-1.5-35B-A3B IQ3_S, offload backend. Before, a ~766-token prompt prefilled in 3.6s and ~1,800 tokens either hung past 300s or killed the worker with "CUDA driver error: device not ready". After: 1,834 tokens in 6.1s, 2,487 in 9.6s, 3,765 in 13.9s, 4,598 in 16.9s. Decode was never affected either way, which is what made this look like a memory problem rather than a head problem: generating 30,000 tokens works fine because decode only ever has one position per sequence to project. Adds GGUFLMHead rather than putting the slice in GGUFLinear, since that class is also every q/k/v/o and MLP projection in the model and must not slice. Beyond this fix, an 8 GB card still runs out of room around 5-6k prompt tokens, but in a different place: the GDN chunked kernel's activations (chunk_gated_delta_rule_fwd -> chunk_fwd_o). That is headroom rather than a defect and --max-prefill-length handles it; at 2048 the same box prefills 6,160 tokens in 23.2s and 7,784 in 27.5s. Worth knowing that the 8192 default does not chunk a typical long prompt at all. --- python/freetoken/layers/gguf.py | 31 +++++++++++++++++++++ python/freetoken/models/qwen3_5_moe/gguf.py | 11 +++++++- python/freetoken/models/qwen3_moe/gguf.py | 11 +++++++- 3 files changed, 51 insertions(+), 2 deletions(-) diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py index 28db4df3..ae0b1b43 100644 --- a/python/freetoken/layers/gguf.py +++ b/python/freetoken/layers/gguf.py @@ -109,6 +109,37 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return out +class GGUFLMHead(GGUFLinear): + """LM head over a native GGUF ``output.weight`` (untied embeddings). + + Identical to ``GGUFLinear`` except that during prefill it keeps only the last position + of each sequence, exactly as ``ParallelLMHead`` (layers/embedding.py) and + ``GGUFTiedLMHead`` (models/gemma4/gguf.py) already do. + + This is not an optimization, it is a memory correctness issue. Logits are + [tokens, vocab], so on a large-vocabulary model the full-prefill tensor is enormous: + Ornith-1.5's vocab is 248,320, which in bf16 is 486 KiB of logits PER TOKEN. A + 1,800-token prompt therefore asks for a single 894 MB allocation, which is more than the + free VRAM left on an 8 GB card after weights and caches, and prefill dies with + "CUDA driver error: device not ready" while decode is completely unaffected. Only the + last position of each sequence is ever sampled, so every other row was computed and + thrown away. + + The dense path never hit this because ``ParallelLMHead`` slices; the bug appears only + when a GGUF checkpoint has untied embeddings and the head is swapped for a generic + quantized Linear, which has no reason to know it is the head. + """ + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from freetoken.core import get_global_ctx + + batch = get_global_ctx().batch + if batch.is_prefill: + indices = batch.attn_metadata.get_last_indices(batch.size) + x = x[indices].contiguous() + return super().forward(x) + + class GGUFMergedLinear(BaseOP): """Merged linear projection with parts that have different quant types. diff --git a/python/freetoken/models/qwen3_5_moe/gguf.py b/python/freetoken/models/qwen3_5_moe/gguf.py index a7dd96db..488fdd99 100644 --- a/python/freetoken/models/qwen3_5_moe/gguf.py +++ b/python/freetoken/models/qwen3_5_moe/gguf.py @@ -952,7 +952,16 @@ def swap_linear(owner, attr, quant_type: int): model.lm_head = GGUFTiedLMHead(embed, qt(-1, "token_embd.weight")) else: - swap_linear(model, "lm_head", qt(-1, "output.weight")) + # NOT swap_linear: a plain GGUFLinear would compute logits for every prefill + # position, and [tokens, vocab] is the largest tensor in the model. See GGUFLMHead. + from freetoken.layers.gguf import GGUFLMHead + + head = model.lm_head + out_features, in_features = head.weight.shape + model.lm_head = GGUFLMHead( + in_features, out_features, qt(-1, "output.weight"), + has_bias=head.bias is not None, + ) __all__ = [ diff --git a/python/freetoken/models/qwen3_moe/gguf.py b/python/freetoken/models/qwen3_moe/gguf.py index a3640d99..86a91806 100644 --- a/python/freetoken/models/qwen3_moe/gguf.py +++ b/python/freetoken/models/qwen3_moe/gguf.py @@ -463,7 +463,16 @@ def swap_linear(owner, attr, quant_type: int): model.lm_head = GGUFTiedLMHead(embed, qt(-1, "token_embd.weight")) else: - swap_linear(model, "lm_head", qt(-1, "output.weight")) + # NOT swap_linear: a plain GGUFLinear would compute logits for every prefill + # position, and [tokens, vocab] is the largest tensor in the model. See GGUFLMHead. + from freetoken.layers.gguf import GGUFLMHead + + head = model.lm_head + out_features, in_features = head.weight.shape + model.lm_head = GGUFLMHead( + in_features, out_features, qt(-1, "output.weight"), + has_bias=head.bias is not None, + ) __all__ = [ From dd48aacb9f6edced7cdb1f2687f39af0078bf5d4 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Tue, 25 Aug 2026 02:00:09 -0700 Subject: [PATCH 29/36] fix(triton): extend-attention tiles overflowed shared memory on sm_75 Two problems, both only reachable on pre-Ampere cards. The head_dim <= 128 branch of _select_extend_tile returned (128, 64) unconditionally, with no fits() check at all, unlike every other branch. At block_d 128 that tile needs 96KB, which is fine on sm_80/sm_89 and fatal on Turing's 64KB hard limit. The shared-memory estimate itself also has the wrong shape for Turing. Ampere+ tiles with mma.sync m16n8k16, Turing with m16n8k8, and Triton stages materially more for the same logical tile as a result. Measured from triton's own OutOfResources reports on a Quadro RTX 6000 (65536 limit): (128,64)@block_d 256 -> 196608 (64,64)@256 -> 131072 (64,32)@256 -> 98304 (128,64)@128 -> 98304 The existing (block_m + 2*block_n)*block_d*2 estimate scales with N where the real cost scales with M, so it under-predicts exactly the tall-M tiles that overflow. A head_dim-256 model on a 64KB card therefore selected (64,32), which the estimate said fit at 65536, and still died needing 98304. The corrected bound is applied ONLY when compute capability < 8. The same tiles measure 114688 for (128,64)@256 on sm_89, so applying it there would shrink tiles that genuinely fit and cost roughly 2x on large head_dim. Ampere+ selection is unchanged, and the existing parametrised tests over A100/H100/ sm_89/unknown budgets all still pass. Turing also needs rungs below (64,32) at head_dim <= 256, so the ladder continues to (32,32) and (32,16) there; Ampere+ keeps its original two-way choice. Verified on the hardware: Ornith-1.5-35B-A3B IQ3_S on a Quadro RTX 6000 (sm_75, 24GB) now serves, where before it died at warmup. 64-69 tok/s with 10240 of 10496 expert slots resident. Correctness spot-checked at temperature 0 against known facts ("The capital city of France is" -> " Paris, ..."). --- python/freetoken/kernel/triton/attention.py | 34 ++++++++++++++++++--- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c2358d84..2de7b549 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -18,7 +18,9 @@ def _optin_smem_bytes(device_index: int) -> int: return int(getattr(props, "shared_memory_per_block_optin", 0)) -def _select_extend_tile(head_dim: int, block_d: int, smem_optin: int) -> tuple[int, int]: +def _select_extend_tile( + head_dim: int, block_d: int, smem_optin: int, pre_ampere: bool = False +) -> tuple[int, int]: """Pick ``(BLOCK_M, BLOCK_N)`` for the extend/prefill kernel, shared-memory aware. Larger tiles run materially faster (~2x for head_dim 512 on H100) but their bf16 @@ -31,12 +33,35 @@ def _select_extend_tile(head_dim: int, block_d: int, smem_optin: int) -> tuple[i budget = smem_optin * 0.8 # headroom for scores/acc/alignment/triton scratch def fits(block_m: int, block_n: int) -> bool: + if pre_ampere: + # Turing stages materially more through shared memory than Ampere+ for the + # same tile: mma.sync m16n8k8 tiles differently from m16n8k16. Measured from + # triton's own OutOfResources reports on an sm_75 Quadro RTX 6000: + # (128,64)@256 -> 196608 (64,64)@256 -> 131072 (64,32)@256 -> 98304 + # (128,64)@128 -> 98304 + # against a 65536 limit. The Ampere+ estimate below under-predicts these by up + # to 2x, which is why a head_dim-256 model picked (64,32) and still died. + # The same tiles on sm_89 measure 114688 for (128,64)@256, so this bound is + # deliberately NOT applied there: it would shrink tiles that actually fit. + return (block_m + block_n) * block_d * 4 <= smem_optin return (block_m + 2 * block_n) * block_d * 2 <= budget if head_dim <= 128: - return 128, 64 - if head_dim <= 256: + # Guarded like every other branch. Unguarded this returned (128,64) unconditionally, + # which needs 96KB at block_d 128 -- fine on sm_80/sm_89, fatal on a 64KB Turing + # card, where warmup died with OutOfResources: Required 98304, limit 65536. return (128, 64) if fits(128, 64) else (64, 32) + if head_dim <= 256: + if fits(128, 64): + return 128, 64 + # Turing needs rungs below (64,32); Ampere+ keeps its original two-way choice so + # behaviour there is unchanged. + if pre_ampere: + for cand in ((64, 32), (32, 32), (32, 16)): + if fits(*cand): + return cand + return 16, 16 + return 64, 32 if head_dim <= 384: return (32, 64) if fits(32, 64) else (32, 32) return (32, 64) if fits(32, 64) else (16, 16) @@ -799,7 +824,8 @@ def extend_paged_attention( # shared memory fits them, shrink on consumer GPUs (sm_89 ~99KB) where the default # 128x64 overflows once head_dim >= 256 (e.g. gemma4: SWA 256, full-attention 512). block_m, block_n = _select_extend_tile( - head_dim, block_d, _optin_smem_bytes(q.device.index) + head_dim, block_d, _optin_smem_bytes(q.device.index), + pre_ampere=torch.cuda.get_device_capability(q.device.index)[0] < 8, ) grid = (qo_indptr.numel() - 1, num_q_heads, triton.cdiv(max_q_len, block_m)) if k_extend is not None or v_extend is not None: From 9952a39dd2a8718efca495940f78376f13f636ac Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Tue, 25 Aug 2026 15:32:04 -0700 Subject: [PATCH 30/36] fix(gguf): MoE GEMV exceeded CUDA's gridDim.z limit on long prompts moe_vec.cuh carries the routed-pair index (token * top_k + slot) in gridDim.z, which CUDA caps at 65535 on every architecture. All 19 launchers computed the grid as (block_num_y, 1, tokens * top_k), so a prefill wider than 65535 / top_k tokens failed the launch with cudaErrorInvalidConfiguration, which torch reports as "CUDA error: invalid argument". At top_k 8 the ceiling is 8191 tokens. fused_experts_gguf issues two launches, gate_up as (top_k, N) and down as (1, N * top_k), and both reach z = N * top_k, so both overflow together. Nothing about this is architecture- or OS-specific; it needs only a long enough prompt and was reported with a 20k-token one. Fixed by launching token-aligned chunks rather than reshaping the grid. Moving the pair index into x would also fit, since gridDim.x allows 2^31-1, but it would cost the locality the kernel is built around: with the row index in x, consecutive blocks walk rows of the SAME expert and reuse its weight pages. Chunking keeps the access pattern and leaves the kernel body untouched. Chunks are whole tokens so the kernel's own token = blockIdx.z / topk arithmetic stays valid against the offset vy / dst / topk_ids pointers. The 19 near-identical launcher bodies now share one moe_vec_launch<> helper. Tests assert the launch succeeds past the ceiling AND that rows below it are bit-identical however many chunks were used. The second is the real check: a wrong per-chunk offset would still launch, look healthy, and emit fluent nonsense. Verified to z = 160000 (20k tokens x top_k 8), and the kernel compiles for sm_75 and sm_89. --- python/freetoken/kernel/csrc/gguf/moe_vec.cuh | 179 ++++++++---------- tests/kernels/test_moe_vec_grid_z.py | 99 ++++++++++ 2 files changed, 183 insertions(+), 95 deletions(-) create mode 100644 tests/kernels/test_moe_vec_grid_z.py diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh index 8cef9e08..dd0550bf 100644 --- a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh +++ b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh @@ -51,8 +51,25 @@ static __global__ void moe_vec_q( } } -template -static void moe_vec_q4_0_q8_1_cuda( + +// CUDA caps gridDim.z at 65535 on every architecture. The routed-pair index +// (token * top_k + slot) is carried in z, so a prefill wider than 65535 / top_k tokens +// overflows it and cudaLaunchKernel fails with cudaErrorInvalidConfiguration, surfaced by +// torch as "CUDA error: invalid argument". At top_k 8 that ceiling is only 8191 tokens, +// which any long prompt crosses; it is not architecture- or OS-specific. +// +// Fixed by launching in token-aligned chunks rather than by reshaping the grid. Moving the +// pair index into x would also fit, but it would cost the locality this kernel is built +// around: with the row index in x, consecutive blocks walk rows of the SAME expert and +// reuse its weight pages. Chunking keeps that intact and leaves the kernel untouched. +// +// Chunks are whole tokens so the kernel's own token = blockIdx.z / topk arithmetic stays +// valid against the offset pointers. +#define MOE_VEC_MAX_GRID_Z 65535 + +template +static void moe_vec_launch( const void* vx, const void* vy, scalar_t* dst, @@ -64,10 +81,36 @@ static void moe_vec_q4_0_q8_1_cuda( const int token_stride, cudaStream_t stream) { const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + // top_k is the router's experts-per-token (<= 128 in practice), so this is >= 511. + const int tokens_per_chunk = MOE_VEC_MAX_GRID_Z / (top_k > 0 ? top_k : 1); + for (int t0 = 0; t0 < tokens; t0 += tokens_per_chunk) { + const int nt = (tokens - t0) < tokens_per_chunk ? (tokens - t0) : tokens_per_chunk; + const dim3 block_nums(block_num_y, 1, nt * top_k); + moe_vec_q + <<>>( + vx, + (const void*)(((const int*)vy) + (size_t)t0 * token_stride), + dst + (size_t)t0 * top_k * nrows, + topk_ids + (size_t)t0 * top_k, + top_k, ncols, nrows, token_stride); + } +} + +template +static void moe_vec_q4_0_q8_1_cuda( + const void* vx, + const void* vy, + scalar_t* dst, + const int* topk_ids, + const int top_k, + const int tokens, + const int ncols, + const int nrows, + const int token_stride, + cudaStream_t stream) { + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -82,11 +125,8 @@ static void moe_vec_q4_1_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -101,11 +141,8 @@ static void moe_vec_q5_0_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -120,11 +157,8 @@ static void moe_vec_q5_1_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -139,11 +173,8 @@ static void moe_vec_q8_0_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -158,11 +189,8 @@ static void moe_vec_q2_K_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -177,11 +205,8 @@ static void moe_vec_q3_K_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -196,11 +221,8 @@ static void moe_vec_q4_K_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -215,11 +237,8 @@ static void moe_vec_q5_K_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -234,11 +253,8 @@ static void moe_vec_q6_K_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -253,11 +269,8 @@ static void moe_vec_iq2_xxs_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -272,11 +285,8 @@ static void moe_vec_iq2_xs_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -291,11 +301,8 @@ static void moe_vec_iq2_s_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -310,11 +317,8 @@ static void moe_vec_iq3_xxs_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -329,11 +333,8 @@ static void moe_vec_iq1_s_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -348,11 +349,8 @@ static void moe_vec_iq1_m_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -367,11 +365,8 @@ static void moe_vec_iq4_nl_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -386,11 +381,8 @@ static void moe_vec_iq4_xs_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } template @@ -405,9 +397,6 @@ static void moe_vec_iq3_s_q8_1_cuda( const int nrows, const int token_stride, cudaStream_t stream) { - const int block_num_y = (nrows + GGML_CUDA_MMV_Y - 1) / GGML_CUDA_MMV_Y; - const dim3 block_nums(block_num_y, 1, tokens * top_k); - const dim3 block_dims(WARP_SIZE, GGML_CUDA_MMV_Y, 1); - moe_vec_q - <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride); + moe_vec_launch( + vx, vy, dst, topk_ids, top_k, tokens, ncols, nrows, token_stride, stream); } diff --git a/tests/kernels/test_moe_vec_grid_z.py b/tests/kernels/test_moe_vec_grid_z.py new file mode 100644 index 00000000..545029be --- /dev/null +++ b/tests/kernels/test_moe_vec_grid_z.py @@ -0,0 +1,99 @@ +"""Fused GGUF MoE past CUDA's gridDim.z ceiling. + +``moe_vec.cuh`` carries the routed-pair index (token * top_k + slot) in ``gridDim.z``, +which CUDA caps at 65535 on every architecture. ``fused_experts_gguf`` issues two launches +-- gate_up as ``(top_k, N)`` and down as ``(1, N * top_k)`` -- and both reach +``z = N * top_k``. So a prefill wider than ``65535 / top_k`` tokens used to fail the launch +with ``cudaErrorInvalidConfiguration``, which torch surfaces as the rather unhelpful +"CUDA error: invalid argument". At top_k 8 that ceiling is 8191 tokens, which any long +prompt crosses; it is not architecture- or OS-specific. Reported against PR #131 with a +20k-token prompt on an RTX 3070. + +The launcher now issues token-aligned chunks. The equality assertion below is the part +that matters: chunking is only correct if each chunk's ``vy`` / ``dst`` / ``topk_ids`` +offsets line up with the kernel's own ``token = blockIdx.z / topk`` arithmetic. A fix that +launches but mis-offsets would pass a smoke test, look entirely healthy, and emit fluent +nonsense in service. +""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +TOP_K = 8 +SLOTS, H, I = 16, 512, 256 +CEIL = 65535 // TOP_K # 8191 tokens + + +def _pack_q4_0(S: int, OUT: int, K: int, dev) -> torch.Tensor: + """Valid Q4_0 rows: per 32-element block, a finite fp16 scale then 16 nibble bytes. + + Random bytes are not a usable bank here -- the 2-byte scale would sometimes decode to + Inf/NaN, and the resulting NaNs would fail the equality check for reasons that have + nothing to do with the grid geometry under test. + """ + nb = K // 32 + nib = torch.randint(0, 256, (S, OUT, nb, 16), dtype=torch.uint8) + scale = (0.02 + 0.03 * torch.rand(S, OUT, nb)).to(torch.float16) + sb = scale.view(torch.uint8).reshape(S, OUT, nb, 2) + return torch.cat([sb, nib], dim=-1).reshape(S, OUT, nb * 18).contiguous().to(dev) + + +@pytest.fixture(scope="module") +def banks(): + dev = torch.device("cuda") + torch.manual_seed(0) + max_t = 2 * CEIL + 16 + return { + "gate_up": _pack_q4_0(SLOTS, 2 * I, H, dev), + "down": _pack_q4_0(SLOTS, H, I, dev), + "ids": torch.randint(0, SLOTS, (max_t, TOP_K), dtype=torch.int32, device=dev), + "w": torch.rand(max_t, TOP_K, device=dev, dtype=torch.float32), + "x": (torch.randn(max_t, H, device=dev, dtype=torch.bfloat16) * 0.5).contiguous(), + } + + +def _run(b, n: int) -> torch.Tensor: + from freetoken.models.gguf.dequant import GGML_Q4_0 + from freetoken.moe.fused_q4_0 import fused_experts_gguf + + return fused_experts_gguf( + b["x"][:n].contiguous(), b["gate_up"], b["down"], + b["w"][:n].contiguous(), b["ids"][:n].contiguous(), + "silu", GGML_Q4_0, + ) + + +@pytest.mark.parametrize( + "n", + [ + CEIL - 1, # 65528: last z that fits + CEIL, # 65528 + 8: first launch that used to fail + CEIL + 1, + 2 * CEIL + 7, # several chunks, deliberately not a chunk multiple + ], +) +def test_moe_vec_launches_past_grid_z_ceiling(banks, n): + """A launch whose z exceeds 65535 must succeed rather than raise 'invalid argument'.""" + out = _run(banks, n) + torch.cuda.synchronize() + assert out.shape == (n, H) + assert torch.isfinite(out).all(), "output contains non-finite values" + + +@pytest.mark.parametrize("n", [CEIL - 1, CEIL, CEIL + 1, 2 * CEIL + 7]) +def test_chunking_does_not_disturb_rows(banks, n): + """Rows below the ceiling must be bit-identical however many chunks were launched. + + This is what catches a wrong per-chunk pointer offset, which a crash test cannot. + """ + ref = _run(banks, 64) + got = _run(banks, n) + torch.cuda.synchronize() + assert torch.equal(got[:64], ref), ( + f"first 64 rows differ at n={n}: chunk offsets are wrong " + f"(max abs diff {(got[:64].float() - ref.float()).abs().max().item()})" + ) From c71bf2e533bf0be5edf619ec069863cb1fba806e Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Tue, 25 Aug 2026 15:50:50 -0700 Subject: [PATCH 31/36] feat(gguf): deepseek4 routed-expert bank loader Ported from qwen3_5_moe/gguf_experts.py; llama.cpp emits the same three stacked tensors for both architectures, so only the arithmetic differs. DeepSeek-V4-Flash is 43 served layers of 256 experts at moe_inter_dim 2048 over dim 4096. Not yet reachable: the deepseek4 architecture is not registered and gguf.py does not exist, so nothing calls this. Landing it separately keeps the verified piece reviewable on its own. Verified against the real 164GB antirez/deepseek-v4-gguf Q4KExperts checkpoint: 43 layers uniform Q4_K on both banks, gate_up (256, 4096, 2304) and down (256, 4096, 1152), each cross-checked against the file's own tensor dims (Q4_K is 144 bytes per 256 elements, so H=4096 -> 16*144 and I=2048 -> 8*144). Two things carried over deliberately. The gate_up fusion reshapes to [E, I, rb] and concatenates on dim=1 so each expert keeps its own gate rows followed by its own up rows; cat(dim=0) then reshape is the version that looks right and lays every expert's gate down before any up, which loads and runs at full speed and emits fluent nonsense. And layers >= num_layers are skipped: the file carries a trailing NextN/MTP block, and counting it makes a uniform checkpoint look mixed. --- .../models/deepseek_v4/gguf_experts.py | 254 ++++++++++++++++++ 1 file changed, 254 insertions(+) create mode 100644 python/freetoken/models/deepseek_v4/gguf_experts.py diff --git a/python/freetoken/models/deepseek_v4/gguf_experts.py b/python/freetoken/models/deepseek_v4/gguf_experts.py new file mode 100644 index 00000000..2ba32b08 --- /dev/null +++ b/python/freetoken/models/deepseek_v4/gguf_experts.py @@ -0,0 +1,254 @@ +"""Routed-expert host banks for a deepseek4 GGUF checkpoint. + +Ported from ``models/qwen3_5_moe/gguf_experts.py``; the two are structurally the same job +because llama.cpp emits the same three stacked tensors for both architectures +(``ffn_gate_exps`` / ``ffn_up_exps`` / ``ffn_down_exps``). What differs is only the +arithmetic: DeepSeek-V4-Flash is 43 served layers of 256 experts at +``moe_inter_dim`` 2048 over ``dim`` 4096, and the checkpoints worth loading are uniformly +Q4_K on all three banks rather than qwen35moe's per-layer mix. + +Only the ROUTED experts come through here. The shared expert +(``ffn_{gate,up,down}_shexp``) is an ordinary quantized Linear that +``convert_deepseek4_to_gguf`` swaps for a ``GGUFLinear``, exactly as qwen3_5_moe handles +its own shared expert -- it is dense per token, so there is nothing to offload. + +A note on which checkpoints reach this code at all. The offload slot pool is ONE +allocation per bank shared by every layer, and ``moe_vec.cuh`` addresses it as +``expert * nrows * (ncols / qk)`` with no padding allowance, so a bank whose ggml type +varies by layer cannot be served. Of the thirteen published +``unsloth/DeepSeek-V4-Flash-0731-GGUF`` variants, eleven mix types across layers and the +two that do not are MXFP4 (ggml type 39), which has no entry in ``BLOCK_SHAPE`` and no +vendored kernel. The ``antirez/deepseek-v4-gguf`` builds are the ones that load: their +Q4KExperts variant is uniformly Q4_K across all 43 layers, verified by reading the file. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from freetoken.models.gguf.dequant import GGML_NAME, GGML_Q4_K, row_bytes + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + + +def gguf_expert_types(model_path: str, num_layers: int) -> dict[str, list[int]]: + """Scan the tensor table and return the per-layer ggml type of each expert bank. + + Returns ``{"gate_up": [...], "down": [...]}``, each a list of ``num_layers`` ggml type + enums. gate and up must agree per layer because they are row-concatenated into one + bank and therefore must share a row stride; a mismatch raises here naming both types. + + ``expert_banks._gguf_banks`` consumes this and is what rejects a bank that is + non-uniform ACROSS layers, with the user-facing message about ``--pure``. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + + gate_types: list[int | None] = [None] * num_layers + up_types: list[int | None] = [None] * num_layers + down_types: list[int | None] = [None] * num_layers + + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if layer >= num_layers: + # The trailing NextN/MTP block. DeepSeek-V4-Flash ships nextn_predict_layers=1, + # so the file carries a block at index num_layers that is not served; counting + # it here would make a uniform checkpoint look mixed. + continue + + if t.name.endswith("ffn_gate_exps.weight"): + gate_types[layer] = t.ggml_type + elif t.name.endswith("ffn_up_exps.weight"): + up_types[layer] = t.ggml_type + elif t.name.endswith("ffn_down_exps.weight"): + down_types[layer] = t.ggml_type + + gate_up_types: list[int] = [] + for layer in range(num_layers): + gate_t, up_t = gate_types[layer], up_types[layer] + if gate_t is None or up_t is None: + raise ValueError( + f"deepseek4 GGUF: layer {layer} is missing routed-expert tensors " + f"(gate={gate_t}, up={up_t}); every layer of this architecture is MoE" + ) + if gate_t != up_t: + raise ValueError( + f"deepseek4 GGUF: layer {layer} has ffn_gate_exps " + f"{GGML_NAME.get(gate_t, gate_t)} but ffn_up_exps " + f"{GGML_NAME.get(up_t, up_t)}; they are row-concatenated into one bank and " + "cannot have different row strides" + ) + gate_up_types.append(gate_t) + + for layer in range(num_layers): + if down_types[layer] is None: + raise ValueError(f"deepseek4 GGUF: layer {layer} is missing ffn_down_exps") + + return {"gate_up": gate_up_types, "down": down_types} + + +def gguf_expert_specs( + config: "ModelConfig", types: dict[str, list[int]] +) -> dict[str, tuple[tuple[int, ...], torch.dtype]]: + """Expert bank shapes as ``{name: (shape, dtype)}`` -- ``alloc_layer_banks``' contract. + + Packed block bytes, in torch order:: + + gate_up (E, 2*I, row_bytes(H, t_gate_up)) uint8 + down (E, H, row_bytes(I, t_down)) uint8 + + One spec per bank rather than per layer, for the slot-pool stride reason in the module + docstring. A non-uniform bank is rejected here instead of being mis-decoded. + """ + E = config.num_experts + H = config.hidden_size + I = config.moe_intermediate_size + out: dict[str, tuple[tuple[int, ...], torch.dtype]] = {} + for name, elems in (("gate_up", H), ("down", I)): + distinct = sorted(set(types[name])) + if len(distinct) != 1: + names = [GGML_NAME.get(t, t) for t in distinct] + raise ValueError( + f"deepseek4 expert bank {name!r} mixes ggml types across layers ({names}); " + "a bank must be uniform because its slot pool is one allocation with one " + "stride" + ) + rb = row_bytes(elems, distinct[0]) + shape = (E, 2 * I, rb) if name == "gate_up" else (E, H, rb) + out[name] = (shape, torch.uint8) + return out + + +def load_gguf_expert_sources( + model_path: str, config: "ModelConfig", *, layer_sink=None +) -> dict[str, list[torch.Tensor]]: + """Per-layer host banks holding the routed experts' native packed block bytes. + + Nothing is dequantized: the bytes handed to the offload cache are the same ones the + kernels decode in the K-loop. + + ``layer_sink`` None (serving) pins each completed layer through an internally owned + ``PinPipeline``; a supplied sink (converter) receives the completion notifications + instead and may release banks, so the returned tensors live only as long as it allows. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks + + types = gguf_expert_types(model_path, config.num_layers) + specs = gguf_expert_specs(config, types) + + L = config.num_layers + E = config.num_experts + H = config.hidden_size + I = config.moe_intermediate_size + + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in hb} + + gate_buf: dict[int, torch.Tensor] = {} + up_buf: dict[int, torch.Tensor] = {} + seen_gate: set[int] = set() + seen_up: set[int] = set() + seen_down: set[int] = set() + + def _load(sink) -> None: + tracker = LayerCompletionTracker(2, hb, sink) if sink is not None else None + + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if layer >= L: + continue # trailing NextN/MTP block, not served + + if t.name.endswith("ffn_gate_exps.weight"): + gate_buf[layer] = t.packed() + seen_gate.add(layer) + elif t.name.endswith("ffn_up_exps.weight"): + up_buf[layer] = t.packed() + seen_up.add(layer) + elif t.name.endswith("ffn_down_exps.weight"): + # torch shape [E, H, I] is ggml dims [I, H, E] with I fastest, so the reader + # returns [E*H, row_bytes(I)] already in expert-major row order. Reshaping + # to [E, H, row_bytes(I)] is a view, not a copy. Note the row_bytes is over + # I (the fastest dim), not over E. + down_rb = specs["down"][0][2] + banks["down"][layer].copy_(t.packed().reshape(E, H, down_rb)) + seen_down.add(layer) + if tracker is not None: + tracker.note(layer) + else: + continue + + if layer in gate_buf and layer in up_buf: + rb = specs["gate_up"][0][2] + # gate and up each arrive as [E*I, row_bytes(H)], and ggml's fastest-first + # dims [H, I, E] make E the slowest axis, so those rows are EXPERT-MAJOR: + # expert e owns rows [e*I, (e+1)*I). + # + # The bank must be [E, 2I, row_bytes(H)] with each expert's own gate rows + # followed by its OWN up rows, so reshape to [E, I, rb] and concatenate on + # the row axis within each expert (dim=1). + # + # cat(dim=0) then reshape(E, 2I, rb) -- the version that looks obviously + # right -- lays every expert's gate down before any up, so expert 0 would + # get its gate rows plus expert 1's gate rows and its up would sit E*I rows + # away. That loads, runs at full speed, and emits fluent nonsense. This + # exact bug cost real debugging time on qwen35moe. + g = gate_buf[layer].reshape(E, I, rb) + u = up_buf[layer].reshape(E, I, rb) + banks["gate_up"][layer].copy_(torch.cat([g, u], dim=1)) + del gate_buf[layer], up_buf[layer] + if tracker is not None: + tracker.note(layer) + + if layer_sink is not None: + _load(layer_sink) + elif torch.cuda.is_available(): + with PinPipeline() as pins: + _load(pins) + else: + _load(None) # CUDA-less: mmap banks stay pageable, never pinned + + want = set(range(L)) + missing_gate, missing_up, missing_down = ( + want - seen_gate, want - seen_up, want - seen_down) + if missing_gate or missing_up or missing_down: + raise ValueError( + f"deepseek4 GGUF is missing routed experts: gate {sorted(missing_gate)}, " + f"up {sorted(missing_up)}, down {sorted(missing_down)}" + ) + + return banks + + +def dummy_gguf_expert_sources(config: "ModelConfig") -> dict[str, list[torch.Tensor]]: + """Random expert banks shaped like ``load_gguf_expert_sources`` output. + + Q4_K throughout, matching the checkpoints that actually load (see module docstring). + """ + from freetoken.moe.host_banks import alloc_layer_banks, pin_banks + + L = config.num_layers + types = {"gate_up": [GGML_Q4_K] * L, "down": [GGML_Q4_K] * L} + specs = gguf_expert_specs(config, types) + + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in hb} + for t in banks["gate_up"] + banks["down"]: + t.random_(0, 256) + if torch.cuda.is_available(): + pin_banks(hb) + return banks + + +__all__ = [ + "gguf_expert_types", + "gguf_expert_specs", + "load_gguf_expert_sources", + "dummy_gguf_expert_sources", +] From 8fb43b5b90732d4d57fed508d686abffbaf86cca Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Tue, 25 Aug 2026 16:04:12 -0700 Subject: [PATCH 32/36] feat(gguf): deepseek4 config from GGUF metadata alone parse_config recovers DeepseekV4Args from the checkpoint's inference/config.json, which ships beside safetensors. A standalone .gguf has no such file, so this rebuilds the same dataclass from KV keys. Nothing falls back to a dataclass default: a silently-defaulted hyperparameter here yields a model that loads and generates confidently wrong text. block_count does NOT relate to the MTP block consistently across architectures, so the served layer count is derived rather than assumed. qwen35moe counts its NextN block inside block_count (Ornith: 41 blocks, blk.40 IS the MTP block, 40 served); deepseek4 does not (43 blocks, blk.0..blk.42 all served, MTP carries no blk tensors). Subtracting nextn_predict_layers unconditionally silently drops the last real layer, which is exactly what the first version did. compress_ratios is the authority instead: one entry per served layer plus MTP. Entries != 0 mark a layer with an attention compressor and == 4 one with the lightning indexer, so the schedule is checked against the tensor table on every load. That check caught the off-by-one immediately (40/20 instead of 41/21) rather than leaving it to surface as degraded output after a 145 GiB load. Verified against antirez/deepseek-v4-gguf Q4KExperts: 43 served layers, 41 compressors, 21 indexers, expert types (Q4_K, Q4_K), score_func sqrtsoftplus from expert_gating_func 4, route_scale 1.5. iter_gguf_weights and convert_deepseek4_to_gguf still to come; the arch is not registered yet, so none of this is reachable. --- python/freetoken/models/deepseek_v4/gguf.py | 257 ++++++++++++++++++++ 1 file changed, 257 insertions(+) create mode 100644 python/freetoken/models/deepseek_v4/gguf.py diff --git a/python/freetoken/models/deepseek_v4/gguf.py b/python/freetoken/models/deepseek_v4/gguf.py new file mode 100644 index 00000000..0004dcb8 --- /dev/null +++ b/python/freetoken/models/deepseek_v4/gguf.py @@ -0,0 +1,257 @@ +"""Serve a deepseek4 GGUF checkpoint. + +Unlike the safetensors path, everything here comes from the GGUF's own metadata. The +reference ``parse_config`` recovers ``DeepseekV4Args`` from the checkpoint's +``inference/config.json``, which ships beside the weights; a standalone .gguf has no such +file, and being self-describing is the point of the format. ``_args_from_gguf`` below +rebuilds the same dataclass from KV keys alone. + +The mapping from GGUF tensor to model parameter was established by reading both sides +rather than by analogy with the qwen adapters, because three tensors do not behave the way +the names suggest: + +* ``attn_output_a`` is Q8_0 in the file but ``attn.wo_a`` is a bare ``nn.Parameter`` in + bfloat16, not a Linear (attention.py: "wo_a: dequantized to bf16, the reference runs a + bf16 grouped-output einsum"). It must be dequantized to dense, and it has no ``.weight`` + suffix. +* the compressor and indexer projections are **F16** in the file, i.e. unquantized. F16 is + in ``GGML_UNQUANTIZED``, so ``fused_mul_mat_gguf`` would take the ``x @ qweight.T`` path + while ``GGUFLinear`` allocates a uint8 buffer. They must land dense on a normal + ``.weight``, never packed. +* ``Indexer.wq_b`` is declared ``Linear(kind="fp8")`` (compress.py), which allocates a + ``.scale`` that no GGUF tensor can fill, because that tensor is F16 here. That Linear is + replaced outright rather than populated. + +Routed experts never pass through this module; they are streamed from the offload cache by +``gguf_experts.load_gguf_expert_sources``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Iterator + +import torch + +from freetoken.models.config import DSV4AttentionGroupConfig, ModelConfig, RotaryConfig + +from .args import DeepseekV4Args + +if TYPE_CHECKING: + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "deepseek4" + +# llama.cpp's expert-gating enum. DeepSeek-V4 scores with sqrt-softplus; 1 and 2 are the +# long-standing softmax/sigmoid values. An unknown id raises rather than silently picking a +# scoring function, because the wrong one routes to the wrong experts and still produces +# fluent text. +_GATING = {1: "softmax", 2: "sigmoid", 4: "sqrtsoftplus"} + + +def _kv(shim: "GgufConfigShim", key: str, default: Any = None) -> Any: + """One ``deepseek4.*`` metadata value. No default means the key is mandatory.""" + full = f"{_ARCH}.{key}" + md = shim if isinstance(shim, dict) else shim.metadata + if full not in md: + if default is None: + raise ValueError( + f"deepseek4 GGUF is missing required metadata key {full!r}; this file does " + f"not carry the config this adapter needs" + ) + return default + return md[full] + + +def _args_from_gguf(shim: "GgufConfigShim") -> DeepseekV4Args: + """Rebuild DeepseekV4Args from GGUF metadata alone. + + Every field is sourced from a key that is actually present in the checkpoint; nothing + is left to the dataclass default, because a silently-defaulted hyperparameter here + produces a model that loads and generates confidently wrong text. + + Cross-checks worth keeping: ``compress_ratios`` carries one entry per layer plus the + MTP layers, entries != 0 mark layers with an attention compressor, and entries == 4 + mark layers with the lightning indexer. Those counts must match the tensor table (41 + and 21 respectively for DeepSeek-V4-Flash), which is what makes this mapping + self-validating rather than merely plausible. + """ + ratios = tuple(int(x) for x in _kv(shim, "attention.compress_ratios")) + swiglu = [float(x) for x in _kv(shim, "swiglu_clamp_exp", [])] + gate_id = int(_kv(shim, "expert_gating_func")) + if gate_id not in _GATING: + raise ValueError( + f"deepseek4 GGUF: unknown expert_gating_func {gate_id}; known values are " + f"{sorted(_GATING)} (routing with the wrong scoring function still generates " + f"fluent text, so this is not defaulted)" + ) + + return DeepseekV4Args( + max_batch_size=1, + max_seq_len=int(_kv(shim, "context_length")), + # The fp8/fp4 reference paths do not apply: a GGUF carries its own block-quantized + # weights and the adapter swaps the quantized projections for GGUF ops. + dtype="bf16", + scale_fmt=None, + expert_dtype=None, + vocab_size=int(_kv(shim, "vocab_size")), + dim=int(_kv(shim, "embedding_length")), + moe_inter_dim=int(_kv(shim, "expert_feed_forward_length")), + n_layers=int(_kv(shim, "block_count")), + n_hash_layers=int(_kv(shim, "hash_layer_count")), + n_mtp_layers=int(_kv(shim, "nextn_predict_layers", 0)), + n_heads=int(_kv(shim, "attention.head_count")), + n_routed_experts=int(_kv(shim, "expert_count")), + n_shared_experts=int(_kv(shim, "expert_shared_count")), + n_activated_experts=int(_kv(shim, "expert_used_count")), + score_func=_GATING[gate_id], + route_scale=float(_kv(shim, "expert_weights_scale")), + swiglu_limit=(swiglu[0] if swiglu else 10.0), + q_lora_rank=int(_kv(shim, "attention.q_lora_rank")), + head_dim=int(_kv(shim, "attention.key_length")), + rope_head_dim=int(_kv(shim, "rope.dimension_count")), + norm_eps=float(_kv(shim, "attention.layer_norm_rms_epsilon")), + o_groups=int(_kv(shim, "attention.output_group_count")), + o_lora_rank=int(_kv(shim, "attention.output_lora_rank")), + window_size=int(_kv(shim, "attention.sliding_window")), + compress_ratios=ratios, + compress_rope_theta=float(_kv(shim, "attention.compress_rope_freq_base")), + original_seq_len=int(_kv(shim, "rope.scaling.original_context_length")), + rope_theta=float(_kv(shim, "rope.freq_base")), + rope_factor=float(_kv(shim, "rope.scaling.factor")), + beta_fast=int(_kv(shim, "rope.scaling.yarn_beta_fast")), + beta_slow=int(_kv(shim, "rope.scaling.yarn_beta_slow")), + index_n_heads=int(_kv(shim, "attention.indexer.head_count")), + index_head_dim=int(_kv(shim, "attention.indexer.key_length")), + index_topk=int(_kv(shim, "attention.indexer.top_k")), + hc_mult=int(_kv(shim, "hyper_connection.count")), + hc_sinkhorn_iters=int(_kv(shim, "hyper_connection.sinkhorn_iterations")), + hc_eps=float(_kv(shim, "hyper_connection.epsilon")), + ) + + + +def _check_schedule(model_path: str, args: DeepseekV4Args, served: int) -> None: + """Cross-check the compress_ratios schedule against the tensor table. + + Cheap (the tensor table is metadata, not weights) and worth doing every load: it is the + difference between finding a layer-count error here and finding it as degraded output + after a 145 GiB load. + """ + from freetoken.models.gguf.reader import gguf_tensor_names + + names = gguf_tensor_names(model_path) + want_compressor = sum(1 for r in args.compress_ratios[:served] if r != 0) + want_indexer = sum(1 for r in args.compress_ratios[:served] if r == 4) + got_compressor = sum( + 1 for i in range(served) if f"blk.{i}.attn_compressor_kv.weight" in names) + got_indexer = sum( + 1 for i in range(served) if f"blk.{i}.indexer.attn_q_b.weight" in names) + + for label, want, got in (("compressor", want_compressor, got_compressor), + ("indexer", want_indexer, got_indexer)): + if want != got: + raise ValueError( + f"deepseek4 GGUF: compress_ratios predicts {want} layers with a {label} " + f"but the file has {got}; the per-layer schedule does not match this " + f"checkpoint (a wrong served-layer count is the usual cause)" + ) + +def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + """ModelConfig for a deepseek4 GGUF, mirroring deepseek_v4/config.py::parse_config. + + The served layer count excludes the trailing MTP/NextN block: ``block_count`` counts it + but it is not part of the forward pass, and treating it as a layer makes a uniform + expert bank look mixed. + """ + args = _args_from_gguf(shim) + model_path = getattr(shim, "model_path", None) + + # How block_count relates to the MTP block is NOT consistent across architectures, so + # it is derived rather than assumed. qwen35moe counts its NextN block inside + # block_count (Ornith: block_count 41, blk.0..blk.40 where blk.40 is the MTP block, 40 + # served). deepseek4 does not (block_count 43, blk.0..blk.42 all served, and the MTP + # layer carries no blk tensors at all). Subtracting n_mtp_layers unconditionally + # silently drops the last real layer here. + # + # compress_ratios is the authority: it carries one entry per served layer plus the MTP + # layers, so the served count falls out of it and is then cross-checked below. + served_layers = len(args.compress_ratios) - args.n_mtp_layers + if served_layers != args.n_layers: + raise ValueError( + f"deepseek4 GGUF: compress_ratios implies {served_layers} served layers " + f"({len(args.compress_ratios)} entries minus {args.n_mtp_layers} MTP) but " + f"block_count is {args.n_layers}; refusing to guess which is right" + ) + args.n_layers = served_layers + + rope_scaling = { + "rope_type": "yarn", + "factor": args.rope_factor, + "beta_fast": args.beta_fast, + "beta_slow": args.beta_slow, + "original_max_position_embeddings": args.original_seq_len, + } + + from .gguf_experts import gguf_expert_types + + types = gguf_expert_types(model_path, served_layers) if model_path else None + expert_types = (types["gate_up"][0], types["down"][0]) if types else None + + # The schedule derived from compress_ratios must match what the file actually contains. + # A compressor exists where ratio != 0 and a lightning indexer where ratio == 4, so + # these counts are an independent check on the layer count above: an off-by-one shows + # up here as a mismatch rather than as a quietly missing layer at serving time. + if model_path: + _check_schedule(model_path, args, served_layers) + + return ModelConfig( + num_layers=served_layers, + num_qo_heads=args.n_heads, + num_kv_heads=1, # MLA: a single shared latent KV head (K == V) + head_dim=args.head_dim, + hidden_size=args.dim, + vocab_size=args.vocab_size, + intermediate_size=args.moe_inter_dim, + hidden_act="silu", + rms_norm_eps=args.norm_eps, + tie_word_embeddings=False, # output.weight is a separate tensor from token_embd + rotary_config=RotaryConfig( + head_dim=args.head_dim, + rotary_dim=args.rope_head_dim, + max_position=args.max_seq_len, + base=args.rope_theta, + scaling=rope_scaling, + ), + num_experts=args.n_routed_experts, + num_experts_per_tok=args.n_activated_experts, + moe_intermediate_size=args.moe_inter_dim, + norm_topk_prob=True, + model_type="deepseek_v4", + architectures=["DeepseekV4ForCausalLM"], + moe_enabled=True, + expert_quant="gguf", + attn_sm_scale=args.head_dim**-0.5, + dsv4_args=args, + gguf_model_path=model_path, + gguf_expert_types=expert_types, + attention_groups=( + DSV4AttentionGroupConfig( + name="dsv4", + layer_ids=tuple(range(served_layers)), + num_kv_heads=1, + head_dim=args.head_dim, + sliding_window=args.window_size, + ), + ), + ) + + +def is_gguf_model(config: ModelConfig) -> bool: + """True when this config came from a GGUF checkpoint (native block-quant path).""" + return getattr(config, "gguf_model_path", None) is not None + + +__all__ = [ + "parse_gguf_config", + "is_gguf_model", +] From 5952308b925e8d5e4b095bf89c600c2485dc95f9 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Tue, 25 Aug 2026 17:01:24 -0700 Subject: [PATCH 33/36] feat(gguf): serve deepseek4 checkpoints Completes the adapter and registers the architecture. Reconciles against the real 164GB antirez/deepseek-v4-gguf Q4KExperts file: 1199 model parameters, 1199 matched, 0 unknown destinations, 0 shape mismatches, 0 unfilled. Four things did not carry over from the qwen adapters: block_count does not relate to the MTP block consistently. qwen35moe counts its NextN block inside it, deepseek4 does not. Subtracting nextn_predict_layers unconditionally dropped the last real layer; the compressor/indexer schedule cross-check caught it as 40/20 instead of 41/21. deepseek_v4 is raw nn.Module and loads via load_state_dict over named_parameters(), unlike the BaseOP-based qwen trees. FreeToken's GGUFLinear holds qweight as a plain attribute, which that loader cannot see, and assigning a non-Module over a Module child raises. Hence GGUFLinearNN / GGUFEmbeddingNN here, carrying the packed bytes in a non-grad uint8 Parameter named .weight so the loader's .to(p.dtype) cast is a no-op. dequant.dequantize's pure-torch path implements only Q4_0 and Q6_K; this checkpoint's attention projections and lm_head are Q8_0, which needs the CUDA kernel. Unquantized F16/F32 are reinterpreted rather than dequantized. Three tensors do not behave as their names suggest, and each would have been a silent mis-wire: attn_output_a is Q8_0 in the file but attn.wo_a is a bare bf16 nn.Parameter (no .weight, never packed); the compressor and indexer projections are F16, i.e. in GGML_UNQUANTIZED, so GGUFLinear's uint8 buffer cannot hold them; and Indexer.wq_b is declared Linear(kind="fp8") with a .scale no GGUF tensor can fill, so it is rebuilt as bf16 rather than populated. tid2eid is a routing index, not a weight: read as int32 and widened, never sent through a float dequant that would round token ids. Registration covers all five points: GGUF_ARCH_TO_REGISTRY, the register.py ModelSpec, _TOKENIZER_ARCH and _STOP_TOKENS, and the package exports. Stop tokens were read from the checkpoint's own vocab (eos id 1 <|end_of_sentence|>, plus <|EOT|> 128805) rather than assumed. Not yet loaded end to end: that needs ~145 GiB of pinned host RAM for the expert banks. --- .../freetoken/models/deepseek_v4/__init__.py | 13 + python/freetoken/models/deepseek_v4/gguf.py | 326 ++++++++++++++++++ python/freetoken/models/gguf/config.py | 1 + python/freetoken/models/gguf/tokenizer.py | 6 + python/freetoken/models/register.py | 6 + 5 files changed, 352 insertions(+) diff --git a/python/freetoken/models/deepseek_v4/__init__.py b/python/freetoken/models/deepseek_v4/__init__.py index b3cb41db..3fb8c91d 100644 --- a/python/freetoken/models/deepseek_v4/__init__.py +++ b/python/freetoken/models/deepseek_v4/__init__.py @@ -15,6 +15,13 @@ from .args import DeepseekV4Args, load_args from .config import parse_config +from .gguf import ( + convert_deepseek4_to_gguf, + is_gguf_model, + iter_gguf_weights, + parse_gguf_config, +) +from .gguf_experts import gguf_expert_types, load_gguf_expert_sources from .model import DeepseekV4ForCausalLM from .weight import iter_weights, load_dsfp4_expert_sources @@ -25,4 +32,10 @@ "DeepseekV4ForCausalLM", "iter_weights", "load_dsfp4_expert_sources", + "parse_gguf_config", + "iter_gguf_weights", + "convert_deepseek4_to_gguf", + "is_gguf_model", + "gguf_expert_types", + "load_gguf_expert_sources", ] diff --git a/python/freetoken/models/deepseek_v4/gguf.py b/python/freetoken/models/deepseek_v4/gguf.py index 0004dcb8..160398db 100644 --- a/python/freetoken/models/deepseek_v4/gguf.py +++ b/python/freetoken/models/deepseek_v4/gguf.py @@ -251,7 +251,333 @@ def is_gguf_model(config: ModelConfig) -> bool: return getattr(config, "gguf_model_path", None) is not None + +class GGUFLinearNN(torch.nn.Module): + """A GGUF-quantized Linear that DSV4's loader can actually fill. + + FreeToken's own ``layers.gguf.GGUFLinear`` is a ``BaseOP`` holding ``qweight`` as a + plain tensor. That works for the qwen models, whose trees are built from BaseOP, but + deepseek_v4 is raw ``nn.Module`` and loads via + ``DeepseekV4ForCausalLM.load_state_dict``, which walks ``named_parameters()`` and + demands a key for every one. A plain attribute is invisible there, and assigning a + non-Module over a Module child raises outright. + + So the packed block bytes live in an ordinary ``nn.Parameter`` -- uint8, requires_grad + False -- named ``weight`` to match the naming the rest of this model uses. The loader's + ``.to(p.dtype)`` cast is then a no-op on uint8, and the tensor arrives byte-for-byte. + """ + + def __init__(self, in_features: int, out_features: int, quant_type: int, + bias: bool = False): + super().__init__() + from freetoken.models.gguf.dequant import row_bytes + + self.in_features = in_features + self.out_features = out_features + self._quant_type = int(quant_type) + self.weight = torch.nn.Parameter( + torch.empty(out_features, row_bytes(in_features, self._quant_type), + dtype=torch.uint8), + requires_grad=False, + ) + if bias: + self.bias = torch.nn.Parameter(torch.empty(out_features), requires_grad=False) + else: + self.register_parameter("bias", None) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from freetoken.layers.gguf import fused_mul_mat_gguf + + out = fused_mul_mat_gguf(x, self.weight, self._quant_type) + return out if self.bias is None else out + self.bias + + +class GGUFEmbeddingNN(torch.nn.Module): + """GGUF-quantized vocab embedding, as an nn.Module for the same reason as above. + + The table is never dequantized whole: only the looked-up rows are gathered in packed + form and dequantized per lookup. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int, quant_type: int): + super().__init__() + from freetoken.models.gguf.dequant import row_bytes + + self.num_embeddings = num_embeddings + self.embedding_dim = embedding_dim + self._quant_type = int(quant_type) + self.weight = torch.nn.Parameter( + torch.empty(num_embeddings, row_bytes(embedding_dim, self._quant_type), + dtype=torch.uint8), + requires_grad=False, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.gguf import ggml_dequantize + + flat = x.flatten() + rows = self.weight.index_select(0, flat) + y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, + torch.bfloat16) + return y.view(*x.shape, self.embedding_dim) + + +def _dense(t, dtype: torch.dtype) -> torch.Tensor: + """A GgufTensor as a dense tensor of its torch shape. + + Two paths, because neither covers everything. Unquantized types (F32/F16/BF16) are + already values, so the packed bytes are simply reinterpreted -- no kernel needed, and + it works without CUDA. Block-quantized types go through the vendored CUDA dequant: + ``dequant.dequantize``'s pure-torch fallback only implements Q4_0 and Q6_K, and this + checkpoint stores its attention projections and lm_head as Q8_0. + """ + from freetoken.models.gguf.dequant import ( + BLOCK_SHAPE, + GGML_BF16, + GGML_F16, + GGML_F32, + GGML_UNQUANTIZED, + ) + + gt = int(t.ggml_type) + raw = t.packed() + if gt in GGML_UNQUANTIZED: + view = {GGML_F32: torch.float32, GGML_F16: torch.float16, + GGML_BF16: torch.bfloat16}[gt] + return raw.reshape(-1).view(view).reshape(t.shape).to(dtype) + + from freetoken.kernel.gguf import ggml_dequantize + + block, type_size = BLOCK_SHAPE[gt] + in_features = t.row_bytes // type_size * block + out = ggml_dequantize(raw.cuda().contiguous(), gt, t.rows, in_features, + torch.bfloat16) + return out.reshape(t.shape).to(dtype) + + +def _to_bf16(t) -> torch.Tensor: + return _dense(t, torch.bfloat16) + + +def _to_f32(t) -> torch.Tensor: + return _dense(t, torch.float32) + + +def _to_i64(t) -> torch.Tensor: + """Read an I32 index table as int64. + + tid2eid is a routing table, not a weight: dequantizing it through a float path would + round large token ids. Reinterpret the raw bytes instead. + """ + return t.packed().reshape(-1).view(torch.int32).reshape(t.shape).to(torch.int64) + + +# suffix -> (destination template, kind). "packed" lands on a GGUFLinear's .weight; +# "bf16"/"f32" are dequantized onto an ordinary parameter. The destination is spelled out +# per tensor rather than derived from the name, because three of them do not follow the +# pattern the names imply (see the module docstring). +_LAYER_MAP: dict[str, tuple[str, str]] = { + "attn_norm.weight": ("attn_norm.weight", "f32"), + "ffn_norm.weight": ("ffn_norm.weight", "f32"), + "attn_q_a.weight": ("attn.wq_a.weight", "packed"), + "attn_q_a_norm.weight": ("attn.q_norm.weight", "f32"), + "attn_q_b.weight": ("attn.wq_b.weight", "packed"), + "attn_kv.weight": ("attn.wkv.weight", "packed"), + "attn_kv_a_norm.weight": ("attn.kv_norm.weight", "f32"), + # wo_a is a bare nn.Parameter in bf16, NOT a Linear: no .weight, never packed. + "attn_output_a.weight": ("attn.wo_a", "bf16"), + "attn_output_b.weight": ("attn.wo_b.weight", "packed"), + "attn_sinks.weight": ("attn.attn_sink", "f32"), + # compressor / indexer projections are F16 in the file. F16 is in GGML_UNQUANTIZED, so + # GGUFLinear cannot hold them -- they must land dense on a normal .weight. + "attn_compressor_kv.weight": ("attn.compressor.wkv.weight", "bf16"), + "attn_compressor_gate.weight": ("attn.compressor.wgate.weight", "bf16"), + "attn_compressor_norm.weight": ("attn.compressor.norm.weight", "f32"), + "attn_compressor_ape.weight": ("attn.compressor.ape", "f32"), + "indexer.attn_q_b.weight": ("attn.indexer.wq_b.weight", "bf16"), + "indexer.proj.weight": ("attn.indexer.weights_proj.weight", "bf16"), + "indexer_compressor_kv.weight": ("attn.indexer.compressor.wkv.weight", "bf16"), + "indexer_compressor_gate.weight": ("attn.indexer.compressor.wgate.weight", "bf16"), + "indexer_compressor_norm.weight": ("attn.indexer.compressor.norm.weight", "f32"), + "indexer_compressor_ape.weight": ("attn.indexer.compressor.ape", "f32"), + "hc_attn_base.weight": ("hc_attn_base", "f32"), + "hc_attn_fn.weight": ("hc_attn_fn", "f32"), + "hc_attn_scale.weight": ("hc_attn_scale", "f32"), + "hc_ffn_base.weight": ("hc_ffn_base", "f32"), + "hc_ffn_fn.weight": ("hc_ffn_fn", "f32"), + "hc_ffn_scale.weight": ("hc_ffn_scale", "f32"), + "ffn_gate_inp.weight": ("ffn.gate.weight", "bf16"), + "exp_probs_b.bias": ("ffn.gate.bias", "f32"), + # DeepSeek names the shared expert gate/up/down; Expert calls them w1/w3/w2. + "ffn_gate_shexp.weight": ("ffn.shared_experts.w1.weight", "packed"), + "ffn_up_shexp.weight": ("ffn.shared_experts.w3.weight", "packed"), + "ffn_down_shexp.weight": ("ffn.shared_experts.w2.weight", "packed"), +} + +_GLOBAL_MAP: dict[str, tuple[str, str]] = { + "output_norm.weight": ("norm.weight", "f32"), + # output.weight is Q8_0 but `head` is a bare bf16 nn.Parameter consumed by F.linear, + # so it is dequantized rather than swapped. deepseek_v4/model.py already slices to the + # last prefill position itself, so it needs no GGUFLMHead. + "output.weight": ("head", "bf16"), + "output_hc_base.weight": ("hc_head_base", "f32"), + "output_hc_fn.weight": ("hc_head_fn", "f32"), + "output_hc_scale.weight": ("hc_head_scale", "f32"), +} + +# Routed experts are streamed from the offload cache, never yielded here. +_EXPERT_SUFFIXES = frozenset( + {"ffn_gate_exps.weight", "ffn_up_exps.weight", "ffn_down_exps.weight"}) + + +def iter_gguf_weights( + model_path: str, + device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """Yield (param_name, tensor) for every non-expert deepseek4 parameter.""" + import re + + from freetoken.models.gguf.reader import iter_gguf_tensors + + assert not include_moe_experts, ( + "deepseek4 GGUF keeps its routed experts in the offload cache; they are loaded by " + "gguf_experts.load_gguf_expert_sources, not by iter_gguf_weights." + ) + assert include_non_moe + + conv = {"packed": lambda t: t.packed(), "bf16": _to_bf16, "f32": _to_f32} + + for t in iter_gguf_tensors(model_path): + name = t.name + m = re.match(r"^blk\.(\d+)\.(.+)$", name) + if m is None: + dest = _GLOBAL_MAP.get(name) + if dest is None: + if name == "token_embd.weight": + yield "embed.weight", t.packed() + continue + raise ValueError( + f"deepseek4 GGUF: unmapped global tensor {name!r}; this checkpoint does " + f"not match the layout this adapter expects" + ) + path, kind = dest + yield path, conv[kind](t) + continue + + layer, suffix = int(m.group(1)), m.group(2) + if suffix in _EXPERT_SUFFIXES: + continue # offload cache + if suffix == "ffn_gate_tid2eid.weight": + # Hash routing table on the first n_hash_layers layers; an index, not a weight. + yield f"layers.{layer}.ffn.gate.tid2eid", _to_i64(t) + continue + dest = _LAYER_MAP.get(suffix) + if dest is None: + raise ValueError( + f"deepseek4 GGUF: unmapped tensor {name!r}; this checkpoint does not match " + f"the layout this adapter expects" + ) + path, kind = dest + yield f"layers.{layer}.{path}", conv[kind](t) + + +def _scan_quant_types(model_path: str) -> dict[tuple[int, str], int]: + """(layer, suffix) -> ggml type, straight from the tensor table. + + A guessed type allocates a wrong-sized packed buffer, so nothing here has a default. + Globals use layer -1. + """ + import re + + from freetoken.models.gguf.reader import iter_gguf_tensors + + out: dict[tuple[int, str], int] = {} + for t in iter_gguf_tensors(model_path): + m = re.match(r"^blk\.(\d+)\.(.+)$", t.name) + if m: + out[(int(m.group(1)), m.group(2))] = int(t.ggml_type) + else: + out[(-1, t.name)] = int(t.ggml_type) + return out + + +def convert_deepseek4_to_gguf(model, config: ModelConfig, *, model_path: str) -> None: + """In place: swap deepseek4's quantized projections + embedding for native GGUF ops. + + Swapped to GGUFLinear (Q8_0 in the checkpoint): attention wq_a / wq_b / wkv / wo_b and + the shared expert's w1 / w2 / w3. + + Deliberately NOT swapped: + * ``attn.wo_a`` is a bare bf16 nn.Parameter, not a Linear; it is dequantized dense. + * the compressor's wkv / wgate and the indexer's weights_proj are already + ``Linear(kind="bf16")`` and their tensors are F16, so they take dense weights. + * ``head`` is a bare bf16 nn.Parameter consumed by F.linear. + + Replaced rather than swapped: ``indexer.wq_b`` is declared ``Linear(kind="fp8")``, + which allocates a ``.scale`` no GGUF tensor can fill because that tensor is F16 here. + It becomes a bf16 Linear so ``.weight`` is bf16 and ``scale`` is None. + """ + from .layers import Linear + + quant = _scan_quant_types(model_path) + + def qt(layer: int, suffix: str) -> int: + key = (layer, suffix) + if key not in quant: + where = suffix if layer < 0 else f"blk.{layer}.{suffix}" + raise ValueError( + f"deepseek4 GGUF {model_path}: expected tensor {where} is absent, so its " + f"quant type cannot be read; this checkpoint does not match the layout " + f"this adapter expects" + ) + return quant[key] + + def swap_linear(owner, attr: str, quant_type: int) -> None: + lin = getattr(owner, attr) + setattr( + owner, attr, + GGUFLinearNN(lin.in_features, lin.out_features, quant_type, + bias=getattr(lin, "bias", None) is not None), + ) + + # DeepseekV4ForCausalLM is an engine wrapper; the parameters live on the inner + # Transformer, and state_dict() names them relative to it (no "_transformer." prefix), + # which is what iter_gguf_weights emits. + root = getattr(model, "_transformer", model) + + root.embed = GGUFEmbeddingNN( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + quant_type=qt(-1, "token_embd.weight"), + ) + + for layer_idx, layer in enumerate(root.layers): + attn = layer.attn + swap_linear(attn, "wq_a", qt(layer_idx, "attn_q_a.weight")) + swap_linear(attn, "wq_b", qt(layer_idx, "attn_q_b.weight")) + swap_linear(attn, "wkv", qt(layer_idx, "attn_kv.weight")) + swap_linear(attn, "wo_b", qt(layer_idx, "attn_output_b.weight")) + + idx = getattr(attn, "indexer", None) + if idx is not None: + # F16 in the file, fp8 in the module: rebuild as bf16 so there is no orphan + # .scale and F.linear is used instead of the block-fp8 GEMM. + old = idx.wq_b + idx.wq_b = Linear(old.in_features, old.out_features, + bias=getattr(old, "bias", None) is not None, kind="bf16") + + shexp = layer.ffn.shared_experts + swap_linear(shexp, "w1", qt(layer_idx, "ffn_gate_shexp.weight")) + swap_linear(shexp, "w3", qt(layer_idx, "ffn_up_shexp.weight")) + swap_linear(shexp, "w2", qt(layer_idx, "ffn_down_shexp.weight")) + + __all__ = [ "parse_gguf_config", + "iter_gguf_weights", + "convert_deepseek4_to_gguf", "is_gguf_model", ] diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 1ab5704c..fd8d16a8 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -25,6 +25,7 @@ # config's expert_count is absent so moe_enabled comes out False. "qwen35": "Qwen35GGUFForCausalLM", "qwen3moe": "Qwen3MoeGGUFForCausalLM", + "deepseek4": "DeepseekV4GGUFForCausalLM", } diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index ddc4433b..e5d4f4d5 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -23,6 +23,8 @@ "qwen35moe": "qwen2", "qwen35": "qwen2", "qwen3moe": "qwen2", + # tokenizer.ggml.model is gpt2 (BPE), pre joyai-llm, 129280 entries. + "deepseek4": "llama", } # Per-arch chat/stop tokens, in preference order: the first one present in the vocab @@ -36,6 +38,10 @@ # Dense sibling: same vocab and same chat markers as the MoE variant. "qwen35": ("<|im_end|>", "<|endoftext|>"), "qwen3moe": ("<|im_end|>", "<|endoftext|>"), + # Read from the vocab: eos id 1 is the document end, <|EOT|> (128805) ends a + # chat turn. <｜User｜> is deliberately not a stop -- the template emits it + # before the model speaks, not after. + "deepseek4": ("<|EOT|>", "<｜end▁of▁sentence｜>"), } diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 48bd93e7..73237ccc 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -129,6 +129,12 @@ class ModelSpec: parse_config="parse_gguf_config", iter_weights="iter_gguf_weights", ), + "DeepseekV4GGUFForCausalLM": ModelSpec( + "freetoken.models.deepseek_v4", + "DeepseekV4ForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + ), "GptOssForCausalLM": ModelSpec( "freetoken.models.gpt_oss", "GptOssForCausalLM", From 7da6b77e3a20eaa310d246650745b05538661b5a Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Tue, 25 Aug 2026 17:11:15 -0700 Subject: [PATCH 34/36] fix(gguf): unquantized projections multiplied raw bytes instead of values GGUFLinear and GGUFEmbedding store every ggml type in a uint8 buffer of row_bytes width, including F32/F16/BF16, where "packed" just means the raw value bytes. fused_mul_mat_gguf's unquantized branch multiplied that byte view directly, so in_features came out as row_bytes -- twice too wide for F16 -- and the call failed with "mat1 and mat2 shapes cannot be multiplied (4x2048 and 4096x248320)". Only reachable when a checkpoint stores a projection unquantized, which is why it had not surfaced: Ornith's output.weight is Q6_K and takes a dequant path. Apodex-1.1-mini ships output.weight as F16 and fails at the first forward. The bytes are now reinterpreted through the type's real dtype before the matmul. The ACTIVATION is cast rather than the weight: converting the weight would copy the whole matrix per call (about 1 GB per forward for a 248k-vocab lm_head) and allocating that during CUDA graph capture fails outright, which is what the first version of this fix did. x is [tokens, hidden], so casting it is negligible, and computing in the stored precision matches what llama.cpp does for these tensors. Verified on Apodex-1.1-mini Q2_K (qwen35moe, 40 layers, 256 experts, F16 output.weight) on an 8GB RTX 4060: serves, and correct at temperature 0 -- "The capital city of France is" -> " Paris", "The largest planet in our solar system is" -> " Jupiter. Its diameter is approximately 142,98[4 km]". --- python/freetoken/layers/gguf.py | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py index ae0b1b43..bc4754e8 100644 --- a/python/freetoken/layers/gguf.py +++ b/python/freetoken/layers/gguf.py @@ -41,6 +41,9 @@ from freetoken.models.gguf.dequant import ( BLOCK_SHAPE, DEQUANT_TYPES, + GGML_BF16, + GGML_F16, + GGML_F32, GGML_NAME, GGML_UNQUANTIZED, MMQ_TYPES, @@ -48,6 +51,14 @@ row_bytes, ) +# ggml type -> the dtype its raw bytes represent. Only the unquantized types appear here; +# everything else goes through a dequant kernel. +_UNQUANTIZED_DTYPE = { + GGML_F32: torch.float32, + GGML_F16: torch.float16, + GGML_BF16: torch.bfloat16, +} + from .base import BaseOP # Below this token count, the MMVQ GEMV kernel wins (matches vLLM's heuristic). @@ -73,7 +84,22 @@ def fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int if x.shape[0] == 0: return x.new_empty((0, out_features)) if qweight_type in GGML_UNQUANTIZED: - return x @ qweight.T + # GGUFLinear/GGUFEmbedding store every type in a uint8 buffer of row_bytes width, + # including the unquantized ones, where "packed" just means the raw F32/F16/BF16 + # bytes. Those must be reinterpreted before the matmul: multiplying the byte view + # directly gives an in_features of row_bytes (2x too wide for F16) and fails with + # "mat1 and mat2 shapes cannot be multiplied". A checkpoint only reaches this path + # when it stores a projection unquantized -- Apodex-1.1-mini ships output.weight as + # F16, which is how this surfaced; models whose lm_head is Q6_K never hit it. + w = qweight + if w.dtype == torch.uint8: + w = w.view(_UNQUANTIZED_DTYPE[qweight_type]) + # Cast the ACTIVATION, not the weight. Converting the weight would copy the whole + # matrix on every call -- about 1 GB per forward for a 248k-vocab lm_head -- and + # allocating that during CUDA graph capture fails outright. x is [tokens, hidden], + # so casting it is negligible, and computing in the stored precision is what + # llama.cpp does for these tensors anyway. + return (x.to(w.dtype) @ w.T).to(x.dtype) if x.shape[0] <= _MMVQ_SAFE and qweight_type in MMVQ_TYPES: return ggml_mul_mat_vec_a8(qweight, x, qweight_type, out_features) if qweight_type in MMQ_TYPES: From 3178a59de15d87cd1294ffd9010b2c6ad42139fa Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Tue, 25 Aug 2026 23:11:51 -0700 Subject: [PATCH 35/36] fix(moe): GGUF checkpoints could never reach the CPU executor Three defects with one root: "gguf" is a container tag, not a weight layout. The checkpoint picks a ggml type per tensor and the concrete CPU format has to be recovered from the bank types, but two call sites tested the tag directly. _cpu_moe_executor_viable compared expert_quant against _WFMT_IDS, which answers False for EVERY GGUF checkpoint. That silently disabled the automatic residency split on hosts where CUDA pinning is quota-capped -- WSL caps it near 40% of RAM (measured: 81.78 GiB of 204). The symptom was not a clear refusal but cudaHostRegister failing partway through the banks, which reads as a memory shortage rather than a dispatch gap. gemm1_dot handled bf16 and q4_0 and then FELL THROUGH to the NVFP4 path, which dereferences scale/global pointers that are null for GGUF banks. An unhandled format therefore segfaulted inside a worker thread with no Python traceback. It now raises through TORCH_CHECK naming the format, same reasoning as the kernel default: guards in #138: an unhandled case that reads null or uninitialised memory is far worse than one that errors. GGUFEmbedding dequantized unconditionally, but the unquantized types are raw value bytes with no dequant kernel at all (ggml_dequantize rejects type 1 outright). DeepSeek-V4 ships token_embd as F16 and died on the first lookup. The gathered rows are now reinterpreted for those types, matching the fix already applied to fused_mul_mat_gguf. --- python/freetoken/engine/engine.py | 15 +++++++++++++++ .../freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 7 +++++++ python/freetoken/layers/gguf.py | 7 ++++++- 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 018e9622..c386e28b 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1144,6 +1144,21 @@ def _cpu_moe_executor_viable(model_config) -> bool: return False expert_quant = getattr(model_config, "expert_quant", "none") fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16") + if fmt == "gguf": + # "gguf" is a container tag, not a layout: the checkpoint picks a ggml type per + # tensor and the concrete CPU format has to be recovered from the bank types. + # Testing the tag against _WFMT_IDS answers False for EVERY GGUF checkpoint, which + # silently disables the automatic residency split on hosts where CUDA pinning is + # quota-capped (WSL caps it near half of RAM). The symptom is not a clear refusal + # but cudaHostRegister failing partway through the banks. + from freetoken.moe.cpu_executor import _GGML_TO_CPU_FMT + + types = getattr(model_config, "gguf_expert_types", None) + if not types: + return False + gate_up, down = int(types[0]), int(types[1]) + # one weight_format serves both banks, so mixed types cannot run on the CPU path + return gate_up == down and gate_up in _GGML_TO_CPU_FMT return fmt == "mxfp4" or fmt in _WFMT_IDS diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 68e6324c..48210b9d 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1650,6 +1650,13 @@ struct CpuMoeExecutor { const uint8_t* w = gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; return q6_k_dot_f32_scalar(w, x, H); // W4A16: bf16 activations, K-quant dequant } + // Anything that reaches here is assumed NVFP4 and dereferences the scale/global + // pointers, which are null for formats that do not have them (the GGUF banks pass 0). + // Falling through with an unhandled format therefore segfaults inside the worker + // thread rather than reporting anything useful, so reject it here instead. + TORCH_CHECK(fmt == WF_NVFP4 || fmt == WF_DSFP4, + "cpu_moe gemm1_dot: unhandled weight_format ", fmt, + " (handled: bf16=0, nvfp4=1, mxfp4=2, dsfp4=3, q4_0=4, q4_k=5, q6_k=6)"); const size_t r = (size_t)e * (2 * I) + row; if (use_vnni) return nvi8dot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py index bc4754e8..5936d34d 100644 --- a/python/freetoken/layers/gguf.py +++ b/python/freetoken/layers/gguf.py @@ -280,7 +280,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: flat = x.flatten() rows = self.qweight.index_select(0, flat) # [n, row_bytes] packed - y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, torch.bfloat16) + if self._quant_type in GGML_UNQUANTIZED: + # Raw value bytes, not blocks: there is no dequant kernel for the unquantized + # types (ggml_dequantize rejects type 1), so reinterpret the gathered rows. + y = rows.view(_UNQUANTIZED_DTYPE[self._quant_type]).to(torch.bfloat16) + else: + y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, torch.bfloat16) y = y.view(*x.shape, self.embedding_dim) if self._embed_scale is not None: if self._embed_scale_t is None: From bb432e8c473c557cdef6517abccf0c129bd8c544 Mon Sep 17 00:00:00 2001 From: Victor Cruz Date: Tue, 25 Aug 2026 23:12:09 -0700 Subject: [PATCH 36/36] feat(gguf): DeepSeek-V4 GGUF serves end to end Five fixes that stood between a reconciled adapter and a generated token, found by loading the real 164GB antirez/deepseek-v4-gguf Q4KExperts checkpoint. convert_deepseek4_to_gguf was never called from the model's __init__, so the fp8 Linears were never swapped and load_state_dict demanded a wq_a.scale no GGUF tensor can fill. The reconciliation test passed because it called the converter explicitly; the engine does not. Third time this hook has been missed in this package (gemma4 and qwen35moe before it). GGUFLinearNN passed its input straight to fused_mul_mat_gguf, which takes [tokens, in_features] and treats dim 0 as the batch. F.linear -- which it replaces -- accepts any number of leading dims, and deepseek_v4 relies on that: its attention passes 3-D tensors. Collapsing one silently reshaped q so the sparse-attention kernel's `b, m, h, d = q.shape` unpack failed. Leading dims are now folded and restored. GGUFEmbeddingNN dequantized unconditionally; token_embd is F16 here, which has no dequant kernel. Reinterpreted instead. The DSV4 sparse-attention kernel exceeded Turing's 64KB shared-memory block limit. BLOCK_T was tuned for sm_120's ~99KB budget, but the dominant cost is not the KV tile: q and acc are [BLOCK_H, D] in fp32, which at BLOCK_H=16 and head_dim 512 is 65536 B on its own -- an entire Turing block before a single KV byte, which is why shrinking BLOCK_T alone left the requirement stuck at 66624. BLOCK_H, BLOCK_T and num_stages are now chosen from the device's opt-in shared memory; sm_80 and above are unchanged. _TOKENIZER_ARCH mapped deepseek4 to "llama". The file says tokenizer.ggml.model = gpt2: the llama converter is sentencepiece-shaped and encodes a space as U+2581, so against a GPT2-BPE vocab every space was silently DROPPED on detokenization ("ThecapitalcityofFranceisParis"). The model was correct throughout; only the detokenizer was wrong, which presents as model damage and is not. Mapped to the qwen2 (GPT2-BPE) entry. Verified on a Quadro RTX 6000 (Turing, sm_75, 24GB) with 204 GiB of WSL RAM: 145 GiB of expert banks split 24 layers GPU-pinned / 19 OS-locked for CPU decode, CPU executor on the Q4_K kernels, ~18GB VRAM. At temperature 0: "The capital city of France is" -> " Paris. The capital city of England is London..." "The largest planet in our solar system is" -> " Jupiter. It is so big that more than 1,300 Earths" Known limitation: CUDA graph capture crashes the worker on this configuration, so this runs with --cuda-graph-max-bs 0. Decode is slower than it should be as a result. Not yet diagnosed. --- .../kernel/triton/dsv4/sparse_attn.py | 93 +++++++++++++++++-- python/freetoken/models/deepseek_v4/gguf.py | 32 ++++++- python/freetoken/models/deepseek_v4/model.py | 14 +++ python/freetoken/models/gguf/tokenizer.py | 7 +- 4 files changed, 129 insertions(+), 17 deletions(-) diff --git a/python/freetoken/kernel/triton/dsv4/sparse_attn.py b/python/freetoken/kernel/triton/dsv4/sparse_attn.py index c22f891e..316824ba 100644 --- a/python/freetoken/kernel/triton/dsv4/sparse_attn.py +++ b/python/freetoken/kernel/triton/dsv4/sparse_attn.py @@ -39,11 +39,82 @@ import triton import triton.language as tl -BLOCK_H = 16 +_BLOCK_H_LARGE, _BLOCK_H_SMALL = 16, 8 +BLOCK_H = _BLOCK_H_LARGE # The gather has exactly ONE tl.load site (the pool base is selected per column), so it stages # a single [BLOCK_T, D] KV tile -- 67968 B at BLOCK_T=32, num_stages=2, which fits the ~99KB # consumer-Blackwell (sm_120, e.g. RTX 5090) budget. (BLOCK_T=64 would need ~103KB.) -BLOCK_T = 32 +# +# 32 does NOT fit every card. Turing (sm_75) caps shared memory at 64KB per block, and the +# same launch there reports Required: 100416 -- the tile scales with the head dim, which is +# 512 on DeepSeek-V4, so the figure above is not a universal constant. Halving the KV tile +# halves the staged bytes and costs iterations, not correctness. +_BLOCK_T_LARGE, _BLOCK_T_SMALL = 32, 16 + + +def _tile_plan(device_index: int | None = None, head_dim: int = 512) -> tuple[int, int, int]: + """(BLOCK_H, BLOCK_T, num_stages) that fit this device's opt-in shared memory. + + The dominant cost is NOT the KV tile: the kernel holds q and acc as [BLOCK_H, D] in + fp32, which at BLOCK_H=16 and head_dim 512 is 16*512*4*2 = 65536 B on its own -- exactly + a Turing block's entire budget, before a single KV byte. That is why shrinking BLOCK_T + alone leaves the requirement stuck at 66624. BLOCK_H is what has to come down on a 64KB + card; halving it costs head-parallelism, not correctness. + """ + try: + props = torch.cuda.get_device_properties(device_index) + optin = int(getattr(props, "shared_memory_per_block_optin", 0)) + except Exception: + optin = 0 + + if optin >= 102400: + return _BLOCK_H_LARGE, _BLOCK_T_LARGE, 2 + + # fp32 q + acc, the fixed floor, then the staged KV tile on top + for block_h in (_BLOCK_H_LARGE, _BLOCK_H_SMALL): + for block_t, stages in ((_BLOCK_T_LARGE, 2), (_BLOCK_T_SMALL, 2), (_BLOCK_T_SMALL, 1)): + need = 2 * block_h * head_dim * 4 + stages * block_t * head_dim * 2 + if need <= optin: + return block_h, block_t, stages + return _BLOCK_H_SMALL, _BLOCK_T_SMALL, 1 + + +def _unused_block_t(device_index: int | None = None) -> int: + """(BLOCK_T, num_stages) that fit this device's opt-in shared memory. + + num_stages=2 double-buffers the KV tile, so it roughly doubles the staged bytes. On a + 64KB card the small tile alone still lands at 66624 B -- about 1KB over -- so the tight + path also drops to a single stage. That costs pipelining, not correctness, and is a + smaller loss than halving the tile again to BLOCK_T=8. + """ + block_t = _block_t(device_index) + if block_t == _BLOCK_T_LARGE: + return block_t, 2 + try: + props = torch.cuda.get_device_properties(device_index) + optin = int(getattr(props, "shared_memory_per_block_optin", 0)) + except Exception: + optin = 0 + return block_t, (2 if optin >= 98304 else 1) + + +def _block_t(device_index: int | None = None) -> int: + """KV tile width that fits this device's opt-in shared memory. + + Queried per device rather than hardcoded: the budget is 64KB on sm_75, ~99KB on sm_89 + and sm_120, and ~164KB on sm_80/sm_90. Falling back to the small tile when the budget + is unknown is the safe direction -- a tile that does not fit fails the launch outright. + """ + try: + props = torch.cuda.get_device_properties(device_index) + except Exception: + return _BLOCK_T_SMALL + optin = int(getattr(props, "shared_memory_per_block_optin", 0)) + # measured requirement at BLOCK_T=32 with head_dim 512; leave the margin triton needs + return _BLOCK_T_LARGE if optin >= 102400 else _BLOCK_T_SMALL + + +BLOCK_T = _BLOCK_T_LARGE MAX_SPLITS = 32 MIN_TILES_PER_SPLIT = 4 @@ -330,7 +401,8 @@ def sparse_attn_paged( n_splits, ) - grid = (m, b, triton.cdiv(h, BLOCK_H)) + block_h, block_t, n_stages = _tile_plan(q.device.index, d) + grid = (m, b, triton.cdiv(h, block_h)) _sparse_attn_paged_kernel[grid]( q, window_pool, cmp_pool, o, sink, idx, cnt, float(softmax_scale), @@ -342,11 +414,11 @@ def sparse_attn_paged( idx.stride(0), idx.stride(1), idx.stride(2), stride_nb, stride_nm, D=d, - BLOCK_H=BLOCK_H, - BLOCK_T=BLOCK_T, + BLOCK_H=block_h, + BLOCK_T=block_t, HAS_COUNTS=has_counts, num_warps=8, - num_stages=2, + num_stages=n_stages, ) return o @@ -355,7 +427,8 @@ def _sparse_attn_paged_splitk( q, window_pool, cmp_pool, sink, idx, cnt, o, b, m, h, d, topk, n_window, softmax_scale, has_counts, stride_nb, stride_nm, n_splits, ): - head_blocks = triton.cdiv(h, BLOCK_H) + block_h, block_t, n_stages = _tile_plan(q.device.index, d) + head_blocks = triton.cdiv(h, block_h) mid_o = torch.empty((b, m, h, n_splits, d), dtype=torch.float32, device=q.device) mid_lse = torch.empty((b, m, h, n_splits), dtype=torch.float32, device=q.device) @@ -371,12 +444,12 @@ def _sparse_attn_paged_splitk( idx.stride(0), idx.stride(1), idx.stride(2), stride_nb, stride_nm, D=d, - BLOCK_H=BLOCK_H, - BLOCK_T=BLOCK_T, + BLOCK_H=block_h, + BLOCK_T=block_t, HAS_COUNTS=has_counts, NUM_SPLITS=n_splits, num_warps=8, - num_stages=2, + num_stages=n_stages, ) _sparse_attn_splitk_merge_kernel[(m, b, h)]( mid_o, mid_lse, o, sink, diff --git a/python/freetoken/models/deepseek_v4/gguf.py b/python/freetoken/models/deepseek_v4/gguf.py index 160398db..962f3c2f 100644 --- a/python/freetoken/models/deepseek_v4/gguf.py +++ b/python/freetoken/models/deepseek_v4/gguf.py @@ -47,6 +47,9 @@ # fluent text. _GATING = {1: "softmax", 2: "sigmoid", 4: "sqrtsoftplus"} +# ggml type -> the dtype its raw bytes represent (unquantized types only). +_UNQ_DTYPE = {0: torch.float32, 1: torch.float16, 30: torch.bfloat16} + def _kv(shim: "GgufConfigShim", key: str, default: Any = None) -> Any: """One ``deepseek4.*`` metadata value. No default means the key is mandatory.""" @@ -288,8 +291,17 @@ def __init__(self, in_features: int, out_features: int, quant_type: int, def forward(self, x: torch.Tensor) -> torch.Tensor: from freetoken.layers.gguf import fused_mul_mat_gguf - out = fused_mul_mat_gguf(x, self.weight, self._quant_type) - return out if self.bias is None else out + self.bias + # fused_mul_mat_gguf takes [tokens, in_features] and treats dim 0 as the batch, so + # leading dims must be folded and restored. F.linear -- which this replaces -- + # accepts any number of leading dims, and deepseek_v4 relies on that: its attention + # passes 3-D tensors, and collapsing one silently reshapes q so the sparse-attention + # kernel's `b, m, h, d = q.shape` unpack fails. + shape = x.shape + flat = x.reshape(-1, shape[-1]) if x.dim() != 2 else x + out = fused_mul_mat_gguf(flat, self.weight, self._quant_type) + if self.bias is not None: + out = out + self.bias + return out if x.dim() == 2 else out.view(*shape[:-1], out.shape[-1]) class GGUFEmbeddingNN(torch.nn.Module): @@ -313,12 +325,22 @@ def __init__(self, num_embeddings: int, embedding_dim: int, quant_type: int): ) def forward(self, x: torch.Tensor) -> torch.Tensor: - from freetoken.kernel.gguf import ggml_dequantize + from freetoken.models.gguf.dequant import GGML_UNQUANTIZED flat = x.flatten() rows = self.weight.index_select(0, flat) - y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, - torch.bfloat16) + if self._quant_type in GGML_UNQUANTIZED: + # Unquantized types are raw value bytes in the uint8 buffer; there is no + # dequant kernel for them (ggml_dequantize rejects type 1 outright), so the + # gathered rows are reinterpreted instead. DeepSeek-V4 ships token_embd as F16. + from freetoken.models.deepseek_v4.gguf import _UNQ_DTYPE + + y = rows.view(_UNQ_DTYPE[self._quant_type]).to(torch.bfloat16) + else: + from freetoken.kernel.gguf import ggml_dequantize + + y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, + torch.bfloat16) return y.view(*x.shape, self.embedding_dim) diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index c00d20cc..4cd92d88 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -223,6 +223,20 @@ def __init__(self, config): self._transformer = Transformer(self._args) self._bound = False + # A GGUF checkpoint carries native block-quantized weights, so the dense/fp8 + # projections have to be swapped for GGUF ops before load_state_dict runs -- that + # walks named_parameters() and demands a key for every one, so an unswapped + # fp8 Linear asks for a .scale no GGUF tensor can fill. Mirrors gemma4/model.py + # and qwen3_5_moe/model.py. + from .gguf import convert_deepseek4_to_gguf, is_gguf_model + + if is_gguf_model(config): + assert config.gguf_model_path is not None, ( + "expert_quant=='gguf' but ModelConfig.gguf_model_path is unset; the " + "per-tensor ggml types can only be read from the file" + ) + convert_deepseek4_to_gguf(self, config, model_path=config.gguf_model_path) + def _ensure_bound(self) -> None: if self._bound: return diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index e5d4f4d5..609d49ed 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -23,8 +23,11 @@ "qwen35moe": "qwen2", "qwen35": "qwen2", "qwen3moe": "qwen2", - # tokenizer.ggml.model is gpt2 (BPE), pre joyai-llm, 129280 entries. - "deepseek4": "llama", + # tokenizer.ggml.model is gpt2 (BPE), pre joyai-llm, 129280 entries. The llama + # converter is sentencepiece-shaped and encodes a space as U+2581; a GPT2 BPE + # vocab uses the Ġ prefix instead, so that mapping silently DROPS every space on + # detokenization ("ThecapitalcityofFranceisParis"). qwen2 is the GPT2-BPE entry. + "deepseek4": "qwen2", } # Per-arch chat/stop tokens, in preference order: the first one present in the vocab