From 535f894a90c506ec24fd4885c4e39415f6324a11 Mon Sep 17 00:00:00 2001 From: Chen Cheng <85005591+ischencheng@users.noreply.github.com> Date: Tue, 3 Nov 2026 16:55:57 -0500 Subject: [PATCH 1/2] feat(gguf): serve deepseek-v4.1 gguf checkpoints on deepseek_v41 llama.cpp 10-25 added the deepseek41 layout. The reader takes the config from the file, undoes the converter (engram gate stored as the product, one per_layer_token_embd tensor, hash constants as KVs) and streams the Engram rows from the file as the HF path does. --- docs/models.md | 9 + python/freetoken/checkpoint/convert.py | 18 +- python/freetoken/kernel/aot_models.py | 2 + python/freetoken/models/deepseek_v4/gguf.py | 22 +- .../freetoken/models/deepseek_v41/__init__.py | 6 +- .../freetoken/models/deepseek_v41/config.py | 26 +- .../models/deepseek_v41/engram_table.py | 27 +- python/freetoken/models/deepseek_v41/gguf.py | 299 ++++++++++++++++++ python/freetoken/models/deepseek_v41/model.py | 29 +- .../freetoken/models/deepseek_v41/weight.py | 3 + python/freetoken/models/gguf/config.py | 1 + python/freetoken/models/gguf/tensors.py | 18 +- python/freetoken/models/register.py | 11 +- tests/models/deepseek_v41/common.py | 161 +++++++++- tests/models/deepseek_v41/harness.py | 18 +- tests/models/deepseek_v41/test_config.py | 35 +- tests/models/deepseek_v41/test_engram.py | 76 ++++- .../models/deepseek_v41/test_forward_smoke.py | 28 +- 18 files changed, 713 insertions(+), 76 deletions(-) create mode 100644 python/freetoken/models/deepseek_v41/gguf.py diff --git a/docs/models.md b/docs/models.md index 9b85e65e73..22c03b1e0b 100644 --- a/docs/models.md +++ b/docs/models.md @@ -153,6 +153,15 @@ unset on PCIe GPUs: wider grids add no bandwidth there. `deepseek4v` mmproj; the text-only 0731 release has no image token and refuses the mmproj. ggml-org's Vision-Exp text GGUF ships a chat template that renders no image placeholder; serve unsloth's release for image requests. +- A DeepSeek-V4.1 GGUF (llama.cpp's `deepseek41` layout, such as + [unsloth/DeepSeek-V4.1-Flash-1017-GGUF](https://huggingface.co/unsloth/DeepSeek-V4.1-Flash-1017-GGUF)) + serves text; its `deepseek41v` mmproj and DSpark draft are not read. The Engram tables + stream from the file's `per_layer_token_embd` tensor, so keep that file on a fast NVMe. + A build that ships the tables as a separate shard needs it linked next to the others under + its shard name, as llama.cpp does. A file converted for an earlier llama.cpp branch has + no `ple.layers` key and is refused, and so are unsloth's UD-Q3_K_XL, UD-Q4_K_XL and UD-Q8_K_XL, + whose MXFP8 tables only unsloth's llama.cpp PR reads. `ft checkpoint` does not convert these + GGUFs: serve the `.gguf`. - DeepSeek-V4.1-Flash serves text and images from the HF checkpoint as shipped (fp8 block-32 dense weights, fp4 experts on the offload cache, 890 B/token global KV in packed fp4/fp8 pools). Its two 98 GiB Engram tables stream from the checkpoint shards on diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index 3c496d8a53..4861437f86 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -219,6 +219,15 @@ def convert_checkpoint( # 1) dense weights (host tensors; load straight to CPU to avoid GPU pressure) _progress("dense", 0, 0) # phase start; per-tensor cumulative bytes follow (total unknown) + # files the model reads directly from the checkpoint dir, not through FTW entries (Qwen3.8-Flash-Next PLE table); first, since a hook may refuse the checkpoint + from freetoken.models.register import _load_attr, get_model_spec + + try: + side_hook = _load_attr(get_model_spec(mc.architectures[0]).module, "ftw_side_files") + except (AttributeError, KeyError, ValueError): + side_hook = None + side_files = side_hook(model_path, out_dir) if side_hook is not None else [] + dense_bytes = 0 from freetoken.models.gguf.reader import gguf_config_source @@ -232,15 +241,6 @@ def convert_checkpoint( dense_bytes += tensor.numel() * tensor.element_size() _progress("dense", dense_bytes, 0) - # files the model reads directly from the checkpoint dir, not through FTW entries (Qwen3.8-Flash-Next PLE table) - from freetoken.models.register import _load_attr, get_model_spec - - try: - side_hook = _load_attr(get_model_spec(mc.architectures[0]).module, "ftw_side_files") - except (AttributeError, KeyError, ValueError): - side_hook = None - side_files = side_hook(model_path, out_dir) if side_hook is not None else [] - # 2) offload expert banks (post-repack) + alpha scales (slow path auto-picks parallel/serial) quant_format = None layer_quant_formats = None diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index 66b4a65ab6..07874149df 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -343,6 +343,8 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int # the routed experts share DSV4's ds_fp4 bank layout (hidden 5120, intermediate 2304). name="deepseek-ai/DeepSeek-V4.1-Flash", architecture="DeepseekV41ForCausalLM", + # llama.cpp deepseek41 GGUF (unsloth 1017 GGUF; text only) + arch_aliases=("DeepseekV41GGUFForCausalLM",), hidden_size=5120, kv_groups=(), top_k=6, diff --git a/python/freetoken/models/deepseek_v4/gguf.py b/python/freetoken/models/deepseek_v4/gguf.py index 29055e51ce..23ace76989 100644 --- a/python/freetoken/models/deepseek_v4/gguf.py +++ b/python/freetoken/models/deepseek_v4/gguf.py @@ -31,10 +31,9 @@ import numpy as np import torch -from freetoken.layers.quantization.ggml_types import is_quantized from freetoken.models.gguf.mmproj import Mmproj, find_mmproj, log_mmproj, mmproj_tensor from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index, gguf_tensor_names, load_gguf_metadata -from freetoken.models.gguf.tensors import check_scales, fused, matrix +from freetoken.models.gguf.tensors import check_scales, dequantized_tensor as _float, fused, matrix from freetoken.utils import init_logger from .args import DeepseekV4Args @@ -398,25 +397,6 @@ def iter_mmproj_weights(mmproj: Mmproj) -> Iterator[tuple[str, torch.Tensor]]: # Weights # ---------------------------------------------------------------------------------------------- -def _float(info: GgufTensorInfo) -> torch.Tensor: - """A float tensor in torch order, dequantizing packed ggml rows: dynamic quants quantize - tensors the module consumes as floats (the mHC mixers, the compressor gates, the router). - Trailing unit dims in ``ne`` are dropped -- some writers ship tensors 4-D (llama.cpp's - loaders try that shape first) and a trailing 1 carries no elements.""" - ne = info.ne - while len(ne) > 1 and ne[-1] == 1: - ne = ne[:-1] - if not is_quantized(info.ggml_type): - from freetoken.models.gguf.tensors import float_tensor - - return float_tensor(info).view(*reversed(ne)) - from freetoken.third_party.gguf import quants - from freetoken.third_party.gguf.constants import GGMLQuantizationType - - flat = quants.dequantize(np.array(info.data()), GGMLQuantizationType(info.ggml_type)) - return torch.from_numpy(flat).view(*reversed(ne)) - - def _tid2eid(info: GgufTensorInfo) -> torch.Tensor: """The hash-route table as the module's int64 [vocab, top_k].""" ne = info.ne diff --git a/python/freetoken/models/deepseek_v41/__init__.py b/python/freetoken/models/deepseek_v41/__init__.py index bcdef078c7..9af7163b7f 100644 --- a/python/freetoken/models/deepseek_v41/__init__.py +++ b/python/freetoken/models/deepseek_v41/__init__.py @@ -8,7 +8,11 @@ """ from .config import parse_config +from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import DeepseekV41ForCausalLM from .weight import ftw_side_files, iter_expert_pieces, iter_weights, iter_vision_weights -__all__ = ["DeepseekV41ForCausalLM", "parse_config", "iter_weights", "iter_vision_weights", "iter_expert_pieces", "ftw_side_files"] +__all__ = [ + "DeepseekV41ForCausalLM", "parse_config", "iter_weights", "iter_vision_weights", "iter_expert_pieces", "ftw_side_files", + "iter_gguf_weights", "parse_gguf_config", "to_gguf", +] diff --git a/python/freetoken/models/deepseek_v41/config.py b/python/freetoken/models/deepseek_v41/config.py index 9e738939c8..981b80f589 100644 --- a/python/freetoken/models/deepseek_v41/config.py +++ b/python/freetoken/models/deepseek_v41/config.py @@ -31,15 +31,8 @@ class VisionConfig: max_wh_ratio: float | None = None -def parse_config(hf_config: Any) -> ModelConfig: - args = DeepseekV41Args.from_hf(hf_config) - text = getattr(hf_config, "text_config", None) or hf_config - max_position = int(getattr(text, "max_position_embeddings", 0) or getattr(hf_config, "max_position_embeddings", 0) or 0) - if max_position <= 0: - max_position = int(args.original_seq_len * args.rope_factor) - vc = getattr(hf_config, "vision_config", None) - vision = VisionConfig(**{f.name: getattr(vc, f.name, f.default) for f in fields(VisionConfig)}) if vc is not None else None - +def model_config_from_args(args: DeepseekV41Args, *, max_position: int, vision: VisionConfig | None = None, image_token_id: int | None = None) -> ModelConfig: + """The ModelConfig of a DeepSeek-V4.1 text tower, from the HF config or a GGUF's metadata.""" return ModelConfig( num_layers=args.n_layers, num_qo_heads=args.n_heads, @@ -78,7 +71,7 @@ def parse_config(hf_config: Any) -> ModelConfig: attn_sm_scale=args.head_dim**-0.5, dsv41_args=args, vision_config=vision, - image_token_id=getattr(hf_config, "image_token_id", None), + image_token_id=image_token_id, attention_groups=( DSV4AttentionGroupConfig( name="dsv41", @@ -92,4 +85,15 @@ def parse_config(hf_config: Any) -> ModelConfig: ) -__all__ = ["parse_config"] +def parse_config(hf_config: Any) -> ModelConfig: + args = DeepseekV41Args.from_hf(hf_config) + text = getattr(hf_config, "text_config", None) or hf_config + max_position = int(getattr(text, "max_position_embeddings", 0) or getattr(hf_config, "max_position_embeddings", 0) or 0) + if max_position <= 0: + max_position = int(args.original_seq_len * args.rope_factor) + vc = getattr(hf_config, "vision_config", None) + vision = VisionConfig(**{f.name: getattr(vc, f.name, f.default) for f in fields(VisionConfig)}) if vc is not None else None + return model_config_from_args(args, max_position=max_position, vision=vision, image_token_id=getattr(hf_config, "image_token_id", None)) + + +__all__ = ["model_config_from_args", "parse_config"] diff --git a/python/freetoken/models/deepseek_v41/engram_table.py b/python/freetoken/models/deepseek_v41/engram_table.py index 3bec75c3f4..7b1c34ab64 100644 --- a/python/freetoken/models/deepseek_v41/engram_table.py +++ b/python/freetoken/models/deepseek_v41/engram_table.py @@ -66,13 +66,15 @@ def _tensor_extent(folder: str, name: str) -> _TensorExtent: @dataclass(frozen=True) class EngramRowSource: - """The fp8 table tensor in its shard (one extent) and its host-resident scales.""" + """One layer's table in one file extent: fp8 rows with host-resident scales, or a GGUF's rows of ``ggml_type``.""" path: str base: int num_rows: int head_dim: int - scales: torch.Tensor # [num_rows, head_dim // 32] uint8 (e8m0 codes), CPU + scales: torch.Tensor | None # [num_rows, head_dim // 32] uint8 (e8m0 codes), CPU; None for a GGUF table + ggml_type: int | None = None + row_bytes: int = 0 # a GGUF row's bytes on disk def engram_row_source(folder: str, layer_id: int) -> EngramRowSource: @@ -105,12 +107,15 @@ def __init__( self.head_dim = source.head_dim self.n_cols = n_cols self.num_rows = source.num_rows - self.row_bytes = FP8_E8M0_B32.row_bytes(source.head_dim) # values | scales + self.ggml_type = source.ggml_type + # staged rows: values | scales for fp8, the file's bytes for a GGUF + self.row_bytes = FP8_E8M0_B32.row_bytes(source.head_dim) if source.ggml_type is None else source.row_bytes self.scales = source.scales self.device = device + stored = source.head_dim if source.ggml_type is None else source.row_bytes self.store = RowStore( paths=[source.path], extent_file=[0], extent_base=[source.base], rows_per_extent=source.num_rows, - row_bytes=source.head_dim, row_stride=source.head_dim, use_io_uring=use_io_uring, + row_bytes=stored, row_stride=stored, use_io_uring=use_io_uring, ) # the graph staging is allocated up front (pinned alloc inside capture is illegal) and # outlives any one capture so a re-capture keeps the same pointers; the eager staging @@ -160,15 +165,18 @@ def stage(self, row_ids: torch.Tensor, *, graph: bool) -> None: ids = row_ids.reshape(-1).to(torch.int64).contiguous() rows = pinned[: n * self.row_bytes].view(n, self.row_bytes) self.store.stage_rows(ids.data_ptr(), n, rows.data_ptr(), self.row_bytes) - rows[:, self.head_dim :].copy_(self.scales.index_select(0, ids)) + if self.scales is not None: + rows[:, self.head_dim :].copy_(self.scales.index_select(0, ids)) def flush(self, *, signal: bool) -> None: self.store.flush(self.flag.data_ptr() if signal else 0) # ----- device (EngramTable protocol) ----- def lookup(self, num_tokens: int) -> torch.Tensor: + from freetoken.kernel import ggml from freetoken.kernel.row_store import wait_reset from freetoken.kernel.triton.dsv41.pack import unpack_rows + from freetoken.layers.quantization.ggml_types import float_dtype stream = torch.cuda.current_stream(self.device) capturing = torch.cuda.is_current_stream_capturing() @@ -180,7 +188,14 @@ def lookup(self, num_tokens: int) -> torch.Tensor: dev[:nbytes].copy_(pinned[:nbytes], non_blocking=True) if not capturing: self._eager_read[self._eager_slot].record(stream) - values = unpack_rows(dev[:nbytes].view(n, self.row_bytes), FP8_E8M0_B32, self.head_dim) + rows = dev[:nbytes].view(n, self.row_bytes) + if self.ggml_type is None: + values = unpack_rows(rows, FP8_E8M0_B32, self.head_dim) + elif (dtype := float_dtype(self.ggml_type)) is not None: + # copy=True: bf16 rows must not alias the staging buffer the next copy overwrites + values = rows.view(dtype).to(torch.bfloat16, copy=True) + else: + values = ggml.dequantize(rows, self.ggml_type, self.head_dim) return values.view(num_tokens, self.n_cols * self.head_dim) diff --git a/python/freetoken/models/deepseek_v41/gguf.py b/python/freetoken/models/deepseek_v41/gguf.py new file mode 100644 index 0000000000..223f52ddab --- /dev/null +++ b/python/freetoken/models/deepseek_v41/gguf.py @@ -0,0 +1,299 @@ +"""DeepSeek-V4.1-Flash from a llama.cpp ``deepseek41`` GGUF (unsloth ``DeepSeek-V4.1-Flash-1017-GGUF``): config from the GGUF metadata, tensor names, the dense state dict and the Engram table. + +The converter (conversion/deepseek41.py) keeps the geometry of the HF checkpoint and changes three things the loader undoes: +the Engram gate is stored as the product ``q_weight * k_weight`` (the model only ever uses the product), all Engram tables sit back to back in one +``per_layer_token_embd`` tensor, and the hash constants and the compressed token map are written as KVs (the runtime has no HF tokenizer to derive the map from). +The DSpark stage and the vision tower are separate files that this reader does not serve; ``exp_probs_b_vl`` in the text file stays unread with them. +""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING, Iterator + +import torch + +from freetoken.layers.quantization.ggml_types import row_bytes +from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index +from freetoken.models.gguf.tensors import check_scales, dequantized_tensor, fused, matrix + +from .args import DeepseekV41Args, Mode +from .config import model_config_from_args +from .engram import EngramHash +from .engram_table import EngramRowSource + +if TYPE_CHECKING: + from freetoken.models.config import ModelConfig + from freetoken.models.gguf.config import GgufConfigShim + +_TABLE = "per_layer_token_embd.weight" + +# engine module leaf (after "model.layers.N.") -> GGUF tensor names (without ".weight") +_LAYER_LEAVES: dict[str, tuple[str, ...]] = { + "attn.wq_a": ("attn_q_a",), + "attn.wq_b": ("attn_q_b",), + "attn.wkv": ("attn_kv",), + "attn.wo_b": ("attn_output_b",), + "attn.compressor.wkv": ("attn_compressor_kv",), + "attn.indexer.wk": ("indexer.attn_k",), + "attn.indexer.wq_b": ("indexer.attn_q_b",), + "attn.indexer.weights_proj": ("indexer.proj",), + "engram.wkv": ("engram_kv",), + "ffn.shared_experts.w1": ("ffn_gate_shexp",), + "ffn.shared_experts.w3": ("ffn_up_shexp",), + "ffn.shared_experts.w2": ("ffn_down_shexp",), + # routed experts (GGUF expert banks), gate / up / down order + "ffn.experts.0.w1": ("ffn_gate_exps",), + "ffn.experts.0.w3": ("ffn_up_exps",), + "ffn.experts.0.w2": ("ffn_down_exps",), +} +_TOP_LEAVES: dict[str, tuple[str, ...]] = { + "model.embed": ("token_embd",), + "head": ("output",), +} +_LAYER_RE = re.compile(r"^model\.layers\.(\d+)\.(.+)$") + + +def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: + """Engine module name -> the GGUF tensors (without ``.weight``) that store it; empty for modules the GGUF has no matrix for.""" + if name in _TOP_LEAVES: + return _TOP_LEAVES[name] + m = _LAYER_RE.match(name) + if m is None: + return () + return tuple(f"blk.{m.group(1)}.{leaf}" for leaf in _LAYER_LEAVES.get(m.group(2), ())) + + +def _args_from_metadata(shim: GgufConfigShim) -> DeepseekV41Args: + md, prefix, path = shim.metadata, shim.model_type, shim.model_path + + def kv(key: str, default=None): + return md.get(f"{prefix}.{key}", default) + + def need(key: str): + val = kv(key) + if val is None: + raise KeyError(f"{path}: missing GGUF metadata key {prefix}.{key}") + return val + + from freetoken.third_party.gguf.constants import ExpertGatingFuncType + + if int(need("expert_gating_func")) != int(ExpertGatingFuncType.SQRTSOFTPLUS): + raise NotImplementedError(f"{path}: expert_gating_func other than sqrtsoftplus is not served") + if int(kv("hash_layer_count", 0) or 0): + raise NotImplementedError(f"{path}: hash-routed layers are not served on DeepSeek-V4.1") + if need("rope.scaling.type") != "yarn": + raise NotImplementedError(f"{path}: rope scaling other than yarn is not served") + head_dim = int(need("attention.key_length")) + if int(need("attention.value_length")) != head_dim: + raise NotImplementedError(f"{path}: attention key and value lengths differ") + if int(kv("attention.head_count_kv", 1) or 1) != 1: + raise NotImplementedError(f"{path}: more than one latent KV head is not served") + # the clamp is per-layer in the GGUF; the model takes one value + clamps = {float(v) for k in ("swiglu_clamp_exp", "swiglu_clamp_shexp") for v in need(k)} + if len(clamps) != 1: + raise NotImplementedError(f"{path}: non-uniform swiglu_clamp_exp / swiglu_clamp_shexp is not served") + + engram: dict = {} + if (layers := kv("ple.layers")) is not None: + ngram, heads = int(need("ple.ngram_size")), int(need("ple.heads_per_ngram")) + cols = (ngram - 1) * heads + sizes = [int(s) for s in need("ple.head_vocab_sizes")] + engram = dict( + engram_layer_ids=[int(i) for i in layers], + engram_num_embeddings=[sum(sizes[i * cols : (i + 1) * cols]) for i in range(len(layers))], + engram_max_ngram_size=ngram, + # llama.cpp reads neither engram.vocab_size nor engram.compressed_vocab_size and some files leave them out; + # the primes draw upward from one below this value, and the token map's ids are dense + engram_vocab_size=sizes[0], + engram_n_heads=heads, + engram_head_dim=int(need("embedding_length_per_layer_input")), + engram_pad_id=int(need("engram.pad_token_id")), + engram_compressed_vocab_size=max(need("engram.token_map")) + 1, + ) + return DeepseekV41Args( + vocab_size=int(shim.vocab_size), + dim=int(need("embedding_length")), + moe_inter_dim=int(need("expert_feed_forward_length")), + n_layers=int(need("block_count")), + n_mtp_layers=0, # the converter leaves the DSpark stage out + n_heads=int(need("attention.head_count")), + n_routed_experts=int(need("expert_count")), + n_shared_experts=int(kv("expert_shared_count", 1) or 1), + n_activated_experts=int(need("expert_used_count")), + score_func="sqrtsoftplus", + norm_topk_prob=bool(need("expert_weights_norm")), + route_scale=float(need("expert_weights_scale")), + swiglu_limit=clamps.pop(), + q_lora_rank=int(need("attention.q_lora_rank")), + head_dim=head_dim, + rope_head_dim=int(need("rope.dimension_count")), + norm_eps=float(need("attention.layer_norm_rms_epsilon")), + o_groups=int(need("attention.output_group_count")), + o_lora_rank=int(need("attention.output_lora_rank")), + window_size=int(need("attention.sliding_window")), + compress_ratios=need("attention.compress_ratios"), + kv_source_layers=need("attention.kv_source_layers"), + index_source_layers=need("attention.indexer.source_layers"), + compress_rope_theta=float(need("attention.compress_rope_freq_base")), + original_seq_len=int(need("rope.scaling.original_context_length")), + rope_theta=float(need("rope.freq_base")), + rope_factor=float(need("rope.scaling.factor")), + beta_fast=int(need("rope.scaling.yarn_beta_fast")), + beta_slow=int(need("rope.scaling.yarn_beta_slow")), + index_n_heads=int(need("attention.indexer.head_count")), + index_head_dim=int(need("attention.indexer.key_length")), + index_topk=int(need("attention.indexer.top_k")), + candidate_source_layer=int(kv("attention.indexer.candidate_source_layer", -1)), + candidate_topk_blocks=int(kv("attention.indexer.candidate_top_k_blocks", 0)), + candidate_block_size=int(kv("attention.indexer.candidate_block_size", 0)), + hc_mult=int(need("hyper_connection.count")), + hc_sinkhorn_iters=int(need("hyper_connection.sinkhorn_iterations")), + hc_eps=float(need("hyper_connection.epsilon")), + **engram, + ) + + +def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: + from dataclasses import replace + + check_scales(shim.model_path) + args = _args_from_metadata(shim) + index = gguf_tensor_index(shim.model_path) + if not args.engram_layer_ids and (f"{shim.model_type}.engram.token_map" in shim.metadata or any(".engram_" in name for name in index)): + raise NotImplementedError(f"{shim.model_path}: Engram tensors without {shim.model_type}.ple.layers, so not llama.cpp's deepseek41 layout (converted for an earlier branch?)") + # a metadata-only GGUF (an FTW dir's source_metadata.gguf) has no tensor table to check + if index: + if args.engram_layer_ids and _TABLE not in index: + raise NotImplementedError(f"{shim.model_path}: no {_TABLE} tensor, the Engram tables llama.cpp's converter writes") + for name, info in index.items(): + if name.endswith(".ffn_gate_exps.weight") and info.ne[-1] != args.n_routed_experts: + raise NotImplementedError( + f"{shim.model_path}: {name} carries {info.ne[-1]} experts but expert_count says {args.n_routed_experts} (per-layer prunes are not served)" + ) + config = model_config_from_args(args, max_position=int(shim.metadata[f"{shim.model_type}.context_length"])) + return replace(config, expert_quant="gguf") + + +# ---------------------------------------------------------------------------------------------- +# Weights +# ---------------------------------------------------------------------------------------------- + + +def iter_gguf_weights( + model_path: str, + device: torch.device, + *, + include_moe_experts: bool, + include_non_moe: bool, + include_vision: bool = True, +) -> Iterator[tuple[str, torch.Tensor]]: + """The dense state dict from the GGUF, in the names the HF loader yields; the routed experts load as offload banks (``freetoken.models.gguf.experts``) and the Engram table streams from disk.""" + if include_moe_experts: + raise ValueError("DeepSeek-V4.1 routed experts are served from the offload cache; run with --moe-strategy offload (include_moe_experts must be False).") + if not include_non_moe: + return + from freetoken.utils import cached_load_hf_config + + config = parse_gguf_config(cached_load_hf_config(model_path)) + args: DeepseekV41Args = config.dsv41_args + index = gguf_tensor_index(model_path) + + def t(name: str) -> GgufTensorInfo: + return index[name if name in index else f"{name}.weight"] + + def bf16(name: str) -> torch.Tensor: + return dequantized_tensor(t(name)).to(torch.bfloat16) + + def linear(dst: str, src: str, k: int) -> Iterator[tuple[str, torch.Tensor]]: + info = t(src) + yield from fused(dst, [info], [matrix(info)], k) + + yield from linear("model.embed", "token_embd", args.dim) + yield from linear("head", "output", args.dim) + yield "model.norm.weight", bf16("output_norm") + for layer in range(args.n_layers): + blk, pre = f"blk.{layer}", f"model.layers.{layer}" + m = f"{pre}.attn" + role = args.roles[layer] + yield f"{pre}.attn_norm.weight", bf16(f"{blk}.attn_norm") + yield f"{pre}.ffn_norm.weight", bf16(f"{blk}.ffn_norm") + for sub in ("attn", "ffn"): + for part in ("fn", "base", "scale"): + yield f"{pre}.hc_{sub}_{part}", dequantized_tensor(t(f"{blk}.hc_{sub}_{part}")) + + yield from linear(f"{m}.wq_a", f"{blk}.attn_q_a", args.dim) + yield f"{m}.q_norm.weight", bf16(f"{blk}.attn_q_a_norm") + yield from linear(f"{m}.wq_b", f"{blk}.attn_q_b", args.q_lora_rank) + yield from linear(f"{m}.wkv", f"{blk}.attn_kv", args.dim) + yield f"{m}.kv_norm.weight", bf16(f"{blk}.attn_kv_a_norm") + # wo_a feeds a grouped bf16 einsum, not the quant method + yield f"{m}.wo_a", bf16(f"{blk}.attn_output_a") + yield from linear(f"{m}.wo_b", f"{blk}.attn_output_b", args.o_groups * args.o_lora_rank) + yield f"{m}.attn_sink", dequantized_tensor(t(f"{blk}.attn_sinks")) + if role.mode is Mode.FULL: + c = f"{blk}.attn_compressor" + yield f"{m}.compressor.norm.weight", bf16(f"{c}_norm") + if role.ratio > 1: + # plain fp32 parameters behind a softmax pooling, not the quant method + yield f"{m}.compressor.wkv", dequantized_tensor(t(f"{c}_kv")) + yield f"{m}.compressor.wgate", dequantized_tensor(t(f"{c}_gate")) + else: + yield from linear(f"{m}.compressor.wkv", f"{c}_kv", args.dim) + yield from linear(f"{m}.indexer.wk", f"{blk}.indexer.attn_k", args.head_dim) + yield f"{m}.indexer.k_norm.weight", bf16(f"{blk}.indexer.k_norm") + if role.has_indexer: + yield from linear(f"{m}.indexer.wq_b", f"{blk}.indexer.attn_q_b", args.q_lora_rank) + yield from linear(f"{m}.indexer.weights_proj", f"{blk}.indexer.proj", args.dim) + + g = f"{pre}.ffn.gate" + yield f"{g}.weight", bf16(f"{blk}.ffn_gate_inp") + yield f"{g}.bias", dequantized_tensor(t(f"{blk}.exp_probs_b.bias")) + yield from linear(f"{pre}.ffn.shared_experts.w1", f"{blk}.ffn_gate_shexp", args.dim) + yield from linear(f"{pre}.ffn.shared_experts.w3", f"{blk}.ffn_up_shexp", args.dim) + yield from linear(f"{pre}.ffn.shared_experts.w2", f"{blk}.ffn_down_shexp", args.moe_inter_dim) + if layer in args.engram_layer_ids: + e = f"{pre}.engram" + yield from linear(f"{e}.wkv", f"{blk}.engram_kv", args.engram_hash_cols * args.engram_head_dim) + # the file holds the product q_weight * k_weight; the layer multiplies them back + gate = dequantized_tensor(t(f"{blk}.engram_gate")).float() + yield f"{e}.q_weight", gate + yield f"{e}.k_weight", torch.ones_like(gate) + + +# ---------------------------------------------------------------------------------------------- +# Engram +# ---------------------------------------------------------------------------------------------- + + +def gguf_engram_tables(model_path: str, args: DeepseekV41Args) -> tuple[EngramHash, list[EngramRowSource]]: + """The hash from the file's token map and one row source per Engram layer: each is a slice of the single ``per_layer_token_embd`` tensor. + + The bucket primes and the multipliers derive from the config as on the HF path, so they are checked against the KVs llama.cpp hashes with.""" + from freetoken.models.gguf.reader import load_gguf_metadata + + md = load_gguf_metadata(model_path) + arch = md["general.architecture"] + hash = EngramHash(args, md[f"{arch}.engram.token_map"], args.engram_compressed_vocab_size) + starts = [sum(hash.rows_per_layer[:i]) for i in range(len(hash.rows_per_layer))] + offsets = (hash.offsets + torch.tensor(starts).unsqueeze(1)).flatten().tolist() + for key, derived in ( + ("ple.head_offsets", offsets), + ("ple.head_vocab_sizes", hash.primes.flatten().tolist()), + ("ple.layer_multipliers", hash.multipliers.flatten().tolist()), + ): + if [int(v) for v in md[f"{arch}.{key}"]] != derived: + raise ValueError(f"{model_path}: {arch}.{key} differs from the Engram hash this build derives from the config") + + info = gguf_tensor_index(model_path)[_TABLE] + head_dim, rows = info.ne + if head_dim != args.engram_head_dim or rows < sum(hash.rows_per_layer): + raise ValueError(f"{model_path}: {_TABLE} is {info.ne}, expected {args.engram_head_dim} values wide and at least {sum(hash.rows_per_layer)} rows") + nbytes = row_bytes(info.ggml_type, head_dim) + return hash, [ + EngramRowSource(info.path, info.offset + start * nbytes, layer_rows, head_dim, None, info.ggml_type, nbytes) + for start, layer_rows in zip(starts, hash.rows_per_layer) + ] + + +__all__ = ["gguf_engram_tables", "iter_gguf_weights", "parse_gguf_config", "to_gguf"] diff --git a/python/freetoken/models/deepseek_v41/model.py b/python/freetoken/models/deepseek_v41/model.py index 8112022f76..90593f9906 100644 --- a/python/freetoken/models/deepseek_v41/model.py +++ b/python/freetoken/models/deepseek_v41/model.py @@ -135,7 +135,7 @@ class Transformer(BaseOP): def __init__(self, config: ModelConfig, *, prefix: str = "model"): args: DeepseekV41Args = config.dsv41_args self.args = args - self.embed = VocabParallelEmbedding(num_embeddings=args.vocab_size, embedding_dim=args.dim) + self.embed = VocabParallelEmbedding(num_embeddings=args.vocab_size, embedding_dim=args.dim, quant_config=config.quant, prefix=f"{prefix}.embed") self.layers = OPList([Block(config, i, prefix=f"{prefix}.layers.{i}") for i in range(args.n_layers)]) self.norm = RMSNorm(args.dim, args.norm_eps) @@ -227,7 +227,7 @@ def engram_layers(self) -> List[EngramLayer]: return [b.engram for b in self.model.layers.op_list if b.engram is not None] def load_host_tables(self, engine_config) -> int: - """Attach the Engram tables (streamed from the checkpoint shards; zeros for dummy weights) + """Attach the Engram tables (streamed from the checkpoint shards or the GGUF; zeros for dummy weights) and take over ``forward_host_ctx`` so every dispatch stages its rows. Returns the pinned host bytes the engine reserves from its pin budget.""" import os @@ -242,24 +242,33 @@ def load_host_tables(self, engine_config) -> int: for layer in layers: layer.attach_table(ZeroEngramTable(layer.width, device)) return 0 - from freetoken.utils import download_hf_weight - from transformers import AutoTokenizer + from freetoken.models.gguf.reader import gguf_config_source - from .engram import EngramHash, build_compressed_token_map from .engram_table import EngramDiskTable, EngramHost, engram_row_source - folder = download_hf_weight(engine_config.model_path) - token_map, compressed_vocab = build_compressed_token_map(AutoTokenizer.from_pretrained(folder, trust_remote_code=True)) - hash = EngramHash(self._args, token_map, compressed_vocab) + if (gguf_src := gguf_config_source(engine_config.model_path)) is not None: + from .gguf import gguf_engram_tables + + hash, sources = gguf_engram_tables(gguf_src, self._args) + else: + from freetoken.utils import download_hf_weight + from transformers import AutoTokenizer + + from .engram import EngramHash, build_compressed_token_map + + folder = download_hf_weight(engine_config.model_path) + token_map, compressed_vocab = build_compressed_token_map(AutoTokenizer.from_pretrained(folder, trust_remote_code=True)) + hash = EngramHash(self._args, token_map, compressed_vocab) + sources = [engram_row_source(folder, layer.layer_id) for layer in layers] graph_rows = max(engine_config.max_running_req, engine_config.cuda_graph_max_bs or 0, 1) tables = [ EngramDiskTable( - engram_row_source(folder, layer.layer_id), hash.n_cols, device, + source, hash.n_cols, device, max_graph_rows=graph_rows, max_extend_tokens=engine_config.max_extend_tokens, use_io_uring=os.getenv("FREETOKEN_ENGRAM_IO_URING", "1") != "0", ) - for layer in layers + for source in sources ] for layer, table in zip(layers, tables): layer.attach_table(table) diff --git a/python/freetoken/models/deepseek_v41/weight.py b/python/freetoken/models/deepseek_v41/weight.py index 962bd19589..52cacb9e7a 100644 --- a/python/freetoken/models/deepseek_v41/weight.py +++ b/python/freetoken/models/deepseek_v41/weight.py @@ -182,9 +182,12 @@ def ftw_side_files(model_path: str, out_dir: str) -> list[str]: loader never reads ``.safetensors`` as weights, so the shard's other tensors ride along unused.""" import shutil + from freetoken.models.gguf.reader import gguf_config_source from freetoken.models.loader import safetensors_weight_map from freetoken.utils import download_hf_weight + if gguf_config_source(model_path) is not None: + raise NotImplementedError(f"{model_path}: an FTW of a DeepSeek-V4.1 GGUF has no Engram table to read; serve the .gguf directly") folder = download_hf_weight(model_path) args = load_args(folder) shard_of = safetensors_weight_map(folder) diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index 122e105684..fe75a0a44e 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -22,6 +22,7 @@ GGUF_ARCH_TO_REGISTRY: dict[str, str] = { "flint": "FlintGGUFForCausalLM", "deepseek4": "DeepseekV4GGUFForCausalLM", + "deepseek41": "DeepseekV41GGUFForCausalLM", "gemma4": "Gemma4GGUFForCausalLM", "gpt-oss": "GptOssGGUFForCausalLM", "glm4moe": "Glm4MoeGGUFForCausalLM", diff --git a/python/freetoken/models/gguf/tensors.py b/python/freetoken/models/gguf/tensors.py index c1c2f2522f..2c69d461b7 100644 --- a/python/freetoken/models/gguf/tensors.py +++ b/python/freetoken/models/gguf/tensors.py @@ -19,6 +19,22 @@ def float_tensor(info: GgufTensorInfo) -> torch.Tensor: return torch.from_numpy(np.array(info.data())).view(_FLOAT_DTYPES[info.ggml_type]).view(*reversed(info.ne)) +def dequantized_tensor(info: GgufTensorInfo) -> torch.Tensor: + """A float tensor in torch order, dequantizing packed ggml rows: dynamic quants quantize tensors the module consumes as floats (the mHC mixers, the compressor gates, the router). + + Trailing unit dims in ``ne`` are dropped: some writers ship tensors 4-D (llama.cpp's loaders try that shape first) and a trailing 1 carries no elements.""" + ne = info.ne + while len(ne) > 1 and ne[-1] == 1: + ne = ne[:-1] + if not is_quantized(info.ggml_type): + return float_tensor(info).view(*reversed(ne)) + from freetoken.third_party.gguf import quants + from freetoken.third_party.gguf.constants import GGMLQuantizationType + + flat = quants.dequantize(np.array(info.data()), GGMLQuantizationType(info.ggml_type)) + return torch.from_numpy(flat).view(*reversed(ne)) + + def packed_rows(info: GgufTensorInfo) -> torch.Tensor: """Packed ggml rows ``[rows, row_bytes]`` of a quantized tensor.""" rows = int(np.prod(info.ne[1:])) @@ -105,4 +121,4 @@ def fused(prefix: str, infos: list[GgufTensorInfo], parts: list[torch.Tensor], k yield f"{prefix}.ggml_scale", scales -__all__ = ["check_scales", "float_tensor", "fused", "matrix", "packed_rows", "table_matrix", "weight_scale"] +__all__ = ["check_scales", "dequantized_tensor", "float_tensor", "fused", "matrix", "packed_rows", "table_matrix", "weight_scale"] diff --git a/python/freetoken/models/register.py b/python/freetoken/models/register.py index 6cf94dcd1c..40ae7af02b 100644 --- a/python/freetoken/models/register.py +++ b/python/freetoken/models/register.py @@ -301,9 +301,18 @@ def _gemma4_vision_unserved(section: Any) -> str | None: checkpoint_roots=(("model.layers", "layers"), ("model.embed", "embed"), ("model.norm", "norm")), packed_modules_mapping=_EXPERTS_W123_PACKED, unquantized_modules=( - "head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.wk", "*.indexer.weights_proj", "vision.*", "aligner.*", + "embed", "head", "*.compressor.wkv", "*.compressor.wgate", "*.indexer.wk", "*.indexer.weights_proj", "vision.*", "aligner.*", ), ), + # llama.cpp deepseek41 GGUF (text; unsloth DeepSeek-V4.1-Flash-1017-GGUF): the same model classes with every tensor in its packed ggml type + "DeepseekV41GGUFForCausalLM": ModelSpec( + "freetoken.models.deepseek_v41", + "DeepseekV41ForCausalLM", + parse_config="parse_gguf_config", + iter_weights="iter_gguf_weights", + # the GGUF scheme reads expert types in (gate, up, down) order + packed_modules_mapping=_EXPERTS_W132_PACKED, + ), "Qwen3_5MoeForConditionalGeneration": ModelSpec( "freetoken.models.qwen3_5_moe", "Qwen3_5MoeForConditionalGeneration", diff --git a/tests/models/deepseek_v41/common.py b/tests/models/deepseek_v41/common.py index d0f8efeef9..0e77a524cb 100644 --- a/tests/models/deepseek_v41/common.py +++ b/tests/models/deepseek_v41/common.py @@ -50,7 +50,7 @@ def tiny_hf_config(**over) -> SimpleNamespace: return SimpleNamespace(architectures=["DeepseekV41ForCausalLM"], model_type="deepseek_v41", text_config=tiny_text_config(**over)) -def checkpoint_tensors(text: dict, seed: int = 0) -> dict[str, torch.Tensor]: +def checkpoint_tensors(text: dict, seed: int = 0, *, random_gate: bool = False) -> dict[str, torch.Tensor]: """Random bf16 weights under the checkpoint's names (the reference ``convert.py`` dialect without quantization: no ``.scale`` tensors). hc / sink / router-bias / engram gates in fp32 as shipped.""" g = torch.Generator().manual_seed(seed) @@ -109,8 +109,8 @@ def rnd(*shape, scale=0.02, dtype=torch.bfloat16): e = f"{p}.engram" cols = (text["engram_max_ngram_size"] - 1) * text["engram_n_heads"] t[f"{e}.wkv.weight"] = rnd(dim * (hc + 1), cols * text["engram_head_dim"]) - t[f"{e}.q_weight"] = torch.ones(hc, dim, dtype=torch.bfloat16) - t[f"{e}.k_weight"] = torch.ones(hc, dim, dtype=torch.bfloat16) + t[f"{e}.q_weight"] = rnd(hc, dim, scale=1.0) if random_gate else torch.ones(hc, dim, dtype=torch.bfloat16) + t[f"{e}.k_weight"] = rnd(hc, dim, scale=1.0) if random_gate else torch.ones(hc, dim, dtype=torch.bfloat16) rows = text["engram_num_embeddings"][text["engram_layer_ids"].index(L)] t[f"{e}.embed.weight"] = rnd(rows, text["engram_head_dim"], scale=1.0).to(torch.float8_e4m3fn) t[f"{e}.embed.scale"] = torch.full((rows, text["engram_head_dim"] // 32), 127, dtype=torch.uint8).view(torch.float8_e8m0fnu) @@ -120,7 +120,7 @@ def rnd(*shape, scale=0.02, dtype=torch.bfloat16): return t -def write_tiny_checkpoint(folder: str, seed: int = 0, *, quantized: bool = False, **over) -> dict[str, torch.Tensor]: +def write_tiny_checkpoint(folder: str, seed: int = 0, *, quantized: bool = False, random_gate: bool = False, **over) -> dict[str, torch.Tensor]: """A loadable checkpoint directory: HF-style config.json + one safetensors shard + index.""" from safetensors.torch import save_file @@ -136,7 +136,7 @@ def write_tiny_checkpoint(folder: str, seed: int = 0, *, quantized: bool = False "weight_block_size": [32, 32], "scale_fmt": "ue8m0", "expert_dtype": "fp4", } json.dump(config, f) - tensors = checkpoint_tensors(text, seed) + tensors = checkpoint_tensors(text, seed, random_gate=random_gate) if quantized: tensors = _quantize_checkpoint(tensors) save_file(tensors, os.path.join(folder, "model-00001-of-00001.safetensors")) @@ -174,3 +174,154 @@ def _quantize_checkpoint(tensors): out[name] = (weight.float() / expanded).to(torch.float8_e4m3fn) out[stem + ".scale"] = scale.to(torch.float8_e8m0fnu) return out + + +_GGUF_ATTN = { + "wq_a.weight": "attn_q_a", "q_norm.weight": "attn_q_a_norm", "wq_b.weight": "attn_q_b", "wkv.weight": "attn_kv", + "kv_norm.weight": "attn_kv_a_norm", "wo_a.weight": "attn_output_a", "wo_b.weight": "attn_output_b", "attn_sink": "attn_sinks", + "compressor.wkv.weight": "attn_compressor_kv", "compressor.wgate.weight": "attn_compressor_gate", "compressor.norm.weight": "attn_compressor_norm", + "indexer.wk.weight": "indexer.attn_k", "indexer.k_norm.weight": "indexer.k_norm", "indexer.wq_b.weight": "indexer.attn_q_b", "indexer.weights_proj.weight": "indexer.proj", +} +_GGUF_LAYER = { + "attn_norm.weight": "attn_norm", "ffn_norm.weight": "ffn_norm", "ffn.gate.weight": "ffn_gate_inp", "ffn.gate.bias": "exp_probs_b.bias", "ffn.gate.bias_vl": "exp_probs_b_vl.bias", + "ffn.shared_experts.w1.weight": "ffn_gate_shexp", "ffn.shared_experts.w3.weight": "ffn_up_shexp", "ffn.shared_experts.w2.weight": "ffn_down_shexp", + "hc_attn_fn": "hc_attn_fn", "hc_attn_base": "hc_attn_base", "hc_attn_scale": "hc_attn_scale", "hc_ffn_fn": "hc_ffn_fn", "hc_ffn_base": "hc_ffn_base", "hc_ffn_scale": "hc_ffn_scale", + "engram.wkv.weight": "engram_kv", +} + + +def _pack_q8_0(rows: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``[n, k]`` floats to ggml Q8_0 blocks of 32 values: ``[n, k // 32 * 34]`` bytes (fp16 scale, then int8s), and the values they decode to.""" + n = rows.shape[0] + blocks = rows.float().view(-1, 32) + d = (blocks.abs().amax(-1, keepdim=True).clamp_min(1e-3) / 127).to(torch.float16) + q = torch.round(blocks / d.float()).clamp(-127, 127).to(torch.int8) + packed = torch.cat([d.view(torch.uint8).view(-1, 2), q.view(torch.uint8)], dim=-1).view(n, -1) + return packed, (q.float() * d.float()).view(n, -1) + + +def _write_gguf(path: str, kvs: list[tuple[str, int, object, int | None]], tensors: list[tuple[str, tuple[int, ...], int, bytes]]) -> None: + """A little-endian GGUF v3 file. ``kvs``: (key, value type, value, array element type); ``tensors``: (name, torch-order shape, ggml type, bytes).""" + import struct + + from freetoken.third_party.gguf import GGUFValueType as K + + fmt = {K.UINT32: " bytes: + raw = text.encode() + return struct.pack(" bytes: + if kind == K.STRING: + return string(val) + if kind == K.ARRAY: + return struct.pack(" dict[int, torch.Tensor]: + """The tiny checkpoint as llama.cpp's converter lays it out (conversion/deepseek41.py), every tensor float except an optional Q8_0 ``table_type`` table; + ``kv_over`` replaces array KVs. Returns each Engram layer's rows as they decode, keyed by layer.""" + import numpy as np + + from freetoken.models.deepseek_v41.args import DeepseekV41Args + from freetoken.models.deepseek_v41.engram import bucket_primes, hash_multipliers + from freetoken.third_party.gguf import ExpertGatingFuncType, GGMLQuantizationType as Q, GGUFValueType as K + + text = tiny_text_config(**over) + tensors = checkpoint_tensors(text, seed, random_gate=random_gate) + args = DeepseekV41Args.from_hf(SimpleNamespace(text_config=text)) + sizes = [p for layer in bucket_primes(args) for per_ngram in layer for p in per_ngram] + rs = text["rope_scaling"] + limit = [float(text["swiglu_limit"])] * text["num_hidden_layers"] + kvs: list[tuple[str, int, object, int | None]] = [("general.architecture", K.STRING, "deepseek41", None)] + + def add_kv(key, kind, value, sub=None): + kvs.append((f"deepseek41.{key}", kind, (kv_over or {}).get(key, value), sub)) + + for key, value in { + "block_count": text["num_hidden_layers"], "context_length": text["max_position_embeddings"], "embedding_length": text["hidden_size"], + "expert_feed_forward_length": text["moe_intermediate_size"], "expert_count": text["n_routed_experts"], "expert_used_count": text["num_experts_per_tok"], + "expert_shared_count": text["n_shared_experts"], "expert_gating_func": int(ExpertGatingFuncType.SQRTSOFTPLUS), "hash_layer_count": 0, + "attention.head_count": text["num_attention_heads"], "attention.head_count_kv": 1, "attention.key_length": text["head_dim"], "attention.value_length": text["head_dim"], + "attention.q_lora_rank": text["q_lora_rank"], "attention.sliding_window": text["sliding_window"], "attention.output_group_count": text["o_groups"], + "attention.output_lora_rank": text["o_lora_rank"], "rope.dimension_count": text["qk_rope_head_dim"], "rope.scaling.original_context_length": rs["original_max_position_embeddings"], + "attention.indexer.head_count": text["index_n_heads"], "attention.indexer.key_length": text["index_head_dim"], "attention.indexer.top_k": text["index_topk"], + "hyper_connection.count": text["hc_mult"], "hyper_connection.sinkhorn_iterations": text["hc_sinkhorn_iters"], + "attention.indexer.candidate_top_k_blocks": text["candidate_topk_blocks"], "attention.indexer.candidate_block_size": text["candidate_block_size"], + "ple.ngram_size": text["engram_max_ngram_size"], "ple.heads_per_ngram": text["engram_n_heads"], "embedding_length_per_layer_input": text["engram_head_dim"], + "engram.pad_token_id": text["engram_pad_token_id"], + }.items(): + add_kv(key, K.UINT32, value) + for key, value in { + "expert_weights_scale": text["routed_scaling_factor"], "attention.layer_norm_rms_epsilon": text["rms_norm_eps"], "rope.freq_base": text["rope_theta"], + "rope.scaling.factor": rs["factor"], "rope.scaling.yarn_beta_fast": rs["beta_fast"], "rope.scaling.yarn_beta_slow": rs["beta_slow"], + "attention.compress_rope_freq_base": text["compress_rope_theta"], "hyper_connection.epsilon": text["hc_eps"], + }.items(): + add_kv(key, K.FLOAT32, float(value)) + add_kv("attention.indexer.candidate_source_layer", K.INT32, text["candidate_source_layer_id"]) + add_kv("expert_weights_norm", K.BOOL, text["norm_topk_prob"]) + add_kv("rope.scaling.type", K.STRING, "yarn") + add_kv("swiglu_clamp_exp", K.ARRAY, limit, K.FLOAT32) + add_kv("swiglu_clamp_shexp", K.ARRAY, limit, K.FLOAT32) + add_kv("attention.compress_ratios", K.ARRAY, text["compress_ratios"], K.UINT32) + add_kv("attention.kv_source_layers", K.ARRAY, text["kv_source_layer_ids"], K.UINT32) + add_kv("attention.indexer.source_layers", K.ARRAY, text["index_source_layer_ids"], K.UINT32) + add_kv("ple.layers", K.ARRAY, text["engram_layer_ids"], K.UINT32) + add_kv("ple.layer_multipliers", K.ARRAY, hash_multipliers(args.engram_layer_ids, args.engram_max_ngram_size, args.engram_compressed_vocab_size).flatten().tolist(), K.UINT64) + add_kv("ple.head_vocab_sizes", K.ARRAY, sizes, K.UINT64) + add_kv("ple.head_offsets", K.ARRAY, np.cumsum([0, *sizes[:-1]]).tolist(), K.UINT64) + add_kv("engram.token_map", K.ARRAY, list(range(text["vocab_size"])), K.UINT32) + + out: list[tuple[str, tuple[int, ...], int, bytes]] = [] + + def add(name: str, t: torch.Tensor) -> None: + kind = Q.BF16 if t.dtype is torch.bfloat16 else Q.F32 + raw = t.view(torch.uint16) if t.dtype is torch.bfloat16 else t.float() + out.append((name, tuple(t.shape), int(kind), raw.numpy().tobytes())) + + def weight_name(gguf_name: str) -> str: + return gguf_name if gguf_name.endswith(".bias") else f"{gguf_name}.weight" + + add("token_embd.weight", tensors["embed.weight"]) + add("output_norm.weight", tensors["norm.weight"]) + add("output.weight", tensors["head.weight"]) + table: dict[int, torch.Tensor] = {} + for L in range(text["num_hidden_layers"]): + attn, layer = f"layers.{L}.attn.", f"layers.{L}." + for prefix, leaves in ((attn, _GGUF_ATTN), (layer, _GGUF_LAYER)): + for leaf, gguf_name in leaves.items(): + if prefix + leaf in tensors: + add(f"blk.{L}.{weight_name(gguf_name)}", tensors[prefix + leaf]) + for proj, gguf_name in (("w1", "ffn_gate_exps"), ("w3", "ffn_up_exps"), ("w2", "ffn_down_exps")): + add(f"blk.{L}.{gguf_name}.weight", torch.stack([tensors[f"{layer}ffn.experts.{e}.{proj}.weight"] for e in range(text["n_routed_experts"])])) + if L in text["engram_layer_ids"]: + add(f"blk.{L}.engram_gate.weight", tensors[f"{layer}engram.q_weight"].float() * tensors[f"{layer}engram.k_weight"].float()) + e = text["engram_layer_ids"].index(L) + rows = sum(sizes[e * args.engram_hash_cols : (e + 1) * args.engram_hash_cols]) + table[L] = (tensors[f"{layer}engram.embed.weight"][:rows].float() * torch.exp2(tensors[f"{layer}engram.embed.scale"][:rows].view(torch.uint8).float() - 127.0)).to(torch.bfloat16) + packed = torch.cat(list(table.values())) + if table_type == "Q8_0": + data, decoded = _pack_q8_0(packed) + out.append(("per_layer_token_embd.weight", (packed.shape[0], packed.shape[1]), int(Q.Q8_0), data.numpy().tobytes())) + table = dict(zip(table, decoded.to(torch.bfloat16).split([t.shape[0] for t in table.values()]))) + else: + add("per_layer_token_embd.weight", packed) + _write_gguf(path, kvs, out) + return table diff --git a/tests/models/deepseek_v41/harness.py b/tests/models/deepseek_v41/harness.py index 83c0ec61ff..ee9a6d077c 100644 --- a/tests/models/deepseek_v41/harness.py +++ b/tests/models/deepseek_v41/harness.py @@ -20,6 +20,7 @@ from freetoken.models import create_model from freetoken.models.deepseek_v41.config import parse_config from freetoken.models.deepseek_v41.engram import ZeroEngramTable +from freetoken.models.deepseek_v41.gguf import iter_gguf_weights, parse_gguf_config from freetoken.models.deepseek_v41.weight import iter_weights from freetoken.utils.hf import cached_load_hf_config from freetoken.utils.torch_utils import torch_dtype @@ -35,30 +36,31 @@ def __init__(self, checkpoint: str, *, max_seq_len: int = 1024, max_running_req: set_tp_info(0, 1) set_rope_device(self.device) hf_config = cached_load_hf_config(checkpoint) - mc = parse_config(hf_config) + gguf = checkpoint.endswith(".gguf") + mc = parse_gguf_config(hf_config) if gguf else parse_config(hf_config) args = mc.dsv41_args args.max_seq_len = max_seq_len args.max_batch_size = max_running_req + 1 args.swa_decoder_replay = swa_decoder_replay self.args = args quant = NoQuantConfig() - if quantized: + offload = quantized or gguf + if offload: from freetoken.models.register import checkpoint_quant_config, get_model_spec - quant = checkpoint_quant_config(checkpoint, hf_config, get_model_spec("DeepseekV41ForCausalLM")) - self.config = replace(mc, moe_strategy="offload" if quantized else "fused", decode_target="gpu", quant=quant) + quant = checkpoint_quant_config(checkpoint, hf_config, get_model_spec(mc.architectures[0] if not gguf else "DeepseekV41GGUFForCausalLM")) + self.config = replace(mc, moe_strategy="offload" if offload else "fused", decode_target="gpu", quant=quant) mc = self.config with torch.device("meta"), torch_dtype(torch.bfloat16): self.model = create_model(self.config) - state = _materialize_loaded_weight_state_dict( - self.model.state_dict(), iter_weights(checkpoint, self.device, include_moe_experts=False), device=self.device, - ) + reader = iter_gguf_weights(checkpoint, self.device, include_moe_experts=False, include_non_moe=True) if gguf else iter_weights(checkpoint, self.device, include_moe_experts=False) + state = _materialize_loaded_weight_state_dict(self.model.state_dict(), reader, device=self.device) self.model.load_state_dict(state) finalize_quant(self.model) from freetoken.moe.expert_banks import load_expert_banks method = self.model.model.layers.op_list[0].ffn.experts.quant_method - if quantized: + if offload: from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache self.banks = load_expert_banks( diff --git a/tests/models/deepseek_v41/test_config.py b/tests/models/deepseek_v41/test_config.py index fab385fb93..a064d755e0 100644 --- a/tests/models/deepseek_v41/test_config.py +++ b/tests/models/deepseek_v41/test_config.py @@ -16,7 +16,7 @@ from freetoken.models.register import get_model_spec from freetoken.utils.torch_utils import torch_dtype -from .common import ENGRAM_LAYERS, INDEX_SOURCES, KV_SOURCES, N_LAYERS, RATIOS, tiny_hf_config, write_tiny_checkpoint +from .common import ENGRAM_LAYERS, INDEX_SOURCES, KV_SOURCES, N_LAYERS, RATIOS, tiny_hf_config, write_tiny_checkpoint, write_tiny_gguf def test_roles_follow_the_dsv41_mode_table(): @@ -167,3 +167,36 @@ def test_fp8_wo_a_layer_rejects_a_dequantized_ftw_layout_clearly(): state2["attn.wo_a"], state2["attn.wo_a_scale"] = w, scale bf16_layer.load_state_dict(state2, prefix="attn") assert torch.equal(bf16_layer.wo_a, w.to(torch.bfloat16)) + + +def test_gguf_metadata_builds_the_checkpoint_config(tmp_path): + from dataclasses import fields + + from freetoken.models.deepseek_v41.gguf import parse_gguf_config + from freetoken.models.gguf.config import build_gguf_shim + + path = str(tmp_path / "tiny.gguf") + write_tiny_gguf(path) + shim = build_gguf_shim(path) + assert shim.architectures == ["DeepseekV41GGUFForCausalLM"] + gguf, hf = parse_gguf_config(shim), parse_config(tiny_hf_config()) + assert gguf.expert_quant == "gguf" and hf.expert_quant == "ds_fp4" + assert gguf.rotary_config == hf.rotary_config and gguf.num_experts == hf.num_experts and gguf.attention_groups == hf.attention_groups + # derived from the file, or not stored in it + skip = {"n_mtp_layers", "engram_num_embeddings", "engram_vocab_size", "dspark_block_size", "dspark_target_layer_ids", "norm_eps", "hc_eps"} + for f in fields(type(gguf.dsv41_args)): + if f.name not in skip: + assert getattr(gguf.dsv41_args, f.name) == getattr(hf.dsv41_args, f.name), f.name + assert gguf.dsv41_args.norm_eps == pytest.approx(hf.dsv41_args.norm_eps) and gguf.dsv41_args.hc_eps == pytest.approx(hf.dsv41_args.hc_eps) + + +def test_gguf_from_an_earlier_branch_is_refused(tmp_path): + from freetoken.models.deepseek_v41.gguf import parse_gguf_config + from freetoken.models.gguf.config import build_gguf_shim + + path = str(tmp_path / "old.gguf") + write_tiny_gguf(path) + shim = build_gguf_shim(path) + shim.metadata = {k: v for k, v in shim.metadata.items() if k != "deepseek41.ple.layers"} + with pytest.raises(NotImplementedError, match="ple.layers"): + parse_gguf_config(shim) diff --git a/tests/models/deepseek_v41/test_engram.py b/tests/models/deepseek_v41/test_engram.py index 51a8329744..f8f74923ce 100644 --- a/tests/models/deepseek_v41/test_engram.py +++ b/tests/models/deepseek_v41/test_engram.py @@ -11,7 +11,7 @@ from freetoken.models.deepseek_v41.args import DeepseekV41Args from freetoken.models.deepseek_v41.engram import EngramHash, bucket_primes, hash_multipliers -from .common import VOCAB, requires_cuda, tiny_hf_config, write_tiny_checkpoint +from .common import ENGRAM_ROWS, VOCAB, requires_cuda, tiny_hf_config, write_tiny_checkpoint, write_tiny_gguf def _args(**over) -> DeepseekV41Args: @@ -123,6 +123,80 @@ def oracle(rows: torch.Tensor) -> torch.Tensor: assert torch.equal(out, oracle(dec_rows)) +@requires_cuda +@pytest.mark.parametrize("table_type", ["BF16", "Q8_0"]) +def test_gguf_tables_read_their_slices_of_the_one_tensor(tmp_path, table_type): + """Two Engram layers back to back in ``per_layer_token_embd``: each table reads its own rows, eagerly and from a captured graph.""" + from freetoken.kernel.row_store import alloc_sync_flag, probe_wait_sync + from freetoken.models.deepseek_v41.engram_table import EngramDiskTable + from freetoken.models.deepseek_v41.gguf import gguf_engram_tables, parse_gguf_config + from freetoken.models.gguf.config import build_gguf_shim + + over = dict(engram_layer_ids=[1, 3], engram_num_embeddings=[ENGRAM_ROWS, ENGRAM_ROWS]) + path = str(tmp_path / "tiny.gguf") + refs = write_tiny_gguf(path, table_type=table_type, **over) + args = parse_gguf_config(build_gguf_shim(path)).dsv41_args + hash, sources = gguf_engram_tables(path, args) + assert [s.num_rows for s in sources] == hash.rows_per_layer and sources[1].base - sources[0].base == hash.rows_per_layer[0] * sources[0].row_bytes + ids = torch.randint(3, VOCAB, (24,), generator=torch.Generator().manual_seed(3)) + rows = hash.row_ids(ids) + wait_sync = probe_wait_sync("auto", torch.device("cuda")) + for l, (layer, source) in enumerate(zip((1, 3), sources)): + table = EngramDiskTable(source, hash.n_cols, torch.device("cuda"), max_graph_rows=4, max_extend_tokens=32) + oracle = refs[layer][rows[:, l].flatten()].view(rows.shape[0], -1).cuda() + table.stage(rows[:, l], graph=False) + table.flush(signal=False) + assert torch.equal(table.lookup(rows.shape[0]), oracle) + table.wait_sync = wait_sync + if wait_sync: + table.flag = alloc_sync_flag() + out = torch.empty(1, hash.n_cols * 32, dtype=torch.bfloat16, device="cuda") + graph, stream = torch.cuda.CUDAGraph(), torch.cuda.Stream() + with torch.cuda.stream(stream): + if not wait_sync: + table.stage(rows[-1:, l], graph=True) + table.flush(signal=False) + with torch.cuda.graph(graph, stream=stream): + out.copy_(table.lookup(1)) + torch.cuda.synchronize() + graph.replay() + table.stage(rows[-1:, l], graph=True) + table.flush(signal=wait_sync) + torch.cuda.synchronize() + assert torch.equal(out, oracle[-1:]) + + +def test_gguf_hash_uses_the_token_map_in_the_file(tmp_path): + """Ids the file's map merges hash to the same rows: nothing but the file knows the map.""" + from freetoken.models.deepseek_v41.gguf import gguf_engram_tables, parse_gguf_config + from freetoken.models.gguf.config import build_gguf_shim + + token_map = [i // 2 for i in range(VOCAB)] + path = str(tmp_path / "merged.gguf") + write_tiny_gguf(path, kv_over={"engram.token_map": token_map}, engram_compressed_vocab_size=max(token_map) + 1) + hash, _ = gguf_engram_tables(path, parse_gguf_config(build_gguf_shim(path)).dsv41_args) + rows = lambda *ids: hash.row_ids(torch.tensor(ids)) + assert torch.equal(rows(10, 20), rows(11, 21)) + assert not torch.equal(rows(10, 20), rows(12, 20)) + + +@pytest.mark.parametrize("key", ["ple.head_offsets", "ple.head_vocab_sizes", "ple.layer_multipliers"]) +def test_gguf_hash_constants_are_checked_against_the_file(tmp_path, key): + from freetoken.models.deepseek_v41.gguf import gguf_engram_tables, parse_gguf_config + from freetoken.models.gguf.config import build_gguf_shim + from freetoken.models.gguf.reader import load_gguf_metadata + + path = str(tmp_path / "good.gguf") + write_tiny_gguf(path) + good = load_gguf_metadata(path)[f"deepseek41.{key}"] + gguf_engram_tables(path, parse_gguf_config(build_gguf_shim(path)).dsv41_args) + # primes also size the table, so a wrong one is caught by the same check through its own config + bad = str(tmp_path / "bad.gguf") + write_tiny_gguf(bad, kv_over={key: [v + 1 if i == 1 else v for i, v in enumerate(good)]}) + with pytest.raises(ValueError, match=key): + gguf_engram_tables(bad, parse_gguf_config(build_gguf_shim(bad)).dsv41_args) + + @requires_cuda def test_eager_staging_does_not_overwrite_an_in_flight_lookup(tmp_path): """The overlap scheduler stages batch k+1 while batch k's lookup copy is still queued behind diff --git a/tests/models/deepseek_v41/test_forward_smoke.py b/tests/models/deepseek_v41/test_forward_smoke.py index e0861e8874..423c73bf00 100644 --- a/tests/models/deepseek_v41/test_forward_smoke.py +++ b/tests/models/deepseek_v41/test_forward_smoke.py @@ -6,7 +6,7 @@ import pytest import torch -from .common import VOCAB, requires_cuda, write_tiny_checkpoint +from .common import VOCAB, requires_cuda, write_tiny_checkpoint, write_tiny_gguf pytestmark = requires_cuda @@ -39,6 +39,32 @@ def test_prefill_and_decode(checkpoint): assert out.shape == (2, VOCAB) and torch.isfinite(out).all() +def test_gguf_forward_matches_the_checkpoint(tmp_path): + """The same bf16 weights through the GGUF reader and the HF reader give the same logits, with non-zero Engram rows and random q / k gate weights so the file's stored product and ``wkv`` count.""" + from .harness import TinyEngine + + class Rows: + def __init__(self, width): + self.rows = torch.randn(600, width, generator=torch.Generator().manual_seed(9)).to(torch.bfloat16).cuda() + + def lookup(self, num_tokens): + return self.rows[:num_tokens] + + folder, path = str(tmp_path / "hf"), str(tmp_path / "tiny.gguf") + write_tiny_checkpoint(folder, random_gate=True) + write_tiny_gguf(path, random_gate=True) + engines = [TinyEngine(c, max_seq_len=1024, max_running_req=2, swa_decoder_replay="exact", engram_table=Rows(192)) for c in (folder, path)] + reqs = [e.new_request(0, _tokens(300, 1)) for e in engines] + logits = [e.prefill([r]) for e, r in zip(engines, reqs)] + assert torch.equal(*logits) + for e, r in zip(engines, reqs): + e.finish_prefill([r]) + # the bf16 fused and ggml decode kernels round differently; allow 2^-7 of the largest logit + for step in range(4): + hf, ggml = [e.decode([r], [10 + step]).float() for e, r in zip(engines, reqs)] + assert (hf - ggml).abs().max() <= 2**-7 * hf.abs().max() + + def test_chunked_prefill_matches_single_shot(checkpoint): """Two 128-aligned chunks must reproduce the single-shot prefill's last-token logits: the compressor carry, the window ring and the compressed rows are all resumed through the pool. The From 86f716eddaa4e6169238a01380793f23bfb566c1 Mon Sep 17 00:00:00 2001 From: Chen Cheng <85005591+ischencheng@users.noreply.github.com> Date: Fri, 13 Nov 2026 16:45:36 -0500 Subject: [PATCH 2/2] fix(gguf): refuse deepseek-v4.1 gguf mismatches at config time Refuse an FTW of the GGUF before the writer creates the output dir and keep the side-file hook after the dense loop as on main. Check the Engram hash KVs and the table in parse_gguf_config, refuse a table with a .scale, and share the clamp and expert_count readers with the deepseek4 GGUF reader. --- python/freetoken/checkpoint/convert.py | 26 ++++-- python/freetoken/models/deepseek_v4/gguf.py | 31 ++----- .../freetoken/models/deepseek_v41/__init__.py | 4 +- python/freetoken/models/deepseek_v41/gguf.py | 87 +++++++++++-------- .../freetoken/models/deepseek_v41/weight.py | 11 ++- python/freetoken/models/gguf/config.py | 29 ++++++- tests/models/deepseek_v41/common.py | 17 +++- tests/models/deepseek_v41/test_config.py | 39 ++++++++- tests/models/deepseek_v41/test_engram.py | 5 +- 9 files changed, 167 insertions(+), 82 deletions(-) diff --git a/python/freetoken/checkpoint/convert.py b/python/freetoken/checkpoint/convert.py index 4861437f86..6dbded318c 100644 --- a/python/freetoken/checkpoint/convert.py +++ b/python/freetoken/checkpoint/convert.py @@ -214,20 +214,21 @@ def convert_checkpoint( from freetoken.utils.progress import byte_bar, count_bar - writer = FTWWriter(out_dir, shard_limit=shard_limit) - n_weight = n_bank = n_alpha = 0 - - # 1) dense weights (host tensors; load straight to CPU to avoid GPU pressure) - _progress("dense", 0, 0) # phase start; per-tensor cumulative bytes follow (total unknown) - # files the model reads directly from the checkpoint dir, not through FTW entries (Qwen3.8-Flash-Next PLE table); first, since a hook may refuse the checkpoint + # a family refuses a source it has no FTW form for before the writer creates out_dir from freetoken.models.register import _load_attr, get_model_spec try: - side_hook = _load_attr(get_model_spec(mc.architectures[0]).module, "ftw_side_files") + check_source = _load_attr(get_model_spec(mc.architectures[0]).module, "ftw_check_source") except (AttributeError, KeyError, ValueError): - side_hook = None - side_files = side_hook(model_path, out_dir) if side_hook is not None else [] + check_source = None + if check_source is not None: + check_source(model_path) + writer = FTWWriter(out_dir, shard_limit=shard_limit) + n_weight = n_bank = n_alpha = 0 + + # 1) dense weights (host tensors; load straight to CPU to avoid GPU pressure) + _progress("dense", 0, 0) # phase start; per-tensor cumulative bytes follow (total unknown) dense_bytes = 0 from freetoken.models.gguf.reader import gguf_config_source @@ -241,6 +242,13 @@ def convert_checkpoint( dense_bytes += tensor.numel() * tensor.element_size() _progress("dense", dense_bytes, 0) + # files the model reads directly from the checkpoint dir, not through FTW entries (Qwen3.8-Flash-Next PLE table) + try: + side_hook = _load_attr(get_model_spec(mc.architectures[0]).module, "ftw_side_files") + except (AttributeError, KeyError, ValueError): + side_hook = None + side_files = side_hook(model_path, out_dir) if side_hook is not None else [] + # 2) offload expert banks (post-repack) + alpha scales (slow path auto-picks parallel/serial) quant_format = None layer_quant_formats = None diff --git a/python/freetoken/models/deepseek_v4/gguf.py b/python/freetoken/models/deepseek_v4/gguf.py index 23ace76989..a786668259 100644 --- a/python/freetoken/models/deepseek_v4/gguf.py +++ b/python/freetoken/models/deepseek_v4/gguf.py @@ -31,6 +31,7 @@ import numpy as np import torch +from freetoken.models.gguf.config import routed_expert_count, uniform_swiglu_clamp from freetoken.models.gguf.mmproj import Mmproj, find_mmproj, log_mmproj, mmproj_tensor from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index, gguf_tensor_names, load_gguf_metadata from freetoken.models.gguf.tensors import check_scales, dequantized_tensor as _float, fused, matrix @@ -126,14 +127,7 @@ def need(key: str): # process nulls vision_n_layers on its copy, and None here is the same answer vc = parse_vision_config(shim) - # the clamp is per-layer in the GGUF (and the shared expert carries its own array); the model - # takes one value, so anything non-uniform is refused instead of silently taking the first - clamp = need("swiglu_clamp_exp") - clamp_vals = list(clamp) if isinstance(clamp, list) else [clamp] - clamp_shexp = kv("swiglu_clamp_shexp") - shexp_vals = list(clamp_shexp) if isinstance(clamp_shexp, list) else [clamp_shexp] if clamp_shexp is not None else [] - if len(set(map(float, clamp_vals + shexp_vals))) != 1: - raise NotImplementedError(f"{shim.model_path}: non-uniform swiglu_clamp_exp / swiglu_clamp_shexp is not served") + swiglu_limit = uniform_swiglu_clamp(shim.model_path, need("swiglu_clamp_exp"), kv("swiglu_clamp_shexp")) from freetoken.third_party.gguf.constants import ExpertGatingFuncType @@ -171,21 +165,8 @@ def need(key: str): num_layers = int(need("block_count")) index = gguf_tensor_index(shim.model_path) - # REAP exports prune experts per layer: as a per-layer expert_count array, or with the - # scalar kept but the expert tensors shrunk. llama.cpp builds every layer from one - # n_expert and refuses both shapes; the model has no per-layer expert count to express. - experts_kv = need("expert_count") - if isinstance(experts_kv, list): - raise NotImplementedError( - f"{shim.model_path}: a per-layer expert_count array (REAP-style prune) is not served" - ) - n_experts = int(experts_kv) - for name, info in index.items(): - if name.endswith(".ffn_gate_exps.weight") and info.ne[-1] != n_experts: - raise NotImplementedError( - f"{shim.model_path}: {name} carries {info.ne[-1]} experts but expert_count says " - f"{n_experts} (REAP-style per-layer prunes are not served)" - ) + # llama.cpp builds every layer from one n_expert; the model has no per-layer expert count to express + n_experts = routed_expert_count(shim.model_path, need("expert_count"), index) args = DeepseekV4Args( max_batch_size=1, vocab_size=int(shim.vocab_size), @@ -194,12 +175,12 @@ def need(key: str): n_layers=num_layers, n_hash_layers=int(kv("hash_layer_count", 0) or 0), n_heads=int(need("attention.head_count")), - n_routed_experts=int(need("expert_count")), + n_routed_experts=n_experts, n_shared_experts=int(kv("expert_shared_count", 1) or 1), n_activated_experts=int(need("expert_used_count")), score_func="sqrtsoftplus", route_scale=float(need("expert_weights_scale")), - swiglu_limit=float(clamp_vals[0]), + swiglu_limit=swiglu_limit, q_lora_rank=int(need("attention.q_lora_rank")), head_dim=head_dim, rope_head_dim=int(need("rope.dimension_count")), diff --git a/python/freetoken/models/deepseek_v41/__init__.py b/python/freetoken/models/deepseek_v41/__init__.py index 9af7163b7f..fabe98d3b1 100644 --- a/python/freetoken/models/deepseek_v41/__init__.py +++ b/python/freetoken/models/deepseek_v41/__init__.py @@ -10,9 +10,9 @@ from .config import parse_config from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf from .model import DeepseekV41ForCausalLM -from .weight import ftw_side_files, iter_expert_pieces, iter_weights, iter_vision_weights +from .weight import ftw_check_source, ftw_side_files, iter_expert_pieces, iter_weights, iter_vision_weights __all__ = [ "DeepseekV41ForCausalLM", "parse_config", "iter_weights", "iter_vision_weights", "iter_expert_pieces", "ftw_side_files", - "iter_gguf_weights", "parse_gguf_config", "to_gguf", + "ftw_check_source", "iter_gguf_weights", "parse_gguf_config", "to_gguf", ] diff --git a/python/freetoken/models/deepseek_v41/gguf.py b/python/freetoken/models/deepseek_v41/gguf.py index 223f52ddab..51179c36d2 100644 --- a/python/freetoken/models/deepseek_v41/gguf.py +++ b/python/freetoken/models/deepseek_v41/gguf.py @@ -14,6 +14,7 @@ import torch from freetoken.layers.quantization.ggml_types import row_bytes +from freetoken.models.gguf.config import routed_expert_count, uniform_swiglu_clamp from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index from freetoken.models.gguf.tensors import check_scales, dequantized_tensor, fused, matrix @@ -64,7 +65,7 @@ def to_gguf(name: str, tensors: dict[str, int]) -> tuple[str, ...]: return tuple(f"blk.{m.group(1)}.{leaf}" for leaf in _LAYER_LEAVES.get(m.group(2), ())) -def _args_from_metadata(shim: GgufConfigShim) -> DeepseekV41Args: +def _args_from_metadata(shim: GgufConfigShim, index: dict[str, GgufTensorInfo]) -> DeepseekV41Args: md, prefix, path = shim.metadata, shim.model_type, shim.model_path def kv(key: str, default=None): @@ -89,10 +90,6 @@ def need(key: str): raise NotImplementedError(f"{path}: attention key and value lengths differ") if int(kv("attention.head_count_kv", 1) or 1) != 1: raise NotImplementedError(f"{path}: more than one latent KV head is not served") - # the clamp is per-layer in the GGUF; the model takes one value - clamps = {float(v) for k in ("swiglu_clamp_exp", "swiglu_clamp_shexp") for v in need(k)} - if len(clamps) != 1: - raise NotImplementedError(f"{path}: non-uniform swiglu_clamp_exp / swiglu_clamp_shexp is not served") engram: dict = {} if (layers := kv("ple.layers")) is not None: @@ -118,13 +115,13 @@ def need(key: str): n_layers=int(need("block_count")), n_mtp_layers=0, # the converter leaves the DSpark stage out n_heads=int(need("attention.head_count")), - n_routed_experts=int(need("expert_count")), + n_routed_experts=routed_expert_count(path, need("expert_count"), index), n_shared_experts=int(kv("expert_shared_count", 1) or 1), n_activated_experts=int(need("expert_used_count")), score_func="sqrtsoftplus", norm_topk_prob=bool(need("expert_weights_norm")), route_scale=float(need("expert_weights_scale")), - swiglu_limit=clamps.pop(), + swiglu_limit=uniform_swiglu_clamp(path, need("swiglu_clamp_exp"), need("swiglu_clamp_shexp")), q_lora_rank=int(need("attention.q_lora_rank")), head_dim=head_dim, rope_head_dim=int(need("rope.dimension_count")), @@ -158,23 +155,55 @@ def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig: from dataclasses import replace check_scales(shim.model_path) - args = _args_from_metadata(shim) index = gguf_tensor_index(shim.model_path) + args = _args_from_metadata(shim, index) if not args.engram_layer_ids and (f"{shim.model_type}.engram.token_map" in shim.metadata or any(".engram_" in name for name in index)): raise NotImplementedError(f"{shim.model_path}: Engram tensors without {shim.model_type}.ple.layers, so not llama.cpp's deepseek41 layout (converted for an earlier branch?)") - # a metadata-only GGUF (an FTW dir's source_metadata.gguf) has no tensor table to check - if index: - if args.engram_layer_ids and _TABLE not in index: - raise NotImplementedError(f"{shim.model_path}: no {_TABLE} tensor, the Engram tables llama.cpp's converter writes") - for name, info in index.items(): - if name.endswith(".ffn_gate_exps.weight") and info.ne[-1] != args.n_routed_experts: - raise NotImplementedError( - f"{shim.model_path}: {name} carries {info.ne[-1]} experts but expert_count says {args.n_routed_experts} (per-layer prunes are not served)" - ) + if args.engram_layer_ids: + _check_engram(shim, args, index) config = model_config_from_args(args, max_position=int(shim.metadata[f"{shim.model_type}.context_length"])) return replace(config, expert_quant="gguf") +def _check_engram(shim: GgufConfigShim, args: DeepseekV41Args, index: dict[str, GgufTensorInfo]) -> None: + """Refuse a file whose Engram hash or table this build would read differently, before any weight loads. + + The bucket primes and the multipliers derive from the config as on the HF path, so they are checked against the KVs llama.cpp hashes with.""" + path, arch = shim.model_path, shim.model_type + hash = _engram_hash(shim.metadata, arch, args) + for key, derived in ( + ("ple.head_offsets", _table_offsets(hash).flatten().tolist()), + ("ple.head_vocab_sizes", hash.primes.flatten().tolist()), + ("ple.layer_multipliers", hash.multipliers.flatten().tolist()), + ): + if [int(v) for v in shim.metadata[f"{arch}.{key}"]] != derived: + raise ValueError(f"{path}: {arch}.{key} differs from the Engram hash this build derives from the config") + # a metadata-only GGUF (an FTW dir's source_metadata.gguf) has no tensor table to check + if not index: + return + if _TABLE not in index: + raise NotImplementedError(f"{path}: no {_TABLE} tensor, the Engram tables llama.cpp's converter writes") + # check_scales lets an NVFP4 tensor's .scale through, and the table reader would decode without it + if _TABLE.replace(".weight", ".scale") in index: + raise NotImplementedError(f"{path}: an Engram table with a per-tensor .scale is not served") + head_dim, rows = index[_TABLE].ne + if head_dim != args.engram_head_dim or rows < sum(hash.rows_per_layer): + raise ValueError(f"{path}: {_TABLE} is {index[_TABLE].ne}, expected {args.engram_head_dim} values wide and at least {sum(hash.rows_per_layer)} rows") + + +def _engram_hash(metadata: dict, arch: str, args: DeepseekV41Args) -> EngramHash: + return EngramHash(args, metadata[f"{arch}.engram.token_map"], args.engram_compressed_vocab_size) + + +def _layer_starts(hash: EngramHash) -> list[int]: + """Each Engram layer's first row in the one table: the layers sit back to back.""" + return [sum(hash.rows_per_layer[:i]) for i in range(len(hash.rows_per_layer))] + + +def _table_offsets(hash: EngramHash) -> torch.Tensor: + return hash.offsets + torch.tensor(_layer_starts(hash)).unsqueeze(1) + + # ---------------------------------------------------------------------------------------------- # Weights # ---------------------------------------------------------------------------------------------- @@ -269,30 +298,16 @@ def linear(dst: str, src: str, k: int) -> Iterator[tuple[str, torch.Tensor]]: def gguf_engram_tables(model_path: str, args: DeepseekV41Args) -> tuple[EngramHash, list[EngramRowSource]]: """The hash from the file's token map and one row source per Engram layer: each is a slice of the single ``per_layer_token_embd`` tensor. - The bucket primes and the multipliers derive from the config as on the HF path, so they are checked against the KVs llama.cpp hashes with.""" + ``args`` comes from ``parse_gguf_config``, which has checked the hash constants and the table against the file.""" from freetoken.models.gguf.reader import load_gguf_metadata md = load_gguf_metadata(model_path) - arch = md["general.architecture"] - hash = EngramHash(args, md[f"{arch}.engram.token_map"], args.engram_compressed_vocab_size) - starts = [sum(hash.rows_per_layer[:i]) for i in range(len(hash.rows_per_layer))] - offsets = (hash.offsets + torch.tensor(starts).unsqueeze(1)).flatten().tolist() - for key, derived in ( - ("ple.head_offsets", offsets), - ("ple.head_vocab_sizes", hash.primes.flatten().tolist()), - ("ple.layer_multipliers", hash.multipliers.flatten().tolist()), - ): - if [int(v) for v in md[f"{arch}.{key}"]] != derived: - raise ValueError(f"{model_path}: {arch}.{key} differs from the Engram hash this build derives from the config") - + hash = _engram_hash(md, md["general.architecture"], args) info = gguf_tensor_index(model_path)[_TABLE] - head_dim, rows = info.ne - if head_dim != args.engram_head_dim or rows < sum(hash.rows_per_layer): - raise ValueError(f"{model_path}: {_TABLE} is {info.ne}, expected {args.engram_head_dim} values wide and at least {sum(hash.rows_per_layer)} rows") - nbytes = row_bytes(info.ggml_type, head_dim) + nbytes = row_bytes(info.ggml_type, args.engram_head_dim) return hash, [ - EngramRowSource(info.path, info.offset + start * nbytes, layer_rows, head_dim, None, info.ggml_type, nbytes) - for start, layer_rows in zip(starts, hash.rows_per_layer) + EngramRowSource(info.path, info.offset + start * nbytes, layer_rows, args.engram_head_dim, None, info.ggml_type, nbytes) + for start, layer_rows in zip(_layer_starts(hash), hash.rows_per_layer) ] diff --git a/python/freetoken/models/deepseek_v41/weight.py b/python/freetoken/models/deepseek_v41/weight.py index 52cacb9e7a..2f8c5376fe 100644 --- a/python/freetoken/models/deepseek_v41/weight.py +++ b/python/freetoken/models/deepseek_v41/weight.py @@ -175,6 +175,14 @@ def engram_table_names(layer_id: int) -> tuple[str, str]: return f"layers.{layer_id}.engram.embed.weight", f"layers.{layer_id}.engram.embed.scale" +def ftw_check_source(model_path: str) -> None: + """``convert_checkpoint`` asks this before writing anything: a GGUF keeps the Engram tables in the .gguf, which an FTW does not carry.""" + from freetoken.models.gguf.reader import gguf_config_source + + if gguf_config_source(model_path) is not None: + raise NotImplementedError(f"{model_path}: an FTW of a DeepSeek-V4.1 GGUF has no Engram table to read; serve the .gguf directly") + + def ftw_side_files(model_path: str, out_dir: str) -> list[str]: """Copy the checkpoint shards holding the Engram tables (fp8 rows + ue8m0 scales) next to an FTW checkpoint. The tables are never FTW entries: ``load_host_tables`` maps them in place from whatever @@ -182,12 +190,9 @@ def ftw_side_files(model_path: str, out_dir: str) -> list[str]: loader never reads ``.safetensors`` as weights, so the shard's other tensors ride along unused.""" import shutil - from freetoken.models.gguf.reader import gguf_config_source from freetoken.models.loader import safetensors_weight_map from freetoken.utils import download_hf_weight - if gguf_config_source(model_path) is not None: - raise NotImplementedError(f"{model_path}: an FTW of a DeepSeek-V4.1 GGUF has no Engram table to read; serve the .gguf directly") folder = download_hf_weight(model_path) args = load_args(folder) shard_of = safetensors_weight_map(folder) diff --git a/python/freetoken/models/gguf/config.py b/python/freetoken/models/gguf/config.py index fe75a0a44e..a450ae7d6e 100644 --- a/python/freetoken/models/gguf/config.py +++ b/python/freetoken/models/gguf/config.py @@ -158,4 +158,31 @@ def build_gguf_shim(model_path: str) -> GgufConfigShim: ) -__all__ = ["GgufConfigShim", "GGUF_ARCH_TO_REGISTRY", "GGUF_MMPROJ_TO_REGISTRY", "build_gguf_shim"] +def uniform_swiglu_clamp(model_path: str, clamp_exp: Any, clamp_shexp: Any = None) -> float: + """The one SwiGLU clamp of the DeepSeek-V4 readers. The GGUF holds it per layer (scalar or array, the shared + expert's optional), the model takes one value, so anything non-uniform is refused instead of taking the first.""" + values = [v for kv in (clamp_exp, clamp_shexp) if kv is not None for v in (kv if isinstance(kv, list) else [kv])] + if len(set(map(float, values))) != 1: + raise NotImplementedError(f"{model_path}: non-uniform swiglu_clamp_exp / swiglu_clamp_shexp is not served") + return float(values[0]) + + +def routed_expert_count(model_path: str, expert_count: Any, index: dict) -> int: + """``expert_count`` when every layer's expert tensors carry it. REAP exports prune experts per layer, as a + per-layer array or with the scalar kept but the tensors shrunk; llama.cpp refuses both, and so does this.""" + if isinstance(expert_count, list): + raise NotImplementedError(f"{model_path}: a per-layer expert_count array (REAP-style prune) is not served") + n_experts = int(expert_count) + for name, info in index.items(): + if name.endswith(".ffn_gate_exps.weight") and info.ne[-1] != n_experts: + raise NotImplementedError( + f"{model_path}: {name} carries {info.ne[-1]} experts but expert_count says " + f"{n_experts} (REAP-style per-layer prunes are not served)" + ) + return n_experts + + +__all__ = [ + "GgufConfigShim", "GGUF_ARCH_TO_REGISTRY", "GGUF_MMPROJ_TO_REGISTRY", "build_gguf_shim", + "routed_expert_count", "uniform_swiglu_clamp", +] diff --git a/tests/models/deepseek_v41/common.py b/tests/models/deepseek_v41/common.py index 0e77a524cb..0f5552add7 100644 --- a/tests/models/deepseek_v41/common.py +++ b/tests/models/deepseek_v41/common.py @@ -235,9 +235,12 @@ def value(kind, val, sub) -> bytes: f.write(blob) -def write_tiny_gguf(path: str, seed: int = 0, *, table_type: str = "BF16", random_gate: bool = False, kv_over: dict | None = None, **over) -> dict[int, torch.Tensor]: +def write_tiny_gguf( + path: str, seed: int = 0, *, table_type: str = "BF16", random_gate: bool = False, kv_over: dict | None = None, + extra_tensors: dict[str, torch.Tensor] | None = None, **over, +) -> dict[int, torch.Tensor]: """The tiny checkpoint as llama.cpp's converter lays it out (conversion/deepseek41.py), every tensor float except an optional Q8_0 ``table_type`` table; - ``kv_over`` replaces array KVs. Returns each Engram layer's rows as they decode, keyed by layer.""" + ``kv_over`` replaces KVs (a list writes an array, a scalar one value), ``extra_tensors`` adds tensors. Returns each Engram layer's rows as they decode, keyed by layer.""" import numpy as np from freetoken.models.deepseek_v41.args import DeepseekV41Args @@ -253,7 +256,13 @@ def write_tiny_gguf(path: str, seed: int = 0, *, table_type: str = "BF16", rando kvs: list[tuple[str, int, object, int | None]] = [("general.architecture", K.STRING, "deepseek41", None)] def add_kv(key, kind, value, sub=None): - kvs.append((f"deepseek41.{key}", kind, (kv_over or {}).get(key, value), sub)) + value = (kv_over or {}).get(key, value) + # an override may turn a per-layer array into one value or back + if isinstance(value, list) and kind != K.ARRAY: + kind, sub = K.ARRAY, kind + elif not isinstance(value, list) and kind == K.ARRAY: + kind, sub = sub, None + kvs.append((f"deepseek41.{key}", kind, value, sub)) for key, value in { "block_count": text["num_hidden_layers"], "context_length": text["max_position_embeddings"], "embedding_length": text["hidden_size"], @@ -323,5 +332,7 @@ def weight_name(gguf_name: str) -> str: table = dict(zip(table, decoded.to(torch.bfloat16).split([t.shape[0] for t in table.values()]))) else: add("per_layer_token_embd.weight", packed) + for name, t in (extra_tensors or {}).items(): + add(name, t) _write_gguf(path, kvs, out) return table diff --git a/tests/models/deepseek_v41/test_config.py b/tests/models/deepseek_v41/test_config.py index a064d755e0..9534d914f5 100644 --- a/tests/models/deepseek_v41/test_config.py +++ b/tests/models/deepseek_v41/test_config.py @@ -16,7 +16,7 @@ from freetoken.models.register import get_model_spec from freetoken.utils.torch_utils import torch_dtype -from .common import ENGRAM_LAYERS, INDEX_SOURCES, KV_SOURCES, N_LAYERS, RATIOS, tiny_hf_config, write_tiny_checkpoint, write_tiny_gguf +from .common import ENGRAM_LAYERS, INDEX_SOURCES, KV_SOURCES, N_LAYERS, RATIOS, requires_cuda, tiny_hf_config, write_tiny_checkpoint, write_tiny_gguf def test_roles_follow_the_dsv41_mode_table(): @@ -200,3 +200,40 @@ def test_gguf_from_an_earlier_branch_is_refused(tmp_path): shim.metadata = {k: v for k, v in shim.metadata.items() if k != "deepseek41.ple.layers"} with pytest.raises(NotImplementedError, match="ple.layers"): parse_gguf_config(shim) + + +def test_gguf_clamp_may_be_one_value_and_a_per_layer_expert_count_is_refused(tmp_path): + from freetoken.models.deepseek_v41.gguf import parse_gguf_config + from freetoken.models.gguf.config import build_gguf_shim + + text = tiny_hf_config().text_config + scalar = str(tmp_path / "scalar.gguf") + write_tiny_gguf(scalar, kv_over={"swiglu_clamp_exp": float(text["swiglu_limit"])}) + assert parse_gguf_config(build_gguf_shim(scalar)).dsv41_args.swiglu_limit == text["swiglu_limit"] + pruned = str(tmp_path / "pruned.gguf") + write_tiny_gguf(pruned, kv_over={"expert_count": [text["n_routed_experts"]] * N_LAYERS}) + with pytest.raises(NotImplementedError, match="per-layer expert_count"): + parse_gguf_config(build_gguf_shim(pruned)) + + +def test_gguf_engram_table_with_a_scale_is_refused(tmp_path, monkeypatch): + """check_scales passes an NVFP4 table's .scale, which the table reader never applies.""" + import freetoken.models.deepseek_v41.gguf as reader + from freetoken.models.gguf.config import build_gguf_shim + + path = str(tmp_path / "scaled.gguf") + write_tiny_gguf(path, extra_tensors={"per_layer_token_embd.scale": torch.ones(1)}) + monkeypatch.setattr(reader, "check_scales", lambda *a, **k: None) + with pytest.raises(NotImplementedError, match="per-tensor .scale"): + reader.parse_gguf_config(build_gguf_shim(path)) + + +@requires_cuda +def test_ftw_conversion_refuses_a_gguf_before_writing(tmp_path): + from freetoken.checkpoint.convert import convert_checkpoint + + path, out = str(tmp_path / "tiny.gguf"), tmp_path / "ftw" + write_tiny_gguf(path) + with pytest.raises(NotImplementedError, match="serve the .gguf directly"): + convert_checkpoint(path, str(out)) + assert not out.exists() diff --git a/tests/models/deepseek_v41/test_engram.py b/tests/models/deepseek_v41/test_engram.py index f8f74923ce..3be6a00fdc 100644 --- a/tests/models/deepseek_v41/test_engram.py +++ b/tests/models/deepseek_v41/test_engram.py @@ -190,11 +190,12 @@ def test_gguf_hash_constants_are_checked_against_the_file(tmp_path, key): write_tiny_gguf(path) good = load_gguf_metadata(path)[f"deepseek41.{key}"] gguf_engram_tables(path, parse_gguf_config(build_gguf_shim(path)).dsv41_args) - # primes also size the table, so a wrong one is caught by the same check through its own config + # primes also size the table, so a wrong one is caught by the same check through its own config; + # the config parse refuses the file, before any dense weight loads bad = str(tmp_path / "bad.gguf") write_tiny_gguf(bad, kv_over={key: [v + 1 if i == 1 else v for i, v in enumerate(good)]}) with pytest.raises(ValueError, match=key): - gguf_engram_tables(bad, parse_gguf_config(build_gguf_shim(bad)).dsv41_args) + parse_gguf_config(build_gguf_shim(bad)) @requires_cuda