Felipe97/llama-cpp-compiled
01.1k
1from __future__ import annotations2 3import json4 5from typing import Any, Callable, Iterable, TYPE_CHECKING6 7import numpy as np8import torch9 10if TYPE_CHECKING:11 from torch import Tensor12 13from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger14 15 16@ModelBase.register("QWenLMHeadModel")17@ModelBase.example("Qwen/Qwen-7B")18class QwenModel(TextModel):19 model_arch = gguf.MODEL_ARCH.QWEN20 21 @staticmethod22 def token_bytes_to_string(b):23 from transformers.convert_slow_tokenizer import bytes_to_unicode24 byte_encoder = bytes_to_unicode()25 return ''.join([byte_encoder[ord(char)] for char in b.decode('latin-1')])26 27 @staticmethod28 def bpe(mergeable_ranks: dict[bytes, int], token: bytes, max_rank: int | None = None) -> list[bytes]:29 parts = [bytes([b]) for b in token]30 while True:31 min_idx = None32 min_rank = None33 for i, pair in enumerate(zip(parts[:-1], parts[1:])):34 rank = mergeable_ranks.get(pair[0] + pair[1])35 if rank is not None and (min_rank is None or rank < min_rank):36 min_idx = i37 min_rank = rank38 if min_rank is None or (max_rank is not None and min_rank >= max_rank):39 break40 assert min_idx is not None41 parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx + 1]] + parts[min_idx + 2:]42 return parts43 44 def set_vocab(self):45 self._set_vocab_qwen()46 47 48@ModelBase.register(49 "Qwen2Model",50 "Qwen2ForCausalLM",51 "Qwen2AudioForConditionalGeneration",52 "KORMoForCausalLM",53 "AudioFlamingo3ForConditionalGeneration",54 "DotsOCRForCausalLM",55)56@ModelBase.example("Qwen/Qwen2.5-7B-Instruct")57class Qwen2Model(TextModel):58 model_arch = gguf.MODEL_ARCH.QWEN259 60 def set_vocab(self):61 try:62 self._set_vocab_sentencepiece()63 except FileNotFoundError:64 self._set_vocab_gpt2()65 66 def set_gguf_parameters(self):67 super().set_gguf_parameters()68 self._try_set_pooling_type()69 70 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:71 if self.hf_arch == "Qwen2Model":72 name = f"model.{name}" # map to Qwen2ForCausalLM tensors73 yield from super().modify_tensors(data_torch, name, bid)74 75 76@ModelBase.register("Qwen2MoeForCausalLM")77@ModelBase.example("Qwen/Qwen1.5-MoE-A2.7B")78class Qwen2MoeModel(TextModel):79 model_arch = gguf.MODEL_ARCH.QWEN2MOE80 81 def set_gguf_parameters(self):82 super().set_gguf_parameters()83 if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:84 self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)85 logger.info(f"gguf: expert feed forward length = {moe_intermediate_size}")86 if (shared_expert_intermediate_size := self.hparams.get('shared_expert_intermediate_size')) is not None:87 self.gguf_writer.add_expert_shared_feed_forward_length(shared_expert_intermediate_size)88 logger.info(f"gguf: expert shared feed forward length = {shared_expert_intermediate_size}")89 90 _experts: list[dict[str, Tensor]] | None = None91 92 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:93 # handle aggregated expert tensors94 # GGUF stores dimensions reversed from PyTorch, so:95 # PyTorch (A,B,C) -> GGUF writes [C,B,A] -> GGML reads ne={C,B,A}96 # Input shapes from HF: (n_expert, n_ff_exp, n_embd) or (n_expert, n_embd, n_ff_exp)97 # Expected GGML ne: {n_embd, n_ff_exp, n_expert} for gate/up, {n_ff_exp, n_embd, n_expert} for down98 if name.endswith("mlp.experts.down_proj") or name.endswith("mlp.experts.down_proj.weight"):99 mapped = f"{name}.weight" if not name.endswith(".weight") else name100 # HF: [n_expert, n_embd, n_ff] -> GGML: {n_ff, n_embd, n_expert}101 yield from super().modify_tensors(data_torch, mapped, bid)102 return103 104 if name.endswith("mlp.experts.gate_up_proj") or name.endswith("mlp.experts.gate_up_proj.weight"):105 if data_torch.ndim < 3 or data_torch.shape[-2] % 2 != 0:106 raise ValueError(f"Unexpected gate_up_proj shape for {name}: {tuple(data_torch.shape)}")107 # HF: [n_expert, 2*n_ff, n_embd] -> split on dim=-2108 n_ff = data_torch.shape[-2] // 2109 gate = data_torch[..., :n_ff, :].contiguous()110 up = data_torch[..., n_ff:, :].contiguous()111 # gate/up: [n_expert, n_ff, n_embd] -> GGML: {n_embd, n_ff, n_expert}112 base_name = name.removesuffix(".weight").removesuffix(".gate_up_proj")113 mapped_gate = f"{base_name}.gate_proj.weight"114 mapped_up = f"{base_name}.up_proj.weight"115 yield from super().modify_tensors(gate, mapped_gate, bid)116 yield from super().modify_tensors(up, mapped_up, bid)117 return118 119 if name.find("experts") != -1:120 n_experts = self.find_hparam(["num_local_experts", "num_experts"])121 assert bid is not None122 123 if self._experts is None:124 self._experts = [{} for _ in range(self.block_count)]125 126 self._experts[bid][name] = data_torch127 128 if len(self._experts[bid]) >= n_experts * 3:129 # merge the experts into a single 3d tensor130 for w_name in ["down_proj", "gate_proj", "up_proj"]:131 datas: list[Tensor] = []132 133 for xid in range(n_experts):134 ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"135 datas.append(self._experts[bid][ename])136 del self._experts[bid][ename]137 138 data_torch = torch.stack(datas, dim=0)139 140 merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"141 142 yield from super().modify_tensors(data_torch, merged_name, bid)143 return144 else:145 return146 147 yield from super().modify_tensors(data_torch, name, bid)148 149 def prepare_tensors(self):150 super().prepare_tensors()151 152 if self._experts is not None:153 # flatten `list[dict[str, Tensor]]` into `list[str]`154 experts = [k for d in self._experts for k in d.keys()]155 if len(experts) > 0:156 raise ValueError(f"Unprocessed experts: {experts}")157 158 159@ModelBase.register("Qwen3ForCausalLM", "Qwen3Model")160@ModelBase.example("Qwen/Qwen3-8B")161class Qwen3Model(Qwen2Model):162 model_arch = gguf.MODEL_ARCH.QWEN3163 164 # extra logic for rerank models165 is_rerank: bool = False166 is_tied_embeddings: bool = False167 token_false_id: int | None = None168 token_true_id: int | None = None169 170 def __init__(self, *args, **kwargs):171 super().__init__(*args, **kwargs)172 173 # track for intern-s1-mini174 hparams = ModelBase.load_hparams(self.dir_model, is_mistral_format=False)175 self.origin_hf_arch = hparams.get('architectures', [None])[0]176 177 if self._is_qwen3_reranker():178 self._find_rerank_config()179 180 def _is_qwen3_reranker(self) -> bool:181 readme_path = self.dir_model / "README.md"182 readme_text = ""183 if readme_path.exists():184 with readme_path.open("r", encoding="utf-8") as f:185 readme_text = f.read()186 187 name_hints = [188 str(self.dir_model.name),189 str(self.hparams.get("_name_or_path", "")),190 str(self.hparams.get("model_type", "")),191 str(self.origin_hf_arch or ""),192 ]193 name_hints = [hint.lower() for hint in name_hints if hint]194 195 if "# qwen3-reranker" in readme_text.lower() or "# qwen3-vl-reranker" in readme_text.lower():196 return True197 198 if any("qwen3-reranker" in hint or "qwen3-vl-reranker" in hint for hint in name_hints):199 return True200 201 return "sequenceclassification" in (self.origin_hf_arch or "").lower()202 203 def set_vocab(self):204 # deal with intern-s1-mini205 if self.origin_hf_arch == 'InternS1ForConditionalGeneration':206 self._set_vocab_interns1()207 return208 209 super().set_vocab()210 211 def _find_rerank_config(self):212 from transformers import AutoTokenizer213 tokenizer = AutoTokenizer.from_pretrained(self.dir_model)214 215 self.is_rerank = True216 self.is_tied_embeddings = self.hparams.get("tie_word_embeddings", False)217 self.token_false_id = tokenizer.convert_tokens_to_ids("no") # ty: ignore[unresolved-attribute, invalid-assignment]218 self.token_true_id = tokenizer.convert_tokens_to_ids("yes") # ty: ignore[unresolved-attribute, invalid-assignment]219 self.sep_token_id = tokenizer.convert_tokens_to_ids("|") # ty: ignore[unresolved-attribute]220 221 assert self.token_false_id is not None and self.token_true_id is not None222 223 def set_gguf_parameters(self):224 super().set_gguf_parameters()225 if self.is_rerank:226 self.gguf_writer.add_pooling_type(gguf.PoolingType.RANK)227 self.gguf_writer.add_classifier_output_labels(["yes", "no"])228 self.gguf_writer.add_chat_template([{229 "name": "rerank",230 "template": "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n"231 "<|im_start|>user\n<Instruct>: Given a web search query, retrieve relevant passages that answer the query\n<Query>: {query}\n<Document>: {document}<|im_end|>\n"232 "<|im_start|>assistant\n<think>\n\n</think>\n\n"233 }])234 235 def _get_cls_out_tensor(self, data_torch: Tensor) -> Tensor:236 # extract "yes" and "no" tokens from the output lm_head tensor237 false_row = data_torch[self.token_false_id]238 true_row = data_torch[self.token_true_id]239 return torch.stack([true_row, false_row], dim=0)240 241 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:242 if self.is_rerank:243 is_tied_head = self.is_tied_embeddings and "embed_tokens" in name244 is_real_head = not self.is_tied_embeddings and "lm_head" in name245 if is_tied_head or is_real_head:246 cls_out_head = (247 gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.CLS_OUT] + ".weight",248 self._get_cls_out_tensor(data_torch),249 )250 yield cls_out_head251 if is_tied_head:252 yield from super().modify_tensors(data_torch, name, bid)253 return254 255 yield from super().modify_tensors(data_torch, name, bid)256 257 258@ModelBase.register("Qwen3MoeForCausalLM")259@ModelBase.example("Qwen/Qwen3-30B-A3B")260class Qwen3MoeModel(Qwen2MoeModel):261 model_arch = gguf.MODEL_ARCH.QWEN3MOE262 263 def __init__(self, *args, **kwargs):264 super().__init__(*args, **kwargs)265 hparams = ModelBase.load_hparams(self.dir_model, False)266 self.origin_hf_arch = hparams.get('architectures', [None])[0]267 268 def set_vocab(self):269 # deal with intern-s1270 if self.origin_hf_arch == 'InternS1ForConditionalGeneration':271 self._set_vocab_interns1()272 return273 274 super().set_vocab()275 276 277class _QwenMtpMixin:278 """Shared MTP wiring for Qwen3-Next and Qwen3.5/3.6 text variants. The HF279 config carries the MTP block under `mtp_num_hidden_layers` (computed from280 the checkpoint when absent, e.g. Qwen3-Next) and the tensors under281 `mtp.*`; we extend block_count, emit the nextn metadata key, and remap282 `mtp.*` to the standard layer-indexed nextn naming so the existing283 tensor_map handles them."""284 285 supports_mtp_export = True286 hparams: dict[str, Any]287 model_arch: gguf.MODEL_ARCH288 gguf_writer: gguf.GGUFWriter289 block_count: int290 tensor_map: gguf.TensorNameMap291 no_mtp: bool292 mtp_only: bool293 _original_block_count: int | None = None294 opt_num_mtp_layers: int = 0295 296 def __init__(self, *args, **kwargs):297 super().__init__(*args, **kwargs)298 self.block_count = self.hparams["num_hidden_layers"]299 if not self.no_mtp:300 n_mtp = self.hparams.get("mtp_num_hidden_layers", 0)301 # Qwen-3-Next doesn't include `mtp_num_hidden_layers` in config.302 if n_mtp == 0:303 assert self.opt_num_mtp_layers != 0304 n_mtp = self.opt_num_mtp_layers305 self.block_count += n_mtp306 self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)307 308 def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:309 hparams = {**self.hparams, **self.hparams.get("text_config", {})}310 key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None)311 type(self)._original_block_count = hparams.get(key)312 type(self).opt_num_mtp_layers = 0313 return super().index_tensors(remote_hf_model_id=remote_hf_model_id) # ty: ignore[unresolved-attribute]314 315 @classmethod316 def filter_tensors(cls, item):317 assert cls._original_block_count is not None318 # TODO: change TextModel to super()319 if (titem := TextModel.filter_tensors(item)) is None:320 return None321 name, gen = titem322 if name.startswith("model.mtp."):323 name = name.replace("model.", "", 1)324 if name.startswith("mtp."):325 if cls.no_mtp:326 return None327 remapper = {328 "fc": "eh_proj",329 "pre_fc_norm_embedding": "enorm",330 "pre_fc_norm_hidden": "hnorm",331 "norm": "shared_head.norm",332 }333 parts = name.split(".", 3)334 if len(parts) == 4 and parts[1] == "layers" and parts[2].isdecimal():335 mtp_idx = int(parts[2])336 name = f"model.layers.{cls._original_block_count + mtp_idx}.{parts[3]}"337 cls.opt_num_mtp_layers = max(cls.opt_num_mtp_layers, mtp_idx + 1)338 elif len(parts) == 3 and parts[1] in remapper:339 name = f"model.layers.{cls._original_block_count}.{remapper[parts[1]]}.{parts[2]}"340 elif cls.mtp_only:341 keep = name in (342 "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",343 "embed_tokens.weight", "norm.weight",344 )345 if not keep:346 return None347 return name, gen348 349 def set_gguf_parameters(self):350 super().set_gguf_parameters() # ty: ignore[unresolved-attribute]351 if self.no_mtp:352 return353 if (n := self.block_count - self.hparams["num_hidden_layers"]) > 0:354 self.gguf_writer.add_nextn_predict_layers(n)355 356 def prepare_metadata(self, vocab_only: bool):357 from_dir = self.fname_out.is_dir()358 super().prepare_metadata(vocab_only=vocab_only) # ty: ignore[unresolved-attribute]359 360 if not self.mtp_only or not from_dir:361 return362 363 output_type: str = self.ftype.name.partition("_")[2] # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]364 fname_default: str = gguf.naming_convention(365 self.metadata.name, self.metadata.basename, self.metadata.finetune, # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]366 self.metadata.version, size_label=None, output_type=output_type, model_type=None) # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]367 self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"368 369 370@ModelBase.register("Qwen3NextForCausalLM")371@ModelBase.example("Qwen/Qwen3-Next-80B-A3B-Instruct")372class Qwen3NextModel(_QwenMtpMixin, Qwen2MoeModel):373 model_arch = gguf.MODEL_ARCH.QWEN3NEXT374 375 def set_gguf_parameters(self):376 super().set_gguf_parameters()377 self.gguf_writer.add_ssm_conv_kernel(self.hparams["linear_conv_kernel_dim"])378 self.gguf_writer.add_ssm_state_size(self.hparams["linear_key_head_dim"])379 self.gguf_writer.add_ssm_group_count(self.hparams["linear_num_key_heads"])380 self.gguf_writer.add_ssm_time_step_rank(self.hparams["linear_num_value_heads"])381 self.gguf_writer.add_ssm_inner_size(self.hparams["linear_value_head_dim"] * self.hparams["linear_num_value_heads"])382 if (layer_types := self.hparams.get("layer_types")) is not None:383 n_layer = self.hparams["num_hidden_layers"]384 if len(layer_types) != n_layer:385 raise ValueError(f"layer_types has {len(layer_types)} entries, expected num_hidden_layers ({n_layer})")386 recurrent = [t == "linear_attention" for t in layer_types]387 recurrent += [False] * (self.block_count - n_layer)388 self.gguf_writer.add_recurrent_layers(recurrent)389 self.gguf_writer.add_full_attention_interval(self.hparams.get("full_attention_interval", 4))390 if (rope_dim := self.hparams.get("head_dim")) is None:391 rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]392 self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.rope_parameters.get("partial_rotary_factor", 0.25)))393 394 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:395 if name.endswith(".A_log"):396 data_torch = -torch.exp(data_torch)397 elif name.endswith(".dt_bias"):398 name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"399 elif "conv1d" in name:400 data_torch = data_torch.squeeze()401 elif name.endswith("norm.weight") and not name.endswith("linear_attn.norm.weight"):402 data_torch = data_torch + 1403 404 if "in_proj_qkvz.weight" in name:405 # original order: [q, k, v, z] * head_count406 # corrected order: [q * head_count, k * head_count, v * head_count, z * head_count]407 head_k_dim = self.hparams["linear_key_head_dim"]408 head_v_dim = self.hparams["linear_value_head_dim"]409 num_v_heads = self.hparams["linear_num_value_heads"]410 num_k_heads = self.hparams["linear_num_key_heads"]411 hidden_size = self.hparams["hidden_size"]412 split_arg_list_qkvz = [413 head_k_dim, # q partition414 head_k_dim, # k partition415 (num_v_heads // num_k_heads * head_v_dim), # v partition416 (num_v_heads // num_k_heads * head_v_dim), # z partition417 ]418 # view as (n_embd, head_count, [q+k+v+z])419 data_torch = data_torch.permute(1, 0).contiguous()420 data_torch = data_torch.view(-1, num_k_heads, sum(split_arg_list_qkvz))421 # split into q, k, v, z422 q, k, v, z = torch.split(data_torch, split_arg_list_qkvz, dim=-1)423 # flatten dim + head_count424 q = q.contiguous().view(hidden_size, -1)425 k = k.contiguous().view(hidden_size, -1)426 v = v.contiguous().view(hidden_size, -1)427 z = z.contiguous().view(hidden_size, -1)428 # stack back429 qkv = torch.cat([q, k, v], dim=-1).permute(1, 0).contiguous()430 z = z.permute(1, 0).contiguous()431 yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_QKV, bid, ".weight"), qkv)432 yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_GATE, bid, ".weight"), z)433 else:434 yield from super().modify_tensors(data_torch, name, bid)435 436 437@ModelBase.register("RND1")438@ModelBase.example("radicalnumerics/RND1-Base-0910")439class RND1Model(Qwen2MoeModel):440 model_arch = gguf.MODEL_ARCH.RND1441 442 def set_gguf_parameters(self):443 super().set_gguf_parameters()444 445 # RND1 specific parameters446 # RND1 uses bidirectional attention447 self.gguf_writer.add_causal_attention(False)448 449 if (mask_token_id := self.hparams.get("mask_token_id")) is not None:450 self.gguf_writer.add_mask_token_id(mask_token_id)451 452 453class _LinearAttentionVReorderBase(Qwen3NextModel):454 model_arch = gguf.MODEL_ARCH.QWEN3NEXT # overridden by subclasses455 """reorders V heads from grouped to tiled order for ggml broadcast456 457 see https://github.com/ggml-org/llama.cpp/pull/19468#discussion_r2786394306458 459 Linear attention may has num_k_heads < num_v_heads. The HF weights store460 V heads grouped by K head: [G0_v0..v{r-1}, G1_v0..v{r-1}, ...].461 ggml binary ops use tiled broadcast: [K0, K1, ..., K0, K1, ...].462 We reorder V heads to tiled order so ggml_repeat can replace the expensive463 interleaved repeat: [G0_v0, G1_v0, ..., G0_v1, G1_v1, ...].464 """465 466 @staticmethod467 def _reorder_v_heads(tensor: Tensor, dim: int, num_k_heads: int, num_v_per_k: int, head_dim: int) -> Tensor:468 """Reorder V heads from grouped (by K head) to tiled order along the given dimension."""469 shape = list(tensor.shape)470 if dim < 0:471 dim += len(shape)472 new_shape = shape[:dim] + [num_k_heads, num_v_per_k, head_dim] + shape[dim + 1:]473 tensor = tensor.reshape(*new_shape)474 perm = list(range(len(new_shape)))475 perm[dim], perm[dim + 1] = perm[dim + 1], perm[dim]476 return tensor.permute(*perm).contiguous().reshape(*shape)477 478 def _transform_nvfp4_weight(self, name: str, weight: Tensor, scale: Tensor) -> tuple[Tensor, Tensor]:479 if not name.endswith((480 ".linear_attn.in_proj_qkv.weight",481 ".linear_attn.in_proj_z.weight",482 ".linear_attn.in_proj_a.weight",483 ".linear_attn.in_proj_b.weight",484 ".linear_attn.out_proj.weight",485 )):486 return weight, scale487 488 num_k_heads = self.hparams["linear_num_key_heads"]489 num_v_heads = self.hparams["linear_num_value_heads"]490 head_k_dim = self.hparams["linear_key_head_dim"]491 head_v_dim = self.hparams["linear_value_head_dim"]492 num_v_per_k = num_v_heads // num_k_heads493 494 def unpack_nibbles(qs: Tensor) -> Tensor:495 lo = torch.bitwise_and(qs, 0x0F)496 hi = torch.bitwise_right_shift(qs, 4)497 return torch.stack((lo, hi), dim=-1).reshape(*qs.shape[:-1], qs.shape[-1] * 2)498 499 def pack_nibbles(codes: Tensor) -> Tensor:500 codes = codes.reshape(*codes.shape[:-1], codes.shape[-1] // 2, 2)501 lo = torch.bitwise_and(codes[..., 0], 0x0F)502 hi = torch.bitwise_left_shift(torch.bitwise_and(codes[..., 1], 0x0F), 4)503 return torch.bitwise_or(lo, hi).contiguous()504 505 def apply_col_perm(qs: Tensor, scales: Tensor, col_perm: Tensor) -> tuple[Tensor, Tensor]:506 assert qs.ndim >= 2507 assert scales.ndim >= 2508 509 k = qs.shape[-1] * 2510 assert col_perm.numel() == k511 assert k % 16 == 0512 513 group_cols = col_perm.reshape(-1, 16)514 group_starts = group_cols[:, 0]515 expected = group_starts.unsqueeze(1) + torch.arange(16, dtype=col_perm.dtype)516 assert torch.equal(group_cols, expected)517 assert torch.all(group_starts % 16 == 0)518 519 group_perm = (group_starts // 16).to(dtype=torch.long)520 expected_groups = torch.arange(scales.shape[-1], dtype=torch.long)521 assert group_perm.numel() == scales.shape[-1]522 assert torch.equal(torch.sort(group_perm).values, expected_groups)523 524 codes = unpack_nibbles(qs)525 codes = codes.index_select(-1, col_perm.to(device=qs.device, dtype=torch.long))526 qs = pack_nibbles(codes)527 scales = scales.index_select(-1, group_perm.to(device=scales.device))528 return qs, scales529 530 def reorder_rows(qs: Tensor, scales: Tensor, head_dim: int) -> tuple[Tensor, Tensor]:531 row_perm = self._reorder_v_heads(532 torch.arange(num_v_heads * head_dim, dtype=torch.long).unsqueeze(-1),533 0, num_k_heads, num_v_per_k, head_dim,534 ).squeeze(-1)535 return (536 qs.index_select(0, row_perm.to(device=qs.device)),537 scales.index_select(0, row_perm.to(device=scales.device)),538 )539 540 if name.endswith(".linear_attn.in_proj_qkv.weight"):541 q_dim = head_k_dim * num_k_heads542 k_dim = head_k_dim * num_k_heads543 q = weight[:q_dim]544 k = weight[q_dim:q_dim + k_dim]545 v = weight[q_dim + k_dim:]546 q_scale = scale[:q_dim]547 k_scale = scale[q_dim:q_dim + k_dim]548 v_scale = scale[q_dim + k_dim:]549 v, v_scale = reorder_rows(v, v_scale, head_v_dim)550 return torch.cat([q, k, v], dim=0), torch.cat([q_scale, k_scale, v_scale], dim=0)551 552 if name.endswith(".linear_attn.in_proj_z.weight"):553 weight, scale = reorder_rows(weight, scale, head_v_dim)554 elif name.endswith((".linear_attn.in_proj_a.weight", ".linear_attn.in_proj_b.weight")):555 weight, scale = reorder_rows(weight, scale, 1)556 elif name.endswith(".linear_attn.out_proj.weight"):557 col_perm = self._reorder_v_heads(558 torch.arange(num_v_heads * head_v_dim, dtype=torch.long).unsqueeze(0),559 1, num_k_heads, num_v_per_k, head_v_dim,560 ).squeeze(0)561 weight, scale = apply_col_perm(weight, scale, col_perm)562 563 return weight, scale564 565 def _repack_nvfp4(self, name: str, weight: Tensor, scale: Tensor, scale2: Tensor, input_scale: Tensor):566 weight, scale = self._transform_nvfp4_weight(name, weight, scale)567 super()._repack_nvfp4(name, weight, scale, scale2, input_scale)568 569 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:570 num_k_heads = self.hparams.get("linear_num_key_heads", 0)571 num_v_heads = self.hparams.get("linear_num_value_heads", 0)572 573 if num_k_heads > 0 and num_v_heads > 0 and num_k_heads != num_v_heads and "linear_attn." in name:574 head_k_dim = self.hparams["linear_key_head_dim"]575 head_v_dim = self.hparams["linear_value_head_dim"]576 num_v_per_k = num_v_heads // num_k_heads577 578 if ".in_proj_qkv." in name:579 # QKV weight: reorder only the V rows580 q_dim = head_k_dim * num_k_heads581 k_dim = head_k_dim * num_k_heads582 q = data_torch[:q_dim]583 k = data_torch[q_dim:q_dim + k_dim]584 v = data_torch[q_dim + k_dim:]585 v = self._reorder_v_heads(v, 0, num_k_heads, num_v_per_k, head_v_dim)586 data_torch = torch.cat([q, k, v], dim=0)587 588 elif ".in_proj_z." in name:589 # Z gate weight: reorder rows (num_v_heads * head_v_dim)590 data_torch = self._reorder_v_heads(data_torch, 0, num_k_heads, num_v_per_k, head_v_dim)591 592 elif ".in_proj_b." in name or ".in_proj_a." in name:593 # Beta/Alpha weight: reorder rows (num_v_heads, head_dim=1)594 data_torch = self._reorder_v_heads(data_torch, 0, num_k_heads, num_v_per_k, 1)595 596 elif ".A_log" in name or ".dt_bias" in name or ".dt_proj" in name:597 # A_log / dt_bias: 1D parameters with num_v_heads elements598 if data_torch.ndim == 1:599 data_torch = self._reorder_v_heads(600 data_torch.unsqueeze(-1), 0, num_k_heads, num_v_per_k, 1601 ).squeeze(-1)602 else:603 data_torch = self._reorder_v_heads(data_torch, -1, num_k_heads, num_v_per_k, 1)604 605 elif ".conv1d" in name:606 # Conv1d kernel: reorder only the V channel portion607 data = data_torch.squeeze()608 qk_channels = head_k_dim * num_k_heads * 2609 qk_part = data[:qk_channels]610 v_part = data[qk_channels:]611 v_part = self._reorder_v_heads(v_part, 0, num_k_heads, num_v_per_k, head_v_dim)612 data_torch = torch.cat([qk_part, v_part], dim=0)613 614 elif ".out_proj." in name:615 # Out projection weight: reorder columns (input dimension)616 data_torch = self._reorder_v_heads(data_torch, 1, num_k_heads, num_v_per_k, head_v_dim)617 618 yield from super().modify_tensors(data_torch, name, bid)619 620 621class _Qwen35MRopeMixin:622 # Qwen3.5 always applies interleaved MRoPE (see Qwen3_5RotaryEmbedding in transformers);623 # the upstream default mrope_section is [11, 11, 10] and llama.cpp's QWEN35 / QWEN35MOE624 # loaders treat qwen35.rope.dimension_sections as required, so make sure it is always625 # written even when a particular checkpoint omits the field in `rope_parameters`.626 _QWEN35_DEFAULT_MROPE_SECTION = [11, 11, 10, 0]627 628 gguf_writer: gguf.GGUFWriter629 rope_parameters: dict630 631 def set_gguf_parameters(self):632 super().set_gguf_parameters() # ty: ignore[unresolved-attribute]633 if "mrope_section" not in self.rope_parameters:634 self.gguf_writer.add_rope_dimension_sections(self._QWEN35_DEFAULT_MROPE_SECTION)635 636 637@ModelBase.register("Qwen3_5ForConditionalGeneration", "Qwen3_5ForCausalLM")638@ModelBase.example("Qwen/Qwen3.5-9B")639class Qwen3_5TextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):640 model_arch = gguf.MODEL_ARCH.QWEN35641 642 643@ModelBase.register("Qwen3_5MoeForConditionalGeneration", "Qwen3_5MoeForCausalLM")644@ModelBase.example("Qwen/Qwen3.5-35B-A3B")645class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):646 model_arch = gguf.MODEL_ARCH.QWEN35MOE647 648 649@ModelBase.register("DFlashDraftModel", "DFlash2DraftModel")650@ModelBase.example("z-lab/Qwen3.5-9B-DFlash")651class DFlashModel(Qwen3Model):652 model_arch = gguf.MODEL_ARCH.DFLASH653 654 def set_vocab(self):655 if self.target_model_dir is None:656 raise ValueError(657 "DFlash draft model requires --target-model-dir to be specified. "658 "Please provide the path to the target model directory containing the tokenizer."659 )660 logger.info(f"DFlash: Using tokenizer from target model: {self.target_model_dir}")661 original_dir = self.dir_model662 self.dir_model = self.target_model_dir663 664 # Reuse the target model's own vocab handler (e.g. Gemma-4 needs its665 # own tokenizer logic, not the Qwen default).666 from . import get_model_class667 with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f:668 target_hparams = json.load(f)669 target_arch = target_hparams["architectures"][0]670 target_cls = get_model_class(target_arch)671 672 if target_cls is not type(self):673 if target_arch == "NemotronHForCausalLM":674 setattr(self, "is_moe", "num_experts_per_tok" in target_hparams)675 target_cls.set_vocab(self) # ty: ignore[unresolved-attribute]676 else:677 super().set_vocab()678 679 self.dir_model = original_dir680 681 mask_token_id = self.hparams.get("dflash_config", {}).get("mask_token_id")682 if mask_token_id is not None:683 self.gguf_writer.add_mask_token_id(mask_token_id)684 685 def set_gguf_parameters(self):686 super().set_gguf_parameters()687 688 dflash_config = self.hparams.get("dflash_config", {})689 block_size = dflash_config.get("block_size", self.hparams.get("block_size", 16))690 self.gguf_writer.add_block_size(block_size)691 692 if "conv_kernel_size" in dflash_config:693 self.gguf_writer.add_conv_kernel_size(int(dflash_config["conv_kernel_size"]))694 self.gguf_writer.add_conv_group_size(int(dflash_config["conv_group_size"]))695 self.gguf_writer.add_selector_rank(int(dflash_config["selector_rank"]))696 self.gguf_writer.add_selector_top_k(int(dflash_config["selector_top_k"]))697 698 output_multiplier = dflash_config.get(699 "output_multiplier", self.hparams.get("output_multiplier")700 )701 if output_multiplier is not None:702 self.gguf_writer.add_logit_scale(float(output_multiplier))703 softcap = dflash_config.get(704 "final_logit_softcapping", self.hparams.get("final_logit_softcapping")705 )706 if softcap is not None and float(softcap) > 0:707 self.gguf_writer.add_final_logit_softcapping(float(softcap))708 embedding_scale = dflash_config.get(709 "input_embedding_scale", self.hparams.get("input_embedding_scale")710 )711 if embedding_scale is not None:712 self.gguf_writer.add_embedding_scale(float(embedding_scale))713 714 target_layer_ids = dflash_config.get("target_layer_ids", [])715 if target_layer_ids:716 extract_layer_ids = [i + 1 for i in target_layer_ids]717 self.gguf_writer.add_target_layers(extract_layer_ids)718 719 use_sliding_window = self.hparams.get("use_sliding_window", False) or dflash_config.get("use_swa", False)720 sliding_window = dflash_config.get("swa_window_size") or self.hparams.get("sliding_window")721 layer_types = self.hparams.get("layer_types")722 if use_sliding_window and sliding_window and layer_types:723 is_swa = [lt == "sliding_attention" for lt in layer_types]724 self.gguf_writer.add_sliding_window(sliding_window)725 self.gguf_writer.add_sliding_window_pattern(is_swa)726 727 causal = self.hparams.get("is_causal")728 if causal is None:729 causal = dflash_config.get("causal")730 if causal is not None:731 self.gguf_writer.add_causal_attention(bool(causal))732 733 # M-RoPE target: the draft ropes on the temporal dim only, so write734 # degenerate sections [n_rot/2, 0, 0, 0]735 if self._target_uses_mrope():736 head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"]737 self.gguf_writer.add_rope_dimension_sections([head_dim // 2, 0, 0, 0])738 739 def _target_uses_mrope(self) -> bool:740 if self.target_model_dir is None:741 return False742 with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f:743 cfg = json.load(f)744 cfg = cfg.get("text_config", cfg)745 rope = cfg.get("rope_parameters") or cfg.get("rope_scaling") or {}746 return "mrope_section" in rope747 748 @classmethod749 def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:750 name, gen = item751 if not name.startswith("model."):752 name = "model." + name753 if "sink" in name and not name.endswith(".weight"):754 name += ".weight"755 return super().filter_tensors((name, gen))756 757 _ROPE_PERMUTE_SUFFIXES = (758 "self_attn.q_proj.weight",759 "self_attn.k_proj.weight",760 "self_attn.q_norm.weight",761 "self_attn.k_norm.weight",762 )763 764 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:765 if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True):766 return767 768 # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd769 if not self.hparams.get("rope_is_neox_style", True) and name.endswith(self._ROPE_PERMUTE_SUFFIXES):770 head_dim = self.hparams["head_dim"]771 shape = data_torch.shape772 data_torch = data_torch.reshape(-1, head_dim // 2, 2, *shape[1:]).transpose(1, 2).reshape(shape)773 774 if name in (775 "model.candidate_selector.predecessor_codebook",776 "model.candidate_selector.successor_codebook",777 ):778 name += ".weight"779 780 yield from super().modify_tensors(data_torch, name, bid)781 782 783@ModelBase.register(784 "Qwen3DSparkModel",785 "DSparkDraftModel",786 "DSparkSpeculator",787 "Lfm2DSparkDraftModel",788 "LingDSparkModel",789)790@ModelBase.example("satgeze/Qwen3.6-27B-DSpark")791class DSparkModel(DFlashModel):792 # DSpark = DFlash + a semi-autoregressive Markov head.793 model_arch = gguf.MODEL_ARCH.DFLASH794 795 def __init__(self, dir_model, *args, **kwargs):796 hparams = kwargs.pop("hparams", None)797 if hparams is None:798 hparams = ModelBase.load_hparams(dir_model, False)799 800 # EAGLE3-style exports use the 1+N bonus-anchor block, DFlash-lineage exports sample from the anchor801 self._sample_from_anchor = hparams.get(802 "sample_from_anchor",803 "transformer_layer_config" not in hparams and "aux_hidden_state_layer_ids" not in hparams)804 if "transformer_layer_config" in hparams:805 hparams = {**hparams, **hparams["transformer_layer_config"]}806 807 super().__init__(dir_model, *args, hparams=hparams, **kwargs)808 809 # normalize both schemas to DFlash's nested dflash_config810 if "aux_hidden_state_layer_ids" in self.hparams:811 self.hparams.setdefault("dflash_config", {812 "mask_token_id": self.hparams.get("mask_token_id"),813 "target_layer_ids": [i - 1 for i in self.hparams["aux_hidden_state_layer_ids"]],814 })815 else:816 self.hparams.setdefault("dflash_config", {817 k: self.hparams[k] for k in ("target_layer_ids", "mask_token_id") if k in self.hparams818 })819 820 if (markov_head_type := self.hparams.get("markov_head_type", "vanilla")) != "vanilla":821 raise ValueError(f"unsupported markov_head_type {markov_head_type!r} (only 'vanilla' is supported)")822 823 n_vocab = self.hparams["vocab_size"]824 self._n_vocab_draft = self.hparams.get("draft_vocab_size") or n_vocab825 if self._n_vocab_draft > n_vocab:826 raise ValueError(f"draft_vocab_size {self._n_vocab_draft} exceeds vocab_size {n_vocab}")827 self._d2t: Tensor | None = None828 829 def set_gguf_parameters(self):830 super().set_gguf_parameters()831 self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor)832 833 # confidence head is optional: vanilla-markov exports ship without it834 has_conf = any("confidence_head.proj" in name for name in self.model_tensors)835 self.gguf_writer.add_has_confidence_head(has_conf)836 837 @classmethod838 def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:839 if item[0] == "t2d": # not used at runtime840 return None841 return super().filter_tensors(item)842 843 _ROPE_PERMUTE_SUFFIXES = (844 "self_attn.q_proj.weight",845 "self_attn.k_proj.weight",846 "self_attn.q_norm.weight",847 "self_attn.k_norm.weight",848 )849 850 def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:851 if name == "model.d2t":852 self._d2t = data_torch853 return854 855 if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith("lm_head.weight"):856 return857 858 # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd859 if not self.hparams.get("rope_is_neox_style", True) and name.endswith(self._ROPE_PERMUTE_SUFFIXES):860 head_dim = self.hparams["head_dim"]861 shape = data_torch.shape862 data_torch = data_torch.reshape(-1, head_dim // 2, 2, *shape[1:]).transpose(1, 2).reshape(shape)863 864 yield from super().modify_tensors(data_torch, name, bid)865 866 def prepare_tensors(self):867 super().prepare_tensors()868 869 n_vocab = self.hparams["vocab_size"]870 if self._n_vocab_draft < n_vocab and self._d2t is None:871 raise ValueError(f"draft_vocab_size {self._n_vocab_draft} < vocab_size {n_vocab} but no d2t table found")872 873 # write d2t as absolute target token ids874 if self._d2t is not None:875 data = LazyTorchTensor.to_eager(self._d2t).to(torch.int64).cpu().numpy().reshape(-1)876 if data.size != self._n_vocab_draft:877 raise ValueError(f"d2t size {data.size} does not match draft_vocab_size {self._n_vocab_draft}")878 data = data + np.arange(data.size, dtype=np.int64)879 if np.any((data < 0) | (data >= n_vocab)):880 raise ValueError(f"d2t target ids out of range for target vocab size {n_vocab}")881 if np.unique(data).size != data.size:882 raise ValueError("d2t contains duplicate target ids")883 logger.info(f"{'d2t,':<30} --> I64, shape = {{{data.size}}}")884 self.gguf_writer.add_tensor("d2t", data, raw_dtype=gguf.GGMLQuantizationType.I64)885