diff --git a/python/freetoken/engine/cache_budget.py b/python/freetoken/engine/cache_budget.py index ab7c0a9f..08019151 100644 --- a/python/freetoken/engine/cache_budget.py +++ b/python/freetoken/engine/cache_budget.py @@ -6,6 +6,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING from freetoken.utils import div_ceil @@ -23,9 +24,97 @@ def expert_bytes_per_slot(sources: dict[str, "list[torch.Tensor]"]) -> int: """ # marlin/b12x gate_up/down alpha scales are fixed [L*E] residency (do not scale # with cache_size), so they are intentionally excluded from the per-slot growth term. - # tensor[0].numel() is the per-row element count (one expert slot); see the matching - # slot-byte idiom in kvcache/linear_state_pool.py and kvcache/dsv4_paged_pool.py. - return sum(t[0][0].numel() * t[0].element_size() for t in sources.values()) + # The GPU slot cache has one stride per bank, chosen from that bank's largest layer + # row. See the matching slot-byte calculation in OffloadMoeCache. + return sum( + max(layer[0].numel() * layer.element_size() for layer in per_layer) + for per_layer in sources.values() + ) + + +@dataclass(frozen=True) +class GeometryPoolPlan: + layer_ids: tuple[int, ...] + row_bytes: tuple[int, ...] + slots: int + + +def plan_geometry_pool_slots( + row_bytes_by_layer: list[tuple[int, ...]], + *, + legacy_cache_size: int, + num_experts: int, + top_k: int, + max_decode_batch: int, +) -> tuple[GeometryPoolPlan, ...] | None: + """Partition fixed max-stride bank arenas into exact-geometry decode pools. + + ``legacy_cache_size`` remains the external budget denomination. Each bank owns + ``legacy_cache_size * max(layer_row_bytes)`` bytes; every planned class must fit + all bank constraints independently. + """ + if not row_bytes_by_layer: + return () + num_banks = len(row_bytes_by_layer[0]) + if num_banks == 0 or any(len(row) != num_banks for row in row_bytes_by_layer): + raise ValueError("every layer must describe the same non-empty bank set") + if ( + legacy_cache_size <= 0 + or num_experts <= 0 + or top_k <= 0 + or max_decode_batch <= 0 + ): + raise ValueError( + "cache size, experts, top_k, and decode batch must be positive" + ) + + grouped: dict[tuple[int, ...], list[int]] = {} + for layer_id, rows in enumerate(row_bytes_by_layer): + if any(value <= 0 for value in rows): + raise ValueError("geometry row bytes must be positive") + grouped.setdefault(tuple(rows), []).append(layer_id) + classes = [(rows, tuple(layer_ids)) for rows, layer_ids in grouped.items()] + budgets = [ + legacy_cache_size * max(rows[bank] for rows in row_bytes_by_layer) + for bank in range(num_banks) + ] + floor = min(num_experts, top_k * max_decode_batch) + slots = [floor] * len(classes) + + def used(bank: int) -> int: + return sum(slots[i] * classes[i][0][bank] for i in range(len(classes))) + + if any(used(bank) > budgets[bank] for bank in range(num_banks)): + return None + + targets = [ + min(len(layer_ids) * num_experts, len(layer_ids) * top_k * max_decode_batch) + for _, layer_ids in classes + ] + caps = [len(layer_ids) * num_experts for _, layer_ids in classes] + + def affordable(index: int) -> bool: + rows = classes[index][0] + return all( + used(bank) + rows[bank] <= budgets[bank] for bank in range(num_banks) + ) + + def fill(limits: list[int]) -> None: + while True: + candidates = [ + i for i in range(len(classes)) if slots[i] < limits[i] and affordable(i) + ] + if not candidates: + return + index = min(candidates, key=lambda i: (slots[i] / limits[i], i)) + slots[index] += 1 + + fill(targets) + fill(caps) + return tuple( + GeometryPoolPlan(layer_ids=layer_ids, row_bytes=rows, slots=slots[index]) + for index, (rows, layer_ids) in enumerate(classes) + ) def net_cache_budget_bytes( diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cd6505d2..38fa8a49 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -613,6 +613,14 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: quant_format=banks.quant_format, decode_target=decode_target, hybrid_max_fetch=config.moe_hybrid_max_fetch, + geometry_pool_top_k=getattr( + config.model_config, "num_experts_per_tok", 0 + ), + geometry_pool_max_batch=max( + config.max_running_req, + config.cuda_graph_max_bs or 0, + 1, + ), ) # before set_bank_sources: the residency validation and the copy plan's skip of non-pinned layers key on the CPU-layer set cache.cpu_layer_ids = cpu_layer_ids @@ -1375,7 +1383,12 @@ def override(attr: str, value: Any): # this is dangerous, use with caution from freetoken.moe.cpu_executor import compiled_extension_supports _act = getattr(model_config, "hidden_act", "silu") - if not _cpu_moe_act_ok: + if bench_fmt == "gguf": + logger.info_rank0( + "benchbw profile recommends hybrid, but GGUF experts do not have a " + "CPU executor; staying on offload" + ) + elif not _cpu_moe_act_ok: logger.info_rank0( f"benchbw profile recommends hybrid, but the CPU MoE executor does not " f"support this model's expert activation " 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/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23e..e675c2b5 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -7,6 +7,7 @@ #include #include +#include #include namespace device { @@ -560,3 +561,179 @@ struct MultiIndexCopyKernel { device.unwrap())(kernel, params); } }; + + +// Multi-bank copy with independent payload, destination stride, and source stride. +// This is required when compact per-layer host rows feed a max-stride GPU slot cache. +struct MultiStridedIndexCopyParams { + const int64_t* __restrict__ dst_ptrs; + const int64_t* __restrict__ src_ptrs; + const int64_t* __restrict__ copy_bytes; + const int64_t* __restrict__ dst_row_strides; + const int64_t* __restrict__ src_row_strides; + const void* __restrict__ dst_indices; + const void* __restrict__ src_indices; + const int64_t* __restrict__ valid_length; + int64_t length; + int num_banks; +}; + +template +__global__ __launch_bounds__(kNumThreads) void fast_index_copy_multi_strided( + const __grid_constant__ MultiStridedIndexCopyParams p +) { + const int b = static_cast(blockIdx.x / kBlocksPerBank); + if (b >= p.num_banks) { + return; + } + const int blk = static_cast(blockIdx.x % kBlocksPerBank); + const auto* src = reinterpret_cast(p.src_ptrs[b]); + auto* dst = reinterpret_cast(p.dst_ptrs[b]); + const int64_t bytes = p.copy_bytes[b]; + const int64_t dst_stride = p.dst_row_strides[b]; + const int64_t src_stride = p.src_row_strides[b]; + const int64_t requested = p.valid_length ? p.valid_length[0] : p.length; + const int64_t n = requested < 0 ? 0 : (requested > p.length ? p.length : requested); + if (bytes <= 0 || (bytes & 15) != 0 || bytes > src_stride || bytes > dst_stride) { + return; + } + const int64_t units = bytes >> 4; + const int64_t total = n * units; + const auto* di = static_cast(p.dst_indices); + const auto* si = static_cast(p.src_indices); + const int64_t grid_stride = static_cast(kBlocksPerBank) * kNumThreads; + for (int64_t u = static_cast(blk) * kNumThreads + threadIdx.x; + u < total; u += grid_stride) { + const int64_t row = u / units; + const int64_t col = (u - row * units) << 4; + const int64_t pd = static_cast(di[row]); + const int64_t ps = static_cast(si[row]); + const uint4 v = *reinterpret_cast(src + ps * src_stride + col); + *reinterpret_cast(dst + pd * dst_stride + col) = v; + } +} + +template +struct MultiStridedIndexCopyKernel { + static void run( + tvm::ffi::TensorView dst_ptrs, + tvm::ffi::TensorView src_ptrs, + tvm::ffi::TensorView copy_bytes, + tvm::ffi::TensorView dst_row_strides, + tvm::ffi::TensorView src_row_strides, + tvm::ffi::TensorView dst_indices, + tvm::ffi::TensorView src_indices, + tvm::ffi::Optional num_indices + ) { + using namespace host; + auto device = SymbolicDevice{}; + auto B = SymbolicSize{"num_banks"}; + auto L = SymbolicSize{"indices length"}; + auto ptr_dtype = SymbolicDType{}; + auto indices_dtype = SymbolicDType{}; + auto num_indices_dtype = SymbolicDType{}; + + TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + .verify(dst_ptrs).verify(src_ptrs).verify(copy_bytes) + .verify(dst_row_strides).verify(src_row_strides); + TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) + .verify(dst_indices).verify(src_indices); + + const int64_t* valid_length = nullptr; + if (num_indices.has_value()) { + TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) + .verify(num_indices.value()); + valid_length = static_cast(num_indices.value().data_ptr()); + } + + const int num_banks = static_cast(B.unwrap()); + const auto params = MultiStridedIndexCopyParams{ + static_cast(dst_ptrs.data_ptr()), + static_cast(src_ptrs.data_ptr()), + static_cast(copy_bytes.data_ptr()), + static_cast(dst_row_strides.data_ptr()), + static_cast(src_row_strides.data_ptr()), + dst_indices.data_ptr(), + src_indices.data_ptr(), + valid_length, + static_cast(L.unwrap()), + num_banks, + }; + const auto use_int32 = indices_dtype.unwrap().bits == 32; + const auto kernel = use_int32 + ? fast_index_copy_multi_strided + : fast_index_copy_multi_strided; + LaunchKernel(static_cast(kBlocksPerBank) * num_banks, kNumThreads, + device.unwrap())(kernel, params); + } +}; + + +// Whole compact GGUF layer -> padded slot cache without a payload-sized CUDA +// staging allocation. ``src_ptr`` is the UVA device alias of registered host +// memory, not necessarily the host virtual address. +struct StridedRowsCopyParams { + uint8_t* __restrict__ dst; + const uint8_t* __restrict__ src; + int64_t rows; + int64_t copy_bytes; + int64_t dst_stride; + int64_t src_stride; +}; + +template +__global__ __launch_bounds__(kNumThreads) void fast_copy_strided_rows( + const __grid_constant__ StridedRowsCopyParams p +) { + const int64_t units = p.copy_bytes >> 4; + const int64_t total = p.rows * units; + const int64_t step = static_cast(gridDim.x) * kNumThreads; + for (int64_t u = static_cast(blockIdx.x) * kNumThreads + threadIdx.x; + u < total; u += step) { + const int64_t row = u / units; + const int64_t col = (u - row * units) << 4; + const uint4 v = *reinterpret_cast(p.src + row * p.src_stride + col); + *reinterpret_cast(p.dst + row * p.dst_stride + col) = v; + } +} + +template +struct StridedRowsCopyKernel { + static void run( + tvm::ffi::TensorView dst, + int64_t src_ptr, + int64_t src_stride, + int64_t copy_bytes, + int64_t rows + ) { + using namespace host; + auto device = SymbolicDevice{}; + auto R = SymbolicSize{"rows"}; + auto D = SymbolicSize{"destination stride"}; + TensorMatcher({R, D}).with_dtype().with_device(device).verify(dst); + assert(rows == static_cast(R.unwrap())); + assert(copy_bytes > 0 && (copy_bytes & 15) == 0); + assert(src_stride >= copy_bytes && (src_stride & 15) == 0); + const int64_t dst_stride = static_cast(D.unwrap()); + assert(dst_stride >= copy_bytes && (dst_stride & 15) == 0); + assert(src_ptr != 0 && (src_ptr & 15) == 0); + assert((reinterpret_cast(dst.data_ptr()) & 15) == 0); + + const auto params = StridedRowsCopyParams{ + static_cast(dst.data_ptr()), + reinterpret_cast(src_ptr), + rows, + copy_bytes, + dst_stride, + src_stride, + }; + const int64_t total_units = rows * (copy_bytes >> 4); + const int64_t wanted = (total_units + kNumThreads - 1) / kNumThreads; + const int64_t capped = wanted < static_cast(kBlocksPerBank) + ? wanted : static_cast(kBlocksPerBank); + const auto blocks = static_cast(capped < 1 ? 1 : capped); + LaunchKernel(blocks, kNumThreads, device.unwrap())( + fast_copy_strided_rows, params + ); + } +}; diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py index 1aaa1303..302f0206 100644 --- a/python/freetoken/kernel/fast_index_copy.py +++ b/python/freetoken/kernel/fast_index_copy.py @@ -185,6 +185,93 @@ def fast_index_copy_multi_jit( module.launch(dst_ptrs, src_ptrs, feat_bytes, dst_indices, src_indices, num_indices) +@lru_cache(maxsize=None) +def _jit_fast_index_copy_multi_strided_module( + *, num_threads: int, blocks_per_bank: int +) -> Module: + args = make_cpp_args(num_threads, blocks_per_bank) + return load_jit( + "fast_index_copy_multi_strided", + *args, + cuda_files=["fast_index_copy.cuh"], + cuda_wrappers=[("launch", f"&MultiStridedIndexCopyKernel<{args}>::run")], + ) + + +def fast_index_copy_multi_strided_jit( + dst_ptrs: torch.Tensor, + src_ptrs: torch.Tensor, + copy_bytes: torch.Tensor, + dst_row_strides: torch.Tensor, + src_row_strides: torch.Tensor, + dst_indices: torch.Tensor, + src_indices: torch.Tensor, + num_indices: torch.Tensor | None = None, + *, + num_threads: int = 1024, + blocks_per_bank: int = 8, +) -> None: + """Copy compact source rows into independently strided destination rows.""" + if _skip_fast_index_copy_enabled(): + return + module = _jit_fast_index_copy_multi_strided_module( + num_threads=num_threads, blocks_per_bank=blocks_per_bank + ) + module.launch( + dst_ptrs, + src_ptrs, + copy_bytes, + dst_row_strides, + src_row_strides, + dst_indices, + src_indices, + num_indices, + ) + + +@lru_cache(maxsize=None) +def _jit_fast_index_copy_rows_strided_module( + *, num_threads: int, blocks_per_bank: int +) -> Module: + args = make_cpp_args(num_threads, blocks_per_bank) + return load_jit( + "fast_index_copy_rows_strided", + *args, + cuda_files=["fast_index_copy.cuh"], + cuda_wrappers=[("launch", f"&StridedRowsCopyKernel<{args}>::run")], + ) + + +def fast_index_copy_rows_strided_jit( + destination: torch.Tensor, + source: torch.Tensor, + *, + num_threads: int = 1024, + blocks_per_bank: int = 8, +) -> None: + """Copy compact pinned-host rows into a wider 2-D CUDA destination.""" + from freetoken.kernel.pinned import device_ptr + + assert destination.is_cuda and destination.dim() == 2 + assert source.device.type == "cpu" and source.dim() == 2 and source.is_contiguous() + assert destination.dtype == source.dtype == torch.uint8 + assert destination.size(0) == source.size(0) + copy_bytes = source.size(1) + assert destination.size(1) >= copy_bytes + if _skip_fast_index_copy_enabled(): + return + module = _jit_fast_index_copy_rows_strided_module( + num_threads=num_threads, blocks_per_bank=blocks_per_bank + ) + module.launch( + destination, + int(device_ptr(source)), + source.stride(0) * source.element_size(), + copy_bytes * source.element_size(), + source.size(0), + ) + + def update_copy_flag_jit(sync_flag: torch.Tensor, delta: int) -> None: assert sync_flag.is_cuda assert sync_flag.numel() == 1 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..40ab620a 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -316,7 +316,7 @@ def _decode_routed( hidden_states, topk_weights, topk_ids, - views=cache.bank_views(), + views=cache.bank_views(layer_id=self.layer_id), n=None, alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, @@ -363,7 +363,7 @@ def _decode_hybrid( hidden_states, gpu_w, gpu_slots, - views=cache.bank_views(), + views=cache.bank_views(layer_id=self.layer_id), n=None, alphas=cache.alphas_for_slots(self.layer_id), is_prefill=False, @@ -531,6 +531,54 @@ def _expert_gemm( return fused_experts_gguf_q4_0( hidden_states, gate_up, down, topk_weights, topk_ids, self.activation ) + if fmt == "gguf": + # GGUF may mix native quantized layers with raw BF16 layers. Quantized + # rows use ggml MMVQ; BF16 rows are byte views over the same padded slot + # caches and go through the regular dense expert kernels. + from freetoken.models.gguf.dequant import GGML_BF16 + + gate_up, down = views + gu_type, dn_type = self.gguf_gate_up_type, self.gguf_down_type + if GGML_BF16 in (gu_type, dn_type): + if gu_type != GGML_BF16 or dn_type != GGML_BF16: + raise ValueError( + "mixed BF16/quantized projections within one GGUF expert layer" + ) + gu_bytes = self.gguf_gate_up_rows * self.hidden_size * 2 + dn_bytes = self.gguf_down_rows * self.intermediate_size * 2 + if gate_up.shape[1] != gu_bytes or down.shape[1] != dn_bytes: + raise ValueError( + "BF16 GGUF expert rows require exact dense cache strides " + f"(gate_up={gate_up.shape[1]}/{gu_bytes}, " + f"down={down.shape[1]}/{dn_bytes})" + ) + dense_gate_up = gate_up.view(torch.bfloat16).view( + gate_up.shape[0], self.gguf_gate_up_rows, self.hidden_size + ) + dense_down = down.view(torch.bfloat16).view( + down.shape[0], self.gguf_down_rows, self.intermediate_size + ) + impl = fused_experts_impl if is_prefill else fused_experts_decode_impl + return impl( + hidden_states, + dense_gate_up, + dense_down, + topk_weights, + topk_ids, + self.activation, + self.apply_router_weight_on_input, + ) + # Mixed-type quantized GGUF experts (per-layer quant types, flat padded + # slot banks): geometry comes from the layer's type attributes. + from freetoken.moe.fused_gguf import fused_experts_gguf + + return fused_experts_gguf( + hidden_states, gate_up, down, topk_weights, topk_ids, self.activation, + gate_up_type=gu_type, + down_type=dn_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..8abc13f9 --- /dev/null +++ b/python/freetoken/models/laguna/gguf.py @@ -0,0 +1,542 @@ +"""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) -> list[tuple[int, int]]: + """Compact flat-slot strides per MoE layer, each 64-byte 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 + def align(n: int) -> int: + return (n + 63) // 64 * 64 + + return [ + (align(2 * I * row_bytes(H, gu)), align(H * row_bytes(I, dn))) + for gu, dn in config.gguf_expert_types + ] + + +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 compact flat ``[E, stride]`` uint8 tensor per MoE layer (bank + index = layer_id - first_k_dense_replace). The GPU slot cache chooses the maximum + per-bank stride, while host RAM stores only each layer's real packed payload. + """ + from freetoken.models.gguf.dequant import row_bytes + from freetoken.models.gguf.reader import iter_gguf_tensors + from freetoken.moe.host_banks import HostBank, LayerCompletionTracker, PinPipeline + + _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 + geometry = _expert_bank_geometry(config) + hb = { + "gate_up": [HostBank((E, gu_stride), torch.uint8) for gu_stride, _ in geometry], + "down": [HostBank((E, dn_stride), torch.uint8) for _, dn_stride in geometry], + } + 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 HostBank, pin_banks + + E = config.num_experts + L = len(config.gguf_expert_types or ()) + assert L, "laguna dummy expert banks need gguf_expert_types" + geometry = _expert_bank_geometry(config) + hb = { + "gate_up": [HostBank((E, gu_stride), torch.uint8) for gu_stride, _ in geometry], + "down": [HostBank((E, dn_stride), torch.uint8) for _, dn_stride in geometry], + } + 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..7a02660a 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -1,6 +1,5 @@ from __future__ import annotations -import math import os from dataclasses import dataclass from typing import Iterator @@ -45,6 +44,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) @@ -93,6 +95,45 @@ MARLIN_MAX_CACHE_SIZE = 992 +class _GeometryPoolState: + """Decode-only LRU state backed by exact-width views into legacy arenas.""" + + def __init__( + self, + *, + num_layers: int, + num_experts: int, + cache_size: int, + device: torch.device, + layer_ids: tuple[int, ...], + row_bytes: tuple[int, ...], + bank_views: tuple[torch.Tensor, ...], + ) -> None: + self.num_layers = num_layers + self.num_experts = num_experts + self.cache_size = cache_size + self.device = device + self.layer_ids = layer_ids + self.row_bytes = row_bytes + self.bank_views = bank_views + self.slot_for_id = torch.full( + (num_layers, num_experts), -1, dtype=torch.int32, device=device + ) + self.id_of_slot = torch.full((cache_size,), -1, dtype=torch.int32, device=device) + self.usage = torch.zeros((cache_size,), dtype=torch.int64, device=device) + self.step = torch.zeros((), dtype=torch.int64, device=device) + self.active_mask = torch.zeros((num_experts,), dtype=torch.int32, device=device) + plan_slots = max(num_experts, cache_size) + self.evict_slots = torch.empty((plan_slots,), dtype=torch.int32, device=device) + self.src_indices = torch.empty((plan_slots,), dtype=torch.int32, device=device) + self.num_indices = torch.zeros((1,), dtype=torch.int64, device=device) + self.lru_stats = torch.zeros((num_layers, N_STATS), dtype=torch.int64, device=device) + self.collect_stats = False + self.copy_dst_ptrs: torch.Tensor | None = None + self.copy_src_ptrs: dict[int, torch.Tensor] = {} + self.copy_feat_bytes: torch.Tensor | None = None + + @dataclass class OffloadMoeCache: num_layers: int @@ -134,6 +175,10 @@ class OffloadMoeCache: # pcie_bw / cpu_bw ratio so the PCIe fetch and the CPU overflow GEMV take equal # time (perfect overlap): fetched : cpu = pcie : cpu - pcie. hybrid_fetch_fraction: float = 0.0 + # Heterogeneous GGUF only. A positive top-k activates exact-geometry decode + # pools carved from the existing max-stride byte arenas. + geometry_pool_top_k: int = 0 + geometry_pool_max_batch: int = 1 def __post_init__(self) -> None: policy_ids = {"lru": 0} @@ -246,6 +291,14 @@ def __post_init__(self) -> None: self._copy_dst_ptrs: torch.Tensor | None = None self._copy_src_ptrs: list[torch.Tensor] | None = None self._copy_feat_bytes: torch.Tensor | None = None + self._copy_payload_bytes: list[torch.Tensor] | None = None + self._copy_src_row_strides: list[torch.Tensor] | None = None + self._copy_dst_row_strides: torch.Tensor | None = None + self.has_heterogeneous_rows = False + self._geometry_pools: list[_GeometryPoolState] = [] + self._geometry_pool_for_layer: dict[int, _GeometryPoolState] = {} + self._pending_geometry_pool: _GeometryPoolState | None = None + self._pending_geometry_prefill = False # The layer whose misses ensure_experts/materialize_layer staged last; consumed # by copy_missing to pick the per-layer source (part of the same pending-copy # state as evict_slots/src_indices/num_indices). @@ -275,6 +328,133 @@ def __post_init__(self) -> None: self.prefill_hit_rows = 0 self.prefill_total_rows = 0 + def _allocate_bank_cache(self, per_layer: list[torch.Tensor]) -> torch.Tensor: + head = per_layer[0] + row_numel = [source[0].numel() for source in per_layer] + if len(set(row_numel)) == 1 and all(source.shape == head.shape for source in per_layer): + shape = (self.cache_size, *head.shape[1:]) + else: + shape = (self.cache_size, max(row_numel)) + return torch.empty(shape, dtype=head.dtype, device=self.device) + + @staticmethod + def _copy_compact_layer( + destination: torch.Tensor, + source: torch.Tensor, + *, + registered_host: bool = False, + ) -> None: + """Copy compact source rows into the leading bytes/elements of padded rows.""" + dst = destination.reshape(destination.shape[0], -1) + src = source.reshape(source.shape[0], -1) + if registered_host and dst.is_cuda and dst.shape[1] != src.shape[1]: + from freetoken.kernel.fast_index_copy import ( + fast_index_copy_rows_strided_jit, + ) + + fast_index_copy_rows_strided_jit(dst, src) + return + dst[:, : src.shape[1]].copy_(src, non_blocking=True) + + def _init_geometry_pools(self) -> None: + self._geometry_pools = [] + self._geometry_pool_for_layer = {} + self._pending_geometry_pool = None + if ( + self.quant_format != "gguf" + or not self.has_heterogeneous_rows + or self.decode_target != "gpu" + or self.geometry_pool_top_k <= 0 + or self._unpinned_layers + ): + return + + from freetoken.engine.cache_budget import plan_geometry_pool_slots + + rows_by_layer = [ + tuple( + self.bank_sources[name][layer_id][0].numel() + * self.bank_sources[name][layer_id].element_size() + for name in self.bank_schema + ) + for layer_id in range(self.num_layers) + ] + plan = plan_geometry_pool_slots( + rows_by_layer, + legacy_cache_size=self.cache_size, + num_experts=self.num_experts, + top_k=self.geometry_pool_top_k, + max_decode_batch=self.geometry_pool_max_batch, + ) + if plan is None: + logger.warning( + "GGUF geometry decode floors do not fit the MoE byte arenas; " + "using the unified max-stride cache" + ) + return + + offsets = [0] * len(self.bank_schema) + for entry in plan: + views = [] + for bank_index, name in enumerate(self.bank_schema): + arena = self.bank_caches[name].reshape(-1) + dtype_bytes = arena.element_size() + row_bytes = entry.row_bytes[bank_index] + if row_bytes % dtype_bytes: + raise ValueError("geometry row bytes must align to the bank dtype") + row_elements = row_bytes // dtype_bytes + pool_elements = entry.slots * row_elements + offset = offsets[bank_index] + view = arena.narrow(0, offset, pool_elements).view(entry.slots, row_elements) + views.append(view) + offsets[bank_index] += pool_elements + pool = _GeometryPoolState( + num_layers=self.num_layers, + num_experts=self.num_experts, + cache_size=entry.slots, + device=self.device, + layer_ids=entry.layer_ids, + row_bytes=entry.row_bytes, + bank_views=tuple(views), + ) + self._geometry_pools.append(pool) + for layer_id in entry.layer_ids: + self._geometry_pool_for_layer[layer_id] = pool + + if self.device.type == "cuda": + from freetoken.kernel.pinned import device_ptr + + for pool in self._geometry_pools: + pool.copy_dst_ptrs = torch.tensor( + [view.data_ptr() for view in pool.bank_views], + dtype=torch.int64, + device=self.device, + ) + pool.copy_feat_bytes = torch.tensor( + pool.row_bytes, dtype=torch.int64, device=self.device + ) + for layer_id in pool.layer_ids: + pool.copy_src_ptrs[layer_id] = torch.tensor( + [ + device_ptr(self.bank_sources[name][layer_id]) + for name in self.bank_schema + ], + dtype=torch.int64, + device=self.device, + ) + self.prefill_hit_d2d = False + detail = ", ".join( + f"layers={len(pool.layer_ids)} slots={pool.cache_size} rows={pool.row_bytes}" + for pool in self._geometry_pools + ) + logger.info("GGUF geometry decode pools: %s", detail) + + def geometry_pool_sizes(self) -> dict[int, int]: + return { + layer_id: pool.cache_size + for layer_id, pool in self._geometry_pool_for_layer.items() + } + def set_bank_sources( self, sources: dict[str, list[torch.Tensor]], @@ -286,7 +466,9 @@ def set_bank_sources( Every bank is a list of ``num_layers`` tensors, one ``[num_experts, ...]`` per layer (independent allocations, so each layer can carry its own host attributes); each slot cache mirrors the bank's row shape and dtype as one - unified GPU pool. The row layouts are produced by the weight loaders / + unified GPU pool. Heterogeneous banks use a flat cache whose stride is the + largest layer row while each host layer remains compact. The row layouts are + produced by the weight loaders / repackers (see ``_BANK_SCHEMAS`` and :mod:`freetoken.moe.nvfp4_backends`) -- the cache machinery is layout-agnostic and just moves rows. @@ -318,24 +500,48 @@ def set_bank_sources( ) self._unpinned_layers = unpinned self.layer_residency = list(residency) + self.has_heterogeneous_rows = False for name in self.bank_schema: per_layer = sources[name] assert len(per_layer) == self.num_layers, (name, len(per_layer)) head = per_layer[0] - for layer_id, source in enumerate(per_layer): - assert source.is_contiguous(), f"bank {name!r} layer {layer_id} must be contiguous" - assert source.size(0) == self.num_experts, (name, layer_id, source.shape) - assert source.shape == head.shape and source.dtype == head.dtype, ( - name, layer_id, source.shape, source.dtype, + if not all(source.size(0) == self.num_experts for source in per_layer): + raise ValueError( + f"bank {name!r} must contain {self.num_experts} experts per layer" + ) + if self.quant_format == "gguf": + if not all( + source.dim() == 2 + and source.dtype == torch.uint8 + and source.is_contiguous() + for source in per_layer + ): + raise ValueError( + f"GGUF bank {name!r} requires 2-D contiguous uint8 rows" + ) + elif not all( + source.shape == head.shape + and source.dtype == head.dtype + and source.is_contiguous() + for source in per_layer + ): + raise ValueError( + f"bank {name!r} requires uniform per-layer shapes and dtypes with " + f"contiguous storage for quant_format={self.quant_format!r}" ) self.bank_sources[name] = list(per_layer) - self.bank_caches[name] = torch.empty( - (self.cache_size, *head.shape[1:]), - dtype=head.dtype, - device=self.device, + self.bank_caches[name] = self._allocate_bank_cache(per_layer) + self.has_heterogeneous_rows |= any( + source.shape != head.shape for source in per_layer[1:] ) self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] self._build_copy_plan() + if self.has_heterogeneous_rows and self.device.type == "cuda" and not self._copy_fused_ok: + raise ValueError( + "heterogeneous GGUF rows require 16-byte-aligned source payloads, " + "strides, and mapped host addresses" + ) + self._init_geometry_pools() if self.prefill_overlap: self._init_prefill_overlap_buffers() @@ -351,23 +557,38 @@ def _build_copy_plan(self) -> None: self._copy_dst_ptrs = None self._copy_src_ptrs = None self._copy_feat_bytes = None + self._copy_payload_bytes = None + self._copy_src_row_strides = None + self._copy_dst_row_strides = None self._copy_dst_ptrs_host: list[int] = [] self._copy_src_ptrs_host: list[list[int]] = [] self._copy_feat_bytes_host: list[int] = [] self._gather_bank_ids: list[int] = [] self._gather_dst_ptrs: torch.Tensor | None = None self._gather_feat_bytes: torch.Tensor | None = None - if not _FUSED_COPY or self.device.type != "cuda" or not self.banks: + if self.device.type != "cuda" or not self.banks: + return + # FREETOKEN_FUSED_COPY disables the uniform multi-bank optimization. Compact + # heterogeneous rows still require the strided correctness path. + if not _FUSED_COPY and not self.has_heterogeneous_rows: return from freetoken.kernel.pinned import device_ptr dst_ptrs, feats = [], [] layer_src_ptrs = [[] for _ in range(self.num_layers)] + layer_payloads = [[] for _ in range(self.num_layers)] + layer_src_strides = [[] for _ in range(self.num_layers)] for per_layer, cache in self.banks: - feat = math.prod(per_layer[0].shape[1:]) * per_layer[0].element_size() + feat = cache[0].numel() * cache.element_size() if feat % 16 != 0 or cache.data_ptr() % 16 != 0: return # leave fused disabled; copy_missing uses the per-bank path for layer_id, source in enumerate(per_layer): + payload = source[0].numel() * source.element_size() + src_stride = source.stride(0) * source.element_size() + if payload > feat or payload % 16 != 0 or src_stride % 16 != 0: + return + layer_payloads[layer_id].append(payload) + layer_src_strides[layer_id].append(src_stride) if layer_id in self._unpinned_layers: # unregistered layer: no device alias exists, and the row is never consumed (CPU decode; pageable prefill) # a 0 placeholder keeps the descriptor shape @@ -388,6 +609,15 @@ def _build_copy_plan(self) -> None: for ptrs in layer_src_ptrs ] self._copy_feat_bytes = torch.tensor(feats, dtype=torch.int64, device=self.device) + self._copy_payload_bytes = [ + torch.tensor(payloads, dtype=torch.int64, device=self.device) + for payloads in layer_payloads + ] + self._copy_src_row_strides = [ + torch.tensor(strides, dtype=torch.int64, device=self.device) + for strides in layer_src_strides + ] + self._copy_dst_row_strides = self._copy_feat_bytes self._copy_dst_ptrs_host = dst_ptrs self._copy_src_ptrs_host = layer_src_ptrs self._copy_feat_bytes_host = feats @@ -439,7 +669,11 @@ def rebuild(self, cache_size: int) -> None: self._prefill_buffer_layer = [None, None] self._prefill_buffer_released = [True, True] self._prefill_buffer_has_release_event = [False, False] - # 2. Drop old GPU tensors (free-before-alloc). + # 2. Drop old GPU tensors (free-before-alloc), including geometry aliases. + self._geometry_pools = [] + self._geometry_pool_for_layer = {} + self._pending_geometry_pool = None + self._pending_geometry_prefill = False self.banks = [] self.bank_caches = {} self.cache_size = cache_size @@ -448,12 +682,10 @@ def rebuild(self, cache_size: int) -> None: torch.cuda.empty_cache() # 3. Reallocate the slot cache from the retained host sources. for name in self.bank_schema: - head = self.bank_sources[name][0] - self.bank_caches[name] = torch.empty( - (cache_size, *head.shape[1:]), dtype=head.dtype, device=self.device - ) + self.bank_caches[name] = self._allocate_bank_cache(self.bank_sources[name]) self.banks = [(self.bank_sources[n], self.bank_caches[n]) for n in self.bank_schema] self._build_copy_plan() # slot caches were reallocated -> refresh fused-copy addrs + self._init_geometry_pools() # 4. Reallocate cache_size-shaped bookkeeping; reset the slot map (cold start). self.slot_for_id.fill_(-1) self.id_of_slot = torch.full((cache_size,), -1, dtype=torch.int32, device=self.device) @@ -552,10 +784,18 @@ def alphas_for_layer(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor] | hi = lo + self.num_experts return self.gate_up_alpha[lo:hi], self.down_alpha[lo:hi] - def bank_views(self, n: int | None = None) -> tuple[torch.Tensor, ...]: - """Per-bank cache views in registration order: the full ``[S]`` slot cache - (decode), or its first ``n`` slots (materialized layer).""" + def bank_views( + self, + n: int | None = None, + *, + layer_id: int | None = None, + ) -> tuple[torch.Tensor, ...]: + """Per-bank decode pool or leading full-layer prefill overlay.""" assert self.banks, "set_bank_sources must register the banks first" + if n is None and self._geometry_pools: + if layer_id is None: + raise ValueError("layer_id is required for geometry decode pools") + return self._geometry_pool_for_layer[layer_id].bank_views if n is None: return tuple(cache for _, cache in self.banks) return tuple(cache[:n] for _, cache in self.banks) @@ -604,6 +844,11 @@ def begin_prefill(self) -> None: return self._prefill_buffer_layer = [None, None] self._prefill_buffer_released = [True, True] + if self._geometry_pools: + from freetoken.moe.offload_kernels import reset_cache + + for pool in self._geometry_pools: + reset_cache(pool) if self.prefill_copy_stream is not None: # Fence this prefill's copy-stream work behind everything already enqueued # on the compute stream. The release/ready events only order against the @@ -641,7 +886,9 @@ def prefetch_prefill_layer(self, layer_id: int) -> None: def copy() -> None: self._invalidate_prefill_buffer(buffer_id) for (per_layer, _), buffer in zip(self.banks, self.prefill_bank_buffers): - buffer[buffer_id].copy_(per_layer[layer_id], non_blocking=True) + self._copy_compact_layer( + buffer[buffer_id], per_layer[layer_id], registered_host=True + ) if self._prefill_hit_d2d_active: self._prefetch_split(layer_id, buffer_id) @@ -670,6 +917,8 @@ def _hit_d2d_usable(self) -> bool: reason = "prefill overlap buffers are not initialized for this device" elif _skip_fast_index_copy_enabled(): reason = "FREETOKEN_SKIP_FAST_INDEX_COPY is set (the hit gather would be a no-op)" + elif self.has_heterogeneous_rows: + reason = "heterogeneous source rows require the full-layer copy path" elif not self._copy_fused_ok: reason = "the fused copy plan is unavailable (bank alignment or FREETOKEN_FUSED_COPY=0)" elif self.cache_size <= 2 * self.num_experts: @@ -807,7 +1056,15 @@ def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) self._pending_src_layer = layer_id self._pending_whole_layer = False - ensure_experts(self, layer_id, expert_ids) + self._pending_geometry_prefill = False + pool = self._geometry_pool_for_layer.get(layer_id) + if pool is not None: + pool.collect_stats = self.collect_stats + self._pending_geometry_pool = pool + ensure_experts(pool, layer_id, expert_ids) + else: + self._pending_geometry_pool = None + ensure_experts(self, layer_id, expert_ids) def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None: """Capped-fetch LRU for the hybrid backend. @@ -831,16 +1088,30 @@ def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None ) def materialize_layer(self, layer_id: int) -> None: - from freetoken.moe.offload_kernels import materialize_layer + from freetoken.moe.offload_kernels import materialize_layer, reset_cache self._pending_src_layer = layer_id self._pending_whole_layer = True + if self._geometry_pools: + for pool in self._geometry_pools: + reset_cache(pool) + self._pending_geometry_pool = None + self._pending_geometry_prefill = True + return + self._pending_geometry_pool = None + self._pending_geometry_prefill = False materialize_layer(self, layer_id) def reset(self) -> None: from freetoken.moe.offload_kernels import reset_cache - reset_cache(self) + if self._geometry_pools: + for pool in self._geometry_pools: + reset_cache(pool) + else: + reset_cache(self) + self._pending_geometry_pool = None + self._pending_geometry_prefill = False # Per-expert recency is not cache_size-shaped, so reset_cache leaves it alone; wipe # it here so a new sequence starts with cold hybrid fetch priorities. self.expert_recency.fill_(-1) @@ -849,6 +1120,8 @@ def reset_stats(self) -> None: self.prefill_hit_rows = 0 self.prefill_total_rows = 0 self.lru_stats.zero_() + for pool in self._geometry_pools: + pool.lru_stats.zero_() self.stat_missing.zero_() self.stat_active.zero_() self.stat_calls.zero_() @@ -883,15 +1156,52 @@ def record_decode_stats_hybrid(self, layer_id: int) -> None: self.stat_steps_layer[layer_id] += 1 def decode_miss_stats(self) -> dict: + lru_stats = self.lru_stats if self.decode_target == "hybrid": active = int(self.stat_active.item()) missing = int(self.stat_missing.item()) calls = int(self.stat_calls.item()) + transferred_by_layer = self.stat_fetched_layer.tolist() + fetched = int(self.stat_fetched.item()) else: - active, missing, calls = (int(x) for x in self.lru_stats.sum(0)) - fetched = int(self.stat_fetched.item()) + if self._geometry_pools: + lru_stats = torch.stack( + [pool.lru_stats for pool in self._geometry_pools] + ).sum(0) + active, missing, calls = (int(x) for x in lru_stats.sum(0)) + missing_by_layer = lru_stats[:, Stat.MISS].tolist() + if self.decode_target == "gpu": + transferred_by_layer = missing_by_layer + elif self.decode_target == "cpu": + transferred_by_layer = [ + 0 if layer_id in self.cpu_layer_ids else rows + for layer_id, rows in enumerate(missing_by_layer) + ] + else: + transferred_by_layer = [0] * self.num_layers + fetched = sum(transferred_by_layer) + bytes_h2d = 0 + if self.bank_sources: + payload_bytes_by_layer = [ + sum( + per_layer[layer_id][0].numel() + * per_layer[layer_id].element_size() + for per_layer in self.bank_sources.values() + ) + for layer_id in range(self.num_layers) + ] + bytes_h2d = sum( + rows * payload_bytes + for rows, payload_bytes in zip( + transferred_by_layer, payload_bytes_by_layer, strict=True + ) + ) return { "layer_calls": calls, + "requested_rows": active, + "miss_rows": missing, + "hit_rows": active - missing, + "bytes_h2d": bytes_h2d, "active_per_layer": (active / calls) if calls else 0.0, "missing_per_layer": (missing / calls) if calls else 0.0, "miss_rate": (missing / active) if active else 0.0, @@ -916,10 +1226,24 @@ def decode_miss_stats_per_layer(self) -> dict: steps = self.stat_steps_layer.tolist() missing = self.stat_missing_layer.tolist() active = self.stat_active_layer.tolist() + fetched = self.stat_fetched_layer.tolist() else: - cols = self.lru_stats.t().tolist() + lru_stats = self.lru_stats + if self._geometry_pools: + lru_stats = torch.stack( + [pool.lru_stats for pool in self._geometry_pools] + ).sum(0) + cols = lru_stats.t().tolist() active, missing, steps = cols[Stat.ACTIVE], cols[Stat.MISS], cols[Stat.CALLS] - fetched = self.stat_fetched_layer.tolist() + if self.decode_target == "gpu": + fetched = missing + elif self.decode_target == "cpu": + fetched = [ + 0 if layer_id in self.cpu_layer_ids else rows + for layer_id, rows in enumerate(missing) + ] + else: + fetched = [0] * self.num_layers per_layer = [] for L in range(self.num_layers): s, m, a, f = steps[L], missing[L], active[L], fetched[L] @@ -969,6 +1293,14 @@ def copy_missing(self) -> None: assert self.banks, "set_bank_sources must register the banks first" layer_id = self._pending_src_layer assert layer_id is not None, "no staged misses (ensure_experts/materialize_layer first)" + if self._pending_geometry_prefill: + for per_layer, cache in self.banks: + self._copy_compact_layer( + cache[: self.num_experts], + per_layer[layer_id], + registered_host=True, + ) + return if layer_id in self._unpinned_layers: if not self._pending_whole_layer: raise RuntimeError( @@ -979,9 +1311,43 @@ def copy_missing(self) -> None: # the only copy a non-pinned layer ever needs is the non-overlap prefill materialize, which schedules the whole layer into slots [0, num_experts) with position == expert id -- a plain synchronous pageable H2D copy # never CUDA-graph captured: prefill is not captured, and decode never reaches this branch (it routes to the CPU executor) for per_layer, cache in self.banks: - cache[: self.num_experts].copy_(per_layer[layer_id]) + self._copy_compact_layer(cache[: self.num_experts], per_layer[layer_id]) + return + pool = self._pending_geometry_pool + if pool is not None and not self._pending_geometry_prefill: + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit + + assert pool.copy_dst_ptrs is not None and pool.copy_feat_bytes is not None + fast_index_copy_multi_jit( + pool.copy_dst_ptrs, + pool.copy_src_ptrs[layer_id], + pool.copy_feat_bytes, + pool.evict_slots, + pool.src_indices, + pool.num_indices, + ) return if self._copy_fused_ok: + assert self._copy_dst_ptrs is not None and self._copy_src_ptrs is not None + if self.has_heterogeneous_rows: + from freetoken.kernel.fast_index_copy import ( + fast_index_copy_multi_strided_jit, + ) + + assert self._copy_payload_bytes is not None + assert self._copy_src_row_strides is not None + assert self._copy_dst_row_strides is not None + fast_index_copy_multi_strided_jit( + self._copy_dst_ptrs, + self._copy_src_ptrs[layer_id], + self._copy_payload_bytes[layer_id], + self._copy_dst_row_strides, + self._copy_src_row_strides[layer_id], + self.evict_slots, + self.src_indices, + self.num_indices, + ) + return from freetoken.kernel.fast_index_copy import fast_index_copy_multi_jit # One launch copies the missing rows for every bank (instead of one launch per @@ -998,6 +1364,12 @@ def copy_missing(self) -> None: ) return + if self.has_heterogeneous_rows: + raise RuntimeError( + "heterogeneous expert rows require the strided fused copy plan " + "(CUDA device, 16-byte aligned rows, and FREETOKEN_FUSED_COPY=1)" + ) + from freetoken.kernel import fast_index_copy_jit for per_layer, cache in self.banks: diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 48923e3b..e30d5c33 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -299,6 +299,21 @@ def shutdown(self) -> None: self.sync_all_ranks() self.engine.shutdown() + def _report_moe_decode_stats(self, finished_reqs: Set[Req]) -> None: + """Log opt-in cumulative, rank-local MoE counters after a completion drain. + + With overlap scheduling, the cumulative total may already include work from the + next launched batch; this is intentionally not a per-request snapshot. + """ + if not finished_reqs or self.config.tp_info.rank != 0: + return + cache = getattr(self.engine, "moe_offload_cache", None) + if cache is None or not cache.collect_stats: + return + logger.info_rank0( + f"MoE decode stats (cumulative, rank-local): {cache.decode_miss_stats()}" + ) + def _process_last_data(self, last_data: ForwardData | None) -> None: if last_data is None: return @@ -388,6 +403,7 @@ def _process_last_data(self, last_data: ForwardData | None) -> None: self.cache_manager.cache_req(req, finished=False) self.finished_reqs = new_finished_reqs + self._report_moe_decode_stats(new_finished_reqs) # Stamp each reply with the post-batch KV page occupancy so the frontend (shell # status bar) can show live KV usage without a separate query. used, total = self._kv_usage_pages() diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index a71b6819..37e94767 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -532,6 +532,12 @@ def _infer_reasoning_parser(model_path: str) -> str | None: choices=["lru"], help="The unified MoE cache eviction policy.", ) + parser.add_argument( + "--moe-collect-stats", + action="store_true", + default=ServerArgs.moe_collect_stats, + help="Collect device-side per-layer MoE decode hit/miss counters.", + ) parser.add_argument( "--moe-cpu-threads", 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/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a164f0b4..850511c1 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -97,6 +97,59 @@ def test_expert_bytes_per_slot_sums_row_bytes_over_banks(): assert expert_bytes_per_slot(sources) == 512 + 256 +def test_geometry_pool_plan_fits_laguna_working_set_in_legacy_arenas(): + from freetoken.engine.cache_budget import plan_geometry_pool_slots + + q4 = (3_538_944, 1_769_472) + bf16 = (12_582_912, 6_291_456) + rows = [q4] * 39 + [bf16] * 8 + plan = plan_geometry_pool_slots( + rows, + legacy_cache_size=298, + num_experts=256, + top_k=10, + max_decode_batch=1, + ) + + assert plan is not None + by_rows = {entry.row_bytes: entry for entry in plan} + assert by_rows[q4].slots >= 390 + assert by_rows[bf16].slots >= 80 + for bank in range(2): + used = sum(entry.slots * entry.row_bytes[bank] for entry in plan) + budget = 298 * max(row[bank] for row in rows) + assert used <= budget + + +def test_geometry_pool_plan_falls_back_when_decode_floors_do_not_fit(): + from freetoken.engine.cache_budget import plan_geometry_pool_slots + + assert ( + plan_geometry_pool_slots( + [(16, 8), (64, 32)], + legacy_cache_size=1, + num_experts=256, + top_k=10, + max_decode_batch=1, + ) + is None + ) + + +def test_expert_bytes_per_slot_uses_each_banks_largest_layer(): + sources = { + "gate_up": [ + torch.zeros(4, 5, dtype=torch.uint8), + torch.zeros(4, 11, dtype=torch.uint8), + ], + "down": [ + torch.zeros(4, 7, dtype=torch.uint8), + torch.zeros(4, 3, dtype=torch.uint8), + ], + } + assert expert_bytes_per_slot(sources) == 11 + 7 + + def test_resolve_auto_applies_ratio_once_and_marlin_cap(): # baseline 1000, weights 100, ratio 0.9 -> budget = 900 - 100 - 0(fixed) = 800 size, pages, overlap = resolve_moe_cache_auto( @@ -407,6 +460,30 @@ def test_adjust_config_defaults_moe_cache_auto_for_auto_resolved_offload_backend assert config.moe_cache_size == 0 # still unresolved -- the scheduler sizes it from VRAM +def test_adjust_config_never_auto_selects_hybrid_for_gguf(monkeypatch): + """A bench profile must not route GGUF banks to the absent CPU executor.""" + from freetoken.engine.engine import _adjust_config + + config = _offload_engine_config() + object.__setattr__(config.model_config, "expert_quant", "gguf") + object.__setattr__(config.model_config, "moe_weight_format", "gguf") + object.__setattr__(config.model_config, "hidden_act", "silu") + monkeypatch.setattr( + "freetoken.moe.bench_profile.load_backend_recommendation", + lambda *_args, **_kwargs: "hybrid", + ) + monkeypatch.setattr("freetoken.engine.engine._profile_gpu", lambda: ("gpu", "uuid")) + monkeypatch.setattr( + "freetoken.moe.cpu_executor.compiled_extension_supports", + lambda _act: True, + ) + + _adjust_config(config) + + assert config.moe_backend == "offload" + assert config.moe_cache_auto is True + + def test_page_table_width_covers_whole_trailing_pages(): # _write_page_table writes WHOLE trailing pages, so the width must reach the last # page's end, not just the next multiple of 32 (DSV4's P=128 exposed the gap). 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_fast_index_copy_strided.py b/tests/kernels/test_fast_index_copy_strided.py new file mode 100644 index 00000000..d43828b5 --- /dev/null +++ b/tests/kernels/test_fast_index_copy_strided.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import pytest +import torch + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fast_index_copy_multi_strided_copies_only_payload_prefixes() -> None: + from freetoken.kernel.fast_index_copy import fast_index_copy_multi_strided_jit + + device = torch.device("cuda") + src0 = torch.arange(4 * 32, dtype=torch.uint8, device=device).view(4, 32) + src1 = ( + (torch.arange(4 * 64, dtype=torch.int32, device=device) % 251) + .to(torch.uint8) + .view(4, 64) + ) + dst0 = torch.full((5, 64), 0xEE, dtype=torch.uint8, device=device) + dst1 = torch.full((5, 80), 0xEE, dtype=torch.uint8, device=device) + + dst_ptrs = torch.tensor( + [dst0.data_ptr(), dst1.data_ptr()], dtype=torch.int64, device=device + ) + src_ptrs = torch.tensor( + [src0.data_ptr(), src1.data_ptr()], dtype=torch.int64, device=device + ) + copy_bytes = torch.tensor([32, 64], dtype=torch.int64, device=device) + dst_row_strides = torch.tensor([64, 80], dtype=torch.int64, device=device) + src_row_strides = torch.tensor([32, 64], dtype=torch.int64, device=device) + dst_indices = torch.tensor([2, 0], dtype=torch.int32, device=device) + src_indices = torch.tensor([1, 3], dtype=torch.int32, device=device) + num_indices = torch.tensor([2], dtype=torch.int64, device=device) + + fast_index_copy_multi_strided_jit( + dst_ptrs, + src_ptrs, + copy_bytes, + dst_row_strides, + src_row_strides, + dst_indices, + src_indices, + num_indices, + ) + torch.cuda.synchronize() + + assert torch.equal(dst0[2, :32], src0[1]) + assert torch.equal(dst0[0, :32], src0[3]) + assert torch.all(dst0[[0, 2], 32:] == 0xEE) + assert torch.all(dst0[[1, 3, 4]] == 0xEE) + assert torch.equal(dst1[2, :64], src1[1]) + assert torch.equal(dst1[0, :64], src1[3]) + assert torch.all(dst1[[0, 2], 64:] == 0xEE) + assert torch.all(dst1[[1, 3, 4]] == 0xEE) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fast_index_copy_rows_strided_avoids_payload_sized_cuda_temporary() -> None: + from freetoken.kernel.fast_index_copy import fast_index_copy_rows_strided_jit + + rows = 16 + payload = 1 << 20 + destination_stride = payload * 2 + source = torch.arange(rows * payload, dtype=torch.int64).view(torch.uint8) + source = source[: rows * payload].reshape(rows, payload).pin_memory() + destination = torch.full( + (rows, destination_stride), 0xEE, dtype=torch.uint8, device="cuda" + ) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + allocated_before = torch.cuda.memory_allocated() + + fast_index_copy_rows_strided_jit(destination, source) + torch.cuda.synchronize() + + temporary_peak = torch.cuda.max_memory_allocated() - allocated_before + assert temporary_peak < 1 << 20 + assert torch.equal(destination[:, :payload].cpu(), source) + assert torch.equal( + destination[:, payload:].cpu(), + torch.full((rows, payload), 0xEE, dtype=torch.uint8), + ) 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..1fca3029 --- /dev/null +++ b/tests/models/test_laguna_weights.py @@ -0,0 +1,212 @@ +"""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_BF16, 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) + expert_type = gguf.GGMLQuantizationType.BF16 if i == L - 1 else q8 + 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), expert_type + ).reshape(E, rows, -1), + raw_dtype=expert_type, + ) + 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 - 2) + ((GGML_BF16, GGML_BF16),) + ) + + +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) + geometry = _expert_bank_geometry(cfg) + assert len(banks["gate_up"]) == L - 1 and len(banks["down"]) == L - 1 + for layer_id, (gu_s, dn_s) in enumerate(geometry): + assert banks["gate_up"][layer_id].shape == (E, gu_s) + assert banks["down"][layer_id].shape == (E, dn_s) + assert banks["gate_up"][layer_id].dtype == torch.uint8 + assert banks["gate_up"][0].shape[1] < banks["gate_up"][-1].shape[1] + # 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/moe/test_offload.py b/tests/moe/test_offload.py index 422ca867..fbce1a1e 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -34,6 +34,477 @@ def _make_layer_and_cache(): return layer, cache +@pytest.mark.parametrize( + "quant_format", + ["bf16", "nvfp4_marlin", "nvfp4_b12x", "ds_fp4", "q4_0"], +) +def test_non_gguf_sources_reject_per_layer_shape_changes_even_with_equal_numel( + quant_format, +): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, + num_experts=2, + cache_size=4, + device=torch.device("cpu"), + quant_format=quant_format, + prefill_overlap=False, + ) + names = cache.bank_schema + sources = {} + for name in names: + sources[name] = [ + torch.zeros((2, 2, 3), dtype=torch.uint8), + torch.zeros((2, 3, 2), dtype=torch.uint8), + ] + + with pytest.raises(ValueError, match="uniform per-layer shapes"): + cache.set_bank_sources(sources) + + +def test_non_gguf_sources_reject_noncontiguous_rows(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, + num_experts=2, + cache_size=4, + device=torch.device("cpu"), + quant_format="bf16", + ) + noncontiguous = torch.zeros((2, 3, 2), dtype=torch.bfloat16).transpose(1, 2) + assert noncontiguous.shape == (2, 2, 3) + assert not noncontiguous.is_contiguous() + sources = { + "gate_up": [ + torch.zeros((2, 2, 3), dtype=torch.bfloat16), + noncontiguous, + ], + "down": [ + torch.zeros((2, 2, 3), dtype=torch.bfloat16), + torch.zeros((2, 2, 3), dtype=torch.bfloat16), + ], + } + + with pytest.raises(ValueError, match="contiguous"): + cache.set_bank_sources(sources) + + +def test_gguf_sources_require_contiguous_2d_uint8_rows(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, + num_experts=2, + cache_size=4, + device=torch.device("cpu"), + quant_format="gguf", + prefill_overlap=False, + ) + bad_rank = { + "gate_up": [torch.zeros((2, 2, 4), dtype=torch.uint8) for _ in range(2)], + "down": [torch.zeros((2, 8), dtype=torch.uint8) for _ in range(2)], + } + with pytest.raises(ValueError, match="2-D contiguous uint8"): + cache.set_bank_sources(bad_rank) + + bad_dtype = { + "gate_up": [torch.zeros((2, 8), dtype=torch.int8) for _ in range(2)], + "down": [torch.zeros((2, 8), dtype=torch.uint8) for _ in range(2)], + } + with pytest.raises(ValueError, match="2-D contiguous uint8"): + cache.set_bank_sources(bad_dtype) + + +def test_gguf_geometry_pools_carve_disjoint_exact_width_views(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, + num_experts=2, + cache_size=8, + device=torch.device("cpu"), + quant_format="gguf", + prefill_overlap=False, + geometry_pool_top_k=1, + geometry_pool_max_batch=1, + ) + sources = { + "gate_up": [ + torch.zeros((2, 4), dtype=torch.uint8), + torch.zeros((2, 8), dtype=torch.uint8), + ], + "down": [ + torch.zeros((2, 2), dtype=torch.uint8), + torch.zeros((2, 4), dtype=torch.uint8), + ], + } + cache.set_bank_sources(sources) + + q_views = cache.bank_views(layer_id=0) + b_views = cache.bank_views(layer_id=1) + assert q_views[0].shape[1:] == (4,) + assert q_views[1].shape[1:] == (2,) + assert b_views[0].shape[1:] == (8,) + assert b_views[1].shape[1:] == (4,) + sizes = cache.geometry_pool_sizes() + assert sizes[0] >= 1 and sizes[1] >= 1 + for bank_index, arena in enumerate(cache.bank_caches.values()): + arena_start = arena.data_ptr() + arena_end = arena_start + arena.numel() * arena.element_size() + ranges = [] + for views in (q_views, b_views): + view = views[bank_index] + start = view.data_ptr() + end = start + view.numel() * view.element_size() + assert arena_start <= start < end <= arena_end + ranges.append((start, end)) + assert ranges[0][1] <= ranges[1][0] or ranges[1][1] <= ranges[0][0] + + old_arena_ptr = cache.bank_caches["gate_up"].data_ptr() + old_sizes = cache.geometry_pool_sizes() + cache.rebuild(10) + assert cache.bank_sources["gate_up"] is not None + assert cache.bank_caches["gate_up"].data_ptr() != old_arena_ptr + assert cache.geometry_pool_sizes()[0] >= old_sizes[0] + assert ( + cache.bank_views(layer_id=0)[0].untyped_storage().data_ptr() + == cache.bank_caches["gate_up"].untyped_storage().data_ptr() + ) + + +def test_heterogeneous_sources_use_compact_host_rows_and_max_gpu_stride(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, + num_experts=2, + cache_size=4, + device=torch.device("cpu"), + quant_format="gguf", + prefill_overlap=True, + ) + small_gu = torch.arange(2 * 16, dtype=torch.uint8).view(2, 16) + large_gu = torch.arange(2 * 32, dtype=torch.uint8).view(2, 32) + small_dn = torch.arange(2 * 8, dtype=torch.uint8).view(2, 8) + large_dn = torch.arange(2 * 16, dtype=torch.uint8).view(2, 16) + sources = { + "gate_up": [small_gu, large_gu], + "down": [small_dn, large_dn], + } + + cache.set_bank_sources(sources) + + assert cache.bank_sources["gate_up"][0].shape == (2, 16) + assert cache.bank_sources["gate_up"][1].shape == (2, 32) + assert cache.bank_caches["gate_up"].shape == (4, 32) + assert cache.bank_caches["down"].shape == (4, 16) + assert cache.has_heterogeneous_rows + + cache.prefetch_prefill_layer(0) + gate_buffer, down_buffer = cache.wait_prefill_layer(0) + assert torch.equal(gate_buffer[:, :16], small_gu) + assert torch.equal(down_buffer[:, :8], small_dn) + + cache.release_prefill_layer(0) + cache.prefetch_prefill_layer(1) + gate_buffer, down_buffer = cache.wait_prefill_layer(1) + assert torch.equal(gate_buffer, large_gu) + assert torch.equal(down_buffer, large_dn) + + cache.rebuild(6) + assert cache.bank_caches["gate_up"].shape == (6, 32) + assert cache.bank_caches["down"].shape == (6, 16) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_geometry_decode_routes_lru_and_copies_into_selected_pool(): + from freetoken.moe.offload_cache import OffloadMoeCache + + def pinned(rows: int, cols: int, offset: int) -> torch.Tensor: + values = (torch.arange(rows * cols, dtype=torch.int32) + offset) % 251 + return values.to(torch.uint8).view(rows, cols).pin_memory() + + cache = OffloadMoeCache( + num_layers=2, + num_experts=4, + cache_size=8, + device=torch.device("cuda"), + quant_format="gguf", + geometry_pool_top_k=1, + geometry_pool_max_batch=1, + ) + sources = { + "gate_up": [pinned(4, 32, 0), pinned(4, 64, 17)], + "down": [pinned(4, 16, 33), pinned(4, 32, 49)], + } + cache.set_bank_sources(sources) + cache.collect_stats = True + + pools = [cache._geometry_pool_for_layer[layer_id] for layer_id in range(2)] + for layer_id in range(2): + ids = torch.tensor([2], dtype=torch.int32, device="cuda") + cache.ensure_experts(layer_id, ids) + cache.copy_missing() + torch.cuda.synchronize() + slot = int(ids.item()) + gate, down = cache.bank_views(layer_id=layer_id) + assert torch.equal(gate[slot].cpu(), sources["gate_up"][layer_id][2]) + assert torch.equal(down[slot].cpu(), sources["down"][layer_id][2]) + assert int(pools[layer_id].slot_for_id[layer_id, 2].item()) == slot + other = pools[1 - layer_id] + assert int(other.slot_for_id[layer_id, 2].item()) == -1 + + stats = cache.decode_miss_stats() + assert stats["layer_calls"] == 2 + assert stats["requested_rows"] == 2 + assert stats["miss_rows"] == 2 + assert stats["hit_rows"] == 0 + assert stats["bytes_h2d"] == (32 + 16) + (64 + 32) + assert stats["miss_rate"] == 1.0 + assert stats["fetched_per_layer"] == 1.0 + assert stats["cpu_per_layer"] == 0.0 + assert stats["fetch_rate"] == 1.0 + per_layer = cache.decode_miss_stats_per_layer()["per_layer"] + assert [entry["steps"] for entry in per_layer] == [1, 1] + assert [entry["fetched_per_step"] for entry in per_layer] == [1.0, 1.0] + + cache.materialize_layer(1) + cache.copy_missing() + torch.cuda.synchronize() + prefill_gate, prefill_down = cache.bank_views(cache.num_experts, layer_id=1) + assert torch.equal(prefill_gate.cpu(), sources["gate_up"][1]) + assert torch.equal(prefill_down.cpu(), sources["down"][1]) + assert all(torch.all(pool.slot_for_id == -1) for pool in pools) + + ids = torch.tensor([2], dtype=torch.int32, device="cuda") + cache.ensure_experts(0, ids) + assert int(pools[0].num_indices.item()) == 1 + cache.copy_missing() + torch.cuda.synchronize() + slot = int(ids.item()) + assert torch.equal( + cache.bank_views(layer_id=0)[0][slot].cpu(), sources["gate_up"][0][2] + ) + + +def test_decode_stats_count_gpu_transfers_in_mixed_cpu_mode(): + from flashlib.kernels.slot_cache import Stat + + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, + num_experts=4, + cache_size=4, + device=torch.device("cpu"), + decode_target="cpu", + ) + cache.cpu_layer_ids = frozenset({1}) + sources = { + "gate_up": [torch.zeros((4, 6), dtype=torch.bfloat16) for _ in range(2)], + "down": [torch.zeros((4, 4), dtype=torch.bfloat16) for _ in range(2)], + } + cache.set_bank_sources(sources) + cache.lru_stats[0, Stat.ACTIVE] = 2 + cache.lru_stats[0, Stat.MISS] = 1 + cache.lru_stats[0, Stat.CALLS] = 1 + cache.lru_stats[1, Stat.ACTIVE] = 3 + cache.lru_stats[1, Stat.MISS] = 2 + cache.lru_stats[1, Stat.CALLS] = 1 + + stats = cache.decode_miss_stats() + + assert stats["requested_rows"] == 5 + assert stats["miss_rows"] == 3 + assert stats["hit_rows"] == 2 + assert stats["bytes_h2d"] == (6 + 4) * 2 + assert stats["fetched_per_layer"] == 0.5 + assert stats["cpu_per_layer"] == 1.0 + assert stats["fetch_rate"] == pytest.approx(1 / 3) + per_layer = cache.decode_miss_stats_per_layer()["per_layer"] + assert [entry["fetched_per_step"] for entry in per_layer] == [1.0, 0.0] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_heterogeneous_unaligned_rows_fail_during_setup(): + from freetoken.moe.offload_cache import OffloadMoeCache + + cache = OffloadMoeCache( + num_layers=2, + num_experts=4, + cache_size=4, + device=torch.device("cuda"), + quant_format="gguf", + ) + sources = { + "gate_up": [ + torch.zeros((4, 30), dtype=torch.uint8, pin_memory=True), + torch.zeros((4, 64), dtype=torch.uint8, pin_memory=True), + ], + "down": [ + torch.zeros((4, 16), dtype=torch.uint8, pin_memory=True), + torch.zeros((4, 32), dtype=torch.uint8, pin_memory=True), + ], + } + with pytest.raises(ValueError, match="16-byte-aligned"): + cache.set_bank_sources(sources) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("disable_uniform_fused_copy", [False, True]) +def test_heterogeneous_copy_missing_uses_source_payload_and_destination_stride( + monkeypatch, disable_uniform_fused_copy +): + import freetoken.moe.offload_cache as offload_cache_module + from freetoken.moe.offload_cache import OffloadMoeCache + + monkeypatch.setattr( + offload_cache_module, "_FUSED_COPY", not disable_uniform_fused_copy + ) + + def pinned(rows: int, cols: int, offset: int) -> torch.Tensor: + t = torch.empty((rows, cols), dtype=torch.uint8, pin_memory=True) + values = (torch.arange(rows * cols, dtype=torch.int32) + offset) % 251 + t.copy_(values.to(torch.uint8).view(rows, cols)) + return t + + cache = OffloadMoeCache( + num_layers=2, + num_experts=4, + cache_size=4, + device=torch.device("cuda"), + quant_format="gguf", + ) + gu0, gu1 = pinned(4, 32, 0), pinned(4, 64, 17) + dn0, dn1 = pinned(4, 16, 33), pinned(4, 32, 49) + cache.set_bank_sources({"gate_up": [gu0, gu1], "down": [dn0, dn1]}) + for bank in cache.bank_caches.values(): + bank.fill_(0xEE) + cache._pending_src_layer = 0 + cache._pending_whole_layer = False + cache.evict_slots[:2] = torch.tensor([3, 1], dtype=torch.int32, device="cuda") + cache.src_indices[:2] = torch.tensor([2, 0], dtype=torch.int32, device="cuda") + cache.num_indices.fill_(2) + + cache.copy_missing() + torch.cuda.synchronize() + + gate_cache = cache.bank_caches["gate_up"] + down_cache = cache.bank_caches["down"] + assert torch.equal(gate_cache[3, :32].cpu(), gu0[2]) + assert torch.equal(gate_cache[1, :32].cpu(), gu0[0]) + assert torch.all(gate_cache[[1, 3], 32:] == 0xEE) + assert torch.equal(down_cache[3, :16].cpu(), dn0[2]) + assert torch.equal(down_cache[1, :16].cpu(), dn0[0]) + assert torch.all(down_cache[[1, 3], 16:] == 0xEE) + + +def test_decode_requests_geometry_bank_views_for_its_layer(monkeypatch): + from freetoken.layers.moe import OffloadMoELayer + + _init_tp() + seen = {} + + class FakeCache: + decode_target = "gpu" + + @staticmethod + def is_cpu_layer(layer_id): + return False + + @staticmethod + def ensure_experts(layer_id, expert_ids): + return None + + @staticmethod + def copy_missing(): + return None + + @staticmethod + def alphas_for_slots(layer_id): + return None + + @staticmethod + def bank_views(*args, **kwargs): + seen.update(kwargs) + return () + + layer = OffloadMoELayer(3, 4, 1, 8, 4) + layer.offload_cache = FakeCache() + expected = torch.zeros((1, 8)) + monkeypatch.setattr(layer, "_expert_gemm", lambda *args, **kwargs: expected) + + result = layer._decode_routed( + torch.zeros((1, 8)), + torch.ones((1, 1)), + torch.zeros((1, 1), dtype=torch.int32), + ) + + assert result is expected + assert seen == {"layer_id": 3} + + +def test_gguf_bf16_layer_reinterprets_padded_slots_for_dense_kernel(monkeypatch): + import freetoken.layers.moe as moe_module + import freetoken.moe.fused_gguf as fused_gguf_module + from freetoken.layers.moe import OffloadMoELayer + from freetoken.moe.offload_cache import OffloadMoeCache + + _init_tp() + layer = OffloadMoELayer( + layer_id=0, + num_experts=2, + top_k=1, + hidden_size=4, + intermediate_size=3, + ) + layer.gguf_gate_up_type = 30 + layer.gguf_down_type = 30 + layer.gguf_gate_up_rows = 6 + layer.gguf_down_rows = 4 + gate_up = torch.arange(2 * 6 * 4, dtype=torch.float32).to(torch.bfloat16) + down = torch.arange(2 * 4 * 3, dtype=torch.float32).to(torch.bfloat16) + views = ( + gate_up.view(torch.uint8).reshape(2, -1), + down.view(torch.uint8).reshape(2, -1), + ) + captured = {} + expected = torch.randn(1, 4) + + def fake_dense(*args): + captured["gate_up"] = args[1] + captured["down"] = args[2] + return expected + + def fail_gguf(*args, **kwargs): + raise AssertionError( + "BF16 GGUF layers must not enter the quantized MMVQ kernel" + ) + + monkeypatch.setattr(moe_module, "fused_experts_impl", fake_dense) + monkeypatch.setattr(fused_gguf_module, "fused_experts_gguf", fail_gguf) + cache = OffloadMoeCache(1, 2, 2, torch.device("cpu"), quant_format="gguf") + + out = layer._expert_gemm( + cache, + torch.randn(1, 4), + torch.ones(1, 1), + torch.zeros(1, 1, dtype=torch.int32), + views=views, + n=2, + alphas=None, + is_prefill=True, + ) + + assert out is expected + assert captured["gate_up"].shape == (2, 6, 4) + assert captured["down"].shape == (2, 4, 3) + assert captured["gate_up"].dtype == torch.bfloat16 + assert captured["down"].dtype == torch.bfloat16 + + def test_dummy_expert_sources_use_moe_layer_count(monkeypatch): from types import SimpleNamespace diff --git a/tests/scheduler/test_moe_stats_reporting.py b/tests/scheduler/test_moe_stats_reporting.py new file mode 100644 index 00000000..6eb42c48 --- /dev/null +++ b/tests/scheduler/test_moe_stats_reporting.py @@ -0,0 +1,36 @@ +from types import SimpleNamespace +from typing import Any + + +def test_moe_decode_stats_log_only_after_request_finishes(monkeypatch): + import freetoken.scheduler.scheduler as scheduler_module + from freetoken.scheduler.scheduler import Scheduler + + expected = { + "requested_rows": 470, + "miss_rows": 80, + "hit_rows": 390, + "bytes_h2d": 1234, + } + cache = SimpleNamespace( + collect_stats=True, + decode_miss_stats=lambda: expected, + ) + scheduler: Any = Scheduler.__new__(Scheduler) + scheduler.engine = SimpleNamespace(moe_offload_cache=cache) + scheduler.config = SimpleNamespace(tp_info=SimpleNamespace(rank=0)) + logged = [] + monkeypatch.setattr(scheduler_module.logger, "info_rank0", logged.append) + + scheduler._report_moe_decode_stats(set()) + assert logged == [] + + scheduler._report_moe_decode_stats({object()}) + assert len(logged) == 1 + assert "cumulative" in logged[0] + assert "rank-local" in logged[0] + assert "'bytes_h2d': 1234" in logged[0] + + scheduler.config.tp_info.rank = 1 + scheduler._report_moe_decode_stats({object()}) + assert len(logged) == 1 diff --git a/tests/server/test_moe_stats_args.py b/tests/server/test_moe_stats_args.py new file mode 100644 index 00000000..ad72ba1c --- /dev/null +++ b/tests/server/test_moe_stats_args.py @@ -0,0 +1,16 @@ +from unittest.mock import patch + +from freetoken.server.args import parse_args + + +class _Config: + def to_dict(self) -> dict: + return {"architectures": ["LlamaForCausalLM"], "torch_dtype": "bfloat16"} + + +def test_moe_collect_stats_flag_is_accepted(): + with patch("freetoken.utils.cached_load_hf_config", lambda _path: _Config()): + args, _ = parse_args( + ["--model", "/models/anon", "--moe-collect-stats"] + ) + assert args.moe_collect_stats is True