Felipe97/llama-cpp-compiled

1from __future__ import annotations2 3import re4from pathlib import Path5from typing import Callable, Iterable, Iterator, TYPE_CHECKING6 7import numpy as np8import torch9 10if TYPE_CHECKING:11    from torch import Tensor12 13from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger14 15from .kimi_linear import KimiLinearModel16 17 18@ModelBase.register("KimiK3ForConditionalGeneration")19@ModelBase.example("moonshotai/Kimi-K3")20class KimiK3Model(TextModel):21    """22    Kimi-K3 text model (KimiLinearForCausalLM under a `language_model.` prefix).23 24    Shares the hybrid MLA + KDA skeleton with kimi-linear, but that converter25    cannot load it: K3 adds cross-layer attention residuals, a latent MoE, the26    situ activation, an MLA output gate and a full-rank KDA gate.27 28    The vision tower and mm_projector are skipped - text only for now.29    """30 31    model_arch = gguf.MODEL_ARCH.KIMI_K332 33    _experts: list[dict[str, Tensor]] | None = None34 35    # `<x>_res_norm.weight` and `<x>_res_proj.weight` are only used as their36    # elementwise product, so they are fused into one [n_embd] vector here.37    # they arrive apart, so buffer the first one and tag it with its kind.38    _res_parts: dict[str, tuple[str, Tensor]]39 40    # HF suffix -> (gguf tensor, per-layer?)41    _RES_FUSIONS = {42        "self_attention_res": (gguf.MODEL_TENSOR.ATTN_RES_SCORE, True),43        "mlp_res":            (gguf.MODEL_TENSOR.FFN_RES_SCORE,   True),44        "output_attn_res":    (gguf.MODEL_TENSOR.OUTPUT_RES_SCORE, False),45    }46 47    # compressed-tensors MXFP4. the `language_model.` prefix is still there, as48    # self.model_tensors is keyed by the raw checkpoint names49    _MXFP4_FORMAT = "mxfp4-pack-quantized"50    _MXFP4_EXPERT_RE = re.compile(51        r"^(?:language_model\.)?model\.layers\.(\d+)"52        r"\.block_sparse_moe\.experts\.(\d+)\.(w[123])\.weight_packed$"53    )54    _MXFP4_PROJ = {55        "w1": gguf.MODEL_TENSOR.FFN_GATE_EXP,56        "w2": gguf.MODEL_TENSOR.FFN_DOWN_EXP,57        "w3": gguf.MODEL_TENSOR.FFN_UP_EXP,58    }59 60    def __init__(self, *args, **kwargs):61        super().__init__(*args, **kwargs)62        self._res_parts = {}63 64    def set_vocab(self):65        # K3 has the same TikToken vocab as K2, so kimi-linear's vocab handling works.66        # borrowed, not inherited: the method only touches TextModel members, and K367        # shares none of kimi-linear's tensor layout.68        KimiLinearModel.set_vocab(self)  # ty: ignore[invalid-argument-type]69 70        # ...but that forces eos to the tokenizer's eos_id, which is [EOS], the71        # document terminator. K3's config says <|end_of_msg|>, the turn terminator;72        # with [EOS] the generation never stops at the end of a turn.73        if (eos := self.hparams.get("eos_token_id")) is not None:74            logger.info(f"restoring configured eos_token_id {eos} (kimi-linear forces the tokenizer's)")75            self.gguf_writer.add_eos_token_id(eos)76 77        # K3 renders chats in python (encoding_k3.py) and ships no jinja template,78        # so add the bundled one when the model has none79        if gguf.SpecialVocab(self.dir_model, load_merges=False).chat_template is None:80            template_path = Path(__file__).parent.parent / "models" / "templates" / "Kimi-K3.jinja"81            logger.info(f"gguf: model has no chat template, using {template_path.name}")82            self.gguf_writer.add_chat_template(template_path.read_text(encoding="utf-8"))83 84    #85    # compressed-tensors MXFP4 -> ggml MXFP486    #87 88    def _is_mxfp4_packed(self) -> bool:89        quant_config = self.hparams.get("quantization_config") or {}90        return (quant_config.get("quant_method") == "compressed-tensors"91                and quant_config.get("format") == self._MXFP4_FORMAT)92 93    def dequant_model(self):94        if not self._is_mxfp4_packed():95            return super().dequant_model()96 97        # skipping base.py's dequant is only safe if the experts are the only98        # quantized tensors, so check it99        stray = [n for n in self.model_tensors100                 if n.endswith(".weight_packed") and not self._MXFP4_EXPERT_RE.match(n)]101        if stray:102            raise NotImplementedError(103                f"{len(stray)} MXFP4 tensor(s) outside the routed experts, e.g. {stray[0]!r}; "104                "only the routed experts have a repack path"105            )106 107    def _mxfp4_expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]):108        """109        One stacked [n_expert, rows, cols] MXFP4 tensor, built lazily.110 111        gguf_writer holds every added tensor until the final write, so building112        this eagerly (like the DeepSeek-V4 path does) keeps all ~1.38 TB of113        experts in memory. lazy means only the tensor being written is resident.114        """115        # meta shapes, so this does not read any weights116        rows, packed_cols = loaders[0][0]().shape117        n_blocks = (packed_cols * 2) // 32118        byte_shape = (len(loaders), rows, n_blocks * 17)119 120        def load(fns: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]) -> np.ndarray:121            out = np.empty(byte_shape, dtype=np.uint8)122            for eid, (packed_fn, scale_fn) in enumerate(fns):123                out[eid] = self.repack_mxfp4_blocks(124                    LazyTorchTensor.to_eager(packed_fn()),125                    LazyTorchTensor.to_eager(scale_fn()),126                )127            return out128 129        # loaders goes through args, not the closure, so that `func` matches130        # LazyBase's single-argument shape131        return gguf.LazyNumpyTensor(132            meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(np.uint8, byte_shape),133            args=(loaders,),134            func=load,135        )136 137    def _write_mxfp4_experts(self) -> None:138        n_experts = self.hparams["num_experts"]139 140        # (bid, wid) -> {expert id: (packed name, scale name)}141        groups: dict[tuple[int, str], dict[int, tuple[str, str]]] = {}142        for name in self.model_tensors:143            m = self._MXFP4_EXPERT_RE.match(name)144            if m is None:145                continue146            bid, eid, wid = int(m.group(1)), int(m.group(2)), m.group(3)147            scale_name = name.removesuffix("_packed") + "_scale"148            if scale_name not in self.model_tensors:149                raise KeyError(f"missing {scale_name} for {name}")150            groups.setdefault((bid, wid), {})[eid] = (name, scale_name)151 152        consumed: list[str] = []153        for (bid, wid), experts in sorted(groups.items()):154            missing = [e for e in range(n_experts) if e not in experts]155            if missing:156                raise KeyError(157                    f"layer {bid} {wid}: {len(missing)} of {n_experts} experts missing, "158                    f"first is {missing[0]}"159                )160            if len(experts) != n_experts:161                raise KeyError(f"layer {bid} {wid}: {len(experts)} experts, expected {n_experts}")162 163            loaders = []164            for eid in range(n_experts):165                packed_name, scale_name = experts[eid]166                loaders.append((self.model_tensors[packed_name], self.model_tensors[scale_name]))167                consumed += [packed_name, scale_name]168 169            data = self._mxfp4_expert_tensor(loaders)170            new_name = self.format_tensor_name(self._MXFP4_PROJ[wid], bid)171            shape = gguf.quant_shape_from_byte_shape(data.shape, gguf.GGMLQuantizationType.MXFP4)172            logger.info(173                f"{new_name}: repacked {n_experts} experts to MXFP4, "174                f"shape = {{{', '.join(str(n) for n in reversed(shape))}}}"175            )176            self.gguf_writer.add_tensor(new_name, data, raw_dtype=gguf.GGMLQuantizationType.MXFP4)177 178        for name in consumed:179            del self.model_tensors[name]180 181    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:182        # not a generator on purpose: base.py chains this with get_tensors(), so the183        # tensors used here must be removed from model_tensors before that starts184        if self._is_mxfp4_packed():185            self._write_mxfp4_experts()186        return ()187 188    def get_tensors(self) -> Iterator[tuple[str, Tensor]]:189        for name, data in super().get_tensors():190            if name.startswith(("vision_tower.", "mm_projector.")):191                continue  # text only192            if name.startswith("language_model."):193                name = name[len("language_model."):]194            yield name, data195 196    def set_gguf_parameters(self):197        # MLA is served as MQA with a single large head, then decompressed198        self.hparams["num_key_value_heads"] = 1199 200        super().set_gguf_parameters()201        self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])202 203        linear_attn_config = self.hparams["linear_attn_config"]204 205        # n_head_kv == 0 marks a KDA (recurrent) layer. the layer lists are 1-indexed,206        # as KimiLinearConfig.is_kda_layer uses (layer_idx + 1)207        full_attn_layers = linear_attn_config["full_attn_layers"]208        n_kv_heads = [209            self.hparams["num_key_value_heads"] if (il + 1) in full_attn_layers else 0210            for il in range(self.hparams["num_hidden_layers"])211        ]212        assert len(n_kv_heads) == self.hparams["num_hidden_layers"]213        self.gguf_writer.add_head_count_kv(n_kv_heads)214 215        # --- KDA ---216        self.gguf_writer.add_ssm_conv_kernel(linear_attn_config["short_conv_kernel_size"])217        self.gguf_writer.add_kda_head_dim(linear_attn_config["head_dim"])218        if (lb := linear_attn_config.get("gate_lower_bound")) is not None:219            self.gguf_writer.add_kda_gate_lower_bound(lb)220 221        # --- MLA ---222        if (q_lora_rank := self.hparams.get("q_lora_rank")) is not None:223            self.gguf_writer.add_q_lora_rank(q_lora_rank)224        kv_lora_rank = self.hparams["kv_lora_rank"]225        self.gguf_writer.add_kv_lora_rank(kv_lora_rank)226 227        qk_nope_head_dim = self.hparams["qk_nope_head_dim"]228        qk_rope_head_dim = self.hparams["qk_rope_head_dim"]229        v_head_dim = self.hparams["v_head_dim"]230        # K3 is nope-only; qk_rope_head_dim still sizes the un-absorbed part of K231        assert self.hparams.get("mla_use_nope"), "K3 MLA is expected to be nope-only"232        self.gguf_writer.add_rope_dimension_count(qk_rope_head_dim)233        # MLA is served as MQA, so the cache holds the compressed latent234        self.gguf_writer.add_key_length(kv_lora_rank + qk_rope_head_dim)235        self.gguf_writer.add_value_length(kv_lora_rank)236        self.gguf_writer.add_key_length_mla(qk_nope_head_dim + qk_rope_head_dim)237        self.gguf_writer.add_value_length_mla(v_head_dim)238 239        # --- MoE ---240        self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"])241        self.gguf_writer.add_expert_shared_count(self.hparams["num_shared_experts"])242        self.gguf_writer.add_leading_dense_block_count(self.hparams["first_k_dense_replace"])243        self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"])244        self.gguf_writer.add_expert_weights_norm(self.hparams["moe_renormalize"])245        assert self.hparams["moe_router_activation_func"] == "sigmoid"246        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)247        # latent MoE: routed experts live in a down-projected space248        if (latent := self.hparams.get("routed_expert_hidden_size")) is not None:249            self.gguf_writer.add_expert_latent_length(latent)250 251        # --- situ activation ---252        assert self.hparams["hidden_act"] == "situ", \253            f"unexpected hidden_act {self.hparams['hidden_act']!r}"254        self.gguf_writer.add_activation_situ_beta(self.hparams["activation_situ_beta"])255        self.gguf_writer.add_activation_situ_linear_beta(self.hparams["activation_situ_linear_beta"])256 257        # --- cross-layer attention residuals ---258        self.gguf_writer.add_attn_res_block_size(self.hparams["attn_res_block_size"])259 260    def prepare_tensors(self):261        super().prepare_tensors()262        if self._experts is not None:263            leftover = [k for d in self._experts for k in d.keys()]264            if leftover:265                raise ValueError(f"Unprocessed experts: {leftover}")266        if self._res_parts:267            raise ValueError(f"Unpaired attention-residual tensors: {sorted(self._res_parts)}")268        if self._is_mxfp4_packed():269            # label the file for what it is; prepare_metadata runs after this270            self._is_mxfp4 = True271            self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE272 273    def _try_fuse_res(self, data_torch: Tensor, name: str, bid: int | None):274        """275        Pair <x>_res_norm.weight with <x>_res_proj.weight and emit their product.276 277        Returns None if this is not a res tensor, [] if buffered until its pair.278        """279        for prefix, (tensor_id, per_layer) in self._RES_FUSIONS.items():280            for kind in ("norm", "proj"):281                if not name.endswith(f"{prefix}_{kind}.weight"):282                    continue283                key = f"{prefix}.{bid}"284                other = self._res_parts.pop(key, None)285                if other is None:286                    self._res_parts[key] = (kind, data_torch)287                    return []288                other_kind, other_data = other289                assert other_kind != kind, f"duplicate {kind} for {key}"290                norm = data_torch if kind == "norm" else other_data291                proj = data_torch if kind == "proj" else other_data292                fused = norm.float().flatten() * proj.float().flatten()293                # ".weight" suffix matches the convention map_tensor_name applies294                new_name = (self.format_tensor_name(tensor_id, bid) if per_layer295                            else gguf.TENSOR_NAMES[tensor_id] + ".weight")296                logger.info(f"fused {prefix}_norm * {prefix}_proj -> {new_name}")297                return [(new_name, fused)]298        return None299 300    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:301        # --- cross-layer attention residuals: fuse norm * proj ---302        fused = self._try_fuse_res(data_torch, name, bid)303        if fused is not None:304            yield from fused305            return306 307        # --- KDA conv1d: HF [d_inner, 1, d_conv] -> ggml ne [d_conv, 1, d_inner, 1] ---308        # GGUF reverses the numpy shape on write, so target numpy (1, d_inner, 1, d_conv).309        # conv_step varies fastest in both layouts, so this is a pure reshape.310        if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")):311            if data_torch.ndim == 3:      # [d_inner, 1, d_conv]312                d_inner, _, d_conv = data_torch.shape313            elif data_torch.ndim == 2:    # [d_inner, d_conv]314                d_inner, d_conv = data_torch.shape315            else:316                raise ValueError(f"unexpected conv1d rank {data_torch.ndim} for {name}")317            data_torch = data_torch.reshape(1, d_inner, 1, d_conv)318 319        # -exp(A_log) is folded here so the graph does not have to320        if name.endswith(".A_log"):321            n_head = self.hparams["num_attention_heads"]322            data_torch = -torch.exp(data_torch.float()[:n_head])323 324        # dt_bias -> the name SSM_DT's mapping expects325        if name.endswith(".dt_bias"):326            name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"327 328        # --- g_proj is two different tensors sharing one HF name ---329        # KDA layers: full-rank gate, [d_inner, n_embd]  (replaces g_a/g_b)330        # MLA layers: output gate,    [n_head*v_head_dim, n_embd]331        # Name-based mapping cannot tell them apart, so resolve by layer type.332        if name.endswith(".self_attn.g_proj.weight"):333            assert bid is not None334            is_kda = (bid + 1) not in self.hparams["linear_attn_config"]["full_attn_layers"]335            tensor_id = gguf.MODEL_TENSOR.SSM_G if is_kda else gguf.MODEL_TENSOR.ATTN_GATE336            yield self.format_tensor_name(tensor_id, bid), data_torch337            return338 339        # --- routed experts: stack per-expert 2D weights into one 3D tensor ---340        if ".block_sparse_moe.experts." in name:341            n_experts = self.hparams["num_experts"]342            assert bid is not None343 344            if self._experts is None:345                self._experts = [{} for _ in range(self.block_count)]346            self._experts[bid][name] = data_torch347 348            if len(self._experts[bid]) < n_experts * 3:349                return350 351            # w1: gate, w2: down, w3: up352            for wid, tensor_id in (("w1", gguf.MODEL_TENSOR.FFN_GATE_EXP),353                                   ("w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP),354                                   ("w3", gguf.MODEL_TENSOR.FFN_UP_EXP)):355                datas = []356                for xid in range(n_experts):357                    ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight"358                    datas.append(self._experts[bid].pop(ename))359                stacked = torch.stack(datas, dim=0)360                yield from super().modify_tensors(stacked, self.format_tensor_name(tensor_id, bid), bid)361            return362 363        # --- MLA absorption: split kv_b into k_b (transposed) and v_b ---364        if name.endswith("kv_b_proj.weight"):365            n_head_kv = self.hparams["num_key_value_heads"]366            v_head_dim = self.hparams["v_head_dim"]367            qk_nope_head_dim = self.hparams["qk_nope_head_dim"]368            assert data_torch.shape[0] == n_head_kv * (v_head_dim + qk_nope_head_dim)369            kv_b = data_torch.view(n_head_kv, v_head_dim + qk_nope_head_dim, data_torch.shape[-1])370            k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1)371            k_b = k_b.transpose(1, 2)372            yield from super().modify_tensors(k_b, name.replace("kv_b_proj", "k_b_proj"), bid)373            yield from super().modify_tensors(v_b, name.replace("kv_b_proj", "v_b_proj"), bid)374            return375 376        yield from super().modify_tensors(data_torch, name, bid)377