diff --git a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu index d88960d5..ecb1f325 100644 --- a/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu +++ b/python/freetoken/kernel/csrc/gguf/gguf_kernel.cu @@ -545,7 +545,8 @@ torch::Tensor ggml_moe_a8_vec( int64_t top_k, int64_t type, int64_t row, - int64_t tokens) { + int64_t tokens, + int64_t expert_stride_bytes) { int col = X.sizes()[1]; const int padded = (col + 512 - 1) / 512 * 512; const at::cuda::OptionalCUDAGuard device_guard(device_of(X)); @@ -568,6 +569,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 3: @@ -581,6 +583,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 6: @@ -594,6 +597,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 7: @@ -607,6 +611,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 8: @@ -620,6 +625,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 10: @@ -633,6 +639,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 11: @@ -646,6 +653,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 12: @@ -659,6 +667,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 13: @@ -672,6 +681,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 14: @@ -685,6 +695,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 16: @@ -698,6 +709,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 17: @@ -711,6 +723,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 18: @@ -724,6 +737,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 19: @@ -737,6 +751,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 20: @@ -750,6 +765,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 21: @@ -763,6 +779,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 22: @@ -776,6 +793,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 23: @@ -789,6 +807,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; case 29: @@ -802,6 +821,7 @@ torch::Tensor ggml_moe_a8_vec( col, row, quant_X.stride(0), + expert_stride_bytes, stream); break; } diff --git a/python/freetoken/kernel/csrc/gguf/moe_vec.cuh b/python/freetoken/kernel/csrc/gguf/moe_vec.cuh index 8cef9e08..e3cae540 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) { const auto row = blockIdx.x * blockDim.y + threadIdx.y; const auto token = blockIdx.z / topk; @@ -27,7 +28,11 @@ 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; + // expert_stride_bytes == 0: dense contiguous banks (original layout). > 0: each + // expert starts at a fixed byte stride (padded banks for mixed-quant models). + 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 +67,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 +87,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 +107,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 +127,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 +147,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 +167,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 +187,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 +207,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 +227,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 +247,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 +267,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 +287,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 +307,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 +327,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 +347,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 +367,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 +387,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 +407,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 +427,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/kernel/gguf.py b/python/freetoken/kernel/gguf.py index 04a16560..40a77ff2 100644 --- a/python/freetoken/kernel/gguf.py +++ b/python/freetoken/kernel/gguf.py @@ -123,9 +123,17 @@ def ggml_moe_a8_vec( quant_type: int, row: int, tokens: int, + expert_stride_bytes: int = 0, ) -> torch.Tensor: - """MMVQ grouped expert GEMV over stacked experts ``weight[E, row, *]``.""" - return _module().ggml_moe_a8_vec(x, weight, topk_ids, top_k, quant_type, row, tokens) + """MMVQ grouped expert GEMV over stacked experts ``weight[E, row, *]``. + + ``expert_stride_bytes`` == 0 assumes dense contiguous banks; > 0 reads each + expert at that fixed byte offset (padded flat banks for mixed-quant models, + where a layer's real payload occupies the leading bytes of each expert slot). + """ + return _module().ggml_moe_a8_vec( + x, weight, topk_ids, top_k, quant_type, row, tokens, expert_stride_bytes + ) def ggml_moe_get_block_size(quant_type: int) -> int: diff --git a/python/freetoken/layers/base.py b/python/freetoken/layers/base.py index b1e939d9..f7dab709 100644 --- a/python/freetoken/layers/base.py +++ b/python/freetoken/layers/base.py @@ -42,7 +42,9 @@ def load_state_dict( if isinstance(param, torch.Tensor): item = state_dict.pop(_concat_prefix(prefix, name)) assert isinstance(item, torch.Tensor) - assert param.shape == item.shape and param.dtype == item.dtype + assert param.shape == item.shape and param.dtype == item.dtype, ( + f"{_concat_prefix(prefix, name)}: model {tuple(param.shape)}/{param.dtype} vs ckpt {tuple(item.shape)}/{item.dtype}" + ) setattr(self, name, item) elif isinstance(param, BaseOP): param.load_state_dict( diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py index ac49b1a5..c6ad4d34 100644 --- a/python/freetoken/layers/gguf.py +++ b/python/freetoken/layers/gguf.py @@ -21,7 +21,15 @@ GGML_F16, GGML_F32, GGML_NAME, + GGML_IQ1_S, + GGML_IQ2_S, + GGML_IQ2_XXS, + GGML_IQ3_XXS, + GGML_IQ4_XS, + GGML_Q3_K, GGML_Q4_0, + GGML_Q4_K, + GGML_Q5_K, GGML_Q6_K, GGML_Q8_0, row_bytes, @@ -32,9 +40,17 @@ # 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} +_MMVQ = { + GGML_Q4_0, GGML_Q8_0, GGML_Q3_K, GGML_Q4_K, GGML_Q5_K, GGML_Q6_K, + GGML_IQ1_S, GGML_IQ2_S, GGML_IQ2_XXS, GGML_IQ3_XXS, GGML_IQ4_XS, +} +# The vendored CUDA MMQ switch covers the standard + K-quants only (no IQ cases); +# IQ types take the dequant fallback for large batches. +_MMQ = {GGML_Q4_0, GGML_Q8_0, GGML_Q3_K, GGML_Q4_K, GGML_Q5_K, GGML_Q6_K} +_DEQUANT = { + GGML_Q4_0, GGML_Q8_0, GGML_Q3_K, GGML_Q4_K, GGML_Q5_K, GGML_Q6_K, + GGML_IQ1_S, GGML_IQ2_S, GGML_IQ2_XXS, GGML_IQ3_XXS, GGML_IQ4_XS, +} # 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 d68d8ded..3b75985d 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -531,6 +531,20 @@ def _expert_gemm( return fused_experts_gguf_q4_0( hidden_states, gate_up, down, topk_weights, topk_ids, self.activation ) + if fmt == "gguf": + # Mixed-type GGUF experts (per-layer quant types, flat padded slot banks): + # same MMVQ path, geometry from the layer's own type attributes (set via + # make_moe_layer extra_attrs by the model, e.g. laguna). + from freetoken.moe.fused_gguf import fused_experts_gguf + + gate_up, down = views + return fused_experts_gguf( + hidden_states, gate_up, down, topk_weights, topk_ids, self.activation, + gate_up_type=self.gguf_gate_up_type, + down_type=self.gguf_down_type, + gate_up_rows=self.gguf_gate_up_rows, + down_rows=self.gguf_down_rows, + ) 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..8d452e4f 100644 --- a/python/freetoken/models/config.py +++ b/python/freetoken/models/config.py @@ -1,6 +1,6 @@ from __future__ import annotations import os -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any, ClassVar, Dict, List, Literal, Tuple, TypeAlias from freetoken.attention.base import AttnType @@ -199,6 +199,7 @@ def _full_group_attn_type(group: FullAttentionGroupConfig) -> AttnType: class ModelConfig: num_layers: int num_qo_heads: int + num_qo_heads_per_layer: tuple[int, ...] | None = field(default=None, kw_only=True) # hybrid models (laguna) vary per layer; None means uniform num_qo_heads. num_kv_heads: int head_dim: int hidden_size: int @@ -273,6 +274,14 @@ class ModelConfig: has_attn_bias: bool = False has_router_bias: bool = False moe_weight_format: str | None = None + # GGUF checkpoints only: ggml quant type of ``token_embd.weight`` (publisher-dependent + # -- Q6_K in Google's QAT release, Q4_0 in Unsloth's). None for non-GGUF checkpoints. + gguf_embed_quant: int | None = None + # Mixed-type GGUF (laguna): (gate_up_type, down_type) ggml ids per MoE layer, + # read from the file's tensor table. None for uniform-quant checkpoints. + gguf_expert_types: tuple[tuple[int, int], ...] | None = None + # source .gguf path; laguna reads per-tensor quant types from it at conversion + gguf_model_path: str | None = None swiglu_limit: float | None = None hidden_act_alpha: float = 1.702 # Full DeepseekV4Args payload for the DSV4-specific machinery (MLA sparse attention, @@ -372,6 +381,12 @@ def is_linear_layer(self, layer_id: int) -> bool: LinearGatedDeltaGroupConfig, ) + def qo_heads(self, layer_id: int) -> int: + """Query-head count for one layer (per-layer override or the uniform count).""" + if self.num_qo_heads_per_layer is not None: + return self.num_qo_heads_per_layer[layer_id] + return self.num_qo_heads + def attn_type_for_layer(self, layer_id: int) -> AttnType: """Canonical per-layer attention-type lookup (the taxonomy is declared top-down on the attention groups; this is the layer-granular view).""" diff --git a/python/freetoken/models/gemma4/gguf.py b/python/freetoken/models/gemma4/gguf.py index 437822b5..7047817d 100644 --- a/python/freetoken/models/gemma4/gguf.py +++ b/python/freetoken/models/gemma4/gguf.py @@ -48,6 +48,18 @@ def _full_rotary_dim(shim: "GgufConfigShim", full_head_dim: int) -> int: return full_head_dim // 4 +def _embed_quant(shim: "GgufConfigShim") -> int: + """ggml quant type of the token embedding table, read off the file. + + Publishers differ (Google's QAT GGUF stores it Q6_K, Unsloth's Q4_0). A + metadata-only GGUF (an FTW dir's source_metadata.gguf) has no tensor table; fall + back to Q6_K, the type llama.cpp's own gemma4 conversion emits. + """ + from freetoken.models.gguf.reader import gguf_tensor_type + + return gguf_tensor_type(shim.model_path, "token_embd.weight") or GGML_Q6_K + + def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: m = shim.metadata @@ -114,6 +126,7 @@ def g(key: str): moe_enabled=True, expert_quant="q4_0", moe_weight_format="q4_0", + gguf_embed_quant=_embed_quant(shim), use_qk_norm=True, attn_sm_scale=1.0, final_logit_softcapping=float(g("final_logit_softcapping")), @@ -230,7 +243,7 @@ def layer_of(name: str) -> int: for t in iter_gguf_tensors(model_path): name = t.name if name == "token_embd.weight": - yield "model.embed_tokens.qweight", t.packed() # Q6_K packed table + yield "model.embed_tokens.qweight", t.packed() # packed table, native quant continue if name == "output_norm.weight": yield "model.norm.weight", _to_bf16(t) @@ -346,6 +359,7 @@ def convert_gemma4_to_gguf(model, config: ModelConfig) -> None: the routed experts (served from the offload cache). """ from freetoken.layers.gguf import GGUFEmbedding, GGUFLinear + embed_quant = config.gguf_embed_quant or GGML_Q6_K def swap_linear(owner, attr, quant_type=GGML_Q4_0): lin = getattr(owner, attr) @@ -360,7 +374,7 @@ def swap_linear(owner, attr, quant_type=GGML_Q4_0): embed = GGUFEmbedding( num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, - quant_type=GGML_Q6_K, + quant_type=embed_quant, embed_scale=config.embedding_scale, ) inner.embed_tokens = embed @@ -372,7 +386,7 @@ def swap_linear(owner, attr, quant_type=GGML_Q4_0): swap_linear(layer.feed_forward.shared_mlp, "down_proj") if config.tie_word_embeddings: - model.lm_head = GGUFTiedLMHead(embed, GGML_Q6_K) + model.lm_head = GGUFTiedLMHead(embed, embed_quant) # -------------------------------------------------------------------------------------- diff --git a/python/freetoken/models/gguf/__init__.py b/python/freetoken/models/gguf/__init__.py index 75e40e4f..ec3f49b4 100644 --- a/python/freetoken/models/gguf/__init__.py +++ b/python/freetoken/models/gguf/__init__.py @@ -5,6 +5,7 @@ gguf_architecture, gguf_config_source, gguf_tensor_names, + gguf_tensor_type, is_gguf_path, iter_gguf_tensors, load_gguf_metadata, @@ -19,6 +20,7 @@ "gguf_architecture", "gguf_config_source", "gguf_tensor_names", + "gguf_tensor_type", "is_gguf_path", "iter_gguf_tensors", "load_gguf_metadata", diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 63b1a18b..a451f2dc 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -18,6 +18,7 @@ # reuses the model classes but a GGUF parse_config / iter_weights). GGUF_ARCH_TO_REGISTRY: dict[str, str] = { "gemma4": "Gemma4GGUFForCausalLM", + "laguna": "LagunaGGUFForCausalLM", } diff --git a/python/freetoken/models/gguf/dequant.py b/python/freetoken/models/gguf/dequant.py index 77c3ea01..3e009a95 100644 --- a/python/freetoken/models/gguf/dequant.py +++ b/python/freetoken/models/gguf/dequant.py @@ -23,7 +23,15 @@ GGML_F16 = 1 GGML_Q4_0 = 2 GGML_Q8_0 = 8 +GGML_Q3_K = 11 +GGML_Q4_K = 12 +GGML_Q5_K = 13 GGML_Q6_K = 14 +GGML_IQ2_XXS = 16 +GGML_IQ3_XXS = 18 +GGML_IQ1_S = 19 +GGML_IQ2_S = 22 +GGML_IQ4_XS = 23 GGML_BF16 = 30 # (block numel, bytes per block) per ggml type. @@ -33,7 +41,15 @@ GGML_BF16: (1, 2), GGML_Q4_0: (32, 18), GGML_Q8_0: (32, 34), + 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_IQ3_XXS: (256, 98), + GGML_IQ1_S: (256, 50), + GGML_IQ2_S: (256, 82), + GGML_IQ4_XS: (256, 136), } GGML_NAME = { @@ -42,7 +58,15 @@ GGML_BF16: "BF16", GGML_Q4_0: "Q4_0", GGML_Q8_0: "Q8_0", + 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_IQ3_XXS: "IQ3_XXS", + GGML_IQ1_S: "IQ1_S", + GGML_IQ2_S: "IQ2_S", + GGML_IQ4_XS: "IQ4_XS", } @@ -115,9 +139,65 @@ 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 the K-/IQ-quants, delegated to gguf-py. + + gguf-py carries numpy decoders for every ggml type (it cannot *quantize* the + K/IQ formats, but decoding is all the reference path needs); porting the bit + math here would only duplicate it. + """ + 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) + + +def dequant_q3_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + return _dequant_gguf_py(raw, out_dtype, GGML_Q3_K) + + +def dequant_q4_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + return _dequant_gguf_py(raw, out_dtype, GGML_Q4_K) + + +def dequant_q5_k(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + return _dequant_gguf_py(raw, out_dtype, GGML_Q5_K) + + +def dequant_iq2_xxs(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + return _dequant_gguf_py(raw, out_dtype, GGML_IQ2_XXS) + + +def dequant_iq3_xxs(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + return _dequant_gguf_py(raw, out_dtype, GGML_IQ3_XXS) + + +def dequant_iq1_s(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + return _dequant_gguf_py(raw, out_dtype, GGML_IQ1_S) + + +def dequant_iq2_s(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + return _dequant_gguf_py(raw, out_dtype, GGML_IQ2_S) + + +def dequant_iq4_xs(raw: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + return _dequant_gguf_py(raw, out_dtype, GGML_IQ4_XS) + + _DEQUANT = { GGML_Q4_0: dequant_q4_0, GGML_Q6_K: dequant_q6_k, + GGML_Q3_K: dequant_q3_k, + GGML_Q4_K: dequant_q4_k, + GGML_Q5_K: dequant_q5_k, + GGML_IQ2_XXS: dequant_iq2_xxs, + GGML_IQ3_XXS: dequant_iq3_xxs, + GGML_IQ1_S: dequant_iq1_s, + GGML_IQ2_S: dequant_iq2_s, + GGML_IQ4_XS: dequant_iq4_xs, } @@ -141,13 +221,13 @@ def dequantize(raw: torch.Tensor, ggml_type: int, out_dtype: torch.dtype) -> tor "GGML_F32", "GGML_F16", "GGML_BF16", - "GGML_Q4_0", + "GGML_Q4_0", "GGML_Q4_K", "GGML_Q5_K", "GGML_IQ2_XXS", "GGML_IQ3_XXS", "GGML_IQ1_S", "GGML_IQ4_XS", "GGML_Q8_0", "GGML_Q6_K", "GGML_NAME", "BLOCK_SHAPE", "row_bytes", "dequant_q4_0", - "dequant_q6_k", + "dequant_q6_k", "dequant_q4_k", "dequant_q5_k", "dequant_iq2_xxs", "dequant_iq3_xxs", "dequant_iq1_s", "dequant_iq4_xs", "dequantize", ] diff --git a/python/freetoken/models/gguf/reader.py b/python/freetoken/models/gguf/reader.py index b950d929..39b30eec 100644 --- a/python/freetoken/models/gguf/reader.py +++ b/python/freetoken/models/gguf/reader.py @@ -175,6 +175,19 @@ def gguf_tensor_names(model_path: str) -> set[str]: return {t.name for t in _reader(model_path).tensors} +def gguf_tensor_type(model_path: str, name: str) -> int | None: + """The ggml quant type of one tensor, or ``None`` if the file has no such tensor. + + Quant choices are per-tensor in GGUF and differ between publishers (e.g. Google's + QAT release stores ``token_embd.weight`` as Q6_K, Unsloth's as Q4_0), so layer + construction reads the type off the file instead of assuming one. + """ + for t in _reader(model_path).tensors: + if t.name == name: + return int(t.tensor_type) + return None + + __all__ = [ "is_gguf_path", "FTW_METADATA_GGUF", @@ -186,4 +199,5 @@ def gguf_tensor_names(model_path: str) -> set[str]: "gguf_architecture", "iter_gguf_tensors", "gguf_tensor_names", + "gguf_tensor_type", ] diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 6d5481c1..fa4395eb 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -13,7 +13,9 @@ from .reader import gguf_architecture, load_gguf_metadata # GGUF architecture -> transformers GGUF tokenizer-converter key. -_TOKENIZER_ARCH = {"gemma4": "gemma4_text"} +# laguna ships a plain gpt2-style BPE (tokenizer.ggml.model = "gpt2"); transformers +# has no "laguna" converter, so route it to the gpt2 one. +_TOKENIZER_ARCH = {"gemma4": "gemma4_text", "laguna": "gpt2"} def load_gguf_tokenizer(model_path: str): @@ -28,9 +30,14 @@ def load_gguf_tokenizer(model_path: str): for k, v in meta.items() if k.startswith("tokenizer.ggml.") } - fast, _extra = convert_gguf_tokenizer(conv_arch, tok_dict) - tokens = tok_dict["tokens"] + # Some converters (gpt2) read .bos_token/.eos_token off the skeleton, which the + # GGUF metadata only carries as ids -- materialize the token strings. + for name in ("bos", "eos"): + tid = tok_dict.get(f"{name}_token_id") + if f"{name}_token" not in tok_dict and tid is not None and int(tid) < len(tokens): + tok_dict[f"{name}_token"] = tokens[int(tid)] + fast, _extra = convert_gguf_tokenizer(conv_arch, tok_dict) def tok_for(id_key: str, default: str) -> str: tid = meta.get(f"tokenizer.ggml.{id_key}") @@ -53,7 +60,10 @@ def tok_for(id_key: str, default: str) -> str: def gguf_eos_token_ids(model_path: str, tokenizer) -> set[int]: - """Stop ids for GGUF generation: the formal plus the chat turn end .""" + """Stop ids for GGUF generation: the formal , the chat turn end , the + GGUF-declared eot, and gemma4's tool-response opener <|tool_response> (the model + emits it right after closing a tool call, so it is a stop id upstream too -- + generation_config.json ships eos_token_id [1, 106, 50]).""" meta = load_gguf_metadata(model_path) tokens = meta["tokenizer.ggml.tokens"] ids: set[int] = set() @@ -64,7 +74,10 @@ def gguf_eos_token_ids(model_path: str, tokenizer) -> set[int]: 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 ("", ""): + eot = meta.get("tokenizer.ggml.eot_token_id") + if eot is not None: + ids.add(int(eot)) + for name in ("", "", "<|tool_response>"): try: ids.add(tokens.index(name)) except ValueError: diff --git a/python/freetoken/models/laguna/__init__.py b/python/freetoken/models/laguna/__init__.py new file mode 100644 index 00000000..3fe907a9 --- /dev/null +++ b/python/freetoken/models/laguna/__init__.py @@ -0,0 +1,15 @@ +from .gguf import ( + dummy_gguf_expert_sources, + iter_gguf_weights, + load_gguf_expert_sources, + parse_gguf_config, +) +from .model import LagunaForCausalLM + +__all__ = [ + "LagunaForCausalLM", + "parse_gguf_config", + "iter_gguf_weights", + "load_gguf_expert_sources", + "dummy_gguf_expert_sources", +] diff --git a/python/freetoken/models/laguna/attention.py b/python/freetoken/models/laguna/attention.py new file mode 100644 index 00000000..7afa9945 --- /dev/null +++ b/python/freetoken/models/laguna/attention.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.nn.functional as F + +from freetoken.attention import AttentionSpec +from freetoken.core import get_global_ctx +from freetoken.layers import BaseOP, LinearReplicated, RMSNorm +from freetoken.layers.rotary import get_rope +from freetoken.models.config import FullAttentionGroupConfig, SWAAttentionGroupConfig +from freetoken.utils import nvtx_annotate + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + + +class LagunaAttention(BaseOP): + def __init__(self, config: ModelConfig, layer_id: int): + self.layer_id = layer_id + group = config.attention_group_for_layer(layer_id) + if not isinstance(group, (FullAttentionGroupConfig, SWAAttentionGroupConfig)): + raise ValueError(f"LagunaAttention does not support {group.kind!r} layers") + + rotary_config = group.rotary_config + self.head_dim = group.head_dim + self.num_kv_heads = group.num_kv_heads + self.num_qo_heads = config.qo_heads(layer_id) + + self.q_dim = self.num_qo_heads * self.head_dim + self.kv_dim = self.num_kv_heads * self.head_dim + + # Separate q/k/v projections (not LinearQKVMerged): laguna GGUFs quantize + # attn_v at a different ggml type than q/k on some layers (XS Q4_K_M), and + # packed rows of different types cannot be fused into one buffer. + self.q_proj = LinearReplicated(config.hidden_size, self.q_dim, has_bias=False) + self.k_proj = LinearReplicated(config.hidden_size, self.kv_dim, has_bias=False) + self.v_proj = LinearReplicated(config.hidden_size, self.kv_dim, has_bias=False) + self.gate_proj = LinearReplicated(config.hidden_size, self.num_qo_heads, has_bias=False) + self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.o_proj = LinearReplicated(self.q_dim, config.hidden_size, has_bias=False) + self.attn_spec = AttentionSpec( + sliding_window=group.sliding_window if isinstance(group, SWAAttentionGroupConfig) else None + ) + + self.rotary = get_rope( + head_dim=self.head_dim, + rotary_dim=rotary_config.rotary_dim, + max_position=rotary_config.max_position, + base=rotary_config.base, + rope_scaling=( + tuple(rotary_config.scaling.items()) + if rotary_config.scaling + else None + ), + ) + + def _apply_rope( + self, + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + positions = positions.reshape(-1) + if positions.device != q.device or positions.dtype != torch.long: + positions = positions.to(device=q.device, dtype=torch.long) + q_view = q.contiguous().view(q.shape[0], -1) + k_view = k.contiguous().view(k.shape[0], -1) + self.rotary.forward(positions, q_view, k_view) + return q_view.view_as(q), k_view.view_as(k) + + @nvtx_annotate("MHA") + def forward(self, x: torch.Tensor) -> torch.Tensor: + ctx = get_global_ctx() + T = x.shape[0] + + gate = F.softplus(self.gate_proj.forward(x).float()) + q_lin = self.q_proj.forward(x) + k_lin = self.k_proj.forward(x) + v_lin = self.v_proj.forward(x) + + q = q_lin.view(T, self.num_qo_heads, self.head_dim) + k = k_lin.view(T, self.num_kv_heads, self.head_dim) + v = v_lin.view(T, self.num_kv_heads, self.head_dim) + + self.q_norm.forward_inplace(q.view(-1, self.num_qo_heads, self.head_dim)) + self.k_norm.forward_inplace(k.view(-1, self.num_kv_heads, self.head_dim)) + + q, k = self._apply_rope(ctx.batch.positions, q, k) + + k = k.reshape(T, self.kv_dim) + v = v.reshape(T, self.kv_dim) + + o = ctx.attn_backend.forward( + q.contiguous(), + k.contiguous(), + v.contiguous(), + self.layer_id, + ctx.batch, + attn_spec=self.attn_spec, + ) + o = o.view(T, self.num_qo_heads, self.head_dim) + o = o * gate.unsqueeze(-1).to(o.dtype) + return self.o_proj.forward(o.reshape(T, self.q_dim)) + + +__all__ = ["LagunaAttention"] diff --git a/python/freetoken/models/laguna/gguf.py b/python/freetoken/models/laguna/gguf.py new file mode 100644 index 00000000..749865ec --- /dev/null +++ b/python/freetoken/models/laguna/gguf.py @@ -0,0 +1,545 @@ +"""Laguna GGUF adapter: ModelConfig from GGUF metadata + native-quant conversion. + +Laguna (poolside) is hybrid full/SWA attention (full at ``il % 4 == 0`` with 48 +query heads, SWA elsewhere with 72), sigmoid-routed MoE with a selection-only +score-correction bias, one shared expert, a per-head softplus attention output +gate, and QK-norm. Reference: llama.cpp ``src/models/laguna.cpp``. + +Unsloth's "Dynamic" GGUFs mix quant types per tensor (Q4_K embed/head, Q5_K/Q6_K +attention + dense, IQ1_S/IQ2_XXS/IQ3_XXS/IQ4_XS expert banks), so conversion +swaps dense projections for :class:`DeferredGGUFLinear`, whose packed buffer is +materialized at load time when each tensor's ggml type is known. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from freetoken.layers import BaseOP +from freetoken.models.config import ( + FullAttentionGroupConfig, + ModelConfig, + RotaryConfig, + SWAAttentionGroupConfig, +) + +if TYPE_CHECKING: + from freetoken.models.gguf.config import GgufConfigShim + + +def _embed_quant(shim: "GgufConfigShim") -> int | None: + from freetoken.models.gguf.reader import gguf_tensor_type + + return gguf_tensor_type(shim.model_path, "token_embd.weight") + + +def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + m = shim.metadata + + def g(key: str): + full_key = f"laguna.{key}" + val = m.get(full_key) + if val is None: + raise KeyError(f"missing GGUF metadata key {full_key}") + return val + + num_layers = int(g("block_count")) + hidden = int(g("embedding_length")) + intermediate = int(g("feed_forward_length")) + context = int(g("context_length")) + + head_counts = tuple(int(h) for h in g("attention.head_count")) + assert len(head_counts) == num_layers, "attention.head_count length != block_count" + num_qo_heads = max(head_counts) + + full_count = min(head_counts) + full_layer_ids = tuple(i for i, c in enumerate(head_counts) if c == full_count) + swa_layer_ids = tuple(i for i in range(num_layers) if i not in full_layer_ids) + assert full_layer_ids == tuple(i for i in range(0, num_layers, 4)), ( + "Laguna full attention layers must be exactly i % 4 == 0" + ) + + num_kv_heads = int(g("attention.head_count_kv")) + head_dim = int(g("attention.key_length")) + assert head_dim == int(g("attention.value_length")), "key_length != value_length" + + full_rotary = RotaryConfig( + head_dim=head_dim, + rotary_dim=int(g("rope.dimension_count")), + max_position=context, + base=float(g("rope.freq_base")), + scaling={ + "rope_type": "yarn", + "factor": float(g("rope.scaling.factor")), + # ggml applies yarn_attn_factor verbatim (1.0 here); without it + # freetoken's yarn would default to 1 + 0.1*ln(factor). + "attention_factor": float(g("rope.scaling.yarn_attn_factor")), + "original_max_position_embeddings": int(g("rope.scaling.original_context_length")), + "beta_fast": float(g("rope.scaling.yarn_beta_fast")), + "beta_slow": float(g("rope.scaling.yarn_beta_slow")), + }, + ) + swa_rotary = RotaryConfig( + head_dim=head_dim, + rotary_dim=int(g("rope.dimension_count_swa")), + max_position=context, + base=float(g("rope.freq_base_swa")), + scaling=None, + ) + + return ModelConfig( + num_layers=num_layers, + num_qo_heads=num_qo_heads, + num_qo_heads_per_layer=head_counts, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + hidden_size=hidden, + vocab_size=int(shim.vocab_size), + intermediate_size=intermediate, + rms_norm_eps=float(g("attention.layer_norm_rms_epsilon")), + rotary_config=full_rotary, + hidden_act="silu", + tie_word_embeddings=bool(shim.tie_word_embeddings), + num_experts=int(g("expert_count")), + num_experts_per_tok=int(g("expert_used_count")), + moe_intermediate_size=int(g("expert_feed_forward_length")), + shared_expert_intermediate_size=int(g("expert_shared_feed_forward_length")), + n_shared_experts=1, + norm_topk_prob=bool(g("expert_weights_norm")), + routed_scaling_factor=float(g("expert_weights_scale")), + # The selection-only e_score_correction bias (blk.N.exp_probs_b) is part of + # the laguna arch itself, not flagged in metadata; has_router_bias (a bias + # on the router *linear*) stays False. + first_k_dense_replace=int(g("leading_dense_block_count")), + use_qk_norm=True, + model_type="laguna", + architectures=list(shim.architectures), + moe_enabled=True, + attention_groups=( + FullAttentionGroupConfig( + name="full", + layer_ids=full_layer_ids, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + rotary_config=full_rotary, + ), + SWAAttentionGroupConfig( + name="swa", + layer_ids=swa_layer_ids, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + rotary_config=swa_rotary, + sliding_window=int(g("attention.sliding_window")), + ), + ), + gguf_embed_quant=_embed_quant(shim), + gguf_model_path=shim.model_path, + # Routed experts: mixed per-layer ggml types (Unsloth Dynamic), served by the + # "gguf" offload bank format; types read from the tensor table when present + # (None on a metadata-only FTW source). + expert_quant="gguf", + moe_weight_format="gguf", + gguf_expert_types=_expert_types(shim), + ) + + +def _expert_types(shim: "GgufConfigShim") -> tuple[tuple[int, int], ...] | None: + """(gate_up, down) ggml type per MoE layer, from the file's tensor table. + + gate and up always share a type in the published files (asserted); a + metadata-only GGUF has no tensor table -> None. + """ + from freetoken.models.gguf.reader import gguf_tensor_names + + if not gguf_tensor_names(shim.model_path): + return None + types = gguf_tensor_types(shim.model_path) + num_layers = int(shim.metadata["laguna.block_count"]) + dense = int(shim.metadata["laguna.leading_dense_block_count"]) + out = [] + for i in range(dense, num_layers): + gu = types[f"blk.{i}.ffn_gate_exps.weight"] + up = types[f"blk.{i}.ffn_up_exps.weight"] + dn = types[f"blk.{i}.ffn_down_exps.weight"] + if gu != up: + raise ValueError(f"blk.{i}: gate/up expert banks have different ggml types") + out.append((gu, dn)) + return tuple(out) + + +# -------------------------------------------------------------------------------------- +# Weight loading: GGUF tensor names -> FreeToken laguna module params. +# -------------------------------------------------------------------------------------- + +# Routed expert banks are consumed by the MoE offload cache, not the state dict. +_EXPERT_SUFFIXES = ("ffn_gate_exps.weight", "ffn_up_exps.weight", "ffn_down_exps.weight") + +# Per-layer 1:1 tensors dequantized to a dense dtype (suffix -> (rel name, dtype)). +# The router gate and selection bias are consumed fp32 (top-k boundary fidelity, +# minimax_m3 precedent); norms load bf16. +_DENSE_MAP = { + "attn_norm.weight": ("input_layernorm.weight", "bf16"), + "attn_q_norm.weight": ("self_attn.q_norm.weight", "bf16"), + "attn_k_norm.weight": ("self_attn.k_norm.weight", "bf16"), + "ffn_norm.weight": ("ffn_norm.weight", "bf16"), + "ffn_gate_inp.weight": ("mlp.gate.weight", "f32"), + "exp_probs_b.bias": ("mlp.e_score_correction_bias", "f32"), +} + +# Packed (native-quant) projections that map 1:1 (suffix -> rel name). q/k/v stay +# separate modules: their ggml types differ within a layer in some laguna files +# (XS Q4_K_M quantizes attn_v as Q6_K on half the layers), so packed rows cannot fuse. +_PACKED_MAP = { + "attn_q.weight": "self_attn.q_proj.qweight", + "attn_k.weight": "self_attn.k_proj.qweight", + "attn_v.weight": "self_attn.v_proj.qweight", + "attn_output.weight": "self_attn.o_proj.qweight", + "attn_gate.weight": "self_attn.gate_proj.qweight", + "ffn_down.weight": "mlp.down_proj.qweight", + "ffn_down_shexp.weight": "mlp.shared_experts.down_proj.qweight", +} +_GATE_UP_SLOTS = { + "ffn_gate.weight": ("mlp.gate_up_proj", "gate"), + "ffn_up.weight": ("mlp.gate_up_proj", "up"), + "ffn_gate_shexp.weight": ("mlp.shared_experts.gate_up_proj", "gate"), + "ffn_up_shexp.weight": ("mlp.shared_experts.gate_up_proj", "up"), +} + + +def _require_tp1(what: str) -> None: + from freetoken.distributed import get_tp_info + + if get_tp_info().size > 1: + raise NotImplementedError(f"laguna GGUF {what} currently supports TP=1 only") + + +def gguf_tensor_types(model_path: str) -> dict[str, int]: + """One pass over the tensor table: name -> ggml type (no tensor data touched).""" + from freetoken.models.gguf.reader import iter_gguf_tensors + + return {t.name: t.ggml_type for t in iter_gguf_tensors(model_path)} + + +def iter_gguf_weights( + model_path: str, + device, + *, + include_moe_experts: bool, + include_non_moe: bool, +): + """Yield (param_name, tensor) for every non-expert laguna param. + + Quantized projections stay packed and are yielded as ``.qweight`` (uint8); + q/k/v and gate/up fuse by concatenating packed rows along the output dim + (valid only when the components share one ggml type -- Unsloth Dynamic keeps + fused groups uniform, enforced here). Norms load bf16; the router gate and + exp_probs_b load fp32. Routed experts go to the offload cache (4b), not here. + """ + import torch + + from freetoken.models.gguf.dequant import dequantize + from freetoken.models.gguf.reader import iter_gguf_tensors + + assert not include_moe_experts, ( + "laguna GGUF experts are loaded into the MoE offload cache, not the state dict" + ) + assert include_non_moe + _require_tp1("weight loading") + + def dense(t, kind): + dtype = torch.float32 if kind == "f32" else torch.bfloat16 + return dequantize(t.packed().reshape(-1), t.ggml_type, dtype).reshape(t.shape) + + gate_up_buf: dict[tuple[int, str], dict[str, tuple]] = {} + + for t in iter_gguf_tensors(model_path): + name = t.name + if name == "token_embd.weight": + yield "model.embed_tokens.qweight", t.packed() + continue + if name == "output_norm.weight": + yield "model.norm.weight", dense(t, "bf16") + continue + if name == "output.weight": + yield "lm_head.qweight", t.packed() + continue + if not name.startswith("blk."): + raise ValueError(f"unmapped laguna GGUF tensor: {name}") + if name.endswith(_EXPERT_SUFFIXES): + continue # routed experts -> offload banks + + layer = int(name.split(".")[1]) + suffix = name.split(".", 2)[2] + base = f"model.layers.{layer}" + + if suffix in _DENSE_MAP: + rel, kind = _DENSE_MAP[suffix] + yield f"{base}.{rel}", dense(t, kind) + elif suffix in _PACKED_MAP: + yield f"{base}.{_PACKED_MAP[suffix]}", t.packed() + elif suffix in _GATE_UP_SLOTS: + rel, slot = _GATE_UP_SLOTS[suffix] + gate_up_buf.setdefault((layer, rel), {})[slot] = (t.packed(), t.ggml_type) + else: + raise ValueError(f"unmapped laguna GGUF tensor: {name}") + + for key in [k for k in gate_up_buf if k[0] == layer]: + gu = gate_up_buf[key] + if len(gu) == 2: + if gu["gate"][1] != gu["up"][1]: + raise ValueError(f"blk.{layer}: mixed ggml types across fused gate/up") + yield f"{base}.{key[1]}.qweight", torch.cat( + [gu["gate"][0], gu["up"][0]], dim=0 + ) + del gate_up_buf[key] + + assert not gate_up_buf, f"incomplete gate_up groups: {sorted(gate_up_buf)}" + + +def is_gguf_model(config: ModelConfig) -> bool: + """True when the model was parsed from a GGUF checkpoint (native-quant path).""" + return config.gguf_embed_quant is not None + + +class DeferredGGUFLinear(BaseOP): + """GGUF linear whose quant type is only known at weight-load time. + + Unsloth Dynamic checkpoints choose the ggml type per tensor, so conversion + cannot size the packed buffer up front; the loader calls :meth:`materialize` + with the tensor's recorded type before copying rows in. + """ + + def __init__(self, in_features: int, out_features: int, has_bias: bool = False): + self.in_features = in_features + self.out_features = out_features + self._quant_type: int | None = None + self.qweight: torch.Tensor | None = None + self.bias = torch.empty(out_features) if has_bias else None + + def materialize(self, quant_type: int) -> None: + from freetoken.models.gguf.dequant import row_bytes + + self._quant_type = quant_type + self.qweight = torch.empty( + self.out_features, row_bytes(self.in_features, quant_type), dtype=torch.uint8 + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from freetoken.layers.gguf import fused_mul_mat_gguf + + assert self.qweight is not None and self._quant_type is not None, ( + "DeferredGGUFLinear used before materialize() -- weight was never loaded" + ) + out = fused_mul_mat_gguf(x, self.qweight, self._quant_type) + if self.bias is not None: + out = out + self.bias + return out + + +class LagunaGGUFLMHead(DeferredGGUFLinear): + def forward(self, x: torch.Tensor) -> torch.Tensor: + from freetoken.core import get_global_ctx + from freetoken.layers.gguf import fused_mul_mat_gguf + batch = get_global_ctx().batch + if batch.is_prefill: + x = x[batch.attn_metadata.get_last_indices(batch.size)].contiguous() + assert self.qweight is not None and self._quant_type is not None + return fused_mul_mat_gguf(x, self.qweight, self._quant_type) + + +def _swap(owner, attr: str) -> None: + # _LinearTPImpl exposes local_*_size (== full sizes at TP=1, which the GGUF + # path requires); the weight shape is the ground truth either way. + old = getattr(owner, attr) + out_features, in_features = old.weight.shape + setattr( + owner, + attr, + DeferredGGUFLinear(in_features, out_features, getattr(old, "bias", None) is not None), + ) + + +def convert_laguna_to_gguf(model, config: ModelConfig) -> None: + """In place: replace laguna's dense projections + embedding + head with GGUF ops. + + Swapped: attention qkv/gate/o, the dense layer's MLP, every shared expert, the + embedding, and the (untied) lm_head. Kept dense: the fp32 MoE router gate and + e_score_correction_bias, all RMSNorms (F32 in the GGUF), and the routed expert + banks (they live on the MoE offload cache). + + When the source ``.gguf`` is known (``config.gguf_model_path``), every swapped + module is materialized here from the file's per-tensor ggml types, so the + packed buffers exist before the engine collects ``model.state_dict()``. + """ + from freetoken.layers.gguf import GGUFEmbedding + + types = gguf_tensor_types(config.gguf_model_path) if config.gguf_model_path else None + + def qt(name: str) -> int | None: + return None if types is None else types.get(name) + + def mat(module: DeferredGGUFLinear, tensor_name: str) -> None: + t = qt(tensor_name) + if t is not None: + module.materialize(t) + + model.model.embed_tokens = GGUFEmbedding( + config.vocab_size, config.hidden_size, config.gguf_embed_quant + ) + for i, layer in enumerate(model.model.layers.op_list): + for attr, tname in ( + ("q_proj", "attn_q.weight"), + ("k_proj", "attn_k.weight"), + ("v_proj", "attn_v.weight"), + ("gate_proj", "attn_gate.weight"), + ("o_proj", "attn_output.weight"), + ): + _swap(layer.self_attn, attr) + mat(getattr(layer.self_attn, attr), f"blk.{i}.{tname}") + mlp = layer.mlp + if hasattr(mlp, "gate_up_proj"): # dense leading layer (LagunaMLP) + _swap(mlp, "gate_up_proj") + _swap(mlp, "down_proj") + mat(mlp.gate_up_proj, f"blk.{i}.ffn_gate.weight") + mat(mlp.down_proj, f"blk.{i}.ffn_down.weight") + else: # LagunaSparseMoeBlock: shared expert only (router stays fp32) + _swap(mlp.shared_experts, "gate_up_proj") + _swap(mlp.shared_experts, "down_proj") + mat(mlp.shared_experts.gate_up_proj, f"blk.{i}.ffn_gate_shexp.weight") + mat(mlp.shared_experts.down_proj, f"blk.{i}.ffn_down_shexp.weight") + # Untied output head (output.weight, Q4_K in the target file). + model.lm_head = LagunaGGUFLMHead(config.hidden_size, config.vocab_size) + mat(model.lm_head, "output.weight") + + + +# -------------------------------------------------------------------------------------- +# Routed expert banks (mixed per-layer ggml types) for the MoE offload cache. +# -------------------------------------------------------------------------------------- + + +def _expert_bank_geometry(config: ModelConfig): + """Uniform flat-slot strides across MoE layers: max payload bytes, 64B aligned.""" + from freetoken.models.gguf.dequant import row_bytes + + assert config.gguf_expert_types, "laguna expert banks need gguf_expert_types" + H, I = config.hidden_size, config.moe_intermediate_size + gu_pay = {gu: 2 * I * row_bytes(H, gu) for gu, _ in config.gguf_expert_types} + dn_pay = {dn: H * row_bytes(I, dn) for _, dn in config.gguf_expert_types} + + def align(n: int) -> int: + return (n + 63) // 64 * 64 + + return align(max(gu_pay.values())), align(max(dn_pay.values())) + + +def load_gguf_expert_sources( + model_path: str, config: ModelConfig, *, layer_sink=None +) -> dict[str, list[torch.Tensor]]: + """Per-MoE-layer host banks of the routed experts' native packed bytes. + + Each bank is one flat ``[E, stride]`` uint8 tensor per MoE layer (bank index = + layer_id - first_k_dense_replace): every expert's real payload occupies the + leading bytes of its padded slot, so all layers share one shape and the ggml + MoE kernels read them via ``expert_stride_bytes``. Mirrors the q4_0 loader's + pin pipeline / layer_sink streaming contract. + """ + from freetoken.models.gguf.dequant import row_bytes + from freetoken.models.gguf.reader import iter_gguf_tensors + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks + + _require_tp1("expert banks") + types = config.gguf_expert_types + assert types, "laguna expert banks need gguf_expert_types (tensor table missing?)" + E = config.num_experts + H, I = config.hidden_size, config.moe_intermediate_size + L = len(types) # MoE layers only + gu_stride, dn_stride = _expert_bank_geometry(config) + + specs = { + "gate_up": ((E, gu_stride), torch.uint8), + "down": ((E, dn_stride), torch.uint8), + } + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in hb} + seen_gu, seen_dn = set(), set() + + def _load(sink) -> None: + tracker = LayerCompletionTracker(2, hb, sink) if sink is not None else None + gu_parts: dict[int, dict[str, torch.Tensor]] = {} + for t in iter_gguf_tensors(model_path): + name = t.name + if not name.startswith("blk.") or not name.endswith(tuple(_EXPERT_SUFFIXES)): + continue + layer = int(name.split(".")[1]) + bank_id = layer - config.first_k_dense_replace + gu_t, dn_t = types[bank_id] + if name.endswith("ffn_down_exps.weight"): + pay = H * row_bytes(I, dn_t) + banks["down"][bank_id][:, :pay].copy_(t.packed().reshape(E, pay)) + seen_dn.add(bank_id) + if tracker is not None: + tracker.note(bank_id) + else: + # gate and up arrive as separate tensors; each expert's slot holds + # gate rows then up rows (the fused gate_up layout the kernel expects). + half = I * row_bytes(H, gu_t) + part = gu_parts.setdefault(bank_id, {}) + part["gate" if "gate" in name else "up"] = t.packed().reshape(E, half) + if len(part) < 2: + continue + dst = banks["gate_up"][bank_id] + dst[:, :half].copy_(part["gate"]) + dst[:, half : 2 * half].copy_(part["up"]) + del gu_parts[bank_id] + seen_gu.add(bank_id) + if tracker is not None: + tracker.note(bank_id) + assert not gu_parts, f"incomplete expert gate/up layers: {sorted(gu_parts)}" + + if layer_sink is not None: + _load(layer_sink) + elif torch.cuda.is_available(): + with PinPipeline() as pins: + _load(pins) + else: + _load(None) + + want = set(range(L)) + assert seen_gu == want and seen_dn == want, ( + f"missing expert layers: gate_up {sorted(want - seen_gu)}, down {sorted(want - seen_dn)}" + ) + return banks + + +def dummy_gguf_expert_sources(config: ModelConfig) -> dict[str, list[torch.Tensor]]: + """Random banks shaped like ``load_gguf_expert_sources`` output.""" + from freetoken.moe.host_banks import alloc_layer_banks, pin_banks + + E = config.num_experts + L = len(config.gguf_expert_types or ()) + assert L, "laguna dummy expert banks need gguf_expert_types" + gu_stride, dn_stride = _expert_bank_geometry(config) + hb = alloc_layer_banks( + {"gate_up": ((E, gu_stride), torch.uint8), "down": ((E, dn_stride), torch.uint8)}, 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__ = [ + "parse_gguf_config", + "iter_gguf_weights", + "gguf_tensor_types", + "is_gguf_model", + "DeferredGGUFLinear", + "LagunaGGUFLMHead", + "convert_laguna_to_gguf", + "load_gguf_expert_sources", + "dummy_gguf_expert_sources", +] diff --git a/python/freetoken/models/laguna/model.py b/python/freetoken/models/laguna/model.py new file mode 100644 index 00000000..c0cd77f3 --- /dev/null +++ b/python/freetoken/models/laguna/model.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import torch +from freetoken.core import get_global_ctx +from freetoken.layers import BaseOP, OPList, ParallelLMHead, RMSNorm, VocabParallelEmbedding +from freetoken.models.blocks import BaseLLMModel +from freetoken.utils import nvtx_annotate + +from .attention import LagunaAttention +from .moe import LagunaMLP, LagunaSparseMoeBlock + + +class LagunaDecoderLayer(BaseOP): + def __init__(self, config, layer_id: int): + self._layer_id = layer_id + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.self_attn = LagunaAttention(config, layer_id) + self.mlp = (LagunaMLP(config.hidden_size, config.intermediate_size) + if layer_id < config.first_k_dense_replace + else LagunaSparseMoeBlock(config, layer_id)) + + @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.self_attn.forward(self.input_layernorm.forward(x)) + x = x + self.mlp.forward(self.ffn_norm.forward(x)) + return x + + +class LagunaModel(BaseOP): + def __init__(self, config): + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = OPList([LagunaDecoderLayer(config, i) for i in range(config.num_layers)]) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + x = self.embed_tokens.forward(input_ids) + for layer in self.layers.op_list: + x = layer.forward(x) + return self.norm.forward(x) + + +class LagunaForCausalLM(BaseLLMModel): + def __init__(self, config): + self.model = LagunaModel(config) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, + tie_word_embeddings=config.tie_word_embeddings, + tied_embedding=None) + super().__init__() + from .gguf import convert_laguna_to_gguf, is_gguf_model + if is_gguf_model(config): + convert_laguna_to_gguf(self, config) + elif config.gguf_model_path: + # GGUF-sourced but no tensor table (a converted FTW dir's metadata-only + # source_metadata.gguf): the per-tensor quant types are unrecoverable, so + # neither module conversion nor expert banks can be built. Refuse loudly + # instead of constructing a dense model that fails weight loading. + raise NotImplementedError( + "laguna FTW conversion is not supported yet -- serve the .gguf file " + "directly (per-tensor quant types live only in its tensor table)" + ) + + def forward(self) -> torch.Tensor: + return self.lm_head.forward(self.model.forward(get_global_ctx().batch.input_ids)) + + +__all__ = ["LagunaDecoderLayer", "LagunaModel", "LagunaForCausalLM"] diff --git a/python/freetoken/models/laguna/moe.py b/python/freetoken/models/laguna/moe.py new file mode 100644 index 00000000..eea90e8e --- /dev/null +++ b/python/freetoken/models/laguna/moe.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Tuple + +import torch +import torch.nn.functional as F + +from freetoken.layers import BaseOP, LinearReplicated, make_moe_layer, silu_and_mul + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + +TopK = Tuple[torch.Tensor, torch.Tensor] + + +class LagunaMLP(BaseOP): + """Plain SwiGLU MLP used by Laguna dense and shared experts.""" + + def __init__(self, hidden_size: int, intermediate_size: int): + self.gate_up_proj = LinearReplicated(hidden_size, 2 * intermediate_size, has_bias=False) + self.down_proj = LinearReplicated(intermediate_size, hidden_size, has_bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up = self.gate_up_proj.forward(x) + del x + y = silu_and_mul(gate_up) + del gate_up + return self.down_proj.forward(y) + + +class LagunaSparseMoeBlock(BaseOP): + """Mixture-of-experts block for Laguna.""" + + def __init__(self, config: ModelConfig, layer_id: int): + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_experts + self.norm_topk_prob = config.norm_topk_prob + self.routed_scaling_factor = config.routed_scaling_factor + + # Router weights are fp32; this is required for exact tie-breaking at the top-k boundary. + self.gate = LinearReplicated(config.hidden_size, config.num_experts, has_bias=False) + self.gate.weight = torch.empty(config.num_experts, config.hidden_size, dtype=torch.float32) + self.e_score_correction_bias = torch.empty(config.num_experts, dtype=torch.float32) + + moe_layer_id = layer_id - config.first_k_dense_replace + # Mixed-type GGUF banks ("gguf" offload format): the kernels need this + # layer's ggml types and output-row geometry (see fused_experts_gguf). + extra_attrs = None + if config.gguf_expert_types is not None: + gu_t, dn_t = config.gguf_expert_types[moe_layer_id] + extra_attrs = { + "gguf_gate_up_type": gu_t, + "gguf_down_type": dn_t, + "gguf_gate_up_rows": 2 * config.moe_intermediate_size, + "gguf_down_rows": config.hidden_size, + } + self.experts = make_moe_layer( + config, + layer_id=moe_layer_id, + renormalize=config.norm_topk_prob, + activation="silu", + extra_attrs=extra_attrs, + ) + self.shared_experts = LagunaMLP( + config.hidden_size, + config.shared_expert_intermediate_size * max(1, config.n_shared_experts), + ) + + def _route(self, hidden_states: torch.Tensor) -> TopK: + logits = F.linear(hidden_states.float(), self.gate.weight) + scores = logits.sigmoid() + scores_for_choice = scores + self.e_score_correction_bias + _, topk_ids = torch.topk(scores_for_choice, self.top_k, dim=-1) + topk_weights = scores.gather(-1, topk_ids) + if self.norm_topk_prob: + topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20) + topk_weights = topk_weights * self.routed_scaling_factor + return topk_weights.to(torch.float32).contiguous(), topk_ids.to(torch.int32).contiguous() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + topk_weights, topk_ids = self._route(hidden_states) + out = self.experts.routed_forward(hidden_states, topk_weights, topk_ids) + out = out + self.shared_experts.forward(hidden_states) + return out.view(num_tokens, hidden_dim) + + +__all__ = ["LagunaMLP", "LagunaSparseMoeBlock"] diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 0c033ca0..d3c3dd77 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -107,6 +107,12 @@ class ModelSpec: parse_config="parse_gguf_config", iter_weights="iter_gguf_weights", ), + "LagunaGGUFForCausalLM": ModelSpec( + "freetoken.models.laguna", + "LagunaForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + ), "GptOssForCausalLM": ModelSpec( "freetoken.models.gpt_oss", "GptOssForCausalLM", diff --git a/python/freetoken/models/weight.py b/python/freetoken/models/weight.py index 6a34f3b9..97eb65d9 100644 --- a/python/freetoken/models/weight.py +++ b/python/freetoken/models/weight.py @@ -342,6 +342,26 @@ def load_q4_0_moe_expert_sources( return loader(model_path, model_config, layer_sink=layer_sink) +def load_gguf_moe_expert_sources( + model_path: str, + model_config, + *, + dummy: bool = False, + layer_sink=None, +) -> dict: + """Load (or fabricate) mixed-type GGUF expert banks (flat padded uint8 slots). + + Per-model: dispatches to the model module's ``load_gguf_expert_sources`` / + ``dummy_gguf_expert_sources`` (laguna is the first user).""" + _config, spec = _spec_for_model_path(model_path) + if dummy: + builder = _model_override(spec, "dummy_gguf_expert_sources") + assert builder is not None, "model defines no dummy_gguf_expert_sources" + return builder(model_config) + loader = _load_attr(spec.module, "load_gguf_expert_sources") + return loader(model_path, model_config, layer_sink=layer_sink) + + def _num_moe_layers(config) -> int: value = getattr(config, "num_moe_layers", None) if value is not None: @@ -406,6 +426,7 @@ def bank(*shape: int, dtype: torch.dtype) -> list[torch.Tensor]: __all__ = [ "load_weight", "load_moe_expert_sources", + "load_gguf_moe_expert_sources", "load_nvfp4_moe_expert_sources", "dummy_moe_expert_sources", "dummy_nvfp4_expert_sources", diff --git a/python/freetoken/moe/expert_banks.py b/python/freetoken/moe/expert_banks.py index 8b6116ba..fcd6f353 100644 --- a/python/freetoken/moe/expert_banks.py +++ b/python/freetoken/moe/expert_banks.py @@ -252,6 +252,28 @@ def _q4_0_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: + if parallel: + raise NotImplementedError( + "parallel reader not implemented for gguf: single packed file (see q4_0)" + ) + if decode_target != "gpu": + raise NotImplementedError( + "mixed-type GGUF experts support the GPU offload backend only " + "(no CPU/hybrid executor format id)" + ) + from freetoken.models.weight import load_gguf_moe_expert_sources + + # Mixed-type GGUF routed experts (laguna): per-layer quant types, flat padded + # [E, stride] uint8 slots so every layer shares one bank shape (kernels read the + # leading payload via expert_stride_bytes). + sink = None if dummy else layer_sink + sources = load_gguf_moe_expert_sources(model_path, model_config, dummy=dummy, layer_sink=sink) + return ExpertBanks( + "gguf", {name: sources[name] for name in _BANK_SCHEMAS["gguf"]}, streamed=sink is not None + ) + + def _dsfp4_banks(model_path, model_config, device, dtype, dummy, parallel=False, workers=8, chunk=_PARALLEL_CHUNK, decode_target="gpu", layer_sink=None) -> ExpertBanks: args = model_config.dsv4_args assert args is not None, "ds_fp4 expert banks require dsv4_args on the model config" @@ -301,6 +323,7 @@ def _model_setup_override(model_config): "nvfp4": _nvfp4_banks, "ds_fp4": _dsfp4_banks, "q4_0": _q4_0_banks, + "gguf": _gguf_banks, } diff --git a/python/freetoken/moe/fused_gguf.py b/python/freetoken/moe/fused_gguf.py new file mode 100644 index 00000000..61de7237 --- /dev/null +++ b/python/freetoken/moe/fused_gguf.py @@ -0,0 +1,94 @@ +"""Grouped expert GEMM over mixed-type GGUF banks (borrowed ggml MoE kernels). + +The generalization of :mod:`freetoken.moe.fused_q4_0` for checkpoints whose +routed-expert quant type varies per layer (Unsloth Dynamic laguna: gate/up +IQ1_S or IQ2_XXS, down IQ3_XXS or IQ4_XS). Because per-expert byte sizes then +differ across layers, the banks are FLAT padded slots -- ``[num_slots, +stride_bytes]`` uint8 with each expert's real payload in the leading bytes -- +and the kernels read them via ``expert_stride_bytes``. Geometry (quant type, +output rows) rides in per-call arguments; MMVQ serves prefill and decode like +the q4_0 path. +""" + +from __future__ import annotations + +import torch + +from freetoken.layers.activation import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul + +_ACT = {"silu": silu_and_mul, "gelu": gelu_and_mul, "gelu_tanh": gelu_tanh_and_mul} + +# moe_vec's CUDA grid puts (tokens * top_k) rows in grid.z, which CUDA caps at +# 65535. Large prefill chunks (e.g. 16384 tokens * top_8 = 131072) exceed that, +# so calls are split into row-count-bounded pieces. The down projection already +# runs at "top_k=1, tokens=num_tokens*top_k" (one row per selected expert), so +# both calls share one chunking helper keyed off total (rows, top_k) pairs. +_MAX_GRID_Z = 65535 + +# Transient memory bound: each call materializes [rows_in_flight, out_rows] plus a +# q8_1 copy of its activations. On a VRAM-tight offload setup (expert cache eats +# everything the KV pool leaves) a 16k-token prefill chunk at top_8 would allocate +# ~1 GiB in one shot and fault asynchronously, so cap rows well below the grid limit. +_MAX_ROWS_IN_FLIGHT = 16384 + + +def _moe_vec_chunked(x, weight, topk_ids, top_k, quant_type, rows, tokens, stride): + from freetoken.kernel.gguf import ggml_moe_a8_vec + + limit = min(_MAX_GRID_Z, _MAX_ROWS_IN_FLIGHT) + if tokens * top_k <= limit: + return ggml_moe_a8_vec(x, weight, topk_ids, top_k, quant_type, rows, tokens, stride) + + chunk = max(1, limit // top_k) + outs = [] + for start in range(0, tokens, chunk): + end = min(start + chunk, tokens) + outs.append( + ggml_moe_a8_vec( + x[start:end], weight, topk_ids[start:end], top_k, quant_type, rows, end - start, stride + ) + ) + return torch.cat(outs, dim=0) + + +def fused_experts_gguf( + hidden_states: torch.Tensor, + gate_up_q: torch.Tensor, # [num_slots, gu_stride] uint8 (flat padded slots) + down_q: torch.Tensor, # [num_slots, dn_stride] uint8 + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: str, + *, + gate_up_type: int, + down_type: int, + gate_up_rows: int, # 2 * intermediate + down_rows: int, # hidden +) -> torch.Tensor: + act_fn = _ACT.get(activation) + if act_fn is None: + raise ValueError(f"unsupported MoE activation {activation!r}") + + num_tokens = hidden_states.shape[0] + top_k = topk_ids.shape[1] + assert gate_up_q.dim() == 2 and down_q.dim() == 2, "gguf banks are flat padded slots" + + gate_up = _moe_vec_chunked( + hidden_states, gate_up_q, topk_ids, top_k, int(gate_up_type), + gate_up_rows, num_tokens, gate_up_q.shape[1], + ) + inter = act_fn(gate_up) + # Down pass: one selected-expert row per (token, k) -- already flat, so it's a + # top_k=1 call over num_tokens*top_k "tokens". topk_ids must flatten the same + # way (row-major [num_tokens, top_k] -> contiguous [num_tokens*top_k, 1]). + flat_ids = topk_ids.reshape(-1, 1) + out = _moe_vec_chunked( + inter, down_q, flat_ids, 1, int(down_type), + down_rows, num_tokens * top_k, down_q.shape[1], + ) + out = out.reshape(num_tokens, top_k, down_rows) * topk_weights.reshape( + num_tokens, top_k, 1 + ).to(out.dtype) + return out.sum(dim=1) + + +__all__ = ["fused_experts_gguf"] diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 6ee76406..aac7c6d9 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"), + # Mixed-type GGUF (laguna): flat padded uint8 slots [E, stride_bytes]; the + # per-layer quant geometry lives on the MoE layer, not the bank shape. + "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) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index a71b6819..c60af269 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -135,6 +135,8 @@ def _infer_tool_call_parser(model_path: str) -> str: " ".join(str(v) for v in text_cfg.get("architectures", []) or []), ] marker = " ".join(candidates).lower() + if "laguna" in marker or "poolside" in marker: + return "poolside_v1" if "gpt_oss" in marker or "gpt-oss" in marker or "gptoss" in marker: return "gpt_oss" # M3 first: its marker also contains the bare "minimax" substring, but the @@ -182,6 +184,8 @@ def _infer_reasoning_parser(model_path: str) -> str | None: " ".join(str(v) for v in text_cfg.get("architectures", []) or []), ] marker = " ".join(candidates).lower() + if "laguna" in marker or "poolside" in marker: + return "poolside_v1" if "gpt_oss" in marker or "gpt-oss" in marker or "gptoss" in marker: return "gpt_oss" if "deepseek" in marker and any( @@ -437,6 +441,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: "deepseekv32", "gemma4", "glm47", + "poolside_v1", "minimax", "minimax_m3", "muse_glimmer", @@ -451,7 +456,7 @@ def _infer_reasoning_parser(model_path: str) -> str | None: type=str, default="auto", choices=[ - "auto", "off", "deepseekv32", "gpt_oss", "qwen3", "glm", + "auto", "off", "deepseekv32", "gpt_oss", "qwen3", "glm", "poolside_v1", "minimax", "minimax_m3", "muse_glimmer", "gemma4", ], help=( diff --git a/python/freetoken/server/function_call_parser.py b/python/freetoken/server/function_call_parser.py index 0804dcc3..3ffbf911 100644 --- a/python/freetoken/server/function_call_parser.py +++ b/python/freetoken/server/function_call_parser.py @@ -1207,6 +1207,8 @@ class Glm47Detector(BaseFormatDetector): Reference: https://github.com/vllm-project/vllm/blob/main/vllm/tool_parsers/glm4_moe_tool_parser.py """ toolcall_opener = "" + _preserve_string_whitespace = False + def __init__(self): super().__init__() self.bot_token = "" @@ -1232,6 +1234,36 @@ def has_tool_call(self, text: str) -> bool: """Check if the text contains a GLM-4.7 format tool call.""" return self.bot_token in text + def _argument_is_string(self, key: str, param_config: Dict) -> bool: + """Whether an XML argument uses the incremental string path. + + Poolside V1 follows JSON Schema literally: only an explicit + ``type: string`` preserves the raw value. GLM keeps its legacy loose + aliases and defaults for backward compatibility. + """ + if self._preserve_string_whitespace: + schema = param_config.get(key, {}) + return isinstance(schema, dict) and schema.get("type") == "string" + return self._schema_param_type(key, param_config, "loose") in ( + "string", + "str", + "enum", + ) + + def _convert_xml_argument( + self, value: str, key: str, param_config: Dict, func_name: str + ): + raw = value.strip() + if self._preserve_string_whitespace: + schema = param_config.get(key, {}) + if isinstance(schema, dict) and "type" in schema: + return self._convert_param_value(raw, key, param_config, func_name) + try: + return json.loads(raw) + except (json.JSONDecodeError, ValueError): + return raw + return self._convert_param_value(raw, key, param_config, func_name) + def _parse_xml_arguments(self, arg_text: str, param_config: Dict | None = None, func_name: str = "") -> dict: """ Parse XML-style arguments into a dictionary. @@ -1249,12 +1281,17 @@ def _parse_xml_arguments(self, arg_text: str, param_config: Dict | None = None, matches = self.func_arg_regex.findall(arg_text) for key, value in matches: key = key.strip() - value = value.strip() if param_config and key in param_config: + if self._preserve_string_whitespace and self._argument_is_string( + key, param_config + ): + args[key] = value + continue # Schema-first: the declared type wins (a string-typed "5" stays "5"). - args[key] = self._convert_param_value(value, key, param_config, func_name) + args[key] = self._convert_xml_argument(value, key, param_config, func_name) continue # Undeclared parameter: legacy loose typing. + value = value.strip() try: parsed_value = json.loads(value) args[key] = parsed_value @@ -1271,7 +1308,9 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult :return: StreamingParseResult with normal_text and parsed calls. """ idx = text.find(self.bot_token) - normal_text = text[:idx].strip() if idx != -1 else text + normal_text = text[:idx] if idx != -1 else text + if idx != -1 and not self._preserve_string_whitespace: + normal_text = normal_text.strip() if self.bot_token not in text: return StreamingParseResult(normal_text=normal_text, calls=[]) @@ -1464,8 +1503,7 @@ def _update_prev() -> None: continue lead = "{" if not self._args_started else "," self._args_started = True - ptype = self._schema_param_type(self._g_key, self._g_config, "loose") - if ptype in ("string", "str", "enum"): + if self._argument_is_string(self._g_key, self._g_config): _emit(lead + json.dumps(self._g_key, ensure_ascii=False) + ':"') self._g_lead_trimmed = False self._g_mode = "pstr" @@ -1475,7 +1513,7 @@ def _update_prev() -> None: continue if mode == "pstr": - if not self._g_lead_trimmed: + if not self._preserve_string_whitespace and not self._g_lead_trimmed: trimmed = buf.lstrip() if trimmed != buf: self._buffer = trimmed @@ -1491,7 +1529,11 @@ def _update_prev() -> None: _emit(self._json_escape_chunk(emit_now)) self._buffer = buf[len(emit_now):] break - tail = buf[:end].rstrip() + tail = ( + buf[:end] + if self._preserve_string_whitespace + else buf[:end].rstrip() + ) _emit(self._json_escape_chunk(tail) + '"') _update_prev() self._buffer = buf[end + len(self._G_VAL_CLOSE):] @@ -1505,7 +1547,9 @@ def _update_prev() -> None: if mode == "pbuf": raw = buf[:end].strip() if self._g_key in self._g_config: - converted = self._convert_param_value(raw, self._g_key, self._g_config, "") + converted = self._convert_xml_argument( + raw, self._g_key, self._g_config, "" + ) else: try: converted = json.loads(raw) @@ -1537,6 +1581,16 @@ def finish_streaming(self) -> str: return residual +class PoolsideV1Detector(Glm47Detector): + """Poolside V1 tool protocol used by Laguna. + + Its envelope matches GLM-4.7, but string arguments are raw text. Preserve + leading/trailing whitespace so source code and patches survive parsing. + """ + + _preserve_string_whitespace = True + + class DeepSeekV32Detector(BaseFormatDetector): """ Detector for DeepSeek V3.2 model function call format using DSML @@ -3532,6 +3586,7 @@ class FunctionCallParser: "minimax_m3": MiniMaxM3Detector, "mistral": MistralDetector, "muse_glimmer": MuseGlimmerDetector, + "poolside_v1": PoolsideV1Detector, "qwen": Qwen25Detector, "qwen25": Qwen25Detector, "qwen3_coder": Qwen3CoderDetector, diff --git a/python/freetoken/server/generation.py b/python/freetoken/server/generation.py index be05d908..0c6ec2ed 100644 --- a/python/freetoken/server/generation.py +++ b/python/freetoken/server/generation.py @@ -351,6 +351,11 @@ def _make_reasoning_parser(spec: GenSpec, state: Any) -> ReasoningParser | None: # GLM's template honors enable_thinking (default on) even with tools; the # generic fallback would force thinking and mislabel disabled output as reasoning. force_reasoning = (spec.chat_template_kwargs or {}).get("enable_thinking") is not False + elif parser_name == "poolside_v1": + # Poolside's template defaults thinking on and opens before generation, + # so the model normally emits only the closer. Explicit false pre-closes the + # block and must leave subsequent answer text visible. + force_reasoning = (spec.chat_template_kwargs or {}).get("enable_thinking") is not False elif parser_name == "gemma4": # Gemma4 defaults thinking off even when tools are present: its template injects an # empty thought channel before generation. Do not let Codex tool definitions make all diff --git a/python/freetoken/server/reasoning_parser.py b/python/freetoken/server/reasoning_parser.py index 6788675c..f205765e 100644 --- a/python/freetoken/server/reasoning_parser.py +++ b/python/freetoken/server/reasoning_parser.py @@ -876,6 +876,7 @@ class ReasoningParser: ReasoningParserEnum: Dict[str, Type[BaseReasoningParser]] = { "deepseekv32": DeepSeekV32ReasoningParser, "gpt_oss": GptOssHarmonyReasoningParser, + "poolside_v1": ThinkReasoningParser, "qwen3": ThinkReasoningParser, "glm": ThinkReasoningParser, "minimax": ThinkReasoningParser, diff --git a/tasks/laguna-handover.md b/tasks/laguna-handover.md new file mode 100644 index 00000000..7ee5e8a9 --- /dev/null +++ b/tasks/laguna-handover.md @@ -0,0 +1,165 @@ +# Laguna GGUF port — handover + +State as of 2026-08-23. Written so a fresh session on a **bigger host** can finish +validating the port against **Laguna-S-2.1** (the XS model was only a stand-in for +bring-up on a 23 GB / 16 GB-VRAM box, which cannot hold S). + +The plan and verified model facts live in `tasks/laguna-todo.md` — read that first, +it is the spec. This file says what is done, what is not, and what will bite you. + +## One-line status + +Phases 1–4 are implemented, reviewed and unit-tested (81 tests green: `tests/models` ++ `tests/kernels/test_gguf_quant_types.py`). Phase 5 (end-to-end validation on real +weights) is **not done**: no token has ever been generated by this port. + +## What works, with evidence + +| Area | Evidence | +|---|---| +| GGML types Q4_K/Q5_K/IQ1_S/IQ2_XXS/IQ3_XXS/IQ4_XS (dequant, mmvq, mmq where it exists, moe_vec) | `tests/kernels/test_gguf_quant_types.py` (27 tests, CUDA vs gguf-py on identical bytes) | +| Config/registry/tokenizer from GGUF metadata | `tests/models/test_laguna_config.py`, committed fixture `tests/fixtures/laguna-s-2.1-metadata.gguf` | +| Attention (per-layer heads, QK-norm, dual rope, softplus gate), MoE router, decoder residuals | `tests/models/test_laguna_modules.py` (stubbed forward vs independent reference) | +| Weight loading, deferred materialization, expert banks | `tests/models/test_laguna_weights.py` (writes a real tiny laguna GGUF with gguf-py) | +| Name-map coverage on the real S file | all 814 tensors mapped or in the expert skip-set; `iter_gguf_weights` yielded 529 params with no duplicates | +| Expert-bank byte path on the real S file | layer 1 / expert 5 gate_up matmul vs gguf-py dequant reference: **0.5 % rel err** | +| S bring-up reached | metadata → weights → expert banks → MoE cache sizing → KV allocation. Died only on host RAM (S needs ~37 GiB pinned; box had 23 GB) | +| XS bring-up reached | same, plus KV alloc of **262144 tokens fp8 = 8.79 GiB** with `--kv-reserve-tokens 262144`; process was killed manually before the ready flip | + +## Validation done on the APEX-Mini model (XS-size) + +Validated end-to-end on `Laguna-XS-2.1-APEX-I-Mini.gguf` (Q3_K/Q4_K/Q5_K/Q6_K/IQ2_S, +12.8 GB) on the 16 GB / 23 GB box: + +- **NIAH 3/3 at 250,054 tokens** -- needle recovered at 10%/50%/90% depth, exact + passcode each time (~433 tok/s prefill). This exercises YaRN over the full window, + the full/SWA split, QK-norm, the softplus gate, and routing, so a llama.cpp token + diff is considered redundant. +- **Decode**: 157-162 tok/s at 64k ctx (8984 expert slots), 21-23 tok/s at ~250k ctx + (2441 slots -- PCIe-bound). 262k KV costs 8.79 GiB fp8, so on 16 GB the 262k and + 64k configs cannot both be "fast". + +## What is NOT done + +1. **S-model validation.** Everything above is XS. S is geometry-identical but 48 + layers / 3072 hidden / 1024 expert, needs >=48 GB RAM (see sizing below) -- run + the same NIAH + decode recipe there. This is the whole point of the handover. +2. **SWA-boundary + eos stop-behaviour** not explicitly tested (short prompts in the + NIAH reasoning block exercised window crossing implicitly, but no dedicated test). +3. **Pre/post-fill perf on IQ types** (IQ2_S, IQ1_S, IQ3_XXS have no MMQ kernel, so + the dense-quant prefill path dequantizes; MoE prefill uses moe_vec). Tune only if + S prefill is slow. +4. **hybrid/cpu MoE backends refuse `gguf`.** No `WFmt` case exists in the compiled + C++ CPU executor, so `--moe-backend hybrid` errors. S host (CPU-bw > PCIe) is where + hybrid matters -- see the full scope in "Things that will bite you" #8 (a C++ SIMD + kernel port, not a Python change). +5. **FTW conversion refused** on purpose (metadata-only GGUF drops per-tensor types). + Serve the `.gguf` directly. +6. **TP=1 only**, text-only. + +## How to run it on the big host + +```bash +ft serve --model /path/Laguna-S-2.1-UD-IQ1_S.gguf \ + --kv-cache-dtype q8_0 --num-tokens 65536 --kv-reserve-tokens 65536 +``` + +- `--kv-reserve-tokens` matters: without it the MoE cache auto-sizer eats the VRAM the + KV pool then needs, and allocation fails late (seen on this box). +- `--kv-cache-dtype q8_0|fp8_e4m3` comes from the *other* (uncommitted) KV-quant + workstream in this tree. With plain `auto` (bf16) the KV cost doubles. + +### Host sizing for S (measured/derived, not guessed) + +- Expert banks, real bytes: **28.8 GiB**. As allocated by this port (uniform padded + slots, see below): **36.9 GiB pinned host RAM**. Budget ≥ 48 GB RAM. +- Dense weights (VRAM): ~3–4 GiB. KV @64k q8_0: ~1.4 GiB (only the 12 full-attention + layers hold full context; the 36 SWA layers are capped at a 512 window). +- VRAM left over becomes the expert LRU cache; more is better, nothing breaks if small. + +## Things that will bite you + +1. **Padding waste (worth fixing).** `OffloadMoeCache.set_bank_sources` asserts every + layer's bank has the same shape, but Laguna's per-layer expert quant types differ, so + `load_gguf_expert_sources` pads every expert slot to the *max* type's size. On S that + is +8.1 GiB of pure waste (28.8 → 36.9). Options, cheapest first: + - group MoE layers by `(gate_up_type, down_type)` and give each group its own cache + (3 groups on S), or + - teach the cache per-layer strides (the kernels already take `expert_stride_bytes`, + so only the host-side bank/copy-plan bookkeeping is in the way). +2. **Mixed types inside one layer.** XS quantizes `attn_v` differently from `attn_q/k` on + half its layers, which is why q/k/v are **separate projections** (not `LinearQKVMerged`). + Do not "optimize" them back into a fused buffer — packed rows of different ggml types + cannot be concatenated. Gate/up fusion *is* still done and is type-checked at load. +3. **Deferred materialization ordering.** The engine collects `model.state_dict()` before + iterating weights, so `convert_laguna_to_gguf` must materialize every + `DeferredGGUFLinear` up front from the file's tensor table (it does — via + `config.gguf_model_path`). If you move that call, loading breaks with missing keys. +4. **`num_qo_heads` is the max (72 on S)**; per-layer counts come from + `config.qo_heads(layer_id)`. Triton allocates decode scratch at the max and gets the + real count from `q.shape[1]` at call time. FlashInfer plans one global head count and + is therefore **not** usable for Laguna; SWA forces the triton backend anyway. +5. **`expert_stride_bytes`** was added to all 19 vendored `moe_vec_*` launchers in + `python/freetoken/kernel/csrc/gguf/`. Value 0 = old dense behaviour, so existing + formats (q4_0/gemma4) are unaffected. If you re-vendor those files from upstream you + will drop this patch. +6. **moe_vec grid.z cap (the crash).** The `moe_vec_q` kernel indexes experts via + `blockIdx.z`, and CUDA caps grid-z at 65535 rows = `tokens*top_k`. `fused_gguf.py` + chunks calls to `min(65535, 16384)` rows for BOTH the grid limit and transient VRAM + use (a 16k-token x top-8 prefill chunk materializes ~1 GiB). The 16384 cap was the + fix for a "CUDA driver error: device not ready" that surfaced asynchronously. +7. **Offload-cache port hygiene.** Orphaned spawn workers hold rendezvous port 1920 with + a bare `multiprocessing` cmdline, so `pkill -f "ft serve"` does NOT kill them; the + next `ft serve` dies EADDRINUSE. Kill by port (`ss -tlnp | grep ':1920 '`) before + restarting, or wait ~60 s for TIME_WAIT to drain. `serve_supervised.sh` in + `.claude/scratch/` does this sweep + drain. +8. **hybrid/cpu fix (for the S host).** To enable `--moe-backend hybrid` for `gguf`: + - Python side: add a `"gguf"` entry to `_WFMT_IDS` in + `python/freetoken/moe/cpu_executor.py` and a `_resolve_gguf_banks` that reads the + flat `[E, stride]` uint8 banks with per-layer `(gate_up_type, down_type)` from + `config.gguf_expert_types`. + - **The hard part (not "~1 file"): the CPU executor's hot path is a compiled C++ + extension** (`python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp`) whose `WFmt` + enum has no `gguf` case. It needs vec-dot kernels for every type laguna uses + (Q3_K, Q4_K, Q5_K, Q6_K, IQ1_S, IQ2_XXS, IQ3_XXS, IQ4_XS — ported from llama.cpp's + ggml-cpu AVX-512 paths) plus per-layer type plumbing from `_resolve_gguf_banks`. + That is a C++ kernel port, several hundred lines of SIMD per family plus the + dispatch -- not a Python fallback (the Python `dequantize` is gguf-py/numpy, + far too slow for the decode GEMV to beat PCIe). + - Scope note: only worth it where CPU-bw > PCIe (i.e. the S host's calibration, NOT + this 16 GB box — its own `ft bench bw` profile recommends `offload` for every + format, so hybrid would lose here regardless). + +## Map of the change + +New: +- `python/freetoken/models/laguna/` — `gguf.py` (config parse, name map, weight iter, + deferred linears, expert-bank loader), `attention.py`, `moe.py`, `model.py` +- `python/freetoken/moe/fused_gguf.py` — mixed-type expert GEMV +- `tests/models/test_laguna_{config,modules,weights}.py`, + `tests/kernels/test_gguf_quant_types.py`, `tests/fixtures/laguna-s-2.1-metadata.gguf` + +Modified: +- `kernel/csrc/gguf/{moe_vec.cuh,gguf_kernel.cu}`, `kernel/gguf.py` — `expert_stride_bytes` +- `models/gguf/dequant.py`, `layers/gguf.py` — the six new ggml types +- `layers/moe.py`, `moe/{expert_banks,offload_cache}.py`, `models/weight.py` — the + `"gguf"` expert-bank format end to end +- `models/config.py` — `num_qo_heads_per_layer` + `qo_heads()`, `gguf_model_path`, + `gguf_expert_types` +- `models/gguf/{config,tokenizer,reader,__init__}.py`, `models/register.py` — laguna + registry entry, gpt2-converter routing, `gguf_tensor_type` +- `models/gemma4/gguf.py`, `layers/base.py` — small shared-infra pickups + +## The other workstream in this repo + +The KV-cache quantization effort (`tasks/todo.md`, `kvcache/quant*.py`, +`kernel/triton/kv_quant.py`, `attention/triton.py`, `engine/*`, `server/args.py`, its +tests) is **separate** but shipped in the commit right after this one, because the two +meet at `--kv-cache-dtype`: the `q8_0` / `fp8_e4m3` flag used in every serve command +above comes from there, and without it laguna's KV falls back to bf16 (2x the bytes). + +Its own status: 53 tests green plus the 33 pre-existing triton-attention tests, but +step 9 of `tasks/todo.md` is open — no needle-in-246k, no perplexity vs bf16, no +measured expert-slot / tok-s gain. So `q8_0` vs `fp8_e4m3` as the default is still +undecided, and on the big host it is worth settling that on the same run that +validates laguna: both questions need one loaded model and a long context. diff --git a/tasks/laguna-todo.md b/tasks/laguna-todo.md new file mode 100644 index 00000000..b074e041 --- /dev/null +++ b/tasks/laguna-todo.md @@ -0,0 +1,169 @@ +# Laguna-S-2.1 GGUF support (unsloth Laguna-S-2.1-UD-IQ1_S.gguf) + +Goal: `ft serve --model Laguna-S-2.1-UD-IQ1_S.gguf --kv-cache-dtype q8_0 --num-tokens 65536` +loads and generates correctly on the RTX 5080 (16 GB), experts on the MoE offload cache. + +## Ground truth (verified from the file header + llama.cpp origin/master) + +Single unsplit GGUF, 33,766,781,984 bytes. **No split-reader work needed for this file.** + +Metadata (`laguna.*`): +- block_count 48, embedding_length 3072, vocab 100352, context_length 262144 +- head_count per layer: `[48,72,72,72] * 12` (48 ⇒ full attention at `il%4==0`, 72 ⇒ SWA) +- head_count_kv 8 (uniform), key/value_length 128, sliding_window 512 +- rope full layers: dim 64 (partial), θ=500000, YaRN factor 32, orig_ctx 8192, + attn_factor 1.0, beta_fast 32, beta_slow 1 +- rope SWA layers: dim 128, θ=10000, **plain rope, no YaRN** (freq_scale 1.0) +- rms_eps 1e-6; leading_dense_block_count 1; expert_count 256, used 10, + expert_ff 1024, shared_expert_ff 1024, weights_norm true, weights_scale 2.5, + gating_func 2 (sigmoid); dense ffn 12288 +- tokenizer: gpt2 BPE, pre="laguna", bos 2, eos 2, eot 24, add_bos true, + chat template embedded +- Tensor quant types by role (per-tensor, "Unsloth Dynamic" = mixed): + - token_embd / output: Q4_K (untied head) + - attn q/k/v/o, attn_gate, dense ffn, shexp gate/up: Q5_K (layer 47: Q6_K) + - ffn_down (dense), down_shexp: Q6_K (one Q8_0) + - expert banks: gate/up IQ1_S (33 layers) or IQ2_XXS (14), down IQ3_XXS (45) or IQ4_XS (2) + - norms, router (ffn_gate_inp), exp_probs_b: F32 +- Layer schedule: layer 0 dense SwiGLU; layers 1..47 MoE (expert bank index = layer_id-1) + +Reference semantics (llama.cpp `src/models/laguna.cpp`, copy in +`/tmp/claude-1000/-home-lucas-ai-FreeToken/61016eb3-71fd-4d7b-97c6-816d5fa0839a/scratchpad/llamacpp-laguna.cpp`): +- QK RMSNorm at head_dim level (weight shape [128]) before rope, Qwen3-style +- attn_gate: `g = softplus(g_proj(pre-norm hidden))`, per-head ([3072, n_head_il]); + multiply attention output per head (broadcast over head_dim) **before** o_proj +- router: fp32 logits → sigmoid → +exp_probs_b bias for top-10 selection only → + gather **unbiased** sigmoid scores → renormalize → ×2.5 → weighted expert sum; + shared expert always added +- pre-attn and pre-ffn RMSNorm only (no post norms); final norm + untied lm_head + +CUDA kernel coverage (verified in `python/freetoken/kernel/csrc/gguf/`): mmvq, mmq +(Q4_K/Q5_K), dequantize and moe_vec for **all** the types above already exist and are +dispatched by ggml type id in `gguf_kernel.cu`. Only the Python-side tables gate them. + +Blockers found in Python side: +- `models/gguf/dequant.py`: only F32/F16/BF16/Q4_0/Q8_0/Q6_K in tables +- `layers/gguf.py`: `_MMVQ/_MMQ/_DEQUANT = {Q4_0, Q8_0, Q6_K}` +- `models/weight.py::load_q4_0_moe_expert_sources`: Q4_0-only expert banks +- `ModelConfig`: scalar `num_qo_heads` (Laguna needs per-layer 48/72) +- GGUF registry: only gemma4 + +Model download running in background task `bp45xxl7b` → +`~/.cache/huggingface/hub/models--unsloth--Laguna-S-2.1-GGUF/...`. +Header fixture: scratchpad `laguna_sparse.gguf` (metadata + tensor table, sparse data). + +## Orchestration + +Implementer: `pc-gpt-5-3-codex-spark` (weak — tasks kept small, exact files/symbols given). +Escalation (user 2026-08-23): on the next Spark failure (context death, usage limit, +botched task), switch the implementer seat to `pc-gpt-5-6-luna` permanently. +Reviewer per phase: `pc-gpt-5-6-terra` `[[effort: medium]]` — findings **with fixes**. +If a phase needs >2 review rounds, the lead (Fable) reviews and fixes directly. + +## Phase 1 — GGML quant-type plumbing (no model code) + +- [x] 1a. `models/gguf/dequant.py`: add ids Q4_K=12, Q5_K=13, IQ2_XXS=16, IQ3_XXS=18, + IQ1_S=19, IQ4_XS=23; BLOCK_SHAPE (256,144)/(256,176)/(256,66)/(256,98)/(256,50)/(256,136); + GGML_NAME entries. Reference dequant: delegate to gguf-py `gguf.quants.dequantize` + (verify installed gguf-py handles these types; else port only what tests need). +- [x] 1b. `layers/gguf.py`: extend `_MMVQ` with all six; `_MMQ` with Q4_K, Q5_K only; + `_DEQUANT` per actual `ggml_dequantize` switch coverage in `gguf_kernel.cu` + (verify each case id before adding). IQ types have no MMQ → prefill path must + fall back to dequant+bf16 matmul; confirm `fused_mul_mat_gguf` already routes that. +- [x] 1c. Tests `tests/kernels/test_gguf_quant_types.py`: per type — CUDA + `ggml_dequantize` vs gguf-py reference on random packed blocks; `mmvq` matvec vs + `F.linear` on dequantized weight; `moe_vec` for IQ1_S/IQ2_XXS/IQ3_XXS/IQ4_XS. +- [x] R1. Terra review round(s) → fixes → tests green. + +Phase 1 notes: gguf-py cannot quantize K/IQ formats, so tests use random +safe-scaled packed bytes with gguf-py dequant as reference; matmul refs model +q8_1 activation quantization; moe_vec asserted bit-exact vs mmvq. MMQ exists +for Q4_K/Q5_K only among the new types (IQ prefill falls back to dequant). +Lead (Fable) fixed tests after 2 implementer rounds. + +## Phase 2 — config, registry, tokenizer + +- [x] 2a. `models/config.py`: add optional per-layer qo-head counts + (`num_qo_heads_per_layer: tuple[int, ...] | None = None`) with accessor defaulting + to scalar `num_qo_heads`; audit consumers that assume the scalar for Q width + (KV geometry is uniform: 8 KV heads × 128 — KV pools unaffected). +- [x] 2b. New `models/laguna/` package: `config.py::parse_gguf_config` building + ModelConfig from the metadata above (full/SWA schedule from head-count array; + per-layer rope params; MoE fields; eos {2,24}); registry entries + (`gguf/config.py::GGUF_ARCH_TO_REGISTRY["laguna"]`, `register.py` ModelSpec + `LagunaGGUFForCausalLM`). +- [x] 2c. `models/gguf/tokenizer.py`: make it arch-generic where gemma-specific + (eos ids from `tokenizer.ggml.eos_token_id` + `eot_token_id`; verify transformers + converts gpt2/"laguna" pre BPE from GGUF; else load via tokenizer.ggml.* fields). +- [x] 2d. Tests: config parse from a metadata dict fixture; registry dispatch; + tokenizer eos/bos/chat-template presence (uses scratchpad header fixture). +- [x] R2. Terra review → fixes. + +Phase 2 notes: yarn scaling needs explicit attention_factor=1.0 (metadata +yarn_attn_factor) or freetoken defaults to ggml-incompatible mscale. Tokenizer +routes laguna->gpt2 converter with bos/eos strings materialized from ids. +Committed metadata-only fixture tests/fixtures/laguna-s-2.1-metadata.gguf (3.5MB). +Registry entry live; model stub raises a clear in-progress error until Phase 3/4. +Review: 1 round (blocker attention_factor + hardening), fixes by Spark. + +## Phase 3 — model modules (torch, quant-agnostic bf16 reference first) + +Phase 3 wiring notes (from Terra's R2 review): FlashInfer plans a single global +qo_head count and LinearQKVMerged/LinearOProj take one Q/O width — attention +modules must be built with `config.qo_heads(layer_id)` and pass it to QKV/O +projections, reshapes, gate logic, and the GQA/tensor-core decision; Triton +decode scratch may keep max-72 allocation but gets the real per-layer count at +invocation. attn backend for SWA is triton-only. + +- [x] 3a. `models/laguna/attention.py`: per-layer head count, QK head-dim RMSNorm, + partial-rope (dim 64 yarn) / full-rope (dim 128 plain) per layer type, per-head + softplus gate (fp32) applied before o_proj. Reuse `layers/rotary.py::get_rope` + (verify partial+yarn support), existing attention backend plumbing (SWA hybrid + pool as in gemma4). +- [x] 3b. `models/laguna/moe.py` + `mlp`: router semantics above; reuse existing MoE + offload machinery (find sigmoid+bias router precedent, e.g. glm/minimax/afmoe + style in repo); shared expert; dense layer 0. Expert bank index = layer_id-1. +- [x] 3c. `models/laguna/model.py`: decoder wiring, final norm, untied head; + `convert_laguna_to_gguf` pass swapping Linear/Embedding → GGUFLinear/GGUFEmbedding + **keeping each tensor's own ggml type** (unlike gemma4's uniform assumption). +- [x] 3d. Unit tests vs handwritten torch reference: one full layer, one SWA layer, + gate math, router top-10 selection/bias/renorm/scale, dense vs MoE layer. +- [x] R3. Terra review — semantics clean; loader-contract findings folded into Phase 4; + test-quality findings fixed in round 1. Module tests: 6, all green. + +## Phase 4 — weight loading + +- [x] 4a. `models/laguna/gguf.py::iter_gguf_weights`: name map (`blk.N.attn_q` etc. → + module params, table in ground truth above); norms/router/bias to f32/bf16; + quantized tensors kept packed with per-tensor type. + R3 contract (blocker): Engine collects model.state_dict() BEFORE iterating + weights, so DeferredGGUFLinear must be materialized earlier — add + `gguf_model_path` to ModelConfig (set from shim.model_path), have + convert_laguna_to_gguf read the tensor table (name→ggml type) and + materialize every swapped module up front using the same name map 4a uses + (share one mapping helper). Also (major): the untied lm_head must select + last-token rows on prefill like ParallelLMHead/GGUFTiedLMHead + (batch.attn_metadata.get_last_indices(batch.size)) — wrap DeferredGGUFLinear + in a LagunaGGUFLMHead that does the gather before the fused matmul. +- [x] 4b. Generalize `models/weight.py::load_q4_0_moe_expert_sources` → type-aware + expert-bank loader (per-layer, per-projection ggml type — gate/up/down differ); + wire type ids through the offload cache to `moe_vec` dispatch. +- [x] 4c. Test: synthetic tiny laguna GGUF (write with gguf-py) loads end-to-end with + dummy weights; every tensor consumed exactly once; unknown tensor → clear error. +- [x] R4. Sol (medium) review — clean except FTW-conversion blocker; laguna FTW + now rejects loudly (documented limitation). Kernel scope independently + cleared by a Terra sub-lens before the reviewer switch. + +## Phase 5 — e2e validation (lead-driven) + +- [ ] 5a. Load real file; `ft serve --kv-cache-dtype q8_0 --num-tokens 65536`; + one-token decode, short prefill; VRAM/expert-cache occupancy sane. +- [ ] 5b. Greedy next-token comparison vs llama.cpp on fixed prompts. +- [ ] 5c. Short needle + SWA-boundary (>512 tok) prompts; eos 2/24 stop behavior. +- [ ] Review section here. + +## Initial limitations + +TP=1 only; FTW conversion unsupported (serve the .gguf directly); text-only; no speculative decoding; split-GGUF variants (BF16 etc.) +out of scope; prefill for IQ expert types goes through moe_vec/dequant fallback +(optimize later if slow). diff --git a/tests/fixtures/laguna-s-2.1-metadata.gguf b/tests/fixtures/laguna-s-2.1-metadata.gguf new file mode 100644 index 00000000..0369262e Binary files /dev/null and b/tests/fixtures/laguna-s-2.1-metadata.gguf differ diff --git a/tests/kernels/test_gguf_quant_types.py b/tests/kernels/test_gguf_quant_types.py new file mode 100644 index 00000000..c35d1c48 --- /dev/null +++ b/tests/kernels/test_gguf_quant_types.py @@ -0,0 +1,172 @@ +"""Coverage tests for the GGUF quant types added for Laguna (Q4_K/Q5_K/IQ*). + +Strategy: gguf-py cannot *quantize* K/IQ formats, but it can *dequantize* any +packed bytes. So every test builds a weight from random-but-safe packed bytes +(fp16 scale fields masked small so the kernels' fp16 intermediates cannot +overflow -- real weights are O(1), random fp16 scales are not) and compares the +CUDA kernels against gguf-py's decode of the SAME bytes. +""" +from __future__ import annotations + +import numpy as np +import pytest +import torch +import torch.nn.functional as F + +if not torch.cuda.is_available(): + pytest.skip("CUDA required", allow_module_level=True) + +import gguf + +from freetoken.models.gguf.dequant import ( + BLOCK_SHAPE, + GGML_IQ1_S, + GGML_IQ2_S, + GGML_IQ2_XXS, + GGML_IQ3_XXS, + GGML_IQ4_XS, + GGML_Q3_K, + GGML_Q4_K, + GGML_Q5_K, + dequantize, +) + +TYPES = [ + GGML_Q3_K, GGML_Q4_K, GGML_Q5_K, + GGML_IQ1_S, GGML_IQ2_S, GGML_IQ2_XXS, GGML_IQ3_XXS, GGML_IQ4_XS, +] + + +def _packed_rows(qtype: int, rows: int, seed: int) -> np.ndarray: + """Random packed rows with every 16-bit field masked to a small positive + fp16 (exponent forced below 1.0) so any scale interpretation stays tiny and + fp16 kernel intermediates cannot overflow. Payload bits keep plenty of + entropy for the sub-block codes.""" + rng = np.random.default_rng(seed) + raw = rng.integers(0, 256, (rows, BLOCK_SHAPE[qtype][1]), dtype=np.uint8) + u16 = raw.view(np.uint16) if raw.shape[1] % 2 == 0 else None + if u16 is None: # odd row_bytes (IQ1_S is 50 -> even; guard anyway) + raw[:, 1::2] &= 0x3B + return raw + u16 &= np.uint16(0x3BFF) # clears sign, caps exponent -> |value| < 1 + return raw + + +def _reference(raw: np.ndarray, qtype: int) -> torch.Tensor: + return torch.from_numpy( + gguf.quants.dequantize(raw, gguf.GGMLQuantizationType(qtype)) + ).float() + + +def _q8_1_activations(x: torch.Tensor) -> torch.Tensor: + """Model the kernels' q8_1 activation quantization (int8 per 32-block with an + fp16 absmax/127 scale) so matmul references carry the same rounding.""" + blocks = x.float().reshape(x.shape[0], -1, 32) + scale = (blocks.abs().amax(dim=-1, keepdim=True) / 127.0).half().float() + q = torch.where(scale > 0, (blocks / scale).round().clamp(-127, 127), blocks) + return (q * scale).reshape(x.shape) + + +def _randn(shape, seed: int) -> torch.Tensor: + g = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(*shape, generator=g, device="cuda", dtype=torch.bfloat16) + + +@pytest.mark.parametrize("qtype", TYPES) +def test_python_reference_matches_gguf(qtype): + """The freetoken reference dequant (used by non-CUDA callers) agrees with + gguf-py on identical bytes.""" + raw = _packed_rows(qtype, rows=2, seed=qtype) + ours = dequantize(torch.from_numpy(raw), qtype, torch.float32).reshape(2, -1) + ref = _reference(raw, qtype).reshape(2, -1) + torch.testing.assert_close(ours, ref, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize("qtype", TYPES) +def test_bytes_cuda_matches_gguf_reference(qtype): + from freetoken.kernel.gguf import ggml_dequantize + + raw = _packed_rows(qtype, rows=4, seed=qtype) + packed = torch.from_numpy(raw).cuda() + got = ggml_dequantize(packed, qtype, 4, BLOCK_SHAPE[qtype][0], torch.float32).cpu() + ref = _reference(raw, qtype).reshape(4, -1) + torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-3) + + +@pytest.mark.parametrize("qtype", TYPES) +def test_mmvq_matches_linear(qtype): + from freetoken.kernel.gguf import ggml_mul_mat_vec_a8 + + block = BLOCK_SHAPE[qtype][0] + rows, cols = 8, 2 * block + raw = _packed_rows(qtype, rows=rows * (cols // block), seed=qtype + 1) + raw = np.ascontiguousarray(raw.reshape(rows, -1)) + packed = torch.from_numpy(raw).cuda() + w = _reference(raw, qtype).reshape(rows, cols).cuda() + x = _randn((1, cols), seed=qtype + 10) + got = ggml_mul_mat_vec_a8(packed, x, qtype, rows).float() + ref = F.linear(_q8_1_activations(x), w) + tol = 5e-3 * ref.abs().max().clamp(min=1.0) + assert (got.reshape(-1) - ref.reshape(-1)).abs().max() <= tol + + +@pytest.mark.parametrize("qtype", [GGML_Q3_K, GGML_Q4_K, GGML_Q5_K]) +def test_mmq_matches_linear(qtype): + from freetoken.kernel.gguf import ggml_mul_mat_a8 + + block = BLOCK_SHAPE[qtype][0] + rows, cols, batch = 8, 2 * block, 8 + raw = _packed_rows(qtype, rows=rows * (cols // block), seed=qtype + 2) + raw = np.ascontiguousarray(raw.reshape(rows, -1)) + packed = torch.from_numpy(raw).cuda() + w = _reference(raw, qtype).reshape(rows, cols).cuda() + x = _randn((batch, cols), seed=qtype + 20) + got = ggml_mul_mat_a8(packed, x, qtype, rows).float() + ref = F.linear(_q8_1_activations(x), w) + tol = 5e-3 * ref.abs().max().clamp(min=1.0) + assert (got - ref).abs().max() <= tol + + +@pytest.mark.parametrize( + "qtype", [GGML_Q3_K, GGML_Q4_K, GGML_IQ2_S, GGML_IQ2_XXS, GGML_IQ3_XXS, GGML_IQ1_S, GGML_IQ4_XS] +) +def test_moe_vec_matches_mmvq(qtype): + """moe_vec shares mmvq's vec_dot; per selected expert it must reproduce the + (reference-validated) mmvq result on that expert's rows bit-exactly.""" + from freetoken.kernel.gguf import ggml_moe_a8_vec, ggml_mul_mat_vec_a8 + + block = BLOCK_SHAPE[qtype][0] + experts, rows, cols, top_k = 4, 4, block, 2 + raw = _packed_rows(qtype, rows=experts * rows, seed=qtype + 3) + bank = np.ascontiguousarray(raw.reshape(experts, rows, -1)) + packed = torch.from_numpy(bank).cuda() + x = _randn((1, cols), seed=qtype + 30) + topk_ids = torch.tensor([[1, 3]], device="cuda", dtype=torch.int32) + out = ggml_moe_a8_vec(x, packed, topk_ids, top_k, qtype, rows, 1) + out = out.reshape(top_k, rows) + for j, e in enumerate([1, 3]): + one = torch.from_numpy(np.ascontiguousarray(bank[e])).cuda() + ref = ggml_mul_mat_vec_a8(one, x, qtype, rows).reshape(rows) + torch.testing.assert_close(out[j], ref, rtol=0.0, atol=0.0) + + +@pytest.mark.parametrize("qtype", [GGML_IQ1_S, GGML_IQ2_S, GGML_IQ3_XXS]) +def test_moe_vec_expert_stride_padded_bank(qtype): + """Mixed-quant banks store each expert's payload in the leading bytes of a + padded flat slot; expert_stride_bytes must reproduce the dense-bank result.""" + from freetoken.kernel.gguf import ggml_moe_a8_vec + + block = BLOCK_SHAPE[qtype][0] + experts, rows, cols, top_k = 4, 4, block, 2 + raw = _packed_rows(qtype, rows=experts * rows, seed=qtype + 40) + dense = torch.from_numpy(np.ascontiguousarray(raw.reshape(experts, rows, -1))).cuda() + payload = rows * raw.shape[1] // 1 # bytes per expert (row_bytes * rows) + payload = rows * dense.shape[2] + stride = payload + 64 # pad each expert slot + flat = torch.zeros(experts, stride, dtype=torch.uint8, device="cuda") + flat[:, :payload] = dense.reshape(experts, payload) + x = _randn((1, cols), seed=qtype + 41) + topk_ids = torch.tensor([[1, 3]], device="cuda", dtype=torch.int32) + ref = ggml_moe_a8_vec(x, dense, topk_ids, top_k, qtype, rows, 1) + got = ggml_moe_a8_vec(x, flat, topk_ids, top_k, qtype, rows, 1, stride) + torch.testing.assert_close(got, ref, rtol=0.0, atol=0.0) diff --git a/tests/models/test_laguna_config.py b/tests/models/test_laguna_config.py new file mode 100644 index 00000000..49f6611d --- /dev/null +++ b/tests/models/test_laguna_config.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +FIXTURE = Path(__file__).resolve().parent.parent / "fixtures" / "laguna-s-2.1-metadata.gguf" + + +def test_laguna_config_parse_and_attention_groups(): + from freetoken.models.gguf.config import build_gguf_shim + from freetoken.models.laguna.gguf import parse_gguf_config + + cfg = parse_gguf_config(build_gguf_shim(str(FIXTURE))) + + assert cfg.num_layers == 48 + assert cfg.num_qo_heads == 72 + assert all(v == 48 for v in cfg.num_qo_heads_per_layer[0::4]) + assert all(v == 72 for i, v in enumerate(cfg.num_qo_heads_per_layer) if i % 4 != 0) + assert cfg.qo_heads(0) == 48 + assert cfg.qo_heads(1) == 72 + assert cfg.num_kv_heads == 8 + assert cfg.head_dim == 128 + assert cfg.hidden_size == 3072 + assert cfg.vocab_size == 100352 + assert cfg.first_k_dense_replace == 1 + assert cfg.num_experts == 256 + assert cfg.num_experts_per_tok == 10 + assert cfg.moe_intermediate_size == 1024 + assert cfg.shared_expert_intermediate_size == 1024 + assert cfg.n_shared_experts == 1 + assert cfg.routed_scaling_factor == 2.5 + assert cfg.norm_topk_prob is True + assert cfg.tie_word_embeddings is False + assert cfg.model_type == "laguna" + assert cfg.moe_enabled is True + assert cfg.use_qk_norm is True + assert cfg.rms_norm_eps == pytest.approx(1e-6) + assert cfg.rotary_config.max_position == 262144 + + assert len(cfg.attention_groups) == 2 + full = cfg.attention_groups[0] + swa = cfg.attention_groups[1] + + assert tuple(full.layer_ids) == tuple(range(0, 48, 4)) + assert tuple(swa.layer_ids) == tuple(i for i in range(0, 48) if i not in set(full.layer_ids)) + assert swa.sliding_window == 512 + + assert full.rotary_config.rotary_dim == 64 + assert full.rotary_config.base == 500000.0 + assert full.rotary_config.scaling is not None + assert full.rotary_config.scaling["rope_type"] == "yarn" + assert full.rotary_config.scaling["factor"] == 32.0 + assert full.rotary_config.scaling["attention_factor"] == 1.0 + assert full.rotary_config.scaling["beta_fast"] == 32.0 + assert full.rotary_config.scaling["beta_slow"] == 1.0 + assert full.rotary_config.scaling["original_max_position_embeddings"] == 8192 + + assert swa.rotary_config.rotary_dim == 128 + assert swa.rotary_config.base == 10000.0 + assert swa.rotary_config.scaling is None + + assert cfg.gguf_embed_quant is None + + +def test_laguna_gguf_arch_registry_map(): + from freetoken.models.gguf.config import GGUF_ARCH_TO_REGISTRY + from freetoken.models import register + + assert GGUF_ARCH_TO_REGISTRY["laguna"] == "LagunaGGUFForCausalLM" + assert "LagunaGGUFForCausalLM" in getattr(register, "_MODEL_REGISTRY") + + +def test_laguna_tokenizer_and_eos_tokens(): + from freetoken.models.gguf.tokenizer import gguf_eos_token_ids, load_gguf_tokenizer + + tok = load_gguf_tokenizer(str(FIXTURE)) + assert tok.eos_token_id == 2 + assert gguf_eos_token_ids(str(FIXTURE), tok) == {2, 24} + assert tok.chat_template + + ids = tok("def foo(): return 1").input_ids + assert tok.decode(ids).endswith("def foo(): return 1") diff --git a/tests/models/test_laguna_modules.py b/tests/models/test_laguna_modules.py new file mode 100644 index 00000000..c6ac2144 --- /dev/null +++ b/tests/models/test_laguna_modules.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +import torch.nn.functional as F + +FIXTURE = Path(__file__).resolve().parent.parent / "fixtures" / "laguna-s-2.1-metadata.gguf" + + +@pytest.fixture(scope="module", autouse=True) +def _tp_one(): + from freetoken.distributed import set_tp_info + + set_tp_info(rank=0, size=1) + + +def _tiny_config(): + from freetoken.models.gguf.config import build_gguf_shim + from freetoken.models.laguna.gguf import parse_gguf_config + from freetoken.models.config import FullAttentionGroupConfig, RotaryConfig, SWAAttentionGroupConfig + + cfg = parse_gguf_config(build_gguf_shim(str(FIXTURE))) + full_rope = replace(cfg.attention_groups[0].rotary_config, head_dim=32, rotary_dim=32) + swa_rope = replace(cfg.attention_groups[1].rotary_config, head_dim=32, rotary_dim=32) + groups = ( + FullAttentionGroupConfig("full", (0, 4), 2, 32, full_rope), + SWAAttentionGroupConfig("swa", (1, 2, 3, 5, 6, 7), 2, 32, swa_rope, 512), + ) + return replace(cfg, num_layers=8, num_qo_heads=6, + num_qo_heads_per_layer=(4, 6, 6, 6, 4, 6, 6, 6), num_kv_heads=2, + head_dim=32, hidden_size=64, intermediate_size=96, moe_intermediate_size=16, + shared_expert_intermediate_size=16, vocab_size=128, num_experts=8, + num_experts_per_tok=3, gguf_embed_quant=None, + gguf_model_path=None, rotary_config=replace(cfg.rotary_config, head_dim=32, rotary_dim=32), + attention_groups=groups) + + +def test_attention_geometry(monkeypatch): + import freetoken.models.laguna.attention as attention_mod + from freetoken.models.laguna.attention import LagunaAttention + monkeypatch.setattr(attention_mod, "get_rope", lambda **kw: SimpleNamespace(rotary_dim=kw["rotary_dim"])) + cfg = _tiny_config(); full = LagunaAttention(cfg, 0); swa = LagunaAttention(cfg, 1) + assert full.num_qo_heads == 4 and full.q_proj.weight.shape == (128, 64) + assert full.k_proj.weight.shape == (64, 64) and full.v_proj.weight.shape == (64, 64) + assert full.rotary.rotary_dim == 32 and full.attn_spec.sliding_window is None + assert swa.num_qo_heads == 6 and swa.attn_spec.sliding_window == 512 + assert swa.gate_proj.weight.shape[0] == 6 + + +def test_attention_forward_applies_gate_independently(monkeypatch): + import freetoken.models.laguna.attention as attention_mod + from freetoken.models.laguna.attention import LagunaAttention + monkeypatch.setattr(attention_mod, "get_rope", lambda **kw: SimpleNamespace(rotary_dim=kw["rotary_dim"], forward=lambda p, q, k: (q, k))) + cfg = _tiny_config(); x = torch.randn(3, 64) + for layer_id, heads in ((0, 4), (1, 6)): + attn = LagunaAttention(cfg, layer_id); seen = {} + qv = torch.randn(3, heads * 32); kv = torch.randn(3, 2 * 32); vv = torch.randn(3, 2 * 32) + gate = torch.randn(3, heads); backend = torch.randn(3, heads, 32, dtype=torch.float16) + def proj(key, out): + class P: + def forward(self, z): seen[key] = z; return out + return P() + class O: + def forward(self, z): seen["o"] = z; return z + class N: + def forward_inplace(self, z): return z + attn.q_proj, attn.k_proj, attn.v_proj = proj("q", qv), proj("k", kv), proj("v", vv) + attn.gate_proj, attn.o_proj = proj("gate", gate), O(); attn.q_norm = attn.k_norm = N() + ctx = SimpleNamespace(batch=SimpleNamespace(positions=torch.arange(3)), attn_backend=SimpleNamespace(forward=lambda *a, **k: backend)) + monkeypatch.setattr(attention_mod, "get_global_ctx", lambda: ctx) + got = attn.forward(x) + assert seen["q"] is x and seen["k"] is x and seen["v"] is x and seen["gate"] is x + expected = backend.view(3, heads, 32) * F.softplus(gate.float()).unsqueeze(-1).to(backend.dtype) + torch.testing.assert_close(seen["o"], expected.reshape(3, heads * 32)); torch.testing.assert_close(got, seen["o"]) + + +def test_router_matches_reference_and_full_forward(): + from freetoken.models.laguna.moe import LagunaSparseMoeBlock + torch.manual_seed(0); blk = LagunaSparseMoeBlock.__new__(LagunaSparseMoeBlock) + blk.top_k, blk.num_experts, blk.norm_topk_prob, blk.routed_scaling_factor = 3, 8, True, 2.5 + blk.gate = SimpleNamespace(weight=torch.randn(8, 64)); blk.e_score_correction_bias = torch.randn(8) + x = torch.randn(5, 64); scores = (x @ blk.gate.weight.T).sigmoid(); ids = torch.topk(scores + blk.e_score_correction_bias, 3, dim=-1).indices + weights = scores.gather(-1, ids); weights = weights / (weights.sum(-1, keepdim=True) + 1e-20) * 2.5 + got_w, got_ids = blk._route(x); torch.testing.assert_close(got_ids, ids.to(torch.int32)); torch.testing.assert_close(got_w, weights) + blk.experts = SimpleNamespace(routed_forward=lambda h, w, i: h * w.sum(-1, keepdim=True)); blk.shared_experts = SimpleNamespace(forward=lambda h: h + 7) + torch.testing.assert_close(blk.forward(x), x * got_w.sum(-1, keepdim=True) + x + 7) + + +def test_decoder_residual_semantics(monkeypatch): + import freetoken.models.laguna.attention as attention_mod + from freetoken.models.laguna.model import LagunaDecoderLayer + monkeypatch.setattr(attention_mod, "get_rope", lambda **kw: SimpleNamespace(rotary_dim=kw["rotary_dim"])) + layer = LagunaDecoderLayer(_tiny_config(), 0); x = torch.randn(3, 64) + layer.input_layernorm = layer.ffn_norm = SimpleNamespace(forward=lambda z: z * 2) + layer.self_attn = SimpleNamespace(forward=lambda z: z + 1); layer.mlp = SimpleNamespace(forward=lambda z: z * 3) + h = x + (x * 2 + 1); torch.testing.assert_close(layer.forward(x), h + h * 2 * 3) + + +def test_laguna_mlp_fused_layout(monkeypatch): + import freetoken.models.laguna.moe as moe_mod + from freetoken.models.laguna.moe import LagunaMLP + monkeypatch.setattr(moe_mod, "silu_and_mul", lambda z: F.silu(z[..., :4]) * z[..., 4:]) + m = LagunaMLP(8, 4); torch.manual_seed(2); m.gate_up_proj.weight.copy_(torch.randn_like(m.gate_up_proj.weight)); m.down_proj.weight.copy_(torch.randn_like(m.down_proj.weight)) + x = torch.randn(3, 8); wg, wu = m.gate_up_proj.weight[:4], m.gate_up_proj.weight[4:] + m.gate_up_proj.forward = lambda z: F.linear(z, m.gate_up_proj.weight); m.down_proj.forward = lambda z: F.linear(z, m.down_proj.weight) + expected = F.linear(F.silu(F.linear(x, wg)) * F.linear(x, wu), m.down_proj.weight) + torch.testing.assert_close(m.forward(x), expected) + + +def test_deferred_gguf_linear_q8(): + from freetoken.models.gguf.dequant import GGML_Q8_0, row_bytes + from freetoken.models.laguna.gguf import DeferredGGUFLinear + layer = DeferredGGUFLinear(64, 32) + with pytest.raises(AssertionError): layer.forward(torch.randn(2, 64)) + layer.materialize(GGML_Q8_0); assert layer.qweight.shape == (32, row_bytes(64, GGML_Q8_0)) + if not torch.cuda.is_available(): pytest.skip("CUDA required for fused GGUF forward") + import gguf + rng = np.random.default_rng(1); weight = rng.standard_normal((32, 64), dtype=np.float32); packed = gguf.quants.quantize(weight, gguf.GGMLQuantizationType.Q8_0) + layer.qweight = torch.from_numpy(np.ascontiguousarray(packed)).cuda(); x = torch.randn(2, 64, device="cuda", dtype=torch.bfloat16); got = layer.forward(x).float(); blocks = x.float().reshape(2, -1, 32); scale = (blocks.abs().amax(dim=-1, keepdim=True) / 127).half().float(); aq = torch.where(scale > 0, (blocks / scale).round().clamp(-127, 127), blocks).mul(scale).reshape_as(x); ref = F.linear(aq, torch.from_numpy(gguf.quants.dequantize(packed, gguf.GGMLQuantizationType.Q8_0)).float().cuda()); assert (got - ref).abs().max() <= 5e-3 * ref.abs().max().clamp(min=1.0) diff --git a/tests/models/test_laguna_weights.py b/tests/models/test_laguna_weights.py new file mode 100644 index 00000000..461e743d --- /dev/null +++ b/tests/models/test_laguna_weights.py @@ -0,0 +1,204 @@ +"""End-to-end Laguna GGUF weight loading over a synthetic tiny checkpoint. + +Writes a real (tiny) laguna GGUF with gguf-py -- Q8_0 quantized projections and +expert banks, F32 norms/router -- then exercises config parsing, the name map, +``iter_gguf_weights`` (every tensor consumed exactly once, fused buffers built), +deferred materialization, and the mixed-type expert-bank loader. +""" +from __future__ import annotations + +import numpy as np +import pytest +import torch + +import gguf + +import freetoken.distributed.info as di +from freetoken.models.gguf.dequant import GGML_Q8_0, row_bytes + +# Tiny geometry: 4 layers (full at 0), heads 4 full / 6 swa, kv 2, head_dim 32. +L, H, FF = 4, 64, 96 +HEADS = [4, 6, 6, 6] +KV, HD = 2, 32 +E, TOPK, I, SHI = 8, 3, 32, 32 +VOCAB = 128 + + +@pytest.fixture(scope="module") +def tiny_gguf(tmp_path_factory): + path = str(tmp_path_factory.mktemp("laguna") / "tiny-laguna.gguf") + w = gguf.GGUFWriter(path, "laguna") + w.add_block_count(L) + w.add_context_length(4096) + w.add_embedding_length(H) + w.add_feed_forward_length(FF) + w.add_head_count(HEADS) + w.add_head_count_kv(KV) + w.add_key_length(HD) + w.add_value_length(HD) + w.add_layer_norm_rms_eps(1e-6) + w.add_sliding_window(512) + w.add_rope_freq_base(500000.0) + w.add_rope_dimension_count(16) + # SWA rope mirrors + yarn keys (raw kv names, mirroring the real file). + w.add_float32("laguna.rope.freq_base_swa", 10000.0) + w.add_uint32("laguna.rope.dimension_count_swa", HD) + w.add_string("laguna.rope.scaling.type", "yarn") + w.add_float32("laguna.rope.scaling.factor", 32.0) + w.add_uint32("laguna.rope.scaling.original_context_length", 8192) + w.add_float32("laguna.rope.scaling.yarn_attn_factor", 1.0) + w.add_float32("laguna.rope.scaling.yarn_beta_fast", 32.0) + w.add_float32("laguna.rope.scaling.yarn_beta_slow", 1.0) + w.add_expert_count(E) + w.add_expert_used_count(TOPK) + w.add_expert_feed_forward_length(I) + w.add_expert_shared_feed_forward_length(SHI) + w.add_bool("laguna.expert_weights_norm", True) + w.add_float32("laguna.expert_weights_scale", 2.5) + w.add_uint32("laguna.expert_gating_func", 2) + w.add_uint32("laguna.leading_dense_block_count", 1) + w.add_uint32("laguna.vocab_size", VOCAB) + # Minimal gpt2 tokenizer metadata so the shim can size the vocab. + w.add_tokenizer_model("gpt2") + w.add_token_list([f"" for i in range(VOCAB)]) + w.add_token_types([1] * VOCAB) + w.add_token_merges([]) + w.add_bos_token_id(2) + w.add_eos_token_id(2) + w.add_uint32("tokenizer.ggml.eot_token_id", 24) + + rng = np.random.default_rng(0) + q8 = gguf.GGMLQuantizationType.Q8_0 + + def quant(name, rows, cols): + data = rng.standard_normal((rows, cols)).astype(np.float32) + w.add_tensor(name, gguf.quants.quantize(data, q8), raw_dtype=q8) + + def f32(name, *shape): + w.add_tensor(name, rng.standard_normal(shape).astype(np.float32)) + + quant("token_embd.weight", VOCAB, H) + quant("output.weight", VOCAB, H) + f32("output_norm.weight", H) + for i in range(L): + p = f"blk.{i}." + nh = HEADS[i] + f32(p + "attn_norm.weight", H) + f32(p + "attn_q_norm.weight", HD) + f32(p + "attn_k_norm.weight", HD) + f32(p + "ffn_norm.weight", H) + quant(p + "attn_q.weight", nh * HD, H) + quant(p + "attn_k.weight", KV * HD, H) + quant(p + "attn_v.weight", KV * HD, H) + quant(p + "attn_output.weight", H, nh * HD) + quant(p + "attn_gate.weight", nh, H) + if i == 0: + quant(p + "ffn_gate.weight", FF, H) + quant(p + "ffn_up.weight", FF, H) + quant(p + "ffn_down.weight", H, FF) + else: + f32(p + "ffn_gate_inp.weight", E, H) + f32(p + "exp_probs_b.bias", E) + quant(p + "ffn_gate_shexp.weight", SHI, H) + quant(p + "ffn_up_shexp.weight", SHI, H) + quant(p + "ffn_down_shexp.weight", H, SHI) + for role, rows, cols in ( + ("ffn_gate_exps", I, H), + ("ffn_up_exps", I, H), + ("ffn_down_exps", H, I), + ): + data = rng.standard_normal((E, rows, cols)).astype(np.float32) + w.add_tensor( + p + role + ".weight", + gguf.quants.quantize(data.reshape(E * rows, cols), q8).reshape(E, rows, -1), + raw_dtype=q8, + ) + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + return path + + +@pytest.fixture(autouse=True) +def _tp1(): + try: + di.get_tp_info() + except RuntimeError: + di.set_tp_info(0, 1) + + +def _config(path): + from freetoken.models.gguf.config import build_gguf_shim + from freetoken.models.laguna.gguf import parse_gguf_config + + return parse_gguf_config(build_gguf_shim(path)) + + +def test_config_and_expert_types(tiny_gguf): + cfg = _config(tiny_gguf) + assert cfg.num_layers == L and cfg.num_qo_heads == 6 + assert cfg.gguf_embed_quant == GGML_Q8_0 + assert cfg.expert_quant == "gguf" and cfg.moe_weight_format == "gguf" + assert cfg.gguf_expert_types == ((GGML_Q8_0, GGML_Q8_0),) * (L - 1) + + +def test_iter_weights_complete_and_fused(tiny_gguf): + from freetoken.models.laguna.gguf import iter_gguf_weights + + got = dict(iter_gguf_weights(tiny_gguf, "cpu", include_moe_experts=False, include_non_moe=True)) + # split q/k/v per layer (mixed-type files forbid fusing) + for i, nh in enumerate(HEADS): + assert got[f"model.layers.{i}.self_attn.q_proj.qweight"].shape == (nh * HD, row_bytes(H, GGML_Q8_0)) + assert got[f"model.layers.{i}.self_attn.k_proj.qweight"].shape == (KV * HD, row_bytes(H, GGML_Q8_0)) + assert got[f"model.layers.{i}.self_attn.v_proj.qweight"].shape == (KV * HD, row_bytes(H, GGML_Q8_0)) + t = got["model.layers.0.mlp.gate_up_proj.qweight"] + assert t.shape == (2 * FF, row_bytes(H, GGML_Q8_0)) + t = got["model.layers.1.mlp.shared_experts.gate_up_proj.qweight"] + assert t.shape == (2 * SHI, row_bytes(H, GGML_Q8_0)) + assert got["model.layers.1.mlp.gate.weight"].dtype == torch.float32 + assert got["model.layers.1.mlp.e_score_correction_bias"].dtype == torch.float32 + assert got["model.layers.0.ffn_norm.weight"].dtype == torch.bfloat16 + assert got["lm_head.qweight"].shape == (VOCAB, row_bytes(H, GGML_Q8_0)) + + +def test_unknown_tensor_rejected(tiny_gguf, tmp_path): + from freetoken.models.laguna.gguf import iter_gguf_weights + + w = gguf.GGUFWriter(str(tmp_path / "bad.gguf"), "laguna") + w.add_block_count(1) + w.add_tensor("blk.0.mystery.weight", np.zeros((4, 4), dtype=np.float32)) + w.write_header_to_file(); w.write_kv_data_to_file(); w.write_tensors_to_file(); w.close() + with pytest.raises(ValueError, match="mystery"): + list(iter_gguf_weights(str(tmp_path / "bad.gguf"), "cpu", + include_moe_experts=False, include_non_moe=True)) + + +def test_expert_bank_loader(tiny_gguf): + from freetoken.models.laguna.gguf import _expert_bank_geometry, load_gguf_expert_sources + + cfg = _config(tiny_gguf) + banks = load_gguf_expert_sources(tiny_gguf, cfg) + gu_s, dn_s = _expert_bank_geometry(cfg) + assert len(banks["gate_up"]) == L - 1 and len(banks["down"]) == L - 1 + for t in banks["gate_up"]: + assert t.shape == (E, gu_s) and t.dtype == torch.uint8 + # payload bytes decode to the source values via gguf-py + half = I * row_bytes(H, GGML_Q8_0) + blob = banks["gate_up"][0][:, : 2 * half] + dec = gguf.quants.dequantize( + np.ascontiguousarray(blob.reshape(E * 2 * I, -1).numpy()), + gguf.GGMLQuantizationType.Q8_0, + ) + assert np.isfinite(dec).all() and dec.std() > 0.5 # real data, not padding + + +def test_deferred_materialization(tiny_gguf): + cfg = _config(tiny_gguf) + from freetoken.models.laguna.gguf import DeferredGGUFLinear + + # conversion materializes from the file's tensor table; emulate on one module + mod = DeferredGGUFLinear(H, 6 * HD) + mod.materialize(GGML_Q8_0) + assert mod.qweight.shape == (6 * HD, row_bytes(H, GGML_Q8_0)) + assert cfg.gguf_model_path == tiny_gguf diff --git a/tests/server/test_parser_auto_selection.py b/tests/server/test_parser_auto_selection.py index 78b1663e..7af30430 100644 --- a/tests/server/test_parser_auto_selection.py +++ b/tests/server/test_parser_auto_selection.py @@ -79,6 +79,27 @@ def test_only_the_families_without_a_tool_format_get_the_generic_fallback(): assert fell_through == NO_DEDICATED_TOOL_FORMAT +def test_laguna_uses_poolside_v1_for_reasoning_and_tools(): + assert _inferred("LagunaGGUFForCausalLM") == ("poolside_v1", "poolside_v1") + + +def test_explicit_poolside_choices_are_accepted(): + config = _Config({"architectures": ["LlamaForCausalLM"], "torch_dtype": "bfloat16"}) + with patch("freetoken.utils.cached_load_hf_config", lambda _path: config): + args, _ = parse_args( + [ + "--model", + ANON_PATH, + "--tool-call-parser", + "poolside_v1", + "--reasoning-parser", + "poolside_v1", + ] + ) + assert args.tool_call_parser == "poolside_v1" + assert args.reasoning_parser == "poolside_v1" + + def test_qwen3_5_is_not_shadowed_by_the_generic_qwen_branch(): """The cascade matches substrings in order, so the specific arm has to come first: a bare ``"qwen" -> qwen25`` reached earlier would swallow every later Qwen and lose its tool format.""" diff --git a/tests/server/test_poolside_v1_parsers.py b/tests/server/test_poolside_v1_parsers.py new file mode 100644 index 00000000..014a0135 --- /dev/null +++ b/tests/server/test_poolside_v1_parsers.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +import json +from types import SimpleNamespace + +from freetoken.server.api_models import Tool +from freetoken.server.function_call_parser import FunctionCallParser +from freetoken.server.generation import GenSpec, _split_reasoning +from freetoken.server.reasoning_parser import ReasoningParser + +TOOLS: list[Tool] = [ + Tool.model_validate( + { + "type": "function", + "function": { + "name": "write_file", + "parameters": { + "type": "object", + "properties": { + "path": {"type": "string"}, + "content": {"type": "string"}, + "mode": {"type": "integer"}, + }, + }, + }, + } + ) +] + + +def _tool_block(content: str = " def f():\n return 1\n") -> str: + return ( + "write_file" + "path/tmp/a.py" + f"content{content}" + "mode420" + "" + ) + + +def test_poolside_v1_reasoning_implicit_open_non_stream() -> None: + parser = ReasoningParser("poolside_v1", force_reasoning=True) + reasoning, content = parser.parse_non_stream( + "I should inspect the file.Here is the answer." + ) + assert reasoning == "I should inspect the file." + assert content == "Here is the answer." + + +def test_poolside_v1_reasoning_stream_marker_can_split_across_chunks() -> None: + parser = ReasoningParser("poolside_v1", force_reasoning=True) + reasoning_parts: list[str] = [] + content_parts: list[str] = [] + for chunk in ["Need ", "the file.Answer", "."]: + reasoning, content = parser.parse_stream_chunk(chunk) + reasoning_parts.append(reasoning) + content_parts.append(content) + reasoning, content = parser.flush() + reasoning_parts.append(reasoning) + content_parts.append(content) + assert "".join(reasoning_parts) == "Need the file." + assert "".join(content_parts) == "Answer." + + +def _poolside_state() -> SimpleNamespace: + return SimpleNamespace(config=SimpleNamespace(reasoning_parser="poolside_v1")) + + +def test_poolside_v1_default_generation_matches_implicit_open_template() -> None: + spec = GenSpec(messages=[], sampling_params=None) + reasoning, content = _split_reasoning( + "Need the file.Answer.", spec, _poolside_state() + ) + assert reasoning == "Need the file." + assert content == "Answer." + + +def test_poolside_v1_explicit_thinking_disable_keeps_visible_content() -> None: + spec = GenSpec( + messages=[], + sampling_params=None, + chat_template_kwargs={"enable_thinking": False}, + ) + reasoning, content = _split_reasoning("Answer.", spec, _poolside_state()) + assert reasoning == "" + assert content == "Answer." + + +def test_poolside_v1_tool_non_stream_preserves_string_whitespace() -> None: + source = " def f():\n return 1\n" + parser = FunctionCallParser(TOOLS, tool_call_parser="poolside_v1") + result = parser.parse_non_stream("Calling it. " + _tool_block(source)) + assert result.normal_text == "Calling it. " + assert len(result.calls) == 1 + assert result.calls[0].name == "write_file" + assert json.loads(result.calls[0].parameters) == { + "path": "/tmp/a.py", + "content": source, + "mode": 420, + } + + +def test_poolside_v1_tool_stream_preserves_string_whitespace_and_types() -> None: + source = " def f():\n return 1\n" + parser = FunctionCallParser(TOOLS, tool_call_parser="poolside_v1") + names: list[str] = [] + arg_fragments: list[str] = [] + normal: list[str] = [] + text = "Calling it. " + _tool_block(source) + for i in range(0, len(text), 3): + visible, calls = parser.parse_stream_chunk(text[i : i + 3]) + normal.append(visible) + for call in calls: + if call.name: + names.append(call.name) + arg_fragments.append(call.parameters) + normal.append(parser.finish_stream()) + assert "".join(normal) == "Calling it. " + assert "".join(names) == "write_file" + assert json.loads("".join(arg_fragments)) == { + "path": "/tmp/a.py", + "content": source, + "mode": 420, + } + + +def _schema_tools() -> list[Tool]: + return [ + Tool.model_validate( + { + "type": "function", + "function": { + "name": "typed", + "parameters": { + "type": "object", + "properties": { + "exact": {"type": "string"}, + "enum_only": {"enum": ["a", "b"]}, + "typeless": {}, + "integer": {"type": "integer"}, + "number": {"type": "number"}, + "boolean": {"type": "boolean"}, + "object": {"type": "object"}, + "array": {"type": "array"}, + }, + }, + }, + } + ) + ] + + +def _schema_block() -> str: + return ( + "typed" + "exact keep me \n" + "enum_only a " + "typeless 5 " + "integer 7 " + "boolean true " + 'object {"x": 1} ' + "array [1, 2] " + "" + ) + + +EXPECTED_SCHEMA_ARGS = { + "exact": " keep me \n", + "enum_only": "a", + "typeless": 5, + "integer": 7, + "boolean": True, + "object": {"x": 1}, + "array": [1, 2], +} + + +def test_poolside_v1_exact_string_schema_semantics_non_stream() -> None: + parser = FunctionCallParser(_schema_tools(), tool_call_parser="poolside_v1") + result = parser.parse_non_stream(_schema_block()) + assert len(result.calls) == 1 + assert json.loads(result.calls[0].parameters) == EXPECTED_SCHEMA_ARGS + + +def test_poolside_v1_exact_string_schema_semantics_one_char_streaming() -> None: + parser = FunctionCallParser(_schema_tools(), tool_call_parser="poolside_v1") + fragments: list[str] = [] + for char in _schema_block(): + _visible, calls = parser.parse_stream_chunk(char) + fragments.extend(call.parameters for call in calls) + parser.finish_stream() + assert json.loads("".join(fragments)) == EXPECTED_SCHEMA_ARGS + + +def _glm_coercion_block() -> str: + return ( + "typed" + "integer1.0" + "number1e3" + "booleanTRUE" + "object{'x': 1}" + "" + ) + + +EXPECTED_GLM_COERCIONS = { + "integer": "1.0", + "number": 1000, + "boolean": True, + "object": {"x": 1}, +} + + +def test_poolside_v1_non_string_schema_uses_glm_coercion_non_stream() -> None: + parser = FunctionCallParser(_schema_tools(), tool_call_parser="poolside_v1") + result = parser.parse_non_stream(_glm_coercion_block()) + args = json.loads(result.calls[0].parameters) + assert args == EXPECTED_GLM_COERCIONS + assert type(args["number"]) is int + + +def test_poolside_v1_non_string_schema_uses_glm_coercion_streaming() -> None: + parser = FunctionCallParser(_schema_tools(), tool_call_parser="poolside_v1") + fragments: list[str] = [] + for char in _glm_coercion_block(): + _visible, calls = parser.parse_stream_chunk(char) + fragments.extend(call.parameters for call in calls) + parser.finish_stream() + args = json.loads("".join(fragments)) + assert args == EXPECTED_GLM_COERCIONS + assert type(args["number"]) is int + + +def test_poolside_v1_preserves_escaped_source_under_one_char_streaming() -> None: + source = ' print("C:\\\\tmp") \n' + parser = FunctionCallParser(TOOLS, tool_call_parser="poolside_v1") + fragments: list[str] = [] + for char in _tool_block(source): + _visible, calls = parser.parse_stream_chunk(char) + fragments.extend(call.parameters for call in calls) + parser.finish_stream() + assert json.loads("".join(fragments))["content"] == source + + +def test_glm47_string_arguments_keep_legacy_trimming() -> None: + source = " padded \n" + parser = FunctionCallParser(TOOLS, tool_call_parser="glm47") + result = parser.parse_non_stream(_tool_block(source)) + assert json.loads(result.calls[0].parameters)["content"] == source.strip() diff --git a/tests/server/test_streaming_model_matrix.py b/tests/server/test_streaming_model_matrix.py index 43922021..90431896 100644 --- a/tests/server/test_streaming_model_matrix.py +++ b/tests/server/test_streaming_model_matrix.py @@ -75,6 +75,9 @@ "glm47": ( "readfilePath/tmp/test_calc.py" ), + "poolside_v1": ( + "readfilePath/tmp/test_calc.py" + ), "gemma4": '<|tool_call>call:read{filePath:<|"|>/tmp/test_calc.py<|"|>}', "minimax": ( '' @@ -107,6 +110,7 @@ "qwen25": ["", ""], "qwen3_coder": ["", "", "", ""], + "poolside_v1": ["", "", ""], "gemma4": ["<|tool_call>", ""], "minimax": ["", "[", ""), "qwen": ("qwen25", "qwen3", "", ""), "glm4.7": ("glm47", "glm", "", ""), + "laguna": ("poolside_v1", "poolside_v1", "", ""), "minimax-m2": ("minimax", "minimax", "", ""), # M3 adaptive mode: the model opens itself (enabled mode pre-opens it # in the template; the parser then runs with force_reasoning=True instead). @@ -428,6 +433,9 @@ def test_reasoning_chat_entrypoint(name): "glm47": ( f"readfilePath{LONG_VALUE}" ), + "poolside_v1": ( + f"readfilePath{LONG_VALUE}" + ), "gemma4": f'<|tool_call>call:read{{filePath:<|"|>{LONG_VALUE}<|"|>}}', "minimax": ( f'{LONG_VALUE}' @@ -664,7 +672,7 @@ def test_empty_arguments_call_emitted_exactly_once(tool, block): @pytest.mark.parametrize( "family", - ["qwen25", "qwen3_coder", "glm47", "gemma4", "minimax", "minimax_m3", + ["qwen25", "qwen3_coder", "glm47", "poolside_v1", "gemma4", "minimax", "minimax_m3", "deepseekv32", "gpt_oss", "muse_glimmer"], ) def test_call_then_trailing_text_in_one_chunk_keeps_order(family):