From eee1723c60c4d7707df48aadb0f1c041dce546d8 Mon Sep 17 00:00:00 2001 From: Poyo-Yon Date: Wed, 16 Sep 2026 22:00:32 +0900 Subject: [PATCH] fix(gemma4): load mixed-quant GGUF checkpoints --- python/freetoken/engine/engine.py | 1 + .../freetoken/kernel/csrc/gguf/gguf_kernel.cu | 19 ++ python/freetoken/kernel/csrc/gguf/moe_vec.cuh | 64 ++++--- python/freetoken/layers/gguf.py | 47 ++++- python/freetoken/layers/moe.py | 25 ++- python/freetoken/models/config.py | 3 + python/freetoken/models/gemma4/gguf.py | 144 ++++++++++++-- python/freetoken/models/gguf/config.py | 2 +- python/freetoken/models/gguf/dequant.py | 105 ++++++++++- python/freetoken/moe/expert_banks.py | 20 +- python/freetoken/moe/fused_q4_0.py | 18 +- python/freetoken/moe/offload_cache.py | 4 + tests/models/test_gemma4_gguf_quant.py | 177 ++++++++++++++++++ 13 files changed, 575 insertions(+), 54 deletions(-) create mode 100644 tests/models/test_gemma4_gguf_quant.py diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 06a74ee5f..4090cdb53 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -690,6 +690,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: hybrid_max_fetch=config.moe_hybrid_max_fetch, layout=layout, max_slots=max_slots, + gguf_quant_types=banks.gguf_quant_types, ) # before set_bank_sources: the residency validation and the copy plan's skip of non-pinned layers key on the CPU-layer set cache.cpu_layer_ids = cpu_layer_ids diff --git a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu index d88960d5f..6348b7ecd 100644 --- a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu +++ b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu @@ -568,6 +568,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 3: @@ -581,6 +582,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 6: @@ -594,6 +596,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 7: @@ -607,6 +610,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 8: @@ -620,6 +624,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 10: @@ -633,6 +638,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 11: @@ -646,6 +652,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 12: @@ -659,6 +666,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 13: @@ -672,6 +680,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 14: @@ -685,6 +694,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 16: @@ -698,6 +708,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 17: @@ -711,6 +722,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 18: @@ -724,6 +736,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 19: @@ -737,6 +750,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 20: @@ -750,6 +764,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 21: @@ -763,6 +778,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 22: @@ -776,6 +792,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 23: @@ -789,6 +806,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; case 29: @@ -802,6 +820,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + W.stride(0), stream); break; } diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh index 8cef9e080..887e553ad 100644 --- a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh +++ b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh @@ -11,7 +11,8 @@ static __global__ void moe_vec_q( const int topk, const int ncols, const int nrows, - const int token_stride) { + const int token_stride, + const int64_t expert_stride_bytes = 0) { const auto row = blockIdx.x * blockDim.y + threadIdx.y; const auto token = blockIdx.z / topk; @@ -27,7 +28,9 @@ static __global__ void moe_vec_q( // partial sum for each thread float tmp = 0.0f; - const block_q_t* x = ((const block_q_t*)vx) + expert * nrows * blocks_per_row; + const block_q_t* x = expert_stride_bytes > 0 + ? (const block_q_t*)((const char*)vx + (size_t)expert * expert_stride_bytes) + : ((const block_q_t*)vx) + expert * nrows * blocks_per_row; const block_q8_1* y = (const block_q8_1*)(((const int*)vy) + token * token_stride); for (auto i = threadIdx.x / (qi / vdr); i < blocks_per_row; i += blocks_per_warp) { @@ -62,12 +65,13 @@ static void moe_vec_q4_0_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -81,12 +85,13 @@ static void moe_vec_q4_1_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -100,12 +105,13 @@ static void moe_vec_q5_0_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -119,12 +125,13 @@ static void moe_vec_q5_1_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -138,12 +145,13 @@ static void moe_vec_q8_0_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -157,12 +165,13 @@ static void moe_vec_q2_K_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -176,12 +185,13 @@ static void moe_vec_q3_K_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -195,12 +205,13 @@ static void moe_vec_q4_K_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -214,12 +225,13 @@ static void moe_vec_q5_K_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -233,12 +245,13 @@ static void moe_vec_q6_K_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -252,12 +265,13 @@ static void moe_vec_iq2_xxs_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -271,12 +285,13 @@ static void moe_vec_iq2_xs_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -290,12 +305,13 @@ static void moe_vec_iq2_s_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -309,12 +325,13 @@ static void moe_vec_iq3_xxs_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -328,12 +345,13 @@ static void moe_vec_iq1_s_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -347,12 +365,13 @@ static void moe_vec_iq1_m_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -366,12 +385,13 @@ static void moe_vec_iq4_nl_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -385,12 +405,13 @@ static void moe_vec_iq4_xs_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } template @@ -404,10 +425,11 @@ static void moe_vec_iq3_s_q8_1_cuda( const int ncols, const int nrows, const int token_stride, + const int64_t expert_stride_bytes, 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); + <<>>(vx, vy, dst, topk_ids, top_k, ncols, nrows, token_stride, expert_stride_bytes); } diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py index ac49b1a5b..81adae718 100644 --- a/python/freetoken/layers/gguf.py +++ b/python/freetoken/layers/gguf.py @@ -20,8 +20,24 @@ GGML_BF16, GGML_F16, GGML_F32, + GGML_IQ1_M, + GGML_IQ1_S, + GGML_IQ2_S, + GGML_IQ2_XS, + GGML_IQ2_XXS, + GGML_IQ3_S, + GGML_IQ3_XXS, + GGML_IQ4_NL, + GGML_IQ4_XS, GGML_NAME, + GGML_Q2_K, + GGML_Q3_K, GGML_Q4_0, + GGML_Q4_1, + GGML_Q4_K, + GGML_Q5_0, + GGML_Q5_1, + GGML_Q5_K, GGML_Q6_K, GGML_Q8_0, row_bytes, @@ -32,9 +48,34 @@ # 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} +_STANDARD_AND_K = { + 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, +} +_IQ = { + 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, +} +_MMVQ = _STANDARD_AND_K | _IQ +# The vendored MMQ switch covers standard and K-quants only. IQ types use +# MMVQ for small batches and the dequant fallback for larger batches. +_MMQ = _STANDARD_AND_K +_DEQUANT = _STANDARD_AND_K | _IQ # Below this token count, the MMVQ GEMV kernel wins (matches vLLM's heuristic). _MMVQ_SAFE = 6 diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index c20945499..f93bc054a 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -423,16 +423,33 @@ def _expert_gemm( hidden_states, topk_weights, topk_ids, view, layer=self, is_prefill=is_prefill ) fmt = cache.quant_format - if fmt == "q4_0": - # Native GGUF Q4_0 experts: dequant-in-kernel grouped GEMV (MMVQ) over the + if fmt in ("q4_0", "gguf"): + # Native GGUF experts: dequant-in-kernel grouped GEMV (MMVQ) over the # streamed packed banks; topk_ids already index the cache slots / layer. + from freetoken.models.gguf.dequant import GGML_Q4_0 from freetoken.moe.fused_q4_0 import fused_experts_gguf_q4_0 gate_up, down = views + gate_up_type, down_type = ( + cache.gguf_quant_types[self.layer_id] + if cache.gguf_quant_types is not None + else (GGML_Q4_0, GGML_Q4_0) + ) return fused_experts_gguf_q4_0( - hidden_states, gate_up, down, topk_weights, topk_ids, self.activation + hidden_states, + gate_up, + down, + topk_weights, + topk_ids, + self.activation, + gate_up_quant_type=gate_up_type, + down_quant_type=down_type, + intermediate_size=self.intermediate_size, + hidden_size=self.hidden_size, ) - raise AssertionError(f"offload experts without a quant method only serve q4_0 banks, got {fmt!r}") + raise AssertionError( + f"offload experts without a quant method only serve native GGUF banks, got {fmt!r}" + ) def make_moe_layer( diff --git a/python/freetoken/models/config.py b/python/freetoken/models/config.py index d7b504a76..3c73c81da 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -314,6 +314,9 @@ class ModelConfig: has_attn_bias: bool = False has_router_bias: bool = False moe_weight_format: str | None = None + # Native GGUF tensor types discovered from the tensor table. Model-specific + # adapters may use this to construct packed layers without dequantizing them. + gguf_quant_types: Any | None = None swiglu_limit: float | None = None hidden_act_alpha: float = 1.702 # Full DeepseekV4Args payload for the DSV4-specific machinery (MLA sparse attention, diff --git a/python/freetoken/models/gemma4/gguf.py b/python/freetoken/models/gemma4/gguf.py index 437822b51..5831a3f6d 100644 --- a/python/freetoken/models/gemma4/gguf.py +++ b/python/freetoken/models/gemma4/gguf.py @@ -6,8 +6,8 @@ ``gemma4.config.parse_config`` -- only the source is GGUF KV metadata instead of a HF config object. transformers' own GGUF->config conversion is rejected by the gemma4 strict dataclass (per-layer ``num_key_value_heads`` array), so we read the -metadata directly. ``expert_quant`` is set to ``"q4_0"`` to route the routed experts -through the native-Q4_0 offload-cache path. +metadata directly. ``expert_quant`` keeps the legacy ``"q4_0"`` provider tag, while +the tensor table records the actual dense and per-layer routed-expert GGML types. """ from __future__ import annotations @@ -48,6 +48,72 @@ def _full_rotary_dim(shim: "GgufConfigShim", full_head_dim: int) -> int: return full_head_dim // 4 +def _gguf_quant_layout(model_path: str, num_layers: int, num_experts: int) -> dict | None: + """Discover native tensor types without touching the tensor payloads. + + A metadata-only FTW source has no tensor table, in which case ``None`` keeps + the historical Q4_0/Q6_K defaults recorded by the converted checkpoint. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + + fields = { + "qkv": ("attn_q.weight", "attn_k.weight", "attn_v.weight"), + "attn_output": ("attn_output.weight",), + "shared_gate_up": ("ffn_gate.weight", "ffn_up.weight"), + "shared_down": ("ffn_down.weight",), + } + per_layer = {name: [set() for _ in range(num_layers)] for name in fields} + expert_gate_up = [None] * num_layers + expert_down = [None] * num_layers + expert_gate_up_bytes = [None] * num_layers + expert_down_bytes = [None] * num_layers + embedding = None + saw_tensor = False + + for tensor in iter_gguf_tensors(model_path): + saw_tensor = True + if tensor.name == "token_embd.weight": + embedding = tensor.ggml_type + continue + if not tensor.name.startswith("blk."): + continue + layer = int(tensor.name.split(".")[1]) + suffix = tensor.name.split(".", 2)[2] + if suffix == "ffn_gate_up_exps.weight": + expert_gate_up[layer] = tensor.ggml_type + expert_gate_up_bytes[layer] = tensor.packed().numel() + elif suffix == "ffn_down_exps.weight": + expert_down[layer] = tensor.ggml_type + expert_down_bytes[layer] = tensor.packed().numel() + else: + for role, suffixes in fields.items(): + if suffix in suffixes: + per_layer[role][layer].add(tensor.ggml_type) + break + + if not saw_tensor: + return None + if embedding is None: + raise ValueError("gemma4 GGUF has no token_embd.weight") + + layout: dict[str, object] = {"embedding": embedding} + for role, values in per_layer.items(): + bad = [i for i, types in enumerate(values) if len(types) != 1] + if bad: + raise ValueError(f"gemma4 GGUF {role} tensors disagree or are missing in layers {bad}") + layout[role] = tuple(next(iter(types)) for types in values) + for role, values in (("expert_gate_up", expert_gate_up), ("expert_down", expert_down)): + missing = [i for i, value in enumerate(values) if value is None] + if missing: + raise ValueError(f"gemma4 GGUF {role} tensors are missing in layers {missing}") + layout[role] = tuple(values) + # Total packed bytes per expert. The loader stores each expert contiguously in + # a fixed-width cache slot and passes the slot stride to the CUDA kernel. + layout["expert_gate_up_bytes"] = tuple(int(n) // num_experts for n in expert_gate_up_bytes) + layout["expert_down_bytes"] = tuple(int(n) // num_experts for n in expert_down_bytes) + return layout + + def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: m = shim.metadata @@ -59,6 +125,7 @@ def g(key: str): num_layers = int(g("block_count")) hidden = int(g("embedding_length")) + num_experts = int(g("expert_count")) num_qo_heads = int(g("attention.head_count")) kv_per_layer = g("attention.head_count_kv") # per-layer list # True -> sliding-window (SWA) layer, False -> full attention. @@ -105,7 +172,7 @@ def g(key: str): rms_norm_eps=float(g("attention.layer_norm_rms_epsilon")), tie_word_embeddings=bool(shim.tie_word_embeddings), rotary_config=full_rotary, - num_experts=int(g("expert_count")), + num_experts=num_experts, num_experts_per_tok=int(g("expert_used_count")), moe_intermediate_size=int(g("expert_feed_forward_length")), norm_topk_prob=True, @@ -114,6 +181,7 @@ def g(key: str): moe_enabled=True, expert_quant="q4_0", moe_weight_format="q4_0", + gguf_quant_types=_gguf_quant_layout(shim.model_path, num_layers, num_experts), use_qk_norm=True, attn_sm_scale=1.0, final_logit_softcapping=float(g("final_logit_softcapping")), @@ -341,7 +409,7 @@ def convert_gemma4_to_gguf(model, config: ModelConfig) -> None: """In place: replace gemma4's dense projections + embedding with native GGUF ops. Quantized in the checkpoint -> swapped: attention qkv/o, shared-MLP gate_up/down - (all Q4_0) and the token embedding (Q6_K, also the tied LM head). Left as dense + and the token embedding (also the tied LM head). Left as dense bf16 (F32 in the GGUF): the router gate, all RMSNorms, the per-layer scalars, and the routed experts (served from the offload cache). """ @@ -356,23 +424,33 @@ def swap_linear(owner, attr, quant_type=GGML_Q4_0): GGUFLinear(in_features, out_features, quant_type, has_bias=lin.bias is not None), ) + layout = config.gguf_quant_types + embedding_type = layout["embedding"] if layout is not None else GGML_Q6_K inner = model.model embed = GGUFEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, - quant_type=GGML_Q6_K, + quant_type=embedding_type, embed_scale=config.embedding_scale, ) inner.embed_tokens = embed - for layer in inner.layers.op_list: - swap_linear(layer.self_attn, "qkv_proj") - swap_linear(layer.self_attn, "o_proj") - swap_linear(layer.feed_forward.shared_mlp, "gate_up_proj") - swap_linear(layer.feed_forward.shared_mlp, "down_proj") + for layer_id, layer in enumerate(inner.layers.op_list): + swap_linear(layer.self_attn, "qkv_proj", layout["qkv"][layer_id] if layout else GGML_Q4_0) + swap_linear(layer.self_attn, "o_proj", layout["attn_output"][layer_id] if layout else GGML_Q4_0) + swap_linear( + layer.feed_forward.shared_mlp, + "gate_up_proj", + layout["shared_gate_up"][layer_id] if layout else GGML_Q4_0, + ) + swap_linear( + layer.feed_forward.shared_mlp, + "down_proj", + layout["shared_down"][layer_id] if layout else GGML_Q4_0, + ) if config.tie_word_embeddings: - model.lm_head = GGUFTiedLMHead(embed, GGML_Q6_K) + model.lm_head = GGUFTiedLMHead(embed, embedding_type) # -------------------------------------------------------------------------------------- @@ -388,6 +466,28 @@ def _q4_0_expert_specs(config: ModelConfig) -> dict[str, tuple[tuple[int, ...], } +def _uses_mixed_gguf_experts(config: ModelConfig) -> bool: + layout = getattr(config, "gguf_quant_types", None) + if layout is None: + return False + return any( + quant_type != GGML_Q4_0 + for role in ("expert_gate_up", "expert_down") + for quant_type in layout[role] + ) + + +def _gguf_expert_specs(config: ModelConfig) -> dict[str, tuple[tuple[int, ...], torch.dtype]]: + if not _uses_mixed_gguf_experts(config): + return _q4_0_expert_specs(config) + layout = config.gguf_quant_types + E = config.num_experts + return { + "gate_up": ((E, max(layout["expert_gate_up_bytes"])), torch.uint8), + "down": ((E, max(layout["expert_down_bytes"])), torch.uint8), + } + + def load_q4_0_expert_sources( model_path: str, config: ModelConfig, *, layer_sink=None ) -> dict[str, list[torch.Tensor]]: @@ -413,7 +513,13 @@ def load_q4_0_expert_sources( L, E = config.num_layers, config.num_experts H, I = config.hidden_size, config.moe_intermediate_size h_bytes, i_bytes = row_bytes(H, GGML_Q4_0), row_bytes(I, GGML_Q4_0) - hb = alloc_layer_banks(_q4_0_expert_specs(config), L) # lazy anon mmaps (unpinned) + mixed = _uses_mixed_gguf_experts(config) + if mixed and layer_sink is not None: + raise NotImplementedError( + "converting mixed-type GGUF expert banks to FTW is not supported yet; " + "serve the source GGUF directly" + ) + hb = alloc_layer_banks(_gguf_expert_specs(config), L) # lazy anon mmaps (unpinned) banks = {name: [b.tensor for b in hb[name]] for name in hb} seen_gu, seen_dn = set(), set() @@ -424,10 +530,18 @@ def _load(sink) -> None: continue layer = int(t.name.split(".")[1]) if t.name.endswith("ffn_gate_up_exps.weight"): - banks["gate_up"][layer].copy_(t.packed().reshape(E, 2 * I, h_bytes)) + packed = t.packed().reshape(E, -1) + if mixed: + banks["gate_up"][layer][:, : packed.shape[1]].copy_(packed) + else: + banks["gate_up"][layer].copy_(packed.reshape(E, 2 * I, h_bytes)) seen_gu.add(layer) elif t.name.endswith("ffn_down_exps.weight"): - banks["down"][layer].copy_(t.packed().reshape(E, H, i_bytes)) + packed = t.packed().reshape(E, -1) + if mixed: + banks["down"][layer][:, : packed.shape[1]].copy_(packed) + else: + banks["down"][layer].copy_(packed.reshape(E, H, i_bytes)) seen_dn.add(layer) else: continue @@ -455,7 +569,7 @@ def dummy_q4_0_expert_sources(config: ModelConfig) -> dict[str, list[torch.Tenso from freetoken.moe.host_banks import alloc_layer_banks, pin_banks L = config.num_layers - hb = alloc_layer_banks(_q4_0_expert_specs(config), L) + hb = alloc_layer_banks(_gguf_expert_specs(config), 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) diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 63b1a18b9..19040ee29 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -21,7 +21,7 @@ } -@dataclass(frozen=True) +@dataclass class GgufConfigShim: architectures: list[str] model_path: str diff --git a/python/freetoken/models/gguf/dequant.py b/python/freetoken/models/gguf/dequant.py index 77c3ea010..818f8d7d5 100644 --- a/python/freetoken/models/gguf/dequant.py +++ b/python/freetoken/models/gguf/dequant.py @@ -1,5 +1,4 @@ -"""GGML block-quant dequantization in pure torch (the formats this repo's GGUF -checkpoints use: Q4_0, Q6_K, plus trivial F32/F16/BF16). +"""GGML block-quant metadata and reference dequantization for native GGUF weights. 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 @@ -8,22 +7,41 @@ (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 +Each dequantizer 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``. +hot Q4_0/Q6_K references mirror ``ggml-quants.c`` directly; other block types +delegate their reference path to gguf-py. """ from __future__ import annotations +from functools import partial + import torch # ggml_type enum values (subset present in these checkpoints). 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. @@ -32,8 +50,24 @@ 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_NAME = { @@ -41,8 +75,24 @@ 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", } @@ -115,10 +165,41 @@ def dequant_q6_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: return y.reshape(-1).to(out_dtype) +def _dequant_gguf_py(raw: torch.Tensor, out_dtype: torch.dtype, ggml_type: int) -> torch.Tensor: + """Reference dequant for additional block types, delegated to gguf-py.""" + import gguf + import numpy as np + + out = gguf.quants.dequantize( + raw.detach().cpu().contiguous().numpy(), gguf.GGMLQuantizationType(ggml_type) + ) + return torch.from_numpy(np.asarray(out)).to(raw.device, out_dtype).reshape(-1) + + _DEQUANT = { GGML_Q4_0: dequant_q4_0, GGML_Q6_K: dequant_q6_k, } +for _ggml_type in ( + 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_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, +): + _DEQUANT[_ggml_type] = partial(_dequant_gguf_py, ggml_type=_ggml_type) def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> torch.Tensor: @@ -142,8 +223,24 @@ def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> tor "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_NAME", "BLOCK_SHAPE", "row_bytes", diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index e3a768b6d..e0ae7b7c4 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -48,6 +48,8 @@ class ExpertBanks: kind: QuantKind | None = None kernel: str | None = None layout: dict | None = None + # Per layer (gate_up ggml type, down ggml type) for native mixed GGUF banks. + gguf_quant_types: tuple[tuple[int, int], ...] | None = None def _dummy_fill(role: str, tensor: torch.Tensor) -> None: @@ -159,8 +161,21 @@ def _q4_0_banks(model_path, model_config, device, dtype, dummy, parallel=False, "(not safetensors), so the common reader doesn't apply -- it needs a GGUF-native " "parallel reader (parse the tensor table, chunked O_DIRECT over the one file)" ) + from freetoken.models.gguf.dequant import GGML_Q4_0 from freetoken.models.weight import load_q4_0_moe_expert_sources + layout = getattr(model_config, "gguf_quant_types", None) + quant_types = None + quant_format = "q4_0" + if layout is not None: + quant_types = tuple(zip(layout["expert_gate_up"], layout["expert_down"])) + if any(pair != (GGML_Q4_0, GGML_Q4_0) for pair in quant_types): + quant_format = "gguf" + if decode_target != "gpu": + raise NotImplementedError( + "mixed-type GGUF experts currently support GPU offload only" + ) + # Native GGUF Q4_0 routed experts: packed block bytes streamed to the GPU and # dequantized inside the borrowed ggml MoE kernels (no bf16 expert copy). Banks are # per-layer HostBanks (pin-after-fill), so conversion streams each completed layer's @@ -168,7 +183,10 @@ def _q4_0_banks(model_path, model_config, device, dtype, dummy, parallel=False, sink = None if dummy else layer_sink sources = load_q4_0_moe_expert_sources(model_path, model_config, dummy=dummy, layer_sink=sink) return ExpertBanks( - "q4_0", {name: sources[name] for name in _BANK_SCHEMAS["q4_0"]}, streamed=sink is not None + quant_format, + {name: sources[name] for name in _BANK_SCHEMAS[quant_format]}, + streamed=sink is not None, + gguf_quant_types=quant_types, ) diff --git a/python/freetoken/moe/fused_q4_0.py b/python/freetoken/moe/fused_q4_0.py index cdab82bf9..48e84ea65 100644 --- a/python/freetoken/moe/fused_q4_0.py +++ b/python/freetoken/moe/fused_q4_0.py @@ -26,6 +26,11 @@ def fused_experts_gguf_q4_0( topk_weights: torch.Tensor, topk_ids: torch.Tensor, activation: str, + *, + gate_up_quant_type: int = GGML_Q4_0, + down_quant_type: int = GGML_Q4_0, + intermediate_size: int | None = None, + hidden_size: int | None = None, ) -> torch.Tensor: from freetoken.kernel.gguf import ggml_moe_a8_vec @@ -34,16 +39,19 @@ def fused_experts_gguf_q4_0( raise ValueError(f"unsupported MoE activation {activation!r}") num_tokens = hidden_states.shape[0] - n2 = gate_up_q.shape[1] # 2 * intermediate - h = down_q.shape[1] # hidden + n2 = 2 * intermediate_size if intermediate_size is not None else gate_up_q.shape[1] + h = hidden_size if hidden_size is not None else down_q.shape[1] top_k = topk_ids.shape[1] - qt = int(GGML_Q4_0) # 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) + gate_up = ggml_moe_a8_vec( + hidden_states, gate_up_q, topk_ids, top_k, int(gate_up_quant_type), 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 7abb9c800..0119222c6 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -45,6 +45,9 @@ # 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 whose tensor type may differ by bank or layer. Slots + # are fixed-width byte rows; gguf_quant_types selects the CUDA decoder. + "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) @@ -144,6 +147,7 @@ class OffloadMoeCache: # bank layout from the expert kernel (a BankSpec per role); when given it replaces the _BANK_SCHEMAS lookup and the slot cap comes from max_slots layout: dict | None = None max_slots: int | None = None + gguf_quant_types: tuple[tuple[int, int], ...] | None = None def __post_init__(self) -> None: policy_ids = {"lru": 0} diff --git a/tests/models/test_gemma4_gguf_quant.py b/tests/models/test_gemma4_gguf_quant.py new file mode 100644 index 000000000..67cf24c3d --- /dev/null +++ b/tests/models/test_gemma4_gguf_quant.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +import copy +from types import SimpleNamespace + +import torch + + +def test_cuda_quant_types_have_packed_row_layouts(): + from freetoken.layers.gguf import _DEQUANT, _MMQ, _MMVQ + from freetoken.models.gguf.dequant import BLOCK_SHAPE + + expected = { + 2: (32, 18), + 3: (32, 20), + 6: (32, 22), + 7: (32, 24), + 8: (32, 34), + 10: (256, 84), + 11: (256, 110), + 12: (256, 144), + 13: (256, 176), + 14: (256, 210), + 16: (256, 66), + 17: (256, 74), + 18: (256, 98), + 19: (256, 50), + 20: (32, 18), + 21: (256, 110), + 22: (256, 82), + 23: (256, 136), + 29: (256, 56), + } + + assert {quant_type: BLOCK_SHAPE[quant_type] for quant_type in expected} == expected + assert _MMVQ == set(expected) + assert _DEQUANT == set(expected) + assert _MMQ == {2, 3, 6, 7, 8, 10, 11, 12, 13, 14} + + +def test_gguf_config_shim_can_be_masked_like_hf_config(): + from freetoken.models.gguf.config import GgufConfigShim + + shim = GgufConfigShim( + architectures=["Gemma4GGUFForCausalLM"], + model_path="model.gguf", + model_type="gemma4", + metadata={}, + vocab_size=4, + tie_word_embeddings=True, + ) + masked = copy.copy(shim) + masked.audio_config = None + assert masked.audio_config is None + + +class _Tensor: + def __init__(self, name: str, ggml_type: int, packed: torch.Tensor): + self.name = name + self.ggml_type = ggml_type + self._packed = packed + + def packed(self) -> torch.Tensor: + return self._packed + + +def test_quant_layout_detects_dense_and_per_layer_expert_types(monkeypatch): + from freetoken.models.gemma4 import gguf + from freetoken.models.gguf import reader + + tensors = [ + _Tensor("token_embd.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.0.attn_q.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.0.attn_k.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.0.attn_v.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.0.attn_output.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.0.ffn_gate.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.0.ffn_up.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.0.ffn_down.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.0.ffn_gate_up_exps.weight", 21, torch.empty(4, 11, dtype=torch.uint8)), + _Tensor("blk.0.ffn_down_exps.weight", 20, torch.empty(4, 13, dtype=torch.uint8)), + _Tensor("blk.1.attn_q.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.1.attn_k.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.1.attn_output.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.1.ffn_gate.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.1.ffn_up.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.1.ffn_down.weight", 8, torch.empty(1, 34, dtype=torch.uint8)), + _Tensor("blk.1.ffn_gate_up_exps.weight", 23, torch.empty(4, 17, dtype=torch.uint8)), + _Tensor("blk.1.ffn_down_exps.weight", 8, torch.empty(4, 19, dtype=torch.uint8)), + ] + monkeypatch.setattr(reader, "iter_gguf_tensors", lambda _path: iter(tensors)) + + layout = gguf._gguf_quant_layout("model.gguf", 2, 4) + + assert layout["embedding"] == 8 + assert layout["qkv"] == (8, 8) + assert layout["shared_gate_up"] == (8, 8) + assert layout["expert_gate_up"] == (21, 23) + assert layout["expert_down"] == (20, 8) + assert layout["expert_gate_up_bytes"] == (11, 17) + assert layout["expert_down_bytes"] == (13, 19) + + +def test_mixed_expert_loader_pads_slots_but_preserves_packed_bytes(monkeypatch): + from freetoken.models.gemma4 import gguf + from freetoken.models.gguf import reader + + config = SimpleNamespace( + num_layers=2, + num_experts=2, + hidden_size=32, + moe_intermediate_size=32, + gguf_quant_types={ + "expert_gate_up": (21, 23), + "expert_down": (20, 8), + "expert_gate_up_bytes": (7, 11), + "expert_down_bytes": (5, 13), + }, + ) + gu0 = torch.arange(14, dtype=torch.uint8).reshape(2, 7) + dn0 = torch.arange(10, dtype=torch.uint8).reshape(2, 5) + gu1 = torch.arange(22, dtype=torch.uint8).reshape(2, 11) + dn1 = torch.arange(26, dtype=torch.uint8).reshape(2, 13) + tensors = [ + _Tensor("blk.0.ffn_gate_up_exps.weight", 21, gu0), + _Tensor("blk.0.ffn_down_exps.weight", 20, dn0), + _Tensor("blk.1.ffn_gate_up_exps.weight", 23, gu1), + _Tensor("blk.1.ffn_down_exps.weight", 8, dn1), + ] + monkeypatch.setattr(reader, "iter_gguf_tensors", lambda _path: iter(tensors)) + monkeypatch.setattr(gguf, "_require_tp1", lambda _what: None) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + banks = gguf.load_q4_0_expert_sources("model.gguf", config) + + assert [tuple(t.shape) for t in banks["gate_up"]] == [(2, 11), (2, 11)] + assert [tuple(t.shape) for t in banks["down"]] == [(2, 13), (2, 13)] + torch.testing.assert_close(banks["gate_up"][0][:, :7], gu0) + torch.testing.assert_close(banks["down"][0][:, :5], dn0) + torch.testing.assert_close(banks["gate_up"][1], gu1) + torch.testing.assert_close(banks["down"][1], dn1) + + +def test_gguf_expert_gemm_passes_independent_quant_types(monkeypatch): + from freetoken.kernel import gguf as kernel + from freetoken.moe import fused_q4_0 + from freetoken.moe.fused_q4_0 import fused_experts_gguf_q4_0 + + calls = [] + + def fake_moe(x, weight, ids, top_k, quant_type, row, tokens): + calls.append((quant_type, tuple(weight.shape), top_k, row, tokens)) + return torch.ones(tokens * top_k, row, dtype=x.dtype) + + monkeypatch.setattr(kernel, "ggml_moe_a8_vec", fake_moe) + monkeypatch.setitem(fused_q4_0._ACT, "silu", lambda x: x[:, : x.shape[1] // 2]) + hidden = torch.ones(1, 4) + gate_up = torch.empty(2, 32, dtype=torch.uint8) + down = torch.empty(2, 48, dtype=torch.uint8) + weights = torch.ones(1, 1) + ids = torch.zeros(1, 1, dtype=torch.int32) + + out = fused_experts_gguf_q4_0( + hidden, + gate_up, + down, + weights, + ids, + "silu", + gate_up_quant_type=21, + down_quant_type=20, + intermediate_size=3, + hidden_size=4, + ) + + assert out.shape == (1, 4) + assert [call[0] for call in calls] == [21, 20]