Brunobkr/llama.cpp_AlgMor24_github

1from __future__ import annotations2 3import json4 5from typing import Any, Callable, Iterable, TYPE_CHECKING6 7import torch8 9if TYPE_CHECKING:10    from torch import Tensor11 12from .base import ModelBase, TextModel, gguf, logger13 14 15@ModelBase.register("QWenLMHeadModel")16class QwenModel(TextModel):17    model_arch = gguf.MODEL_ARCH.QWEN18 19    @staticmethod20    def token_bytes_to_string(b):21        from transformers.convert_slow_tokenizer import bytes_to_unicode22        byte_encoder = bytes_to_unicode()23        return ''.join([byte_encoder[ord(char)] for char in b.decode('latin-1')])24 25    @staticmethod26    def bpe(mergeable_ranks: dict[bytes, int], token: bytes, max_rank: int | None = None) -> list[bytes]:27        parts = [bytes([b]) for b in token]28        while True:29            min_idx = None30            min_rank = None31            for i, pair in enumerate(zip(parts[:-1], parts[1:])):32                rank = mergeable_ranks.get(pair[0] + pair[1])33                if rank is not None and (min_rank is None or rank < min_rank):34                    min_idx = i35                    min_rank = rank36            if min_rank is None or (max_rank is not None and min_rank >= max_rank):37                break38            assert min_idx is not None39            parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx + 1]] + parts[min_idx + 2:]40        return parts41 42    def set_vocab(self):43        self._set_vocab_qwen()44 45 46@ModelBase.register(47    "Qwen2Model",48    "Qwen2ForCausalLM",49    "Qwen2AudioForConditionalGeneration",50    "KORMoForCausalLM",51    "AudioFlamingo3ForConditionalGeneration",52    "DotsOCRForCausalLM",53)54class Qwen2Model(TextModel):55    model_arch = gguf.MODEL_ARCH.QWEN256 57    def set_vocab(self):58        try:59            self._set_vocab_sentencepiece()60        except FileNotFoundError:61            self._set_vocab_gpt2()62 63    def set_gguf_parameters(self):64        super().set_gguf_parameters()65        self._try_set_pooling_type()66 67    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:68        if self.hf_arch == "Qwen2Model":69            name = f"model.{name}"  # map to Qwen2ForCausalLM tensors70        yield from super().modify_tensors(data_torch, name, bid)71 72 73@ModelBase.register("Qwen2MoeForCausalLM")74class Qwen2MoeModel(TextModel):75    model_arch = gguf.MODEL_ARCH.QWEN2MOE76 77    def set_gguf_parameters(self):78        super().set_gguf_parameters()79        if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None:80            self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size)81            logger.info(f"gguf: expert feed forward length = {moe_intermediate_size}")82        if (shared_expert_intermediate_size := self.hparams.get('shared_expert_intermediate_size')) is not None:83            self.gguf_writer.add_expert_shared_feed_forward_length(shared_expert_intermediate_size)84            logger.info(f"gguf: expert shared feed forward length = {shared_expert_intermediate_size}")85 86    _experts: list[dict[str, Tensor]] | None = None87 88    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:89        # handle aggregated expert tensors90        # GGUF stores dimensions reversed from PyTorch, so:91        # PyTorch (A,B,C) -> GGUF writes [C,B,A] -> GGML reads ne={C,B,A}92        # Input shapes from HF: (n_expert, n_ff_exp, n_embd) or (n_expert, n_embd, n_ff_exp)93        # Expected GGML ne: {n_embd, n_ff_exp, n_expert} for gate/up, {n_ff_exp, n_embd, n_expert} for down94        if name.endswith("mlp.experts.down_proj") or name.endswith("mlp.experts.down_proj.weight"):95            mapped = f"{name}.weight" if not name.endswith(".weight") else name96            # HF: [n_expert, n_embd, n_ff] -> GGML: {n_ff, n_embd, n_expert}97            yield from super().modify_tensors(data_torch, mapped, bid)98            return99 100        if name.endswith("mlp.experts.gate_up_proj") or name.endswith("mlp.experts.gate_up_proj.weight"):101            if data_torch.ndim < 3 or data_torch.shape[-2] % 2 != 0:102                raise ValueError(f"Unexpected gate_up_proj shape for {name}: {tuple(data_torch.shape)}")103            # HF: [n_expert, 2*n_ff, n_embd] -> split on dim=-2104            n_ff = data_torch.shape[-2] // 2105            gate = data_torch[..., :n_ff, :].contiguous()106            up = data_torch[..., n_ff:, :].contiguous()107            # gate/up: [n_expert, n_ff, n_embd] -> GGML: {n_embd, n_ff, n_expert}108            base_name = name.removesuffix(".weight").removesuffix(".gate_up_proj")109            mapped_gate = f"{base_name}.gate_proj.weight"110            mapped_up = f"{base_name}.up_proj.weight"111            yield from super().modify_tensors(gate, mapped_gate, bid)112            yield from super().modify_tensors(up, mapped_up, bid)113            return114 115        if name.find("experts") != -1:116            n_experts = self.find_hparam(["num_local_experts", "num_experts"])117            assert bid is not None118 119            if self._experts is None:120                self._experts = [{} for _ in range(self.block_count)]121 122            self._experts[bid][name] = data_torch123 124            if len(self._experts[bid]) >= n_experts * 3:125                # merge the experts into a single 3d tensor126                for w_name in ["down_proj", "gate_proj", "up_proj"]:127                    datas: list[Tensor] = []128 129                    for xid in range(n_experts):130                        ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"131                        datas.append(self._experts[bid][ename])132                        del self._experts[bid][ename]133 134                    data_torch = torch.stack(datas, dim=0)135 136                    merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight"137 138                    yield from super().modify_tensors(data_torch, merged_name, bid)139                return140            else:141                return142 143        yield from super().modify_tensors(data_torch, name, bid)144 145    def prepare_tensors(self):146        super().prepare_tensors()147 148        if self._experts is not None:149            # flatten `list[dict[str, Tensor]]` into `list[str]`150            experts = [k for d in self._experts for k in d.keys()]151            if len(experts) > 0:152                raise ValueError(f"Unprocessed experts: {experts}")153 154 155@ModelBase.register("Qwen3ForCausalLM", "Qwen3Model")156class Qwen3Model(Qwen2Model):157    model_arch = gguf.MODEL_ARCH.QWEN3158 159    # extra logic for rerank models160    is_rerank: bool = False161    is_tied_embeddings: bool = False162    token_false_id: int | None = None163    token_true_id: int | None = None164 165    def __init__(self, *args, **kwargs):166        super().__init__(*args, **kwargs)167 168        # track for intern-s1-mini169        hparams = ModelBase.load_hparams(self.dir_model, is_mistral_format=False)170        self.origin_hf_arch = hparams.get('architectures', [None])[0]171 172        if self._is_qwen3_reranker():173            self._find_rerank_config()174 175    def _is_qwen3_reranker(self) -> bool:176        readme_path = self.dir_model / "README.md"177        readme_text = ""178        if readme_path.exists():179            with readme_path.open("r", encoding="utf-8") as f:180                readme_text = f.read()181 182        name_hints = [183            str(self.dir_model.name),184            str(self.hparams.get("_name_or_path", "")),185            str(self.hparams.get("model_type", "")),186            str(self.origin_hf_arch or ""),187        ]188        name_hints = [hint.lower() for hint in name_hints if hint]189 190        if "# qwen3-reranker" in readme_text.lower() or "# qwen3-vl-reranker" in readme_text.lower():191            return True192 193        if any("qwen3-reranker" in hint or "qwen3-vl-reranker" in hint for hint in name_hints):194            return True195 196        return "sequenceclassification" in (self.origin_hf_arch or "").lower()197 198    def set_vocab(self):199        # deal with intern-s1-mini200        if self.origin_hf_arch == 'InternS1ForConditionalGeneration':201            self._set_vocab_interns1()202            return203 204        super().set_vocab()205 206    def _find_rerank_config(self):207        from transformers import AutoTokenizer208        tokenizer = AutoTokenizer.from_pretrained(self.dir_model)209 210        self.is_rerank = True211        self.is_tied_embeddings = self.hparams.get("tie_word_embeddings", False)212        self.token_false_id = tokenizer.convert_tokens_to_ids("no")  # ty: ignore[unresolved-attribute, invalid-assignment]213        self.token_true_id = tokenizer.convert_tokens_to_ids("yes")  # ty: ignore[unresolved-attribute, invalid-assignment]214        self.sep_token_id = tokenizer.convert_tokens_to_ids("|")  # ty: ignore[unresolved-attribute]215 216        assert self.token_false_id is not None and self.token_true_id is not None217 218    def set_gguf_parameters(self):219        super().set_gguf_parameters()220        if self.is_rerank:221            self.gguf_writer.add_pooling_type(gguf.PoolingType.RANK)222            self.gguf_writer.add_classifier_output_labels(["yes", "no"])223            self.gguf_writer.add_chat_template([{224                "name": "rerank",225                "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"226                            "<|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"227                            "<|im_start|>assistant\n<think>\n\n</think>\n\n"228            }])229 230    def _get_cls_out_tensor(self, data_torch: Tensor) -> Tensor:231        # extract "yes" and "no" tokens from the output lm_head tensor232        false_row = data_torch[self.token_false_id]233        true_row = data_torch[self.token_true_id]234        return torch.stack([true_row, false_row], dim=0)235 236    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:237        if self.is_rerank:238            is_tied_head = self.is_tied_embeddings and "embed_tokens" in name239            is_real_head = not self.is_tied_embeddings and "lm_head" in name240            if is_tied_head or is_real_head:241                cls_out_head = (242                    gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.CLS_OUT] + ".weight",243                    self._get_cls_out_tensor(data_torch),244                )245                yield cls_out_head246                if is_tied_head:247                    yield from super().modify_tensors(data_torch, name, bid)248                return249 250        yield from super().modify_tensors(data_torch, name, bid)251 252 253@ModelBase.register("Qwen3MoeForCausalLM")254class Qwen3MoeModel(Qwen2MoeModel):255    model_arch = gguf.MODEL_ARCH.QWEN3MOE256 257    def __init__(self, *args, **kwargs):258        super().__init__(*args, **kwargs)259        hparams = ModelBase.load_hparams(self.dir_model, False)260        self.origin_hf_arch = hparams.get('architectures', [None])[0]261 262    def set_vocab(self):263        # deal with intern-s1264        if self.origin_hf_arch == 'InternS1ForConditionalGeneration':265            self._set_vocab_interns1()266            return267 268        super().set_vocab()269 270 271class _QwenMtpMixin:272    """Shared MTP wiring for Qwen3-Next and Qwen3.5/3.6 text variants. The HF273    config carries the MTP block under `mtp_num_hidden_layers` (computed from274    the checkpoint when absent, e.g. Qwen3-Next) and the tensors under275    `mtp.*`; we extend block_count, emit the nextn metadata key, and remap276    `mtp.*` to the standard layer-indexed nextn naming so the existing277    tensor_map handles them."""278 279    supports_mtp_export = True280    hparams: dict[str, Any]281    model_arch: gguf.MODEL_ARCH282    gguf_writer: gguf.GGUFWriter283    block_count: int284    tensor_map: gguf.TensorNameMap285    no_mtp: bool286    mtp_only: bool287    _original_block_count: int | None = None288    opt_num_mtp_layers: int = 0289 290    def __init__(self, *args, **kwargs):291        super().__init__(*args, **kwargs)292        self.block_count = self.hparams["num_hidden_layers"]293        if not self.no_mtp:294            n_mtp = self.hparams.get("mtp_num_hidden_layers", 0)295            # Qwen-3-Next doesn't include `mtp_num_hidden_layers` in config.296            if n_mtp == 0:297                assert self.opt_num_mtp_layers != 0298                n_mtp = self.opt_num_mtp_layers299            self.block_count += n_mtp300        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)301 302    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:303        hparams = {**self.hparams, **self.hparams.get("text_config", {})}304        key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None)305        type(self)._original_block_count = hparams.get(key)306        type(self).opt_num_mtp_layers = 0307        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)  # ty: ignore[unresolved-attribute]308 309    @classmethod310    def filter_tensors(cls, item):311        assert cls._original_block_count is not None312        # TODO: change TextModel to super()313        if (titem := TextModel.filter_tensors(item)) is None:314            return None315        name, gen = titem316        if name.startswith("model.mtp."):317            name = name.replace("model.", "", 1)318        if name.startswith("mtp."):319            if cls.no_mtp:320                return None321            remapper = {322                "fc":                    "eh_proj",323                "pre_fc_norm_embedding": "enorm",324                "pre_fc_norm_hidden":    "hnorm",325                "norm":                  "shared_head.norm",326            }327            parts = name.split(".", 3)328            if len(parts) == 4 and parts[1] == "layers" and parts[2].isdecimal():329                mtp_idx = int(parts[2])330                name = f"model.layers.{cls._original_block_count + mtp_idx}.{parts[3]}"331                cls.opt_num_mtp_layers = max(cls.opt_num_mtp_layers, mtp_idx + 1)332            elif len(parts) == 3 and parts[1] in remapper:333                name = f"model.layers.{cls._original_block_count}.{remapper[parts[1]]}.{parts[2]}"334        elif cls.mtp_only:335            keep = name in (336                "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",337                "embed_tokens.weight", "norm.weight",338            )339            if not keep:340                return None341        return name, gen342 343    def set_gguf_parameters(self):344        super().set_gguf_parameters()  # ty: ignore[unresolved-attribute]345        if self.no_mtp:346            return347        if (n := self.block_count - self.hparams["num_hidden_layers"]) > 0:348            self.gguf_writer.add_nextn_predict_layers(n)349 350    def prepare_metadata(self, vocab_only: bool):351        from_dir = self.fname_out.is_dir()352        super().prepare_metadata(vocab_only=vocab_only)  # ty: ignore[unresolved-attribute]353 354        if not self.mtp_only or not from_dir:355            return356 357        output_type: str = self.ftype.name.partition("_")[2]  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]358        fname_default: str = gguf.naming_convention(359            self.metadata.name, self.metadata.basename, self.metadata.finetune,                  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]360            self.metadata.version, size_label=None, output_type=output_type, model_type=None)    # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]361        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"362 363 364@ModelBase.register("Qwen3NextForCausalLM")365class Qwen3NextModel(_QwenMtpMixin, Qwen2MoeModel):366    model_arch = gguf.MODEL_ARCH.QWEN3NEXT367 368    def set_gguf_parameters(self):369        super().set_gguf_parameters()370        self.gguf_writer.add_ssm_conv_kernel(self.hparams["linear_conv_kernel_dim"])371        self.gguf_writer.add_ssm_state_size(self.hparams["linear_key_head_dim"])372        self.gguf_writer.add_ssm_group_count(self.hparams["linear_num_key_heads"])373        self.gguf_writer.add_ssm_time_step_rank(self.hparams["linear_num_value_heads"])374        self.gguf_writer.add_ssm_inner_size(self.hparams["linear_value_head_dim"] * self.hparams["linear_num_value_heads"])375        self.gguf_writer.add_full_attention_interval(self.hparams.get("full_attention_interval", 4))376        if (rope_dim := self.hparams.get("head_dim")) is None:377            rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]378        self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.rope_parameters.get("partial_rotary_factor", 0.25)))379 380    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:381        if name.endswith(".A_log"):382            data_torch = -torch.exp(data_torch)383        elif name.endswith(".dt_bias"):384            name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"385        elif "conv1d" in name:386            data_torch = data_torch.squeeze()387        elif name.endswith("norm.weight") and not name.endswith("linear_attn.norm.weight"):388            data_torch = data_torch + 1389 390        if "in_proj_qkvz.weight" in name:391            # original order:  [q, k, v, z] * head_count392            # corrected order: [q * head_count, k * head_count, v * head_count, z * head_count]393            head_k_dim = self.hparams["linear_key_head_dim"]394            head_v_dim = self.hparams["linear_value_head_dim"]395            num_v_heads = self.hparams["linear_num_value_heads"]396            num_k_heads = self.hparams["linear_num_key_heads"]397            hidden_size = self.hparams["hidden_size"]398            split_arg_list_qkvz = [399                head_k_dim, # q partition400                head_k_dim, # k partition401                (num_v_heads // num_k_heads * head_v_dim), # v partition402                (num_v_heads // num_k_heads * head_v_dim), # z partition403            ]404            # view as (n_embd, head_count, [q+k+v+z])405            data_torch = data_torch.permute(1, 0).contiguous()406            data_torch = data_torch.view(-1, num_k_heads, sum(split_arg_list_qkvz))407            # split into q, k, v, z408            q, k, v, z = torch.split(data_torch, split_arg_list_qkvz, dim=-1)409            # flatten dim + head_count410            q = q.contiguous().view(hidden_size, -1)411            k = k.contiguous().view(hidden_size, -1)412            v = v.contiguous().view(hidden_size, -1)413            z = z.contiguous().view(hidden_size, -1)414            # stack back415            qkv = torch.cat([q, k, v], dim=-1).permute(1, 0).contiguous()416            z = z.permute(1, 0).contiguous()417            yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_QKV,  bid, ".weight"), qkv)418            yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_GATE, bid, ".weight"), z)419        else:420            yield from super().modify_tensors(data_torch, name, bid)421 422 423@ModelBase.register("RND1")424class RND1Model(Qwen2MoeModel):425    model_arch = gguf.MODEL_ARCH.RND1426 427    def set_gguf_parameters(self):428        super().set_gguf_parameters()429 430        # RND1 specific parameters431        # RND1 uses bidirectional attention432        self.gguf_writer.add_causal_attention(False)433 434        if (mask_token_id := self.hparams.get("mask_token_id")) is not None:435            self.gguf_writer.add_mask_token_id(mask_token_id)436 437 438class _LinearAttentionVReorderBase(Qwen3NextModel):439    model_arch = gguf.MODEL_ARCH.QWEN3NEXT  # overridden by subclasses440    """reorders V heads from grouped to tiled order for ggml broadcast441 442    see https://github.com/ggml-org/llama.cpp/pull/19468#discussion_r2786394306443 444    Linear attention may has num_k_heads < num_v_heads. The HF weights store445    V heads grouped by K head: [G0_v0..v{r-1}, G1_v0..v{r-1}, ...].446    ggml binary ops use tiled broadcast: [K0, K1, ..., K0, K1, ...].447    We reorder V heads to tiled order so ggml_repeat can replace the expensive448    interleaved repeat: [G0_v0, G1_v0, ..., G0_v1, G1_v1, ...].449    """450 451    @staticmethod452    def _reorder_v_heads(tensor: Tensor, dim: int, num_k_heads: int, num_v_per_k: int, head_dim: int) -> Tensor:453        """Reorder V heads from grouped (by K head) to tiled order along the given dimension."""454        shape = list(tensor.shape)455        if dim < 0:456            dim += len(shape)457        new_shape = shape[:dim] + [num_k_heads, num_v_per_k, head_dim] + shape[dim + 1:]458        tensor = tensor.reshape(*new_shape)459        perm = list(range(len(new_shape)))460        perm[dim], perm[dim + 1] = perm[dim + 1], perm[dim]461        return tensor.permute(*perm).contiguous().reshape(*shape)462 463    def _transform_nvfp4_weight(self, name: str, weight: Tensor, scale: Tensor) -> tuple[Tensor, Tensor]:464        if not name.endswith((465            ".linear_attn.in_proj_qkv.weight",466            ".linear_attn.in_proj_z.weight",467            ".linear_attn.in_proj_a.weight",468            ".linear_attn.in_proj_b.weight",469            ".linear_attn.out_proj.weight",470        )):471            return weight, scale472 473        num_k_heads = self.hparams["linear_num_key_heads"]474        num_v_heads = self.hparams["linear_num_value_heads"]475        head_k_dim = self.hparams["linear_key_head_dim"]476        head_v_dim = self.hparams["linear_value_head_dim"]477        num_v_per_k = num_v_heads // num_k_heads478 479        def unpack_nibbles(qs: Tensor) -> Tensor:480            lo = torch.bitwise_and(qs, 0x0F)481            hi = torch.bitwise_right_shift(qs, 4)482            return torch.stack((lo, hi), dim=-1).reshape(*qs.shape[:-1], qs.shape[-1] * 2)483 484        def pack_nibbles(codes: Tensor) -> Tensor:485            codes = codes.reshape(*codes.shape[:-1], codes.shape[-1] // 2, 2)486            lo = torch.bitwise_and(codes[..., 0], 0x0F)487            hi = torch.bitwise_left_shift(torch.bitwise_and(codes[..., 1], 0x0F), 4)488            return torch.bitwise_or(lo, hi).contiguous()489 490        def apply_col_perm(qs: Tensor, scales: Tensor, col_perm: Tensor) -> tuple[Tensor, Tensor]:491            assert qs.ndim >= 2492            assert scales.ndim >= 2493 494            k = qs.shape[-1] * 2495            assert col_perm.numel() == k496            assert k % 16 == 0497 498            group_cols = col_perm.reshape(-1, 16)499            group_starts = group_cols[:, 0]500            expected = group_starts.unsqueeze(1) + torch.arange(16, dtype=col_perm.dtype)501            assert torch.equal(group_cols, expected)502            assert torch.all(group_starts % 16 == 0)503 504            group_perm = (group_starts // 16).to(dtype=torch.long)505            expected_groups = torch.arange(scales.shape[-1], dtype=torch.long)506            assert group_perm.numel() == scales.shape[-1]507            assert torch.equal(torch.sort(group_perm).values, expected_groups)508 509            codes = unpack_nibbles(qs)510            codes = codes.index_select(-1, col_perm.to(device=qs.device, dtype=torch.long))511            qs = pack_nibbles(codes)512            scales = scales.index_select(-1, group_perm.to(device=scales.device))513            return qs, scales514 515        def reorder_rows(qs: Tensor, scales: Tensor, head_dim: int) -> tuple[Tensor, Tensor]:516            row_perm = self._reorder_v_heads(517                torch.arange(num_v_heads * head_dim, dtype=torch.long).unsqueeze(-1),518                0, num_k_heads, num_v_per_k, head_dim,519            ).squeeze(-1)520            return (521                qs.index_select(0, row_perm.to(device=qs.device)),522                scales.index_select(0, row_perm.to(device=scales.device)),523            )524 525        if name.endswith(".linear_attn.in_proj_qkv.weight"):526            q_dim = head_k_dim * num_k_heads527            k_dim = head_k_dim * num_k_heads528            q = weight[:q_dim]529            k = weight[q_dim:q_dim + k_dim]530            v = weight[q_dim + k_dim:]531            q_scale = scale[:q_dim]532            k_scale = scale[q_dim:q_dim + k_dim]533            v_scale = scale[q_dim + k_dim:]534            v, v_scale = reorder_rows(v, v_scale, head_v_dim)535            return torch.cat([q, k, v], dim=0), torch.cat([q_scale, k_scale, v_scale], dim=0)536 537        if name.endswith(".linear_attn.in_proj_z.weight"):538            weight, scale = reorder_rows(weight, scale, head_v_dim)539        elif name.endswith((".linear_attn.in_proj_a.weight", ".linear_attn.in_proj_b.weight")):540            weight, scale = reorder_rows(weight, scale, 1)541        elif name.endswith(".linear_attn.out_proj.weight"):542            col_perm = self._reorder_v_heads(543                torch.arange(num_v_heads * head_v_dim, dtype=torch.long).unsqueeze(0),544                1, num_k_heads, num_v_per_k, head_v_dim,545            ).squeeze(0)546            weight, scale = apply_col_perm(weight, scale, col_perm)547 548        return weight, scale549 550    def _repack_nvfp4(self, name: str, weight: Tensor, scale: Tensor, scale2: Tensor, input_scale: Tensor):551        weight, scale = self._transform_nvfp4_weight(name, weight, scale)552        super()._repack_nvfp4(name, weight, scale, scale2, input_scale)553 554    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:555        num_k_heads = self.hparams.get("linear_num_key_heads", 0)556        num_v_heads = self.hparams.get("linear_num_value_heads", 0)557 558        if num_k_heads > 0 and num_v_heads > 0 and num_k_heads != num_v_heads and "linear_attn." in name:559            head_k_dim = self.hparams["linear_key_head_dim"]560            head_v_dim = self.hparams["linear_value_head_dim"]561            num_v_per_k = num_v_heads // num_k_heads562 563            if ".in_proj_qkv." in name:564                # QKV weight: reorder only the V rows565                q_dim = head_k_dim * num_k_heads566                k_dim = head_k_dim * num_k_heads567                q = data_torch[:q_dim]568                k = data_torch[q_dim:q_dim + k_dim]569                v = data_torch[q_dim + k_dim:]570                v = self._reorder_v_heads(v, 0, num_k_heads, num_v_per_k, head_v_dim)571                data_torch = torch.cat([q, k, v], dim=0)572 573            elif ".in_proj_z." in name:574                # Z gate weight: reorder rows (num_v_heads * head_v_dim)575                data_torch = self._reorder_v_heads(data_torch, 0, num_k_heads, num_v_per_k, head_v_dim)576 577            elif ".in_proj_b." in name or ".in_proj_a." in name:578                # Beta/Alpha weight: reorder rows (num_v_heads, head_dim=1)579                data_torch = self._reorder_v_heads(data_torch, 0, num_k_heads, num_v_per_k, 1)580 581            elif ".A_log" in name or ".dt_bias" in name or ".dt_proj" in name:582                # A_log / dt_bias: 1D parameters with num_v_heads elements583                if data_torch.ndim == 1:584                    data_torch = self._reorder_v_heads(585                        data_torch.unsqueeze(-1), 0, num_k_heads, num_v_per_k, 1586                    ).squeeze(-1)587                else:588                    data_torch = self._reorder_v_heads(data_torch, -1, num_k_heads, num_v_per_k, 1)589 590            elif ".conv1d" in name:591                # Conv1d kernel: reorder only the V channel portion592                data = data_torch.squeeze()593                qk_channels = head_k_dim * num_k_heads * 2594                qk_part = data[:qk_channels]595                v_part = data[qk_channels:]596                v_part = self._reorder_v_heads(v_part, 0, num_k_heads, num_v_per_k, head_v_dim)597                data_torch = torch.cat([qk_part, v_part], dim=0)598 599            elif ".out_proj." in name:600                # Out projection weight: reorder columns (input dimension)601                data_torch = self._reorder_v_heads(data_torch, 1, num_k_heads, num_v_per_k, head_v_dim)602 603        yield from super().modify_tensors(data_torch, name, bid)604 605 606class _Qwen35MRopeMixin:607    # Qwen3.5 always applies interleaved MRoPE (see Qwen3_5RotaryEmbedding in transformers);608    # the upstream default mrope_section is [11, 11, 10] and llama.cpp's QWEN35 / QWEN35MOE609    # loaders treat qwen35.rope.dimension_sections as required, so make sure it is always610    # written even when a particular checkpoint omits the field in `rope_parameters`.611    _QWEN35_DEFAULT_MROPE_SECTION = [11, 11, 10, 0]612 613    gguf_writer: gguf.GGUFWriter614    rope_parameters: dict615 616    def set_gguf_parameters(self):617        super().set_gguf_parameters()  # ty: ignore[unresolved-attribute]618        if "mrope_section" not in self.rope_parameters:619            self.gguf_writer.add_rope_dimension_sections(self._QWEN35_DEFAULT_MROPE_SECTION)620 621 622@ModelBase.register("Qwen3_5ForConditionalGeneration", "Qwen3_5ForCausalLM")623class Qwen3_5TextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):624    model_arch = gguf.MODEL_ARCH.QWEN35625 626 627@ModelBase.register("Qwen3_5MoeForConditionalGeneration", "Qwen3_5MoeForCausalLM")628class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):629    model_arch = gguf.MODEL_ARCH.QWEN35MOE630 631 632@ModelBase.register("DFlashDraftModel")633class DFlashModel(Qwen3Model):634    model_arch = gguf.MODEL_ARCH.DFLASH635 636    def set_vocab(self):637        if self.target_model_dir is None:638            raise ValueError(639                "DFlash draft model requires --target-model-dir to be specified. "640                "Please provide the path to the target model directory containing the tokenizer."641            )642        logger.info(f"DFlash: Using tokenizer from target model: {self.target_model_dir}")643        original_dir = self.dir_model644        self.dir_model = self.target_model_dir645 646        # Reuse the target model's own vocab handler (e.g. Gemma-4 needs its647        # own tokenizer logic, not the Qwen default).648        from . import get_model_class649        with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f:650            target_arch = json.load(f)["architectures"][0]651        target_cls = get_model_class(target_arch)652 653        if target_cls is not type(self):654            target_cls.set_vocab(self)  # ty: ignore[unresolved-attribute]655        else:656            super().set_vocab()657 658        self.dir_model = original_dir659 660        mask_token_id = self.hparams.get("dflash_config", {}).get("mask_token_id")661        if mask_token_id is not None:662            self.gguf_writer.add_mask_token_id(mask_token_id)663 664    def set_gguf_parameters(self):665        super().set_gguf_parameters()666 667        block_size = self.hparams.get("block_size", 16)668        self.gguf_writer.add_block_size(block_size)669        dflash_config = self.hparams.get("dflash_config", {})670 671        target_layer_ids = dflash_config.get("target_layer_ids", [])672        if target_layer_ids:673            extract_layer_ids = [i + 1 for i in target_layer_ids]674            self.gguf_writer.add_target_layers(extract_layer_ids)675 676        use_sliding_window = self.hparams.get("use_sliding_window", False)677        sliding_window = self.hparams.get("sliding_window")678        layer_types = self.hparams.get("layer_types")679        if use_sliding_window and sliding_window and layer_types:680            is_swa = [lt == "sliding_attention" for lt in layer_types]681            self.gguf_writer.add_sliding_window(sliding_window)682            self.gguf_writer.add_sliding_window_pattern(is_swa)683 684    @classmethod685    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:686        name, gen = item687        if not name.startswith("model."):688            name = "model." + name689        return super().filter_tensors((name, gen))690 691 692@ModelBase.register("Qwen3DSparkModel")693class DSparkModel(DFlashModel):694    # DSpark = DFlash + a semi-autoregressive Markov head695    model_arch = gguf.MODEL_ARCH.DFLASH696 697    def __init__(self, *args, **kwargs):698        super().__init__(*args, **kwargs)699        # normalize the flat DeepSpec schema to DFlash's nested dflash_config700        self.hparams.setdefault("dflash_config", {701            k: self.hparams[k] for k in ("target_layer_ids", "mask_token_id") if k in self.hparams702        })703 704    @classmethod705    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:706        name, gen = item707        if name.endswith(("embed_tokens.weight", "lm_head.weight")):708            return None709        return super().filter_tensors((name, gen))710