diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cd6505d2..c386e28b 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -10,7 +10,6 @@ from freetoken.attention import AttnType, attention_backend_info, create_attention_backend from freetoken.core import Batch, Context, Req, set_global_ctx from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info -from freetoken.gpu_select import gpu_identity from freetoken.layers import set_rope_device from freetoken.models import create_model, load_weight from freetoken.moe import create_moe_backend, is_offload_moe_backend @@ -295,11 +294,10 @@ def __init__(self, config: EngineConfig): assert not torch.cuda.is_initialized() set_tp_info(rank=config.tp_info.rank, size=config.tp_info.size) _ensure_expandable_segments() # before the first CUDA allocation below - - from freetoken.gpu_select import bind_assigned_gpu - - self.device = bind_assigned_gpu(config.tp_info.rank) _adjust_config(config) + + self.device = torch.device(f"cuda:{config.tp_info.rank}") + torch.cuda.set_device(self.device) torch.manual_seed(42) self.stream = torch.cuda.Stream() torch.cuda.set_stream(self.stream) @@ -611,6 +609,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: prefill_overlap=config.moe_prefill_overlap, prefill_hit_d2d=config.moe_prefill_hit_d2d, quant_format=banks.quant_format, + gguf_expert_types=banks.gguf_expert_types, decode_target=decode_target, hybrid_max_fetch=config.moe_hybrid_max_fetch, ) @@ -651,10 +650,8 @@ def _resolve_hybrid_fetch(self, config: EngineConfig, cache) -> None: return # explicit fixed cap from freetoken.moe.bench_profile import load_hybrid_fetch_fraction - gpu_name, gpu_uuid = _profile_gpu(self.device.index) - fraction = load_hybrid_fetch_fraction( - cache.quant_format, gpu_name=gpu_name, gpu_uuid=gpu_uuid - ) + gpu_name = torch.cuda.get_device_name(self.device) if torch.cuda.is_available() else None + fraction = load_hybrid_fetch_fraction(cache.quant_format, gpu_name=gpu_name) if fraction is None: cache.hybrid_max_fetch = 1 logger.warning_rank0( @@ -997,14 +994,6 @@ def shutdown(self) -> None: destroy_distributed() -def _profile_gpu(index: "int | None" = None) -> Tuple[str | None, str | None]: - """(name, uuid) of visible device ``index`` (default: the current, i.e. bound, device); (None, None) without CUDA.""" - if not torch.cuda.is_available(): - return None, None - ident = gpu_identity(torch.cuda.current_device() if index is None else index) - return ident["name"], ident["uuid"] - - def _ensure_expandable_segments() -> None: """Default the CUDA allocator to expandable segments. @@ -1155,6 +1144,21 @@ def _cpu_moe_executor_viable(model_config) -> bool: return False expert_quant = getattr(model_config, "expert_quant", "none") fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16") + if fmt == "gguf": + # "gguf" is a container tag, not a layout: the checkpoint picks a ggml type per + # tensor and the concrete CPU format has to be recovered from the bank types. + # Testing the tag against _WFMT_IDS answers False for EVERY GGUF checkpoint, which + # silently disables the automatic residency split on hosts where CUDA pinning is + # quota-capped (WSL caps it near half of RAM). The symptom is not a clear refusal + # but cudaHostRegister failing partway through the banks. + from freetoken.moe.cpu_executor import _GGML_TO_CPU_FMT + + types = getattr(model_config, "gguf_expert_types", None) + if not types: + return False + gate_up, down = int(types[0]), int(types[1]) + # one weight_format serves both banks, so mixed types cannot run on the CPU path + return gate_up == down and gate_up in _GGML_TO_CPU_FMT return fmt == "mxfp4" or fmt in _WFMT_IDS @@ -1370,8 +1374,8 @@ def override(attr: str, value: Any): # this is dangerous, use with caution bench_fmt = expert_quant if expert_quant != "none" else (moe_wfmt or "bf16") from freetoken.moe.bench_profile import load_backend_recommendation - gpu_name, gpu_uuid = _profile_gpu() - if load_backend_recommendation(bench_fmt, gpu_name=gpu_name, gpu_uuid=gpu_uuid) == "hybrid": + gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None + if load_backend_recommendation(bench_fmt, gpu_name=gpu_name) == "hybrid": from freetoken.moe.cpu_executor import compiled_extension_supports _act = getattr(model_config, "hidden_act", "silu") diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 880e8637..48210b9d 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1212,7 +1212,151 @@ q4dot_fn select_q4dot() { return q4_0_dot_i8_scalar; } -enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4 }; +// ----------------------------- Q4_K and Q6_K (W4A16) ---------------------------- +// Correctness-first scalar K-quant expert dot kernels: dequantize blocks and accumulate +// against bf16 activation rows. AVX2/VNNI optimizations are a deliberate follow-up. +// Reference: llama.cpp ggml-quants.c, ggml-cuda/convert.cu, models/gguf/dequant.py. + +// Helper: extract 6-bit scale and min from Q4_K's packed scales array. +// Q4_K packs scales and mins using 6 bits each, packed into 12 bytes for 256 elements. +inline void get_scale_min_k4(int j, const uint8_t* q, uint8_t& scale, uint8_t& minv) { + if (j < 4) { + scale = q[j] & 63; + minv = q[j + 4] & 63; + } else { + scale = (q[j + 4] & 0x0F) | ((q[j - 4] >> 6) << 4); + minv = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); + } +} + +float q4_k_dot_f32_scalar(const uint8_t* w, const bf16_t* x, int K) { + // Q4_K, 256-element super-blocks of 144 bytes: half2 dm | scales[12] | qs[128]. + // dm.x scales the quants, dm.y scales the per-sub-block minimum that is SUBTRACTED. + // + // Two things here are easy to get wrong and both produce fluent-looking garbage rather + // than an obvious failure, so this mirrors dequantize_block_q4_K in gguf/dequantize.cuh + // element for element: + // * the quants are UNSIGNED 0..15. There is no -8 bias; that is Q4_0's encoding. Q4_K + // centres the range with the per-sub-block min instead. + // * within a 64-element group the low nibble of byte l is element l and the HIGH nibble + // is element l+32, not l+1. The two nibbles of a byte are 32 apart, and they carry + // different scale/min pairs (is+0 vs is+1). + float acc = 0.0f; + const int nb = K / 256; + + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 144; + + uint16_t dh_scale, dh_min; + std::memcpy(&dh_scale, blk, sizeof(uint16_t)); + std::memcpy(&dh_min, blk + 2, sizeof(uint16_t)); + const float dall = fp16_to_f32(dh_scale); + const float dmin = fp16_to_f32(dh_min); + + const uint8_t* scales = blk + 4; // 12 bytes of packed 6-bit scales + mins + const uint8_t* qs = blk + 16; // 128 bytes of 4-bit quants + const bf16_t* xb = x + (size_t)256 * b; + + for (int il = 0; il < 4; ++il) { // four 64-element groups + const int is = 2 * il; + + uint8_t sc, m; + get_scale_min_k4(is + 0, scales, sc, m); + const float d1 = dall * sc; + const float m1 = dmin * m; + get_scale_min_k4(is + 1, scales, sc, m); + const float d2 = dall * sc; + const float m2 = dmin * m; + + for (int ir = 0; ir < 8; ++ir) { + const uint8_t* q = qs + 32 * il + 4 * ir; + const bf16_t* y = xb + 64 * il + 4 * ir; + for (int l = 0; l < 4; ++l) { + acc += (d1 * (float)(q[l] & 0xF) - m1) * bf16_to_f32(y[l]); + acc += (d2 * (float)(q[l] >> 4) - m2) * bf16_to_f32(y[l + 32]); + } + } + } + } + + return acc; +} + +float q6_k_dot_f32_scalar(const uint8_t* w, const bf16_t* x, int K) { + // Q6_K: 256-element blocks, 210 bytes each. + // Block layout: ql[128] | qh[64] | scales[16] | d (fp16) + // ql: lower 4 bits of 6-bit quant values (128 bytes) + // qh: upper 2 bits of 6-bit quant values, packed (64 bytes, 2 bits per element) + // scales: 16 int8 sub-scales (8 per 128-element half) + // d: fp16 block scale + // Dequant: q in [-32,31], w = d * scales[is] * (q - 32) + // Reference: ggml-cuda/convert.cu::dequantize_block_q6_K and models/gguf/dequant.py::dequant_q6_k + + float acc = 0.0f; + const int nb = K / 256; // number of Q6_K blocks + + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 210; + + const uint8_t* ql = blk; // 128 bytes (lower 4 bits) + const uint8_t* qh = blk + 128; // 64 bytes (upper 2 bits) + const int8_t* scales = (const int8_t*)(blk + 192); // 16 int8 sub-scales + + uint16_t dh; + std::memcpy(&dh, blk + 208, sizeof(uint16_t)); + const float d = fp16_to_f32(dh); + + const int elem_base = 256 * b; + + // Process in two 128-element halves + for (int h = 0; h < 2; ++h) { + const uint8_t* ql_h = ql + 64 * h; // 64 bytes for this half + const uint8_t* qh_h = qh + 32 * h; // 32 bytes for this half + const int8_t* sc_h = scales + 8 * h; // 8 scales for this half + const int h_base = elem_base + 128 * h; + + // Process 4 groups of 32 elements, each group uses 2 of the 8 scales + for (int g = 0; g < 4; ++g) { + const int8_t* sc_g = sc_h + 2 * g; // 2 scales for this group + const int qh_bits_base = 2 * g; // Starting bit position in qh bytes + + // Group 1: ql_h[0:32] low nibbles, qh high bits [0:2] + // Group 2: ql_h[32:64] low nibbles, qh high bits [2:4] + // Group 3: ql_h[0:32] high nibbles, qh high bits [4:6] + // Group 4: ql_h[32:64] high nibbles, qh high bits [6:8] + const bool use_hi_nibble = (g >= 2); + const int ql_offset = (g % 2 == 1) ? 32 : 0; + + for (int l = 0; l < 32; ++l) { + const uint8_t ql_val = ql_h[ql_offset + l]; + const uint8_t qh_val = qh_h[l]; + + // Extract the 6-bit quant value: 4 bits from ql, 2 bits from qh + const int q_lo = use_hi_nibble ? (ql_val >> 4) : (ql_val & 0x0F); + // Shift is per-GROUP only. An extra term in l here silently corrupts the + // upper half of every group: see dequantize_block_q6_K, which reads + // (qh >> 2*g) & 3 for all 32 lanes of the group. + const int q_hi = (qh_val >> qh_bits_base) & 3; + const int q = (q_lo | (q_hi << 4)) - 32; + + // Select scale: use sc_g[0] for first 16 elements, sc_g[1] for next 16 + const int sc_idx = l / 16; + const int scale = sc_g[sc_idx]; + + // Calculate element index + const int elem_offset = 32 * g + l; + const int elem_idx = h_base + elem_offset; + + acc += d * scale * q * bf16_to_f32(x[elem_idx]); + } + } + } + } + + return acc; +} + +enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4, WF_Q4_K = 5, WF_Q6_K = 6 }; // Each ctor pointer arg is the address of a CPU int64 array of length // num_layers (one base address per layer, built by cpu_executor.py's @@ -1259,7 +1403,7 @@ struct CpuMoeExecutor { // it to a captured GPU elementwise kernel removes it while keeping the official // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py). bool input_prequant = false; - // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I). + // K-quant packed-row byte strides (Q4_0: H/32*18, Q4_K: H/256*144, Q6_K: H/256*210 for gate_up over K=H). int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0; float e2m1_lut[16]; float e4m3_lut[256]; @@ -1381,6 +1525,18 @@ struct CpuMoeExecutor { q4_gu_row_bytes = (H / 32) * 18; // K = H (gate_up rows) q4_dn_row_bytes = (I / 32) * 18; // K = I (down rows) } + if (weight_format == WF_Q4_K) { + if (H % 256 != 0 || I % 256 != 0) + throw std::runtime_error("Q4_K CPU MoE requires H and I to be multiples of 256"); + q4_gu_row_bytes = (H / 256) * 144; // K = H (gate_up rows) + q4_dn_row_bytes = (I / 256) * 144; // K = I (down rows) + } + if (weight_format == WF_Q6_K) { + if (H % 256 != 0 || I % 256 != 0) + throw std::runtime_error("Q6_K CPU MoE requires H and I to be multiples of 256"); + q4_gu_row_bytes = (H / 256) * 210; // K = H (gate_up rows) + q4_dn_row_bytes = (I / 256) * 210; // K = I (down rows) + } isa = c.name; // nvfp4 (AVX-VNNI only): W4A8 int8 decode when the CPU supports it. q4_0 is always // W4A8 (activations pre-quantized to Q8_0); select_q4dot picks VPDPBUSD / VPMADDUBSW @@ -1486,6 +1642,21 @@ struct CpuMoeExecutor { gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; return q4dot(w, xi8, xas, H); // W4A8: int8 activations (Q8_0), scale in xas } + if (fmt == WF_Q4_K) { + const uint8_t* w = gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; + return q4_k_dot_f32_scalar(w, x, H); // W4A16: bf16 activations, K-quant dequant + } + if (fmt == WF_Q6_K) { + const uint8_t* w = gu_packed_l + ((size_t)e * (2 * I) + row) * (size_t)q4_gu_row_bytes; + return q6_k_dot_f32_scalar(w, x, H); // W4A16: bf16 activations, K-quant dequant + } + // Anything that reaches here is assumed NVFP4 and dereferences the scale/global + // pointers, which are null for formats that do not have them (the GGUF banks pass 0). + // Falling through with an unhandled format therefore segfaults inside the worker + // thread rather than reporting anything useful, so reject it here instead. + TORCH_CHECK(fmt == WF_NVFP4 || fmt == WF_DSFP4, + "cpu_moe gemm1_dot: unhandled weight_format ", fmt, + " (handled: bf16=0, nvfp4=1, mxfp4=2, dsfp4=3, q4_0=4, q4_k=5, q6_k=6)"); const size_t r = (size_t)e * (2 * I) + row; if (use_vnni) return nvi8dot(gu_packed_l + r * (size_t)(H / 2), gu_scale_l + r * (size_t)(H / 16), @@ -1508,6 +1679,14 @@ struct CpuMoeExecutor { const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes; return q4dot(w, gi8, gas, I); // W4A8: int8 activations (Q8_0), scale in gas } + if (fmt == WF_Q4_K) { + const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes; + return q4_k_dot_f32_scalar(w, g, I); // W4A16: bf16 activations, K-quant dequant + } + if (fmt == WF_Q6_K) { + const uint8_t* w = dn_packed_l + ((size_t)e * H + row) * (size_t)q4_dn_row_bytes; + return q6_k_dot_f32_scalar(w, g, I); // W4A16: bf16 activations, K-quant dequant + } const size_t r = (size_t)e * H + row; if (use_vnni) return nvi8dot(dn_packed_l + r * (size_t)(I / 2), dn_scale_l + r * (size_t)(I / 16), diff --git a/python/freetoken/kernel/triton/dsv4/sparse_attn.py b/python/freetoken/kernel/triton/dsv4/sparse_attn.py index c22f891e..316824ba 100644 --- a/python/freetoken/kernel/triton/dsv4/sparse_attn.py +++ b/python/freetoken/kernel/triton/dsv4/sparse_attn.py @@ -39,11 +39,82 @@ import triton import triton.language as tl -BLOCK_H = 16 +_BLOCK_H_LARGE, _BLOCK_H_SMALL = 16, 8 +BLOCK_H = _BLOCK_H_LARGE # The gather has exactly ONE tl.load site (the pool base is selected per column), so it stages # a single [BLOCK_T, D] KV tile -- 67968 B at BLOCK_T=32, num_stages=2, which fits the ~99KB # consumer-Blackwell (sm_120, e.g. RTX 5090) budget. (BLOCK_T=64 would need ~103KB.) -BLOCK_T = 32 +# +# 32 does NOT fit every card. Turing (sm_75) caps shared memory at 64KB per block, and the +# same launch there reports Required: 100416 -- the tile scales with the head dim, which is +# 512 on DeepSeek-V4, so the figure above is not a universal constant. Halving the KV tile +# halves the staged bytes and costs iterations, not correctness. +_BLOCK_T_LARGE, _BLOCK_T_SMALL = 32, 16 + + +def _tile_plan(device_index: int | None = None, head_dim: int = 512) -> tuple[int, int, int]: + """(BLOCK_H, BLOCK_T, num_stages) that fit this device's opt-in shared memory. + + The dominant cost is NOT the KV tile: the kernel holds q and acc as [BLOCK_H, D] in + fp32, which at BLOCK_H=16 and head_dim 512 is 16*512*4*2 = 65536 B on its own -- exactly + a Turing block's entire budget, before a single KV byte. That is why shrinking BLOCK_T + alone leaves the requirement stuck at 66624. BLOCK_H is what has to come down on a 64KB + card; halving it costs head-parallelism, not correctness. + """ + try: + props = torch.cuda.get_device_properties(device_index) + optin = int(getattr(props, "shared_memory_per_block_optin", 0)) + except Exception: + optin = 0 + + if optin >= 102400: + return _BLOCK_H_LARGE, _BLOCK_T_LARGE, 2 + + # fp32 q + acc, the fixed floor, then the staged KV tile on top + for block_h in (_BLOCK_H_LARGE, _BLOCK_H_SMALL): + for block_t, stages in ((_BLOCK_T_LARGE, 2), (_BLOCK_T_SMALL, 2), (_BLOCK_T_SMALL, 1)): + need = 2 * block_h * head_dim * 4 + stages * block_t * head_dim * 2 + if need <= optin: + return block_h, block_t, stages + return _BLOCK_H_SMALL, _BLOCK_T_SMALL, 1 + + +def _unused_block_t(device_index: int | None = None) -> int: + """(BLOCK_T, num_stages) that fit this device's opt-in shared memory. + + num_stages=2 double-buffers the KV tile, so it roughly doubles the staged bytes. On a + 64KB card the small tile alone still lands at 66624 B -- about 1KB over -- so the tight + path also drops to a single stage. That costs pipelining, not correctness, and is a + smaller loss than halving the tile again to BLOCK_T=8. + """ + block_t = _block_t(device_index) + if block_t == _BLOCK_T_LARGE: + return block_t, 2 + try: + props = torch.cuda.get_device_properties(device_index) + optin = int(getattr(props, "shared_memory_per_block_optin", 0)) + except Exception: + optin = 0 + return block_t, (2 if optin >= 98304 else 1) + + +def _block_t(device_index: int | None = None) -> int: + """KV tile width that fits this device's opt-in shared memory. + + Queried per device rather than hardcoded: the budget is 64KB on sm_75, ~99KB on sm_89 + and sm_120, and ~164KB on sm_80/sm_90. Falling back to the small tile when the budget + is unknown is the safe direction -- a tile that does not fit fails the launch outright. + """ + try: + props = torch.cuda.get_device_properties(device_index) + except Exception: + return _BLOCK_T_SMALL + optin = int(getattr(props, "shared_memory_per_block_optin", 0)) + # measured requirement at BLOCK_T=32 with head_dim 512; leave the margin triton needs + return _BLOCK_T_LARGE if optin >= 102400 else _BLOCK_T_SMALL + + +BLOCK_T = _BLOCK_T_LARGE MAX_SPLITS = 32 MIN_TILES_PER_SPLIT = 4 @@ -330,7 +401,8 @@ def sparse_attn_paged( n_splits, ) - grid = (m, b, triton.cdiv(h, BLOCK_H)) + block_h, block_t, n_stages = _tile_plan(q.device.index, d) + grid = (m, b, triton.cdiv(h, block_h)) _sparse_attn_paged_kernel[grid]( q, window_pool, cmp_pool, o, sink, idx, cnt, float(softmax_scale), @@ -342,11 +414,11 @@ def sparse_attn_paged( idx.stride(0), idx.stride(1), idx.stride(2), stride_nb, stride_nm, D=d, - BLOCK_H=BLOCK_H, - BLOCK_T=BLOCK_T, + BLOCK_H=block_h, + BLOCK_T=block_t, HAS_COUNTS=has_counts, num_warps=8, - num_stages=2, + num_stages=n_stages, ) return o @@ -355,7 +427,8 @@ def _sparse_attn_paged_splitk( q, window_pool, cmp_pool, sink, idx, cnt, o, b, m, h, d, topk, n_window, softmax_scale, has_counts, stride_nb, stride_nm, n_splits, ): - head_blocks = triton.cdiv(h, BLOCK_H) + block_h, block_t, n_stages = _tile_plan(q.device.index, d) + head_blocks = triton.cdiv(h, block_h) mid_o = torch.empty((b, m, h, n_splits, d), dtype=torch.float32, device=q.device) mid_lse = torch.empty((b, m, h, n_splits), dtype=torch.float32, device=q.device) @@ -371,12 +444,12 @@ def _sparse_attn_paged_splitk( idx.stride(0), idx.stride(1), idx.stride(2), stride_nb, stride_nm, D=d, - BLOCK_H=BLOCK_H, - BLOCK_T=BLOCK_T, + BLOCK_H=block_h, + BLOCK_T=block_t, HAS_COUNTS=has_counts, NUM_SPLITS=n_splits, num_warps=8, - num_stages=2, + num_stages=n_stages, ) _sparse_attn_splitk_merge_kernel[(m, b, h)]( mid_o, mid_lse, o, sink, diff --git a/python/freetoken/layers/gguf.py b/python/freetoken/layers/gguf.py index ac49b1a5..5936d34d 100644 --- a/python/freetoken/layers/gguf.py +++ b/python/freetoken/layers/gguf.py @@ -1,5 +1,6 @@ """Native-GGUF quantized layers: weights stay in their packed block layout and are -dequantized *inside* the borrowed llama.cpp kernels (no bf16 copy ever materialized). +dequantized *inside* the borrowed llama.cpp CUDA kernels -- either fused into the matmul +(MMVQ/MMQ) or, for types with no MMQ kernel, by an explicit ``ggml_dequantize`` pass. Mirrors vLLM/sglang's ``GGUFLinearMethod`` / ``GGUFEmbeddingMethod`` dispatch, ported onto FreeToken's ``BaseOP``. FreeToken keeps fused projections (qkv, gate_up) as a @@ -8,6 +9,28 @@ share an input dim, hence the same ``row_bytes``), so a fused layer is still one ``[out, row_bytes]`` qweight -- no per-shard padding bookkeeping needed. +**Merged vs. plain fused projections**: + +When all output parts share the same quant type (the common case in gemma4), a plain +``GGUFLinear`` with concatenated packed rows is valid and efficient -- one kernel launch +dequantizes and multiplies. When parts use different quant types (as in Ornith's IQ3_M +checkpoint, where qkv_proj mixes IQ3_S and Q4_K), row_bytes differs per part, so torch.cat +would produce garbage. ``GGUFMergedLinear`` instead materializes the output of each part +separately via ``fused_mul_mat_gguf`` and concatenates the results along dim=-1 (equivalent +to the GEMM because all parts read the same input: ``cat([x @ W1.T, x @ W2.T]) == x @ cat([W1, W2], 0).T``). + +**Matmul dispatch strategy** (4-tier, per fused_mul_mat_gguf): + +1. **Unquantized (F32, F16, BF16)**: straight torch matmul ``x @ qweight.T``. +2. **Small-batch quantized (batch <= 6, MMVQ types)**: GEMV kernel via ``ggml_mul_mat_vec_a8``. +3. **Large-batch standard quants (MMQ types: Q4_0, Q4_1, Q5_0, Q5_1, Q8_0, K-quants)**: MMQ kernel + via ``ggml_mul_mat_a8``. +4. **Large-batch I-quants (IQ2_XXS, IQ2_XS, IQ3_XXS, IQ1_S, IQ4_NL, IQ3_S, IQ2_S, IQ4_XS, IQ1_M)**: + I-quants have MMVQ and dequant kernels but NO MMQ kernel. Prefill therefore falls back to + ``ggml_dequantize`` + plain torch matmul. This materializes a transient BF16 copy of the weight + (cost: ``out_features * in_features * 2 bytes``), which is a real tradeoff for memory-bound + prefill on large I-quant weights. + TP is assumed to be 1 (the gemma4 GGUF path restricts to TP=1, like the HF path). """ @@ -17,31 +40,40 @@ from freetoken.models.gguf.dequant import ( BLOCK_SHAPE, + DEQUANT_TYPES, GGML_BF16, GGML_F16, GGML_F32, GGML_NAME, - GGML_Q4_0, - GGML_Q6_K, - GGML_Q8_0, + GGML_UNQUANTIZED, + MMQ_TYPES, + MMVQ_TYPES, row_bytes, ) -from .base import BaseOP +# ggml type -> the dtype its raw bytes represent. Only the unquantized types appear here; +# everything else goes through a dequant kernel. +_UNQUANTIZED_DTYPE = { + GGML_F32: torch.float32, + GGML_F16: torch.float16, + GGML_BF16: torch.bfloat16, +} -# 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} +from .base import BaseOP # Below this token count, the MMVQ GEMV kernel wins (matches vLLM's heuristic). _MMVQ_SAFE = 6 def fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor: - """y = x @ dequant(qweight).T, dispatched by batch size and quant type.""" + """y = x @ dequant(qweight).T, dispatched by batch size and quant type. + + Dispatch order: + 1. Unquantized (F32/F16/BF16): plain torch matmul + 2. Small-batch quantized (batch <= 6, in MMVQ_TYPES): GEMV kernel + 3. Large-batch standard quants (in MMQ_TYPES): MMQ kernel + 4. Large-batch with I-quants (in DEQUANT_TYPES but not MMQ_TYPES): dequant + torch matmul + """ from freetoken.kernel.gguf import ( ggml_dequantize, ggml_mul_mat_a8, @@ -51,13 +83,28 @@ def fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int out_features = qweight.shape[0] if x.shape[0] == 0: return x.new_empty((0, out_features)) - if qweight_type in _UNQUANTIZED: - return x @ qweight.T - if x.shape[0] <= _MMVQ_SAFE and qweight_type in _MMVQ: + if qweight_type in GGML_UNQUANTIZED: + # GGUFLinear/GGUFEmbedding store every type in a uint8 buffer of row_bytes width, + # including the unquantized ones, where "packed" just means the raw F32/F16/BF16 + # bytes. Those must be reinterpreted before the matmul: multiplying the byte view + # directly gives an in_features of row_bytes (2x too wide for F16) and fails with + # "mat1 and mat2 shapes cannot be multiplied". A checkpoint only reaches this path + # when it stores a projection unquantized -- Apodex-1.1-mini ships output.weight as + # F16, which is how this surfaced; models whose lm_head is Q6_K never hit it. + w = qweight + if w.dtype == torch.uint8: + w = w.view(_UNQUANTIZED_DTYPE[qweight_type]) + # Cast the ACTIVATION, not the weight. Converting the weight would copy the whole + # matrix on every call -- about 1 GB per forward for a 248k-vocab lm_head -- and + # allocating that during CUDA graph capture fails outright. x is [tokens, hidden], + # so casting it is negligible, and computing in the stored precision is what + # llama.cpp does for these tensors anyway. + return (x.to(w.dtype) @ w.T).to(x.dtype) + if x.shape[0] <= _MMVQ_SAFE and qweight_type in MMVQ_TYPES: return ggml_mul_mat_vec_a8(qweight, x, qweight_type, out_features) - if qweight_type in _MMQ: + if qweight_type in MMQ_TYPES: return ggml_mul_mat_a8(qweight, x, qweight_type, out_features) - if qweight_type in _DEQUANT: + if qweight_type in DEQUANT_TYPES: block, type_size = BLOCK_SHAPE[qweight_type] in_features = qweight.shape[1] // type_size * block weight = ggml_dequantize(qweight, qweight_type, out_features, in_features, x.dtype) @@ -88,6 +135,123 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return out +class GGUFLMHead(GGUFLinear): + """LM head over a native GGUF ``output.weight`` (untied embeddings). + + Identical to ``GGUFLinear`` except that during prefill it keeps only the last position + of each sequence, exactly as ``ParallelLMHead`` (layers/embedding.py) and + ``GGUFTiedLMHead`` (models/gemma4/gguf.py) already do. + + This is not an optimization, it is a memory correctness issue. Logits are + [tokens, vocab], so on a large-vocabulary model the full-prefill tensor is enormous: + Ornith-1.5's vocab is 248,320, which in bf16 is 486 KiB of logits PER TOKEN. A + 1,800-token prompt therefore asks for a single 894 MB allocation, which is more than the + free VRAM left on an 8 GB card after weights and caches, and prefill dies with + "CUDA driver error: device not ready" while decode is completely unaffected. Only the + last position of each sequence is ever sampled, so every other row was computed and + thrown away. + + The dense path never hit this because ``ParallelLMHead`` slices; the bug appears only + when a GGUF checkpoint has untied embeddings and the head is swapped for a generic + quantized Linear, which has no reason to know it is the head. + """ + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from freetoken.core import get_global_ctx + + batch = get_global_ctx().batch + if batch.is_prefill: + indices = batch.attn_metadata.get_last_indices(batch.size) + x = x[indices].contiguous() + return super().forward(x) + + +class GGUFMergedLinear(BaseOP): + """Merged linear projection with parts that have different quant types. + + Used when fusing output-parallel projections (qkv, gate_up) whose parts use different + quantization types. Unlike GGUFLinear (which concatenates packed rows along dim 0 and + requires all parts to share row_bytes), GGUFMergedLinear materializes the output of + each part separately via fused_mul_mat_gguf, then concatenates the results. + + Mathematically equivalent to a single GEMM, since all parts read the same input x: + cat([x @ W1.T, x @ W2.T]) == x @ cat([W1, W2], 0).T + (source: llama.cpp's iq*_m mixed-quant strategy). + """ + + def __init__( + self, + in_features: int, + output_sizes: list[int], + quant_types: list[int], + has_bias: bool = False, + ): + """Initialize a merged linear projection. + + Args: + in_features: Input feature dimension (shared by all parts). + output_sizes: List of output sizes for each part; must all be > 0. + quant_types: List of GGML quant types, one per part; must match output_sizes length. + has_bias: Whether to allocate a bias term. + + Raises: + ValueError: If output_sizes and quant_types lengths do not match, or if any output_size <= 0. + NotImplementedError: If any quant_type is not supported (not in MMVQ_TYPES or GGML_UNQUANTIZED). + """ + if len(output_sizes) != len(quant_types): + raise ValueError( + f"output_sizes length {len(output_sizes)} != quant_types length {len(quant_types)}" + ) + if not all(o > 0 for o in output_sizes): + raise ValueError(f"all output_sizes must be > 0, got {output_sizes}") + + # Validate each quant type is supported. + for qt in quant_types: + if qt not in MMVQ_TYPES and qt not in GGML_UNQUANTIZED: + raise NotImplementedError( + f"quant type {GGML_NAME.get(qt, qt)} not in MMVQ_TYPES or GGML_UNQUANTIZED" + ) + + self.in_features = in_features + self.output_sizes = output_sizes + self.out_features = sum(output_sizes) + self._quant_types = quant_types + self.part_names = [] + + # Allocate packed weight buffers: one named tensor per part (qweight_0, qweight_1, ...). + # Named (not underscore-prefixed) so they are discovered by state_dict. + for i, (out_size, qt) in enumerate(zip(output_sizes, quant_types)): + name = f"qweight_{i}" + self.part_names.append(name) + setattr( + self, + name, + torch.empty(out_size, row_bytes(in_features, qt), dtype=torch.uint8), + ) + + self.bias = torch.empty(self.out_features) if has_bias else None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass: compute each part's output and concatenate along dim=-1. + + Args: + x: Input tensor of shape [..., in_features]. + + Returns: + Tensor of shape [..., out_features] with parts concatenated along dim=-1. + """ + parts = [] + for name, qt in zip(self.part_names, self._quant_types): + qweight = getattr(self, name) + part_out = fused_mul_mat_gguf(x, qweight, qt) + parts.append(part_out) + + out = torch.cat(parts, dim=-1) + if self.bias is not None: + out = out + self.bias + return out + + class GGUFEmbedding(BaseOP): """Vocab embedding stored as a native GGUF block-quantized table. @@ -116,7 +280,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: flat = x.flatten() rows = self.qweight.index_select(0, flat) # [n, row_bytes] packed - y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, torch.bfloat16) + if self._quant_type in GGML_UNQUANTIZED: + # Raw value bytes, not blocks: there is no dequant kernel for the unquantized + # types (ggml_dequantize rejects type 1), so reinterpret the gathered rows. + y = rows.view(_UNQUANTIZED_DTYPE[self._quant_type]).to(torch.bfloat16) + else: + y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, torch.bfloat16) y = y.view(*x.shape, self.embedding_dim) if self._embed_scale is not None: if self._embed_scale_t is None: @@ -125,4 +294,46 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y -__all__ = ["GGUFLinear", "GGUFEmbedding", "fused_mul_mat_gguf"] +def gguf_merged_or_plain( + in_features: int, + output_sizes: list[int], + quant_types: list[int], + has_bias: bool = False, +) -> GGUFLinear | GGUFMergedLinear: + """Choose between GGUFLinear (uniform quant types) and GGUFMergedLinear (mixed types). + + When all output parts share the same quant type (the uniform case, common in gemma4), + return a GGUFLinear with concatenated packed rows -- valid and cheaper since row_bytes + is identical per part (one kernel launch instead of N). + + When quant types differ (the mixed case, produced by llama.cpp's IQ*_M / Q*_K_M), + return a GGUFMergedLinear to avoid torch.cat garbage from misaligned row_bytes. + + Args: + in_features: Input feature dimension. + output_sizes: List of output sizes for each part. + quant_types: List of GGML quant types, one per part. + has_bias: Whether to allocate a bias term. + + Returns: + GGUFLinear if all quant types are identical, else GGUFMergedLinear. + """ + if len(set(quant_types)) == 1: + # Uniform case: all parts use the same quant type. + # Concatenate packed rows (they share row_bytes) into a single [sum(output_sizes), row_bytes] weight. + out_features = sum(output_sizes) + qt = quant_types[0] + lin = GGUFLinear(in_features, out_features, qt, has_bias=has_bias) + return lin + else: + # Mixed case: parts use different quant types. + return GGUFMergedLinear(in_features, output_sizes, quant_types, has_bias=has_bias) + + +__all__ = [ + "GGUFLinear", + "GGUFMergedLinear", + "GGUFEmbedding", + "fused_mul_mat_gguf", + "gguf_merged_or_plain", +] diff --git a/python/freetoken/models/deepseek_v4/__init__.py b/python/freetoken/models/deepseek_v4/__init__.py index b3cb41db..3fb8c91d 100644 --- a/python/freetoken/models/deepseek_v4/__init__.py +++ b/python/freetoken/models/deepseek_v4/__init__.py @@ -15,6 +15,13 @@ from .args import DeepseekV4Args, load_args from .config import parse_config +from .gguf import ( + convert_deepseek4_to_gguf, + is_gguf_model, + iter_gguf_weights, + parse_gguf_config, +) +from .gguf_experts import gguf_expert_types, load_gguf_expert_sources from .model import DeepseekV4ForCausalLM from .weight import iter_weights, load_dsfp4_expert_sources @@ -25,4 +32,10 @@ "DeepseekV4ForCausalLM", "iter_weights", "load_dsfp4_expert_sources", + "parse_gguf_config", + "iter_gguf_weights", + "convert_deepseek4_to_gguf", + "is_gguf_model", + "gguf_expert_types", + "load_gguf_expert_sources", ] diff --git a/python/freetoken/models/deepseek_v4/gguf.py b/python/freetoken/models/deepseek_v4/gguf.py new file mode 100644 index 00000000..160398db --- /dev/null +++ b/python/freetoken/models/deepseek_v4/gguf.py @@ -0,0 +1,583 @@ +"""Serve a deepseek4 GGUF checkpoint. + +Unlike the safetensors path, everything here comes from the GGUF's own metadata. The +reference ``parse_config`` recovers ``DeepseekV4Args`` from the checkpoint's +``inference/config.json``, which ships beside the weights; a standalone .gguf has no such +file, and being self-describing is the point of the format. ``_args_from_gguf`` below +rebuilds the same dataclass from KV keys alone. + +The mapping from GGUF tensor to model parameter was established by reading both sides +rather than by analogy with the qwen adapters, because three tensors do not behave the way +the names suggest: + +* ``attn_output_a`` is Q8_0 in the file but ``attn.wo_a`` is a bare ``nn.Parameter`` in + bfloat16, not a Linear (attention.py: "wo_a: dequantized to bf16, the reference runs a + bf16 grouped-output einsum"). It must be dequantized to dense, and it has no ``.weight`` + suffix. +* the compressor and indexer projections are **F16** in the file, i.e. unquantized. F16 is + in ``GGML_UNQUANTIZED``, so ``fused_mul_mat_gguf`` would take the ``x @ qweight.T`` path + while ``GGUFLinear`` allocates a uint8 buffer. They must land dense on a normal + ``.weight``, never packed. +* ``Indexer.wq_b`` is declared ``Linear(kind="fp8")`` (compress.py), which allocates a + ``.scale`` that no GGUF tensor can fill, because that tensor is F16 here. That Linear is + replaced outright rather than populated. + +Routed experts never pass through this module; they are streamed from the offload cache by +``gguf_experts.load_gguf_expert_sources``. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Iterator + +import torch + +from freetoken.models.config import DSV4AttentionGroupConfig, ModelConfig, RotaryConfig + +from .args import DeepseekV4Args + +if TYPE_CHECKING: + from freetoken.models.gguf.config import GgufConfigShim + +_ARCH = "deepseek4" + +# llama.cpp's expert-gating enum. DeepSeek-V4 scores with sqrt-softplus; 1 and 2 are the +# long-standing softmax/sigmoid values. An unknown id raises rather than silently picking a +# scoring function, because the wrong one routes to the wrong experts and still produces +# fluent text. +_GATING = {1: "softmax", 2: "sigmoid", 4: "sqrtsoftplus"} + + +def _kv(shim: "GgufConfigShim", key: str, default: Any = None) -> Any: + """One ``deepseek4.*`` metadata value. No default means the key is mandatory.""" + full = f"{_ARCH}.{key}" + md = shim if isinstance(shim, dict) else shim.metadata + if full not in md: + if default is None: + raise ValueError( + f"deepseek4 GGUF is missing required metadata key {full!r}; this file does " + f"not carry the config this adapter needs" + ) + return default + return md[full] + + +def _args_from_gguf(shim: "GgufConfigShim") -> DeepseekV4Args: + """Rebuild DeepseekV4Args from GGUF metadata alone. + + Every field is sourced from a key that is actually present in the checkpoint; nothing + is left to the dataclass default, because a silently-defaulted hyperparameter here + produces a model that loads and generates confidently wrong text. + + Cross-checks worth keeping: ``compress_ratios`` carries one entry per layer plus the + MTP layers, entries != 0 mark layers with an attention compressor, and entries == 4 + mark layers with the lightning indexer. Those counts must match the tensor table (41 + and 21 respectively for DeepSeek-V4-Flash), which is what makes this mapping + self-validating rather than merely plausible. + """ + ratios = tuple(int(x) for x in _kv(shim, "attention.compress_ratios")) + swiglu = [float(x) for x in _kv(shim, "swiglu_clamp_exp", [])] + gate_id = int(_kv(shim, "expert_gating_func")) + if gate_id not in _GATING: + raise ValueError( + f"deepseek4 GGUF: unknown expert_gating_func {gate_id}; known values are " + f"{sorted(_GATING)} (routing with the wrong scoring function still generates " + f"fluent text, so this is not defaulted)" + ) + + return DeepseekV4Args( + max_batch_size=1, + max_seq_len=int(_kv(shim, "context_length")), + # The fp8/fp4 reference paths do not apply: a GGUF carries its own block-quantized + # weights and the adapter swaps the quantized projections for GGUF ops. + dtype="bf16", + scale_fmt=None, + expert_dtype=None, + vocab_size=int(_kv(shim, "vocab_size")), + dim=int(_kv(shim, "embedding_length")), + moe_inter_dim=int(_kv(shim, "expert_feed_forward_length")), + n_layers=int(_kv(shim, "block_count")), + n_hash_layers=int(_kv(shim, "hash_layer_count")), + n_mtp_layers=int(_kv(shim, "nextn_predict_layers", 0)), + n_heads=int(_kv(shim, "attention.head_count")), + n_routed_experts=int(_kv(shim, "expert_count")), + n_shared_experts=int(_kv(shim, "expert_shared_count")), + n_activated_experts=int(_kv(shim, "expert_used_count")), + score_func=_GATING[gate_id], + route_scale=float(_kv(shim, "expert_weights_scale")), + swiglu_limit=(swiglu[0] if swiglu else 10.0), + q_lora_rank=int(_kv(shim, "attention.q_lora_rank")), + head_dim=int(_kv(shim, "attention.key_length")), + rope_head_dim=int(_kv(shim, "rope.dimension_count")), + norm_eps=float(_kv(shim, "attention.layer_norm_rms_epsilon")), + o_groups=int(_kv(shim, "attention.output_group_count")), + o_lora_rank=int(_kv(shim, "attention.output_lora_rank")), + window_size=int(_kv(shim, "attention.sliding_window")), + compress_ratios=ratios, + compress_rope_theta=float(_kv(shim, "attention.compress_rope_freq_base")), + original_seq_len=int(_kv(shim, "rope.scaling.original_context_length")), + rope_theta=float(_kv(shim, "rope.freq_base")), + rope_factor=float(_kv(shim, "rope.scaling.factor")), + beta_fast=int(_kv(shim, "rope.scaling.yarn_beta_fast")), + beta_slow=int(_kv(shim, "rope.scaling.yarn_beta_slow")), + index_n_heads=int(_kv(shim, "attention.indexer.head_count")), + index_head_dim=int(_kv(shim, "attention.indexer.key_length")), + index_topk=int(_kv(shim, "attention.indexer.top_k")), + hc_mult=int(_kv(shim, "hyper_connection.count")), + hc_sinkhorn_iters=int(_kv(shim, "hyper_connection.sinkhorn_iterations")), + hc_eps=float(_kv(shim, "hyper_connection.epsilon")), + ) + + + +def _check_schedule(model_path: str, args: DeepseekV4Args, served: int) -> None: + """Cross-check the compress_ratios schedule against the tensor table. + + Cheap (the tensor table is metadata, not weights) and worth doing every load: it is the + difference between finding a layer-count error here and finding it as degraded output + after a 145 GiB load. + """ + from freetoken.models.gguf.reader import gguf_tensor_names + + names = gguf_tensor_names(model_path) + want_compressor = sum(1 for r in args.compress_ratios[:served] if r != 0) + want_indexer = sum(1 for r in args.compress_ratios[:served] if r == 4) + got_compressor = sum( + 1 for i in range(served) if f"blk.{i}.attn_compressor_kv.weight" in names) + got_indexer = sum( + 1 for i in range(served) if f"blk.{i}.indexer.attn_q_b.weight" in names) + + for label, want, got in (("compressor", want_compressor, got_compressor), + ("indexer", want_indexer, got_indexer)): + if want != got: + raise ValueError( + f"deepseek4 GGUF: compress_ratios predicts {want} layers with a {label} " + f"but the file has {got}; the per-layer schedule does not match this " + f"checkpoint (a wrong served-layer count is the usual cause)" + ) + +def parse_gguf_config(shim: "GgufConfigShim") -> ModelConfig: + """ModelConfig for a deepseek4 GGUF, mirroring deepseek_v4/config.py::parse_config. + + The served layer count excludes the trailing MTP/NextN block: ``block_count`` counts it + but it is not part of the forward pass, and treating it as a layer makes a uniform + expert bank look mixed. + """ + args = _args_from_gguf(shim) + model_path = getattr(shim, "model_path", None) + + # How block_count relates to the MTP block is NOT consistent across architectures, so + # it is derived rather than assumed. qwen35moe counts its NextN block inside + # block_count (Ornith: block_count 41, blk.0..blk.40 where blk.40 is the MTP block, 40 + # served). deepseek4 does not (block_count 43, blk.0..blk.42 all served, and the MTP + # layer carries no blk tensors at all). Subtracting n_mtp_layers unconditionally + # silently drops the last real layer here. + # + # compress_ratios is the authority: it carries one entry per served layer plus the MTP + # layers, so the served count falls out of it and is then cross-checked below. + served_layers = len(args.compress_ratios) - args.n_mtp_layers + if served_layers != args.n_layers: + raise ValueError( + f"deepseek4 GGUF: compress_ratios implies {served_layers} served layers " + f"({len(args.compress_ratios)} entries minus {args.n_mtp_layers} MTP) but " + f"block_count is {args.n_layers}; refusing to guess which is right" + ) + args.n_layers = served_layers + + rope_scaling = { + "rope_type": "yarn", + "factor": args.rope_factor, + "beta_fast": args.beta_fast, + "beta_slow": args.beta_slow, + "original_max_position_embeddings": args.original_seq_len, + } + + from .gguf_experts import gguf_expert_types + + types = gguf_expert_types(model_path, served_layers) if model_path else None + expert_types = (types["gate_up"][0], types["down"][0]) if types else None + + # The schedule derived from compress_ratios must match what the file actually contains. + # A compressor exists where ratio != 0 and a lightning indexer where ratio == 4, so + # these counts are an independent check on the layer count above: an off-by-one shows + # up here as a mismatch rather than as a quietly missing layer at serving time. + if model_path: + _check_schedule(model_path, args, served_layers) + + return ModelConfig( + num_layers=served_layers, + num_qo_heads=args.n_heads, + num_kv_heads=1, # MLA: a single shared latent KV head (K == V) + head_dim=args.head_dim, + hidden_size=args.dim, + vocab_size=args.vocab_size, + intermediate_size=args.moe_inter_dim, + hidden_act="silu", + rms_norm_eps=args.norm_eps, + tie_word_embeddings=False, # output.weight is a separate tensor from token_embd + rotary_config=RotaryConfig( + head_dim=args.head_dim, + rotary_dim=args.rope_head_dim, + max_position=args.max_seq_len, + base=args.rope_theta, + scaling=rope_scaling, + ), + num_experts=args.n_routed_experts, + num_experts_per_tok=args.n_activated_experts, + moe_intermediate_size=args.moe_inter_dim, + norm_topk_prob=True, + model_type="deepseek_v4", + architectures=["DeepseekV4ForCausalLM"], + moe_enabled=True, + expert_quant="gguf", + attn_sm_scale=args.head_dim**-0.5, + dsv4_args=args, + gguf_model_path=model_path, + gguf_expert_types=expert_types, + attention_groups=( + DSV4AttentionGroupConfig( + name="dsv4", + layer_ids=tuple(range(served_layers)), + num_kv_heads=1, + head_dim=args.head_dim, + sliding_window=args.window_size, + ), + ), + ) + + +def is_gguf_model(config: ModelConfig) -> bool: + """True when this config came from a GGUF checkpoint (native block-quant path).""" + return getattr(config, "gguf_model_path", None) is not None + + + +class GGUFLinearNN(torch.nn.Module): + """A GGUF-quantized Linear that DSV4's loader can actually fill. + + FreeToken's own ``layers.gguf.GGUFLinear`` is a ``BaseOP`` holding ``qweight`` as a + plain tensor. That works for the qwen models, whose trees are built from BaseOP, but + deepseek_v4 is raw ``nn.Module`` and loads via + ``DeepseekV4ForCausalLM.load_state_dict``, which walks ``named_parameters()`` and + demands a key for every one. A plain attribute is invisible there, and assigning a + non-Module over a Module child raises outright. + + So the packed block bytes live in an ordinary ``nn.Parameter`` -- uint8, requires_grad + False -- named ``weight`` to match the naming the rest of this model uses. The loader's + ``.to(p.dtype)`` cast is then a no-op on uint8, and the tensor arrives byte-for-byte. + """ + + def __init__(self, in_features: int, out_features: int, quant_type: int, + bias: bool = False): + super().__init__() + from freetoken.models.gguf.dequant import row_bytes + + self.in_features = in_features + self.out_features = out_features + self._quant_type = int(quant_type) + self.weight = torch.nn.Parameter( + torch.empty(out_features, row_bytes(in_features, self._quant_type), + dtype=torch.uint8), + requires_grad=False, + ) + if bias: + self.bias = torch.nn.Parameter(torch.empty(out_features), requires_grad=False) + else: + self.register_parameter("bias", None) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from freetoken.layers.gguf import fused_mul_mat_gguf + + out = fused_mul_mat_gguf(x, self.weight, self._quant_type) + return out if self.bias is None else out + self.bias + + +class GGUFEmbeddingNN(torch.nn.Module): + """GGUF-quantized vocab embedding, as an nn.Module for the same reason as above. + + The table is never dequantized whole: only the looked-up rows are gathered in packed + form and dequantized per lookup. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int, quant_type: int): + super().__init__() + from freetoken.models.gguf.dequant import row_bytes + + self.num_embeddings = num_embeddings + self.embedding_dim = embedding_dim + self._quant_type = int(quant_type) + self.weight = torch.nn.Parameter( + torch.empty(num_embeddings, row_bytes(embedding_dim, self._quant_type), + dtype=torch.uint8), + requires_grad=False, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + from freetoken.kernel.gguf import ggml_dequantize + + flat = x.flatten() + rows = self.weight.index_select(0, flat) + y = ggml_dequantize(rows, self._quant_type, flat.shape[0], self.embedding_dim, + torch.bfloat16) + return y.view(*x.shape, self.embedding_dim) + + +def _dense(t, dtype: torch.dtype) -> torch.Tensor: + """A GgufTensor as a dense tensor of its torch shape. + + Two paths, because neither covers everything. Unquantized types (F32/F16/BF16) are + already values, so the packed bytes are simply reinterpreted -- no kernel needed, and + it works without CUDA. Block-quantized types go through the vendored CUDA dequant: + ``dequant.dequantize``'s pure-torch fallback only implements Q4_0 and Q6_K, and this + checkpoint stores its attention projections and lm_head as Q8_0. + """ + from freetoken.models.gguf.dequant import ( + BLOCK_SHAPE, + GGML_BF16, + GGML_F16, + GGML_F32, + GGML_UNQUANTIZED, + ) + + gt = int(t.ggml_type) + raw = t.packed() + if gt in GGML_UNQUANTIZED: + view = {GGML_F32: torch.float32, GGML_F16: torch.float16, + GGML_BF16: torch.bfloat16}[gt] + return raw.reshape(-1).view(view).reshape(t.shape).to(dtype) + + from freetoken.kernel.gguf import ggml_dequantize + + block, type_size = BLOCK_SHAPE[gt] + in_features = t.row_bytes // type_size * block + out = ggml_dequantize(raw.cuda().contiguous(), gt, t.rows, in_features, + torch.bfloat16) + return out.reshape(t.shape).to(dtype) + + +def _to_bf16(t) -> torch.Tensor: + return _dense(t, torch.bfloat16) + + +def _to_f32(t) -> torch.Tensor: + return _dense(t, torch.float32) + + +def _to_i64(t) -> torch.Tensor: + """Read an I32 index table as int64. + + tid2eid is a routing table, not a weight: dequantizing it through a float path would + round large token ids. Reinterpret the raw bytes instead. + """ + return t.packed().reshape(-1).view(torch.int32).reshape(t.shape).to(torch.int64) + + +# suffix -> (destination template, kind). "packed" lands on a GGUFLinear's .weight; +# "bf16"/"f32" are dequantized onto an ordinary parameter. The destination is spelled out +# per tensor rather than derived from the name, because three of them do not follow the +# pattern the names imply (see the module docstring). +_LAYER_MAP: dict[str, tuple[str, str]] = { + "attn_norm.weight": ("attn_norm.weight", "f32"), + "ffn_norm.weight": ("ffn_norm.weight", "f32"), + "attn_q_a.weight": ("attn.wq_a.weight", "packed"), + "attn_q_a_norm.weight": ("attn.q_norm.weight", "f32"), + "attn_q_b.weight": ("attn.wq_b.weight", "packed"), + "attn_kv.weight": ("attn.wkv.weight", "packed"), + "attn_kv_a_norm.weight": ("attn.kv_norm.weight", "f32"), + # wo_a is a bare nn.Parameter in bf16, NOT a Linear: no .weight, never packed. + "attn_output_a.weight": ("attn.wo_a", "bf16"), + "attn_output_b.weight": ("attn.wo_b.weight", "packed"), + "attn_sinks.weight": ("attn.attn_sink", "f32"), + # compressor / indexer projections are F16 in the file. F16 is in GGML_UNQUANTIZED, so + # GGUFLinear cannot hold them -- they must land dense on a normal .weight. + "attn_compressor_kv.weight": ("attn.compressor.wkv.weight", "bf16"), + "attn_compressor_gate.weight": ("attn.compressor.wgate.weight", "bf16"), + "attn_compressor_norm.weight": ("attn.compressor.norm.weight", "f32"), + "attn_compressor_ape.weight": ("attn.compressor.ape", "f32"), + "indexer.attn_q_b.weight": ("attn.indexer.wq_b.weight", "bf16"), + "indexer.proj.weight": ("attn.indexer.weights_proj.weight", "bf16"), + "indexer_compressor_kv.weight": ("attn.indexer.compressor.wkv.weight", "bf16"), + "indexer_compressor_gate.weight": ("attn.indexer.compressor.wgate.weight", "bf16"), + "indexer_compressor_norm.weight": ("attn.indexer.compressor.norm.weight", "f32"), + "indexer_compressor_ape.weight": ("attn.indexer.compressor.ape", "f32"), + "hc_attn_base.weight": ("hc_attn_base", "f32"), + "hc_attn_fn.weight": ("hc_attn_fn", "f32"), + "hc_attn_scale.weight": ("hc_attn_scale", "f32"), + "hc_ffn_base.weight": ("hc_ffn_base", "f32"), + "hc_ffn_fn.weight": ("hc_ffn_fn", "f32"), + "hc_ffn_scale.weight": ("hc_ffn_scale", "f32"), + "ffn_gate_inp.weight": ("ffn.gate.weight", "bf16"), + "exp_probs_b.bias": ("ffn.gate.bias", "f32"), + # DeepSeek names the shared expert gate/up/down; Expert calls them w1/w3/w2. + "ffn_gate_shexp.weight": ("ffn.shared_experts.w1.weight", "packed"), + "ffn_up_shexp.weight": ("ffn.shared_experts.w3.weight", "packed"), + "ffn_down_shexp.weight": ("ffn.shared_experts.w2.weight", "packed"), +} + +_GLOBAL_MAP: dict[str, tuple[str, str]] = { + "output_norm.weight": ("norm.weight", "f32"), + # output.weight is Q8_0 but `head` is a bare bf16 nn.Parameter consumed by F.linear, + # so it is dequantized rather than swapped. deepseek_v4/model.py already slices to the + # last prefill position itself, so it needs no GGUFLMHead. + "output.weight": ("head", "bf16"), + "output_hc_base.weight": ("hc_head_base", "f32"), + "output_hc_fn.weight": ("hc_head_fn", "f32"), + "output_hc_scale.weight": ("hc_head_scale", "f32"), +} + +# Routed experts are streamed from the offload cache, never yielded here. +_EXPERT_SUFFIXES = frozenset( + {"ffn_gate_exps.weight", "ffn_up_exps.weight", "ffn_down_exps.weight"}) + + +def iter_gguf_weights( + model_path: str, + device, + *, + include_moe_experts: bool, + include_non_moe: bool, +) -> Iterator[tuple[str, torch.Tensor]]: + """Yield (param_name, tensor) for every non-expert deepseek4 parameter.""" + import re + + from freetoken.models.gguf.reader import iter_gguf_tensors + + assert not include_moe_experts, ( + "deepseek4 GGUF keeps its routed experts in the offload cache; they are loaded by " + "gguf_experts.load_gguf_expert_sources, not by iter_gguf_weights." + ) + assert include_non_moe + + conv = {"packed": lambda t: t.packed(), "bf16": _to_bf16, "f32": _to_f32} + + for t in iter_gguf_tensors(model_path): + name = t.name + m = re.match(r"^blk\.(\d+)\.(.+)$", name) + if m is None: + dest = _GLOBAL_MAP.get(name) + if dest is None: + if name == "token_embd.weight": + yield "embed.weight", t.packed() + continue + raise ValueError( + f"deepseek4 GGUF: unmapped global tensor {name!r}; this checkpoint does " + f"not match the layout this adapter expects" + ) + path, kind = dest + yield path, conv[kind](t) + continue + + layer, suffix = int(m.group(1)), m.group(2) + if suffix in _EXPERT_SUFFIXES: + continue # offload cache + if suffix == "ffn_gate_tid2eid.weight": + # Hash routing table on the first n_hash_layers layers; an index, not a weight. + yield f"layers.{layer}.ffn.gate.tid2eid", _to_i64(t) + continue + dest = _LAYER_MAP.get(suffix) + if dest is None: + raise ValueError( + f"deepseek4 GGUF: unmapped tensor {name!r}; this checkpoint does not match " + f"the layout this adapter expects" + ) + path, kind = dest + yield f"layers.{layer}.{path}", conv[kind](t) + + +def _scan_quant_types(model_path: str) -> dict[tuple[int, str], int]: + """(layer, suffix) -> ggml type, straight from the tensor table. + + A guessed type allocates a wrong-sized packed buffer, so nothing here has a default. + Globals use layer -1. + """ + import re + + from freetoken.models.gguf.reader import iter_gguf_tensors + + out: dict[tuple[int, str], int] = {} + for t in iter_gguf_tensors(model_path): + m = re.match(r"^blk\.(\d+)\.(.+)$", t.name) + if m: + out[(int(m.group(1)), m.group(2))] = int(t.ggml_type) + else: + out[(-1, t.name)] = int(t.ggml_type) + return out + + +def convert_deepseek4_to_gguf(model, config: ModelConfig, *, model_path: str) -> None: + """In place: swap deepseek4's quantized projections + embedding for native GGUF ops. + + Swapped to GGUFLinear (Q8_0 in the checkpoint): attention wq_a / wq_b / wkv / wo_b and + the shared expert's w1 / w2 / w3. + + Deliberately NOT swapped: + * ``attn.wo_a`` is a bare bf16 nn.Parameter, not a Linear; it is dequantized dense. + * the compressor's wkv / wgate and the indexer's weights_proj are already + ``Linear(kind="bf16")`` and their tensors are F16, so they take dense weights. + * ``head`` is a bare bf16 nn.Parameter consumed by F.linear. + + Replaced rather than swapped: ``indexer.wq_b`` is declared ``Linear(kind="fp8")``, + which allocates a ``.scale`` no GGUF tensor can fill because that tensor is F16 here. + It becomes a bf16 Linear so ``.weight`` is bf16 and ``scale`` is None. + """ + from .layers import Linear + + quant = _scan_quant_types(model_path) + + def qt(layer: int, suffix: str) -> int: + key = (layer, suffix) + if key not in quant: + where = suffix if layer < 0 else f"blk.{layer}.{suffix}" + raise ValueError( + f"deepseek4 GGUF {model_path}: expected tensor {where} is absent, so its " + f"quant type cannot be read; this checkpoint does not match the layout " + f"this adapter expects" + ) + return quant[key] + + def swap_linear(owner, attr: str, quant_type: int) -> None: + lin = getattr(owner, attr) + setattr( + owner, attr, + GGUFLinearNN(lin.in_features, lin.out_features, quant_type, + bias=getattr(lin, "bias", None) is not None), + ) + + # DeepseekV4ForCausalLM is an engine wrapper; the parameters live on the inner + # Transformer, and state_dict() names them relative to it (no "_transformer." prefix), + # which is what iter_gguf_weights emits. + root = getattr(model, "_transformer", model) + + root.embed = GGUFEmbeddingNN( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + quant_type=qt(-1, "token_embd.weight"), + ) + + for layer_idx, layer in enumerate(root.layers): + attn = layer.attn + swap_linear(attn, "wq_a", qt(layer_idx, "attn_q_a.weight")) + swap_linear(attn, "wq_b", qt(layer_idx, "attn_q_b.weight")) + swap_linear(attn, "wkv", qt(layer_idx, "attn_kv.weight")) + swap_linear(attn, "wo_b", qt(layer_idx, "attn_output_b.weight")) + + idx = getattr(attn, "indexer", None) + if idx is not None: + # F16 in the file, fp8 in the module: rebuild as bf16 so there is no orphan + # .scale and F.linear is used instead of the block-fp8 GEMM. + old = idx.wq_b + idx.wq_b = Linear(old.in_features, old.out_features, + bias=getattr(old, "bias", None) is not None, kind="bf16") + + shexp = layer.ffn.shared_experts + swap_linear(shexp, "w1", qt(layer_idx, "ffn_gate_shexp.weight")) + swap_linear(shexp, "w3", qt(layer_idx, "ffn_up_shexp.weight")) + swap_linear(shexp, "w2", qt(layer_idx, "ffn_down_shexp.weight")) + + +__all__ = [ + "parse_gguf_config", + "iter_gguf_weights", + "convert_deepseek4_to_gguf", + "is_gguf_model", +] diff --git a/python/freetoken/models/deepseek_v4/gguf_experts.py b/python/freetoken/models/deepseek_v4/gguf_experts.py new file mode 100644 index 00000000..2ba32b08 --- /dev/null +++ b/python/freetoken/models/deepseek_v4/gguf_experts.py @@ -0,0 +1,254 @@ +"""Routed-expert host banks for a deepseek4 GGUF checkpoint. + +Ported from ``models/qwen3_5_moe/gguf_experts.py``; the two are structurally the same job +because llama.cpp emits the same three stacked tensors for both architectures +(``ffn_gate_exps`` / ``ffn_up_exps`` / ``ffn_down_exps``). What differs is only the +arithmetic: DeepSeek-V4-Flash is 43 served layers of 256 experts at +``moe_inter_dim`` 2048 over ``dim`` 4096, and the checkpoints worth loading are uniformly +Q4_K on all three banks rather than qwen35moe's per-layer mix. + +Only the ROUTED experts come through here. The shared expert +(``ffn_{gate,up,down}_shexp``) is an ordinary quantized Linear that +``convert_deepseek4_to_gguf`` swaps for a ``GGUFLinear``, exactly as qwen3_5_moe handles +its own shared expert -- it is dense per token, so there is nothing to offload. + +A note on which checkpoints reach this code at all. The offload slot pool is ONE +allocation per bank shared by every layer, and ``moe_vec.cuh`` addresses it as +``expert * nrows * (ncols / qk)`` with no padding allowance, so a bank whose ggml type +varies by layer cannot be served. Of the thirteen published +``unsloth/DeepSeek-V4-Flash-0731-GGUF`` variants, eleven mix types across layers and the +two that do not are MXFP4 (ggml type 39), which has no entry in ``BLOCK_SHAPE`` and no +vendored kernel. The ``antirez/deepseek-v4-gguf`` builds are the ones that load: their +Q4KExperts variant is uniformly Q4_K across all 43 layers, verified by reading the file. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from freetoken.models.gguf.dequant import GGML_NAME, GGML_Q4_K, row_bytes + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + + +def gguf_expert_types(model_path: str, num_layers: int) -> dict[str, list[int]]: + """Scan the tensor table and return the per-layer ggml type of each expert bank. + + Returns ``{"gate_up": [...], "down": [...]}``, each a list of ``num_layers`` ggml type + enums. gate and up must agree per layer because they are row-concatenated into one + bank and therefore must share a row stride; a mismatch raises here naming both types. + + ``expert_banks._gguf_banks`` consumes this and is what rejects a bank that is + non-uniform ACROSS layers, with the user-facing message about ``--pure``. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + + gate_types: list[int | None] = [None] * num_layers + up_types: list[int | None] = [None] * num_layers + down_types: list[int | None] = [None] * num_layers + + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if layer >= num_layers: + # The trailing NextN/MTP block. DeepSeek-V4-Flash ships nextn_predict_layers=1, + # so the file carries a block at index num_layers that is not served; counting + # it here would make a uniform checkpoint look mixed. + continue + + if t.name.endswith("ffn_gate_exps.weight"): + gate_types[layer] = t.ggml_type + elif t.name.endswith("ffn_up_exps.weight"): + up_types[layer] = t.ggml_type + elif t.name.endswith("ffn_down_exps.weight"): + down_types[layer] = t.ggml_type + + gate_up_types: list[int] = [] + for layer in range(num_layers): + gate_t, up_t = gate_types[layer], up_types[layer] + if gate_t is None or up_t is None: + raise ValueError( + f"deepseek4 GGUF: layer {layer} is missing routed-expert tensors " + f"(gate={gate_t}, up={up_t}); every layer of this architecture is MoE" + ) + if gate_t != up_t: + raise ValueError( + f"deepseek4 GGUF: layer {layer} has ffn_gate_exps " + f"{GGML_NAME.get(gate_t, gate_t)} but ffn_up_exps " + f"{GGML_NAME.get(up_t, up_t)}; they are row-concatenated into one bank and " + "cannot have different row strides" + ) + gate_up_types.append(gate_t) + + for layer in range(num_layers): + if down_types[layer] is None: + raise ValueError(f"deepseek4 GGUF: layer {layer} is missing ffn_down_exps") + + return {"gate_up": gate_up_types, "down": down_types} + + +def gguf_expert_specs( + config: "ModelConfig", types: dict[str, list[int]] +) -> dict[str, tuple[tuple[int, ...], torch.dtype]]: + """Expert bank shapes as ``{name: (shape, dtype)}`` -- ``alloc_layer_banks``' contract. + + Packed block bytes, in torch order:: + + gate_up (E, 2*I, row_bytes(H, t_gate_up)) uint8 + down (E, H, row_bytes(I, t_down)) uint8 + + One spec per bank rather than per layer, for the slot-pool stride reason in the module + docstring. A non-uniform bank is rejected here instead of being mis-decoded. + """ + E = config.num_experts + H = config.hidden_size + I = config.moe_intermediate_size + out: dict[str, tuple[tuple[int, ...], torch.dtype]] = {} + for name, elems in (("gate_up", H), ("down", I)): + distinct = sorted(set(types[name])) + if len(distinct) != 1: + names = [GGML_NAME.get(t, t) for t in distinct] + raise ValueError( + f"deepseek4 expert bank {name!r} mixes ggml types across layers ({names}); " + "a bank must be uniform because its slot pool is one allocation with one " + "stride" + ) + rb = row_bytes(elems, distinct[0]) + shape = (E, 2 * I, rb) if name == "gate_up" else (E, H, rb) + out[name] = (shape, torch.uint8) + return out + + +def load_gguf_expert_sources( + model_path: str, config: "ModelConfig", *, layer_sink=None +) -> dict[str, list[torch.Tensor]]: + """Per-layer host banks holding the routed experts' native packed block bytes. + + Nothing is dequantized: the bytes handed to the offload cache are the same ones the + kernels decode in the K-loop. + + ``layer_sink`` None (serving) pins each completed layer through an internally owned + ``PinPipeline``; a supplied sink (converter) receives the completion notifications + instead and may release banks, so the returned tensors live only as long as it allows. + """ + from freetoken.models.gguf.reader import iter_gguf_tensors + from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline, alloc_layer_banks + + types = gguf_expert_types(model_path, config.num_layers) + specs = gguf_expert_specs(config, types) + + L = config.num_layers + E = config.num_experts + H = config.hidden_size + I = config.moe_intermediate_size + + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in hb} + + gate_buf: dict[int, torch.Tensor] = {} + up_buf: dict[int, torch.Tensor] = {} + seen_gate: set[int] = set() + seen_up: set[int] = set() + seen_down: set[int] = set() + + def _load(sink) -> None: + tracker = LayerCompletionTracker(2, hb, sink) if sink is not None else None + + for t in iter_gguf_tensors(model_path): + if not t.name.startswith("blk."): + continue + layer = int(t.name.split(".")[1]) + if layer >= L: + continue # trailing NextN/MTP block, not served + + if t.name.endswith("ffn_gate_exps.weight"): + gate_buf[layer] = t.packed() + seen_gate.add(layer) + elif t.name.endswith("ffn_up_exps.weight"): + up_buf[layer] = t.packed() + seen_up.add(layer) + elif t.name.endswith("ffn_down_exps.weight"): + # torch shape [E, H, I] is ggml dims [I, H, E] with I fastest, so the reader + # returns [E*H, row_bytes(I)] already in expert-major row order. Reshaping + # to [E, H, row_bytes(I)] is a view, not a copy. Note the row_bytes is over + # I (the fastest dim), not over E. + down_rb = specs["down"][0][2] + banks["down"][layer].copy_(t.packed().reshape(E, H, down_rb)) + seen_down.add(layer) + if tracker is not None: + tracker.note(layer) + else: + continue + + if layer in gate_buf and layer in up_buf: + rb = specs["gate_up"][0][2] + # gate and up each arrive as [E*I, row_bytes(H)], and ggml's fastest-first + # dims [H, I, E] make E the slowest axis, so those rows are EXPERT-MAJOR: + # expert e owns rows [e*I, (e+1)*I). + # + # The bank must be [E, 2I, row_bytes(H)] with each expert's own gate rows + # followed by its OWN up rows, so reshape to [E, I, rb] and concatenate on + # the row axis within each expert (dim=1). + # + # cat(dim=0) then reshape(E, 2I, rb) -- the version that looks obviously + # right -- lays every expert's gate down before any up, so expert 0 would + # get its gate rows plus expert 1's gate rows and its up would sit E*I rows + # away. That loads, runs at full speed, and emits fluent nonsense. This + # exact bug cost real debugging time on qwen35moe. + g = gate_buf[layer].reshape(E, I, rb) + u = up_buf[layer].reshape(E, I, rb) + banks["gate_up"][layer].copy_(torch.cat([g, u], dim=1)) + del gate_buf[layer], up_buf[layer] + if tracker is not None: + tracker.note(layer) + + if layer_sink is not None: + _load(layer_sink) + elif torch.cuda.is_available(): + with PinPipeline() as pins: + _load(pins) + else: + _load(None) # CUDA-less: mmap banks stay pageable, never pinned + + want = set(range(L)) + missing_gate, missing_up, missing_down = ( + want - seen_gate, want - seen_up, want - seen_down) + if missing_gate or missing_up or missing_down: + raise ValueError( + f"deepseek4 GGUF is missing routed experts: gate {sorted(missing_gate)}, " + f"up {sorted(missing_up)}, down {sorted(missing_down)}" + ) + + return banks + + +def dummy_gguf_expert_sources(config: "ModelConfig") -> dict[str, list[torch.Tensor]]: + """Random expert banks shaped like ``load_gguf_expert_sources`` output. + + Q4_K throughout, matching the checkpoints that actually load (see module docstring). + """ + from freetoken.moe.host_banks import alloc_layer_banks, pin_banks + + L = config.num_layers + types = {"gate_up": [GGML_Q4_K] * L, "down": [GGML_Q4_K] * L} + specs = gguf_expert_specs(config, types) + + hb = alloc_layer_banks(specs, L) + banks = {name: [b.tensor for b in hb[name]] for name in hb} + for t in banks["gate_up"] + banks["down"]: + t.random_(0, 256) + if torch.cuda.is_available(): + pin_banks(hb) + return banks + + +__all__ = [ + "gguf_expert_types", + "gguf_expert_specs", + "load_gguf_expert_sources", + "dummy_gguf_expert_sources", +] diff --git a/python/freetoken/models/deepseek_v4/model.py b/python/freetoken/models/deepseek_v4/model.py index c00d20cc..4cd92d88 100644 --- a/python/freetoken/models/deepseek_v4/model.py +++ b/python/freetoken/models/deepseek_v4/model.py @@ -223,6 +223,20 @@ def __init__(self, config): self._transformer = Transformer(self._args) self._bound = False + # A GGUF checkpoint carries native block-quantized weights, so the dense/fp8 + # projections have to be swapped for GGUF ops before load_state_dict runs -- that + # walks named_parameters() and demands a key for every one, so an unswapped + # fp8 Linear asks for a .scale no GGUF tensor can fill. Mirrors gemma4/model.py + # and qwen3_5_moe/model.py. + from .gguf import convert_deepseek4_to_gguf, is_gguf_model + + if is_gguf_model(config): + assert config.gguf_model_path is not None, ( + "expert_quant=='gguf' but ModelConfig.gguf_model_path is unset; the " + "per-tensor ggml types can only be read from the file" + ) + convert_deepseek4_to_gguf(self, config, model_path=config.gguf_model_path) + def _ensure_bound(self) -> None: if self._bound: return diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 63b1a18b..41069b68 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", + "deepseek4": "DeepseekV4GGUFForCausalLM", } diff --git a/python/freetoken/models/gguf/tokenizer.py b/python/freetoken/models/gguf/tokenizer.py index 6d5481c1..5b0d2b59 100644 --- a/python/freetoken/models/gguf/tokenizer.py +++ b/python/freetoken/models/gguf/tokenizer.py @@ -13,7 +13,13 @@ from .reader import gguf_architecture, load_gguf_metadata # GGUF architecture -> transformers GGUF tokenizer-converter key. -_TOKENIZER_ARCH = {"gemma4": "gemma4_text"} +_TOKENIZER_ARCH = { + "gemma4": "gemma4_text", + # tokenizer.ggml.model is gpt2 (BPE), pre joyai-llm, 129280 entries. The llama + # converter is sentencepiece-shaped and encodes a space as U+2581; against a + # GPT2-BPE vocab that silently DROPS every space on detokenization. + "deepseek4": "qwen2", +} def load_gguf_tokenizer(model_path: str): diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 0c033ca0..b605762c 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", ), + "DeepseekV4GGUFForCausalLM": ModelSpec( + "freetoken.models.deepseek_v4", + "DeepseekV4ForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + ), "GptOssForCausalLM": ModelSpec( "freetoken.models.gpt_oss", "GptOssForCausalLM",