char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Georgi Gerganovconvert : skip bias_vl tensor in DeepSeek-V4 DSpark conversion (#28294)d646c9d15

char/aliceai-t5-moe
50.2 KiB1104 linesraw
1from __future__ import annotations
2
3import json
4import re
5from pathlib import Path
6
7from typing import Any, Callable, Iterable, TYPE_CHECKING
8
9import numpy as np
10import torch
11
12if TYPE_CHECKING:
13    from torch import Tensor
14
15from .base import LazyTorchTensor, MmprojModel, ModelBase, TextModel, gguf, logger
16
17from .qwen import QwenModel
18
19
20@ModelBase.register("DeepseekOCRForCausalLM")
21@ModelBase.example("deepseek-ai/DeepSeek-OCR")
22class DeepseekOCRVisionModel(MmprojModel):
23    # HF dynamic_preprocess() max_num, which differs per model
24    preproc_max_tiles = 9
25
26    def __init__(self, *args, **kwargs):
27        super().__init__(*args, **kwargs)
28        self.clip_projector_type = gguf.VisionProjectorType.DEEPSEEKOCR
29
30    def set_gguf_parameters(self):
31        super().set_gguf_parameters()
32        hparams = self.hparams
33        self.gguf_writer.add_clip_projector_type(self.clip_projector_type)
34        # default values below are taken from HF tranformers code
35        self.gguf_writer.add_vision_attention_layernorm_eps(hparams.get("layer_norm_eps", 1e-6))
36        self.gguf_writer.add_vision_use_gelu(True)
37        # calculate proj_scale_factor (used by tinygemma3 test model)
38        image_seq_length = self.preprocessor_config.get("image_seq_length", 256)
39        n_per_side = int(image_seq_length ** 0.5)
40        image_size = self.hparams["image_size"]
41        patch_size = self.hparams["patch_size"]
42        proj_scale_factor = (image_size // patch_size) // n_per_side
43        if proj_scale_factor > 0 and proj_scale_factor != 4:
44            # we only need to write this if it's not the default value
45            # in this case, we are converting a test model
46            self.gguf_writer.add_vision_projector_scale_factor(proj_scale_factor)
47        # @bluebread: there's no window_size in config but just add it here anyway
48        self.gguf_writer.add_vision_window_size(self.hparams.get("window_size", 14))
49
50        self.gguf_writer.add_vision_preproc_min_tiles(2)
51        self.gguf_writer.add_vision_preproc_max_tiles(self.preproc_max_tiles)
52
53        # SAM configuration
54        sam_hparams = hparams['sam']
55        self.gguf_writer.add_vision_sam_layers_count(sam_hparams['layers'])
56        self.gguf_writer.add_vision_sam_embedding_length(sam_hparams['width'])
57        self.gguf_writer.add_vision_sam_head_count(sam_hparams['heads'])
58
59    def get_vision_config(self) -> dict[str, Any]:
60        vision_config: dict[str, Any] | None = self.global_config.get("vision_config")
61
62        if not vision_config:
63            raise ValueError("DeepseekOCR model requires 'vision_config' in the model configuration, but it was not found")
64
65        vision_config['sam'] = vision_config['width']['sam_vit_b']
66        if vision_config['width'].get('clip-l-14-224') is not None:
67            vision_config.update(vision_config['width']['clip-l-14-224'])
68        if isinstance(vision_config['width'], int):
69            vision_config['hidden_size'] = vision_config['width']
70        if vision_config.get('heads') is not None:
71            vision_config['num_heads'] = vision_config['heads']
72            vision_config['intermediate_size'] = vision_config['heads'] * 4
73
74        return vision_config
75
76    def tensor_force_quant(self, name, new_name, bid, n_dims):
77        for nq_name in ('.embeddings.', 'pos_embed', '.rel_pos_h', '.rel_pos_w', '.neck.', '.net_'):
78            if nq_name in name:
79                return gguf.GGMLQuantizationType.F32
80        return super().tensor_force_quant(name, new_name, bid, n_dims)
81
82    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
83        if name.endswith("view_seperator"):
84            data_torch = data_torch.unsqueeze(0)
85        yield from super().modify_tensors(data_torch, name, bid)
86
87    @classmethod
88    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
89        name, gen = item
90
91        # Only process vision-related tensors, skip language model tensors
92        # Vision components: sam_model, vision_model, projector, image_newline, view_seperator
93        # Language model components to skip: lm_head, embed_tokens, layers, norm
94        if name.startswith(("lm_head.", "model.embed_tokens.", "model.layers.", "model.norm.")):
95            return None
96
97        if name.endswith("pos_embed") or name.endswith("rel_pos_h") or name.endswith("rel_pos_w"):
98            name += ".weight"
99
100        return super().filter_tensors((name, gen))
101
102
103@ModelBase.register("UnlimitedOCRForCausalLM")
104@ModelBase.example("baidu/Unlimited-OCR")
105class UnlimitedOCRVisionModel(DeepseekOCRVisionModel):
106    preproc_max_tiles = 32
107
108
109@ModelBase.register("DeepseekOCR2ForCausalLM")
110@ModelBase.example("deepseek-ai/DeepSeek-OCR-2")
111class DeepseekOCR2VisionModel(DeepseekOCRVisionModel):
112    preproc_max_tiles = 6
113
114    def __init__(self, *args, **kwargs):
115        super().__init__(*args, **kwargs)
116        self.clip_projector_type = gguf.VisionProjectorType.DEEPSEEKOCR2
117
118    def set_gguf_parameters(self):
119        # the vision tower's qwen2 encoder is built from fixed defaults,
120        # see build_qwen2_decoder_as_encoder() in deepencoderv2.py
121        if self.hparams.get("patch_size") is None:
122            self.hparams["patch_size"] = 16
123        if self.hparams.get("intermediate_size") is None:
124            self.hparams["intermediate_size"] = 4864
125        if self.hparams.get("num_attention_heads") is None:
126            self.hparams["num_attention_heads"] = 14
127        super().set_gguf_parameters()
128        # qwen2 encoder is GQA: 14 Q heads, 2 KV heads
129        self.gguf_writer.add_vision_head_count_kv(2)
130
131    def get_vision_config(self) -> dict[str, Any]:
132        vision_config = super().get_vision_config()
133        vision_config['hidden_size'] = vision_config['width']['qwen2-0-5b']['dim']
134        if vision_config.get('layers') is None:
135            vision_config['layers'] = 24
136        return vision_config
137
138
139@ModelBase.register("DeepseekForCausalLM")
140@ModelBase.example("deepseek-ai/deepseek-moe-16b-chat")
141class DeepseekModel(TextModel):
142    model_arch = gguf.MODEL_ARCH.DEEPSEEK
143
144    def set_vocab(self):
145        try:
146            self._set_vocab_sentencepiece()
147        except FileNotFoundError:
148            self._set_vocab_gpt2()
149
150    def set_gguf_parameters(self):
151        super().set_gguf_parameters()
152        hparams = self.hparams
153        if (rope_dim := hparams.get("head_dim")) is None:
154            rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"]
155
156        self.gguf_writer.add_rope_dimension_count(rope_dim)
157        self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE)
158        self.gguf_writer.add_leading_dense_block_count(hparams["first_k_dense_replace"])
159        self.gguf_writer.add_vocab_size(hparams["vocab_size"])
160        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])
161        self.gguf_writer.add_expert_weights_scale(1.0)
162        self.gguf_writer.add_expert_count(hparams["n_routed_experts"])
163        self.gguf_writer.add_expert_shared_count(hparams["n_shared_experts"])
164
165    _experts: list[dict[str, Tensor]] | None = None
166
167    @staticmethod
168    def permute(weights: Tensor, n_head: int, n_head_kv: int | None):
169        if n_head_kv is not None and n_head != n_head_kv:
170            n_head = n_head_kv
171        return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:])
172                .swapaxes(1, 2)
173                .reshape(weights.shape))
174
175    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
176        n_head = self.hparams["num_attention_heads"]
177        n_kv_head = self.hparams.get("num_key_value_heads")
178
179        if name.endswith(("q_proj.weight", "q_proj.bias")):
180            data_torch = DeepseekModel.permute(data_torch, n_head, n_head)
181        if name.endswith(("k_proj.weight", "k_proj.bias")):
182            data_torch = DeepseekModel.permute(data_torch, n_head, n_kv_head)
183
184        # process the experts separately
185        if name.find("mlp.experts") != -1:
186            n_experts = self.hparams["n_routed_experts"]
187            assert bid is not None
188
189            if self._experts is None:
190                self._experts = [{} for _ in range(self.block_count)]
191
192            self._experts[bid][name] = data_torch
193
194            if len(self._experts[bid]) >= n_experts * 3:
195                # merge the experts into a single 3d tensor
196                for w_name in ["down_proj", "gate_proj", "up_proj"]:
197                    datas: list[Tensor] = []
198
199                    for xid in range(n_experts):
200                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
201                        datas.append(self._experts[bid][ename])
202                        del self._experts[bid][ename]
203
204                    data_torch = torch.stack(datas, dim=0)
205
206                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
207
208                    yield from super().modify_tensors(data_torch, merged_name, bid)
209                return
210            else:
211                return
212
213        yield from super().modify_tensors(data_torch, name, bid)
214
215    def prepare_tensors(self):
216        super().prepare_tensors()
217
218        if self._experts is not None:
219            # flatten `list[dict[str, Tensor]]` into `list[str]`
220            experts = [k for d in self._experts for k in d.keys()]
221            if len(experts) > 0:
222                raise ValueError(f"Unprocessed experts: {experts}")
223
224
225@ModelBase.register(
226    "DeepseekV2ForCausalLM",
227    "DeepseekV3ForCausalLM",
228    "DeepseekOCRForCausalLM",
229    "UnlimitedOCRForCausalLM",
230    "KimiVLForConditionalGeneration",
231    "KimiK25ForConditionalGeneration",
232    "YoutuForCausalLM",
233    "YoutuVLForConditionalGeneration",
234)
235@ModelBase.example("deepseek-ai/DeepSeek-V2-Lite", "deepseek-ai/DeepSeek-V3")
236class DeepseekV2Model(TextModel):
237    model_arch = gguf.MODEL_ARCH.DEEPSEEK2
238
239    # TODO @ngxson : remove this when we support MTP for deepseek models
240    skip_mtp = True
241
242    merge_expert = True
243
244    def __init__(self, *args, **kwargs):
245        super().__init__(*args, **kwargs)
246        hparams: dict = ModelBase.load_hparams(self.dir_model, is_mistral_format=False)
247        self.origin_hf_arch = hparams.get('architectures', [None])[0]
248
249        # special handling for Deepseek OCR
250        if self.origin_hf_arch in ("DeepseekOCRForCausalLM", "DeepseekOCR2ForCausalLM", "UnlimitedOCRForCausalLM"):
251            self.model_arch = gguf.MODEL_ARCH.DEEPSEEK2OCR
252            self.gguf_writer.arch = gguf.MODEL_ARCH_NAMES[self.model_arch]
253            self.gguf_writer.add_architecture()
254            # default jinja template
255            self.gguf_writer.add_chat_template("{% for m in messages %}{{m['content']}}{% endfor %}")
256
257    @classmethod
258    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
259        name, _ = item
260        # DeepSeek-OCR vision encoder (SAM + DeepSeek-OCR-2 qwen2 tower)
261        if "sam_model" in name or "qwen2_model" in name:
262            return None
263        return super().filter_tensors(item)
264
265    def set_vocab(self):
266        try:
267            self._set_vocab_gpt2()
268            return
269        except Exception:
270            pass
271
272        from transformers import AutoTokenizer
273        tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)
274        tokpre = self.get_vocab_base_pre(tokenizer)
275
276        if tokpre == "kimi-k2":
277            # Build merges list using the approach similar to HunYuanMoE
278            merges = []
279            vocab = {}
280            mergeable_ranks = tokenizer.model._mergeable_ranks  # ty: ignore[unresolved-attribute]
281            for token, rank in mergeable_ranks.items():
282                vocab[QwenModel.token_bytes_to_string(token)] = rank
283                if len(token) == 1:
284                    continue
285                merged = QwenModel.bpe(mergeable_ranks, token, max_rank=rank)
286                if len(merged) == 2:
287                    merges.append(' '.join(map(QwenModel.token_bytes_to_string, merged)))
288
289            # Build token list
290            vocab_size = self.hparams["vocab_size"]
291            special_tokens = tokenizer.special_tokens  # ty: ignore[unresolved-attribute]
292            reverse_vocab = {id_ : encoded_tok for encoded_tok, id_ in {**vocab, **special_tokens}.items()}
293            tokens: list[str] = []
294            toktypes: list[int] = []
295
296            for i in range(vocab_size):
297                if i not in reverse_vocab:
298                    tokens.append(f"[PAD{i}]")
299                    toktypes.append(gguf.TokenType.UNUSED)
300                else:
301                    token = reverse_vocab[i]
302                    tokens.append(token)
303                    if i in special_tokens.values():
304                        toktypes.append(gguf.TokenType.CONTROL)
305                    else:
306                        toktypes.append(gguf.TokenType.NORMAL)
307
308            self.gguf_writer.add_tokenizer_model("gpt2")
309            self.gguf_writer.add_tokenizer_pre(tokpre)
310            self.gguf_writer.add_token_list(tokens)
311            self.gguf_writer.add_token_types(toktypes)
312            self.gguf_writer.add_token_merges(merges)
313
314            special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=False)
315            special_vocab.add_to_gguf(self.gguf_writer)
316        else:
317            raise NotImplementedError(f"Deepseek pre-tokenizer {tokpre!r} is not supported yet!")
318
319    def set_gguf_parameters(self):
320        is_ocr = (self.model_arch == gguf.MODEL_ARCH.DEEPSEEK2OCR)
321
322        if is_ocr:
323            self.hparams['rope_theta'] = self.hparams.get('rope_theta', 10000.0)
324        else:
325            # note: deepseek2 using MLA converts into MQA (ie: GQA with 1 group)
326            self.hparams["num_key_value_heads"] = 1
327
328        self.hparams['rms_norm_eps'] = self.hparams.get('rms_norm_eps', 1e-6)
329
330        super().set_gguf_parameters()
331        hparams = self.hparams
332
333        # first_k_dense_replace: number of leading layers using dense FFN instead of MoE
334        # For non-MoE models (like Youtu), set to n_layer to use dense FFN for all layers
335        # For MoE models (like DeepSeek-V2), this is the number of leading non-MoE layers
336        has_moe = hparams.get("n_routed_experts") is not None
337        first_k_dense_replace = hparams.get("first_k_dense_replace")
338        if first_k_dense_replace is None:
339            # Default: if no MoE, all layers are dense; if MoE, none are dense
340            first_k_dense_replace = hparams["num_hidden_layers"] if not has_moe else 0
341        self.gguf_writer.add_leading_dense_block_count(first_k_dense_replace)
342        kv_lora_rank = hparams.get("kv_lora_rank", 512)
343        self.gguf_writer.add_vocab_size(hparams["vocab_size"])
344        if "q_lora_rank" in hparams and hparams["q_lora_rank"] is not None:
345            self.gguf_writer.add_q_lora_rank(hparams["q_lora_rank"])
346
347        # note: deepseek2 using MLA converts into MQA with larger heads, then decompresses to MHA
348        if not is_ocr:
349            self.gguf_writer.add_kv_lora_rank(kv_lora_rank)
350            self.gguf_writer.add_key_length(kv_lora_rank + hparams["qk_rope_head_dim"])
351            self.gguf_writer.add_value_length(kv_lora_rank)
352            self.gguf_writer.add_key_length_mla(hparams["qk_nope_head_dim"] + hparams["qk_rope_head_dim"])
353            self.gguf_writer.add_value_length_mla(hparams["v_head_dim"])
354
355        # MoE parameters (required by C++ code for DEEPSEEK2 arch)
356        # For non-MoE models like Youtu, use intermediate_size as expert_feed_forward_length
357        moe_intermediate_size = self.find_hparam(["moe_intermediate_size", "intermediate_size"], optional=False)
358        self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)
359
360        if (n_routed_experts := hparams.get("n_routed_experts")) is not None:
361            self.gguf_writer.add_expert_count(n_routed_experts)
362
363        # expert_shared_count is required by C++ code, default to 0 for non-MoE models
364        n_shared_experts = hparams.get("n_shared_experts", 0)
365        self.gguf_writer.add_expert_shared_count(n_shared_experts)
366
367        # When not set, C++ code will use scale_w = false to skip the no-op scaling
368        if (routed_scaling_factor := hparams.get("routed_scaling_factor")) is not None:
369            self.gguf_writer.add_expert_weights_scale(routed_scaling_factor)
370
371        if (norm_topk_prob := hparams.get("norm_topk_prob")) is not None and norm_topk_prob:
372            self.gguf_writer.add_expert_weights_norm(norm_topk_prob)
373
374        self.gguf_writer.add_rope_dimension_count(hparams["qk_rope_head_dim"])
375
376        # Unlimited-OCR sliding window; written for metadata, the decoder ignores it (full MHA)
377        if is_ocr:
378            sliding_window = hparams.get("sliding_window_size") or hparams.get("sliding_window")
379            if sliding_window:
380                self.gguf_writer.add_sliding_window(sliding_window)
381
382        if (rope_mscale_all := self.rope_parameters.get("mscale_all_dim")) is not None:
383            # [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX]
384            # note: for legacy reasons, this is not consistent with the other usages of self.gguf_writer.add_rope_scaling_yarn_log_mul
385            # ref https://github.com/ggml-org/llama.cpp/pull/17945
386            self.gguf_writer.add_rope_scaling_yarn_log_mul(0.1 * rope_mscale_all)
387
388    _experts: list[dict[str, Tensor]] | None = None
389
390    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
391        # skip lm_head.weight if tie_word_embeddings is True
392        if self.hparams.get("tie_word_embeddings", False):
393            if name == "lm_head.weight" or name == "model.lm_head.weight":
394                logger.info("Skipping tied output layer 'lm_head.weight' (will use token_embd.weight)")
395                return
396
397        # skip Multi-Token Prediction (MTP) layers
398        if self.skip_mtp:
399            block_count = self.hparams["num_hidden_layers"]
400            match = re.match(r"model.layers.(\d+)", name)
401            if match and int(match.group(1)) >= block_count:
402                return
403
404        # process the experts separately
405        if self.merge_expert and name.find("mlp.experts") != -1:
406            n_experts = self.hparams["n_routed_experts"]
407            assert bid is not None
408
409            if self._experts is None:
410                self._experts = [{} for _ in range(self.block_count)]
411
412            self._experts[bid][name] = data_torch
413
414            if len(self._experts[bid]) >= n_experts * 3:
415                # merge the experts into a single 3d tensor
416                for w_name in ["down_proj", "gate_proj", "up_proj"]:
417                    datas: list[Tensor] = []
418
419                    for xid in range(n_experts):
420                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
421                        datas.append(self._experts[bid][ename])
422                        del self._experts[bid][ename]
423
424                    data_torch = torch.stack(datas, dim=0)
425
426                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
427
428                    yield from super().modify_tensors(data_torch, merged_name, bid)
429                return
430            else:
431                return
432
433        # note: MLA with the absorption optimization, needs these two split and k_b_proj transposed
434        if name.endswith("kv_b_proj.weight"):
435            name_kb = name.replace("kv_b_proj", "k_b_proj")
436            name_vb = name.replace("kv_b_proj", "v_b_proj")
437
438            n_head_kv = self.hparams["num_key_value_heads"]
439            v_head_dim = self.hparams["v_head_dim"]
440            qk_nope_head_dim = self.hparams["qk_nope_head_dim"]
441
442            assert data_torch.shape[0] == n_head_kv * (v_head_dim + qk_nope_head_dim)
443
444            kv_b = data_torch.view(n_head_kv, v_head_dim + qk_nope_head_dim, data_torch.shape[-1])
445            k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1)
446            k_b = k_b.transpose(1, 2)
447
448            yield from super().modify_tensors(k_b, name_kb, bid)
449            yield from super().modify_tensors(v_b, name_vb, bid)
450            return
451
452        yield from super().modify_tensors(data_torch, name, bid)
453
454    def prepare_tensors(self):
455        super().prepare_tensors()
456
457        if self._experts is not None:
458            # flatten `list[dict[str, Tensor]]` into `list[str]`
459            experts = [k for d in self._experts for k in d.keys()]
460            if len(experts) > 0:
461                raise ValueError(f"Unprocessed experts: {experts}")
462
463
464@ModelBase.register("DeepseekV32ForCausalLM")
465@ModelBase.example("deepseek-ai/DeepSeek-V3.2-Exp")
466class DeepseekV32Model(DeepseekV2Model):
467    model_arch = gguf.MODEL_ARCH.DEEPSEEK32
468    skip_mtp = False
469    supports_mtp_export = True
470    _n_main_layers: int | None = None
471
472    def __init__(self, *args, **kwargs):
473        super().__init__(*args, **kwargs)
474        self.block_count = self.hparams["num_hidden_layers"]
475        if not self.no_mtp:
476            self.block_count += self.hparams.get("num_nextn_predict_layers", 0)
477        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
478
479    def index_tensors(self, remote_hf_model_id: str | None = None):
480        type(self)._n_main_layers = self.hparams["num_hidden_layers"]
481        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
482
483    @classmethod
484    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
485        if (titem := super().filter_tensors(item)) is None:
486            return None
487        name, gen = titem
488
489        # DeepSeek V3.2 appends the NextN/MTP block past num_hidden_layers
490        # (model.layers.61 -> blk.61 in the 62-block file).
491        assert cls._n_main_layers is not None
492        is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._n_main_layers
493
494        # --no-mtp: drop the appended NextN block entirely.
495        if is_mtp and cls.no_mtp:
496            return None
497        # --mtp: keep ONLY NextN-block tensors plus the shared embeddings/
498        # norm/lm_head (so the resulting GGUF carries just the draft head).
499        if cls.mtp_only and not is_mtp and name not in (
500            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
501        ):
502            return None
503
504        return name, gen
505
506    def set_vocab(self):
507        from transformers import AutoTokenizer
508        tokenizer = AutoTokenizer.from_pretrained(self.dir_model)
509        assert getattr(tokenizer, "add_bos_token", False), "Change value of add_bos_token to true in tokenizer_config.json file."
510        self._set_vocab_gpt2()
511
512    def set_gguf_parameters(self):
513        super().set_gguf_parameters()
514
515        # NextN/MTP prediction layers
516        if not self.no_mtp and (num_nextn_predict_layers := self.hparams.get("num_nextn_predict_layers")) is not None:
517            self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
518
519        # DSA indexer parameters
520        self.gguf_writer.add_indexer_head_count(self.hparams["index_n_heads"])
521        self.gguf_writer.add_indexer_key_length(self.hparams["index_head_dim"])
522        self.gguf_writer.add_indexer_top_k(self.hparams["index_topk"])
523
524
525@ModelBase.register("DeepseekV4ForCausalLM")
526@ModelBase.example("deepseek-ai/DeepSeek-V4-Flash-Base")
527class DeepseekV4Model(TextModel):
528    model_arch = gguf.MODEL_ARCH.DEEPSEEK4
529    supports_mtp_export = True
530    _skipped_mtp_tensors = 0
531    _dsv4_main_layers: int | None = None
532    _dsv4_nextn_layers: int = 0
533
534    def __init__(self, *args, **kwargs):
535        type(self)._skipped_mtp_tensors = 0
536        super().__init__(*args, **kwargs)
537
538        with open(self.dir_model / "config.json", "r", encoding="utf-8") as f:
539            raw_hparams = json.load(f)
540        for key, value in raw_hparams.items():
541            self.hparams.setdefault(key, value)
542
543        # workaround for special rope_parameters (main/compress) in transformers 5.x
544        if self.rope_parameters.get("full_attention", self.rope_parameters).get("rope_type") is None:
545            if (rope_scaling := raw_hparams.get("rope_scaling")) is not None:
546                if "rope_type" not in rope_scaling and (rope_type := rope_scaling.get("type")) is not None:
547                    rope_scaling["rope_type"] = rope_type
548                self.rope_parameters.update(**rope_scaling)
549
550        self.block_count = self.hparams["num_hidden_layers"]
551        if self.mtp_only:
552            self.block_count += self.hparams.get("num_nextn_predict_layers", 0)
553        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
554
555        self._dsv4_fp8_dequantized: set[str] = set()
556        self._dsv4_bf16_tensors: set[str] = set()
557        self._dsv4_f32_tensors: set[str] = set()
558        self._dsv4_mxfp4_generated = False
559        self._collect_source_dtypes()
560
561        if type(self)._skipped_mtp_tensors:
562            logger.info("Skipping %d DeepSeek-V4 MTP tensor(s) for conversion v0", type(self)._skipped_mtp_tensors)
563
564        # add a default chat template; if the model has a built-in template, it will be overridden later
565        model_id_hint = self.remote_hf_model_id or self.dir_model.name
566        is_0731 = "0731" in model_id_hint
567        template_name = "deepseek-ai-DeepSeek-V4-Flash-0731.jinja" if is_0731 else "deepseek-ai-DeepSeek-V4.jinja"
568        template_path = Path(__file__).parent.parent / "models" / "templates" / template_name
569        if template_path.is_file():
570            with open(template_path, "r", encoding="utf-8") as f:
571                self.gguf_writer.add_chat_template(f.read())
572
573    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
574        type(self)._dsv4_main_layers = self.hparams["num_hidden_layers"]
575        type(self)._dsv4_nextn_layers = self.hparams.get("num_nextn_predict_layers", 0)
576        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
577
578    @classmethod
579    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
580        name, gen = item
581        if name.startswith(("aligner.", "image_")):
582            return None
583        if name.startswith("mtp."):
584            if not cls.mtp_only:
585                cls._skipped_mtp_tensors += 1
586                return None
587
588            assert cls._dsv4_main_layers is not None
589            parts = name.split(".", 2)
590            if len(parts) < 3 or not parts[1].isdecimal():
591                raise ValueError(f"Unexpected DeepSeek-V4 MTP tensor {name!r}")
592
593            mtp_idx = int(parts[1])
594            if mtp_idx >= cls._dsv4_nextn_layers:
595                raise ValueError(f"Unexpected DeepSeek-V4 MTP layer {mtp_idx}")
596
597            bid = cls._dsv4_main_layers + mtp_idx
598            suffix = parts[2]
599            root_hc_head = {
600                "hc_head_fn",
601                "hc_head_base",
602                "hc_head_scale",
603            }
604            if suffix in root_hc_head:
605                name = suffix
606            elif suffix in (
607                "e_proj.weight", "e_proj.scale",
608                "h_proj.weight", "h_proj.scale",
609            ):
610                name = f"layers.{bid}.nextn.{suffix}"
611            elif suffix == "enorm.weight":
612                name = f"layers.{bid}.nextn.enorm.weight"
613            elif suffix == "hnorm.weight":
614                name = f"layers.{bid}.nextn.hnorm.weight"
615            elif suffix == "norm.weight":
616                name = f"layers.{bid}.nextn.shared_head_norm.weight"
617            else:
618                name = f"layers.{bid}.{suffix}"
619            return name, gen
620
621        if cls.mtp_only:
622            keep = name in (
623                "embed.weight",
624                "norm.weight",
625                "head.weight",
626                "head.scale",
627            )
628            if not keep:
629                return None
630
631        return super().filter_tensors((name, gen))
632
633    @staticmethod
634    def _float8_dtypes() -> tuple[torch.dtype, ...]:
635        return tuple(
636            dtype for dtype in (
637                getattr(torch, "float8_e4m3fn", None),
638                getattr(torch, "float8_e5m2", None),
639            ) if dtype is not None
640        )
641
642    @staticmethod
643    def _e8m0_to_float(scale: Tensor) -> Tensor:
644        torch_float8_e8m0 = getattr(torch, "float8_e8m0fnu", None)
645        if torch_float8_e8m0 is not None and scale.dtype == torch_float8_e8m0:
646            return scale.float()
647
648        bits = scale.view(torch.uint8).float()
649        return torch.exp2(bits - 127.0)
650
651    def _collect_source_dtypes(self) -> None:
652        for name, gen in self.model_tensors.items():
653            dtype = gen().dtype
654            if dtype == torch.bfloat16:
655                self._dsv4_bf16_tensors.add(name)
656            elif dtype == torch.float32:
657                self._dsv4_f32_tensors.add(name)
658
659    def set_gguf_parameters(self):
660        super().set_gguf_parameters()
661        hparams = self.hparams
662
663        self.gguf_writer.add_rope_dimension_count(hparams["qk_rope_head_dim"])
664        self.gguf_writer.add_q_lora_rank(hparams["q_lora_rank"])
665        self.gguf_writer.add_sliding_window(hparams["sliding_window"])
666
667        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])
668        self.gguf_writer.add_expert_shared_count(hparams["n_shared_experts"])
669        self.gguf_writer.add_expert_weights_scale(hparams["routed_scaling_factor"])
670        self.gguf_writer.add_expert_weights_norm(hparams["norm_topk_prob"])
671        self.gguf_writer.add_swiglu_clamp_exp([hparams["swiglu_limit"]] * self.block_count)
672        self.gguf_writer.add_swiglu_clamp_shexp([hparams["swiglu_limit"]] * self.block_count)
673
674        self.gguf_writer.add_indexer_head_count(hparams["index_n_heads"])
675        self.gguf_writer.add_indexer_key_length(hparams["index_head_dim"])
676        self.gguf_writer.add_indexer_top_k(hparams["index_topk"])
677
678        self.gguf_writer.add_attention_output_group_count(hparams["o_groups"])
679        self.gguf_writer.add_attention_output_lora_rank(hparams["o_lora_rank"])
680        self.gguf_writer.add_attention_compress_ratios(hparams["compress_ratios"])
681        self.gguf_writer.add_attention_compress_rope_freq_base(hparams["compress_rope_theta"])
682        self.gguf_writer.add_hyper_connection_count(hparams["hc_mult"])
683        self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hparams["hc_sinkhorn_iters"])
684        self.gguf_writer.add_hyper_connection_epsilon(hparams["hc_eps"])
685        self.gguf_writer.add_hash_layer_count(hparams["num_hash_layers"])
686        if self.model_arch == gguf.MODEL_ARCH.DEEPSEEK4:
687            self.gguf_writer.add_embedding_length_out(hparams["hidden_size"] * hparams["hc_mult"])
688        if self.mtp_only and (num_nextn_predict_layers := hparams.get("num_nextn_predict_layers", 0)) > 0:
689            self.gguf_writer.add_nextn_predict_layers(num_nextn_predict_layers)
690
691    def dequant_model(self):
692        fp8_dtypes = self._float8_dtypes()
693        tensors_to_remove: list[str] = []
694
695        def dequant_fp8_weight(weight: Tensor, scale: Tensor) -> Tensor:
696            out_features, in_features = weight.shape
697            scale_f = self._e8m0_to_float(scale)
698            scale_f = scale_f.repeat_interleave(128, 0)[:out_features]
699            scale_f = scale_f.repeat_interleave(128, 1)[:, :in_features]
700            return weight.float() * scale_f
701
702        for name in list(self.model_tensors.keys()):
703            if not name.endswith(".scale"):
704                continue
705            weight_name = name.removesuffix(".scale") + ".weight"
706            if weight_name not in self.model_tensors:
707                continue
708
709            weight = self.model_tensors[weight_name]
710            scale = self.model_tensors[name]
711            if weight().dtype not in fp8_dtypes:
712                continue
713
714            self.model_tensors[weight_name] = lambda w=weight, s=scale: dequant_fp8_weight(w(), s())
715            self._dsv4_fp8_dequantized.add(weight_name)
716            tensors_to_remove.append(name)
717
718        for name in tensors_to_remove:
719            del self.model_tensors[name]
720
721    def _write_mxfp4_expert_tensor(self, bid: int, proj: str, tensor_key: gguf.MODEL_TENSOR) -> list[str]:
722        n_experts = self.hparams["n_routed_experts"]
723        data: np.ndarray | None = None
724        consumed: list[str] = []
725
726        for eid in range(n_experts):
727            weight_name = f"layers.{bid}.ffn.experts.{eid}.{proj}.weight"
728            scale_name = f"layers.{bid}.ffn.experts.{eid}.{proj}.scale"
729            if weight_name not in self.model_tensors or scale_name not in self.model_tensors:
730                raise KeyError(f"Missing routed expert tensors for {weight_name}")
731
732            weight = LazyTorchTensor.to_eager(self.model_tensors[weight_name]())
733            scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]())
734            packed = self.repack_mxfp4_blocks(weight, scale)
735            if data is None:
736                data = np.empty((n_experts, *packed.shape), dtype=packed.dtype)
737            data[eid] = packed
738            consumed.extend((weight_name, scale_name))
739
740        assert data is not None
741        new_name = self.format_tensor_name(tensor_key, bid)
742        shape = gguf.quant_shape_from_byte_shape(data.shape, gguf.GGMLQuantizationType.MXFP4)
743        logger.info(f"{new_name}: repacked routed experts to MXFP4, shape = {{{', '.join(str(n) for n in reversed(shape))}}}")
744        self.gguf_writer.add_tensor(new_name, data, raw_dtype=gguf.GGMLQuantizationType.MXFP4)
745
746        return consumed
747
748    def _write_hash_routing_tensors(self) -> list[str]:
749        consumed: list[str] = []
750
751        for bid in range(self.hparams["num_hash_layers"]):
752            name = f"layers.{bid}.ffn.gate.tid2eid"
753            if name not in self.model_tensors:
754                raise KeyError(f"Missing hash routing tensor {name}")
755
756            data_torch = LazyTorchTensor.to_eager(self.model_tensors[name]())
757            data = data_torch.to(torch.int32).cpu().numpy()
758            new_name = self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_TID2EID, bid, ".weight")
759            logger.info(f"{new_name}: converted hash routing table to I32, shape = {{{', '.join(str(n) for n in reversed(data.shape))}}}")
760            self.gguf_writer.add_tensor(new_name, data)
761            consumed.append(name)
762
763        return consumed
764
765    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
766        if self._dsv4_mxfp4_generated:
767            return ()
768
769        consumed: list[str] = []
770        main_layers = self.hparams["num_hidden_layers"]
771        if not self.mtp_only:
772            consumed.extend(self._write_hash_routing_tensors())
773        elif self.hparams["num_hash_layers"] > 0:
774            for bid in range(self.hparams["num_hash_layers"]):
775                name = f"layers.{bid}.ffn.gate.tid2eid"
776                if name in self.model_tensors:
777                    consumed.extend(self._write_hash_routing_tensors())
778                    break
779
780        for bid in range(self.block_count):
781            if self.mtp_only and bid < main_layers:
782                continue
783            consumed.extend(self._write_mxfp4_expert_tensor(bid, "w1", gguf.MODEL_TENSOR.FFN_GATE_EXP))
784            consumed.extend(self._write_mxfp4_expert_tensor(bid, "w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP))
785            consumed.extend(self._write_mxfp4_expert_tensor(bid, "w3", gguf.MODEL_TENSOR.FFN_UP_EXP))
786
787        for bid in range(main_layers, self.block_count):
788            e_name = f"layers.{bid}.nextn.e_proj.weight"
789            h_name = f"layers.{bid}.nextn.h_proj.weight"
790            if e_name not in self.model_tensors and h_name not in self.model_tensors:
791                continue
792            if e_name not in self.model_tensors or h_name not in self.model_tensors:
793                raise KeyError(f"Missing DeepSeek-V4 MTP e/h projection pair for block {bid}")
794
795            e_proj = LazyTorchTensor.to_eager(self.model_tensors[e_name]())
796            h_proj = LazyTorchTensor.to_eager(self.model_tensors[h_name]())
797            yield (f"layers.{bid}.nextn.eh_proj.weight", torch.cat((e_proj, h_proj), dim=1).contiguous())
798            consumed.extend((e_name, h_name))
799
800        for name in consumed:
801            del self.model_tensors[name]
802
803        self._dsv4_mxfp4_generated = True
804        return ()
805
806    def _format_dsv4_tensor_name(self, key: gguf.MODEL_TENSOR, bid: int | None, suffix: str = ".weight") -> str:
807        return self.format_tensor_name(key, bid, suffix)
808
809    def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]:
810        root_map: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
811            "embed.weight": (gguf.MODEL_TENSOR.TOKEN_EMBD, ".weight"),
812            "norm.weight": (gguf.MODEL_TENSOR.OUTPUT_NORM, ".weight"),
813            "head.weight": (gguf.MODEL_TENSOR.OUTPUT, ".weight"),
814            "hc_head_fn": (gguf.MODEL_TENSOR.HC_HEAD_FN, ".weight"),
815            "hc_head_base": (gguf.MODEL_TENSOR.HC_HEAD_BASE, ".weight"),
816            "hc_head_scale": (gguf.MODEL_TENSOR.HC_HEAD_SCALE, ".weight"),
817        }
818        if name in root_map:
819            return root_map[name]
820
821        match = re.match(r"layers\.(\d+)\.(.+)$", name)
822        if match is None:
823            raise ValueError(f"Unsupported DeepSeek-V4 tensor {name!r}")
824
825        layer = int(match.group(1))
826        if bid != layer:
827            raise ValueError(f"Tensor {name!r} parsed bid {bid} but layer name has {layer}")
828
829        layer_map: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
830            "hc_attn_fn": (gguf.MODEL_TENSOR.HC_ATTN_FN, ".weight"),
831            "hc_attn_base": (gguf.MODEL_TENSOR.HC_ATTN_BASE, ".weight"),
832            "hc_attn_scale": (gguf.MODEL_TENSOR.HC_ATTN_SCALE, ".weight"),
833            "hc_ffn_fn": (gguf.MODEL_TENSOR.HC_FFN_FN, ".weight"),
834            "hc_ffn_base": (gguf.MODEL_TENSOR.HC_FFN_BASE, ".weight"),
835            "hc_ffn_scale": (gguf.MODEL_TENSOR.HC_FFN_SCALE, ".weight"),
836            "attn.attn_sink": (gguf.MODEL_TENSOR.ATTN_SINKS, ".weight"),
837            "attn.wq_a.weight": (gguf.MODEL_TENSOR.ATTN_Q_A, ".weight"),
838            "attn.wq_b.weight": (gguf.MODEL_TENSOR.ATTN_Q_B, ".weight"),
839            "attn.q_norm.weight": (gguf.MODEL_TENSOR.ATTN_Q_A_NORM, ".weight"),
840            "attn.wkv.weight": (gguf.MODEL_TENSOR.ATTN_KV, ".weight"),
841            "attn.kv_norm.weight": (gguf.MODEL_TENSOR.ATTN_KV_NORM, ".weight"),
842            "attn.wo_a.weight": (gguf.MODEL_TENSOR.ATTN_OUT_A, ".weight"),
843            "attn.wo_b.weight": (gguf.MODEL_TENSOR.ATTN_OUT_B, ".weight"),
844            "attn.compressor.ape": (gguf.MODEL_TENSOR.ATTN_COMPRESSOR_APE, ".weight"),
845            "attn.compressor.wkv.weight": (gguf.MODEL_TENSOR.ATTN_COMPRESSOR_WKV, ".weight"),
846            "attn.compressor.wgate.weight": (gguf.MODEL_TENSOR.ATTN_COMPRESSOR_WGATE, ".weight"),
847            "attn.compressor.norm.weight": (gguf.MODEL_TENSOR.ATTN_COMPRESSOR_NORM, ".weight"),
848            "attn.indexer.wq_b.weight": (gguf.MODEL_TENSOR.INDEXER_ATTN_Q_B, ".weight"),
849            "attn.indexer.weights_proj.weight": (gguf.MODEL_TENSOR.INDEXER_PROJ, ".weight"),
850            "attn.indexer.compressor.ape": (gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_APE, ".weight"),
851            "attn.indexer.compressor.wkv.weight": (gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_WKV, ".weight"),
852            "attn.indexer.compressor.wgate.weight": (gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE, ".weight"),
853            "attn.indexer.compressor.norm.weight": (gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_NORM, ".weight"),
854            "attn_norm.weight": (gguf.MODEL_TENSOR.ATTN_NORM, ".weight"),
855            "ffn_norm.weight": (gguf.MODEL_TENSOR.FFN_NORM, ".weight"),
856            "ffn.gate.weight": (gguf.MODEL_TENSOR.FFN_GATE_INP, ".weight"),
857            "ffn.gate.bias": (gguf.MODEL_TENSOR.FFN_EXP_PROBS_B, ".bias"),
858            "ffn.gate.bias_vl": (gguf.MODEL_TENSOR.FFN_EXP_PROBS_B_VL, ".bias"),
859            "ffn.gate.tid2eid": (gguf.MODEL_TENSOR.FFN_GATE_TID2EID, ".weight"),
860            "ffn.shared_experts.w1.weight": (gguf.MODEL_TENSOR.FFN_GATE_SHEXP, ".weight"),
861            "ffn.shared_experts.w2.weight": (gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, ".weight"),
862            "ffn.shared_experts.w3.weight": (gguf.MODEL_TENSOR.FFN_UP_SHEXP, ".weight"),
863            "nextn.eh_proj.weight": (gguf.MODEL_TENSOR.NEXTN_EH_PROJ, ".weight"),
864            "nextn.enorm.weight": (gguf.MODEL_TENSOR.NEXTN_ENORM, ".weight"),
865            "nextn.hnorm.weight": (gguf.MODEL_TENSOR.NEXTN_HNORM, ".weight"),
866            "nextn.shared_head_norm.weight": (gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ".weight"),
867            "nextn.embed_tokens.weight": (gguf.MODEL_TENSOR.NEXTN_EMBED_TOKENS, ".weight"),
868            "nextn.shared_head_head.weight": (gguf.MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, ".weight"),
869        }
870
871        tensor_name = match.group(2)
872        if tensor_name in layer_map:
873            return layer_map[tensor_name]
874
875        if re.match(r"ffn\.experts\.\d+\.w[123]\.(weight|scale)$", tensor_name):
876            return gguf.MODEL_TENSOR.FFN_GATE_EXP, ".weight"
877
878        raise ValueError(f"Unsupported DeepSeek-V4 tensor {name!r}")
879
880    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
881        if re.match(r"layers\.\d+\.ffn\.experts\.\d+\.w[123]\.(weight|scale)$", name):
882            return []
883
884        # hash layers route text tokens via tid2eid and image tokens via bias_vl; gate.bias is unused
885        if name.endswith(".ffn.gate.bias") and bid is not None and bid < self.hparams["num_hash_layers"]:
886            return []
887
888        tensor_key, suffix = self._map_dsv4_tensor_name(name, bid)
889        if tensor_key == gguf.MODEL_TENSOR.FFN_GATE_TID2EID:
890            return []
891
892        return [(self._format_dsv4_tensor_name(tensor_key, bid, suffix), data_torch)]
893
894    def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
895        del bid  # unused
896
897        if name in self._dsv4_fp8_dequantized and n_dims >= 2:
898            return gguf.GGMLQuantizationType.Q8_0
899        if new_name.endswith(".nextn.eh_proj.weight"):
900            return gguf.GGMLQuantizationType.Q8_0
901        if name in self._dsv4_f32_tensors:
902            return gguf.GGMLQuantizationType.F32
903        if name in self._dsv4_bf16_tensors and n_dims >= 2:
904            return gguf.GGMLQuantizationType.BF16
905
906        return False
907
908    def prepare_metadata(self, vocab_only: bool):
909        from_dir = self.fname_out.is_dir()
910        super().prepare_metadata(vocab_only=vocab_only)
911
912        if not self.mtp_only or not from_dir:
913            return
914
915        output_type: str = self.ftype.name.partition("_")[2]
916        fname_default: str = gguf.naming_convention(
917            self.metadata.name, self.metadata.basename, self.metadata.finetune,
918            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
919        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
920
921    def prepare_tensors(self):
922        super().prepare_tensors()
923        self._is_mxfp4 = True
924        self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE
925
926
927@ModelBase.register("DeepseekV4DSparkModel")
928@ModelBase.example("deepseek-ai/DeepSeek-V4-Flash-DSpark")
929class DeepseekV4DSparkModel(DeepseekV4Model):
930    model_arch = gguf.MODEL_ARCH.DFLASH
931
932    _DSPARK_ROOT_MAP: dict[str, tuple[gguf.MODEL_TENSOR, str]] = {
933        "main_proj.weight": (gguf.MODEL_TENSOR.FC, ".weight"),
934        "main_norm.weight": (gguf.MODEL_TENSOR.ENC_OUTPUT_NORM, ".weight"),
935        "markov_head.markov_w1.weight": (gguf.MODEL_TENSOR.DSPARK_MARKOV_W1, ".weight"),
936        "markov_head.markov_w2.weight": (gguf.MODEL_TENSOR.DSPARK_MARKOV_W2, ".weight"),
937        "confidence_head.proj.weight": (gguf.MODEL_TENSOR.DSPARK_CONF_PROJ, ".weight"),
938    }
939
940    def __init__(self, *args, **kwargs):
941        super().__init__(*args, **kwargs)
942
943        self.block_count = 1 + max(
944            int(match.group(1)) for name in self.model_tensors
945            if (match := re.match(r"layers\.(\d+)\.", name))
946        )
947        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
948
949        self.hparams["compress_ratios"] = [0] * self.block_count
950        self.hparams["num_hash_layers"] = 0
951
952    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
953        if remote_hf_model_id is None:
954            return super().index_tensors()
955
956        with open(self.dir_model / "model.safetensors.index.json", "r", encoding="utf-8") as f:
957            weight_map = json.load(f)["weight_map"]
958
959        part_names = sorted({
960            part_name for name, part_name in weight_map.items()
961            if name.startswith("mtp.")
962        })
963        tensors: dict[str, Callable[[], Tensor]] = {}
964
965        for part_name in part_names:
966            from huggingface_hub import hf_hub_download
967
968            logger.info("gguf: caching remote DSpark part '%s'", part_name)
969            part_path = Path(hf_hub_download(repo_id=remote_hf_model_id, filename=part_name))
970            with gguf.utility.SafetensorsLocal(part_path) as model_part:
971                for name in model_part:
972                    data = model_part[name]
973                    data_gen = lambda data=data: LazyTorchTensor.from_local_tensor(data)  # noqa: E731
974                    if titem := self.filter_tensors((name, data_gen)):
975                        tensor_name, tensor_gen = titem
976                        tensors[tensor_name] = tensor_gen
977
978        return tensors
979
980    @classmethod
981    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
982        name, gen = item
983        if not name.startswith("mtp."):
984            return None
985        return super().filter_tensors((cls._rekey_mtp_tensor_name(name), gen))
986
987    @staticmethod
988    def _rekey_mtp_tensor_name(name: str) -> str:
989        match = re.match(r"mtp\.(\d+)\.(.+)$", name)
990        if match is None:
991            raise ValueError(f"Unexpected DSpark tensor {name!r}")
992
993        stage, rest = match.group(1), match.group(2)
994        root_names = (
995            "main_proj.scale",
996            "norm.weight",
997            "hc_head_fn",
998            "hc_head_base",
999            "hc_head_scale",
1000        )
1001        if rest in DeepseekV4DSparkModel._DSPARK_ROOT_MAP or rest in root_names:
1002            return rest
1003        return f"layers.{stage}.{rest}"
1004
1005    def _map_dsv4_tensor_name(self, name: str, bid: int | None) -> tuple[gguf.MODEL_TENSOR, str]:
1006        if name in self._DSPARK_ROOT_MAP:
1007            return self._DSPARK_ROOT_MAP[name]
1008        return super()._map_dsv4_tensor_name(name, bid)
1009
1010    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
1011        # the DFlash draft uses the plain exp-probs bias (ffn.gate.bias -> FFN_EXP_PROBS_B);
1012        # the mtmd-only hash routing tensors (bias_vl, tid2eid) are not part of the DFLASH arch
1013        if name.endswith(".ffn.gate.bias_vl"):
1014            return
1015        yield from super().modify_tensors(data_torch, name, bid)
1016
1017    def set_vocab(self):
1018        if self.target_model_dir is None:
1019            raise ValueError("DeepSeek-V4 DSpark requires --target-model-dir with the target tokenizer")
1020
1021        original_dir = self.dir_model
1022        try:
1023            self.dir_model = self.target_model_dir
1024            super().set_vocab()
1025        finally:
1026            self.dir_model = original_dir
1027
1028        self.gguf_writer.add_mask_token_id(self.hparams["dspark_noise_token_id"])
1029
1030    def set_gguf_parameters(self):
1031        super().set_gguf_parameters()
1032
1033        self.gguf_writer.add_block_size(self.hparams["dspark_block_size"])
1034        self.gguf_writer.add_target_layers([layer + 1 for layer in self.hparams["dspark_target_layer_ids"]])
1035
1036
1037@ModelBase.register("DeepseekV4ForCausalLM")
1038@ModelBase.example("deepseek-ai/DeepSeek-V4-Flash-Vision-Exp")
1039class DeepseekV4FlashVisionModel(MmprojModel):
1040    def __init__(self, *args, **kwargs):
1041        super().__init__(*args, **kwargs)
1042        assert self.hparams_vision is not None
1043        # no preprocessor_config.json in the repo; normalization is (x/255 - 0.5) / 0.5
1044        # ref: inference/image_processor.py (load_image)
1045        self.preprocessor_config = {
1046            "image_mean": [0.5, 0.5, 0.5],
1047            "image_std":  [0.5, 0.5, 0.5],
1048            **self.preprocessor_config,
1049        }
1050
1051    def get_vision_config(self) -> dict[str, Any] | None:
1052        cfg = self.global_config
1053        if cfg.get("vision_n_layers", 0) == 0:
1054            raise ValueError("DeepseekV4FlashVisionModel requires vision_n_layers > 0 in the model config")
1055        return {
1056            "num_hidden_layers":   cfg["vision_n_layers"],
1057            "hidden_size":         cfg["vision_dim"],
1058            "num_attention_heads": cfg["vision_n_heads"],
1059            "intermediate_size":   cfg["vision_inter_dim"],
1060            "patch_size":          cfg["vision_patch_size"],
1061            # dynamic resolution; only used for compat / warmup
1062            "image_size":          cfg["vision_patch_size"] * cfg["vision_downsample_ratio"] * 16,
1063            "rope_theta":          cfg.get("vision_rope_theta", 10000.0),
1064            "downsample_ratio":    cfg["vision_downsample_ratio"],
1065            "min_pixels":          cfg["vision_min_pixels"],
1066        }
1067
1068    def set_gguf_parameters(self):
1069        super().set_gguf_parameters()
1070        assert self.hparams_vision is not None
1071        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.DEEPSEEK4V)
1072        # vision RMSNorm eps is the pytorch default, NOT the LLM's rms_norm_eps (1e-20)
1073        # ref: inference/vision.py (RMSNorm)
1074        self.gguf_writer.add_vision_attention_layernorm_eps(1e-6)
1075        self.gguf_writer.add_vision_use_silu(True) # SwiGLU MLP
1076        self.gguf_writer.add_vision_projector_scale_factor(self.hparams_vision["downsample_ratio"])
1077        self.gguf_writer.add_vision_min_pixels(self.hparams_vision["min_pixels"])
1078        # hardcoded on the C++ side (see PROJECTOR_TYPE_DEEPSEEK4V in clip.cpp)
1079        # if future models use different values, add GGUF keys for those
1080        assert self.global_config["vision_max_n_token"] == 384
1081        assert self.global_config["vision_max_wh_ratio"] == 8
1082
1083    @classmethod
1084    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
1085        name, _ = item
1086        if not (name.startswith(("vision.", "aligner.", "image_"))):
1087            return None
1088        return super().filter_tensors(item)
1089
1090    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
1091        assert self.hparams_vision is not None
1092        if name == "vision.patch_embed.proj.weight":
1093            # nn.Linear over flattened (3, p, p) patches == conv2d weight
1094            p = self.hparams_vision["patch_size"]
1095            data_torch = data_torch.reshape(data_torch.shape[0], 3, p, p)
1096
1097        if ".mlp.w1." in name:
1098            # fused SwiGLU gate+up
1099            gate, up = data_torch.chunk(2, dim=0)
1100            yield from super().modify_tensors(gate, name.replace("w1", "w1_gate"), bid)
1101            yield from super().modify_tensors(up, name.replace("w1", "w1_up"), bid)
1102            return
1103
1104        yield from super().modify_tensors(data_torch, name, bid)