From 901f9b3a7fb71701c8ab4bffec82239d3bbf4ae1 Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Mon, 27 Jul 2026 19:49:13 +0200 Subject: [PATCH 01/21] model: add Kimi-K3 text model Hybrid KDA (linear) + MLA (full) attention as in Kimi-Linear-48B, plus five things that architecture does not have: 1. cross-layer residual attention (attn_res_block_size) 2. latent MoE (routed experts run at n_expert_latent) 3. situ activation (replaces SwiGLU everywhere) 4. MLA output gate (sigmoid gate before o_proj) 5. full-rank KDA gate (single ssm_g instead of ssm_g_a/ssm_g_b) K3's text_config reports KimiLinearForCausalLM - the older 48B architecture - so get_model_architecture routes on the top-level name instead. The KDA decay gate has two forms, selected by linear_attn_config's gate_lower_bound. It is not a clamp: when set it swaps the activation entirely (fla/ops/kda/gate.py), from -exp(A_log)*softplus(x) to lower_bound*sigmoid(exp(A_log)*x). K3 sets it to -5.0; kimi-linear leaves it unset, so that path is unchanged. Cross-layer residuals reuse ggml_dsv4_hc_pre for the weighted sum. That op is CPU + CUDA only, so Metal/Vulkan will fall back per-node until those kernels exist. The routed experts ship as compressed-tensors "mxfp4-pack-quantized". That is bit-compatible with ggml's MXFP4 - same E2M1 code assignment, same E8M0 scale byte, only the nibble positions within a block differ - so they are repacked rather than dequantized, losslessly and without a ~5.5 TB bf16 round-trip. The repack is built lazily because gguf_writer holds every added tensor until the final write. DeepSeek-V4 was already doing the identical bit-shuffling, so it now shares the helper. Verified against Moonshot's own code path (transformers + fla's Triton KDA kernels) on a tiny model exercising every K3-specific feature. Final-position logits vs the fp32 reference: 6.7e-05 rel / corr 1.00000000 for both the chunked and the recurrent delta-net path. MXFP4 blocks dequantize to the source weights with 0.0e+00 error. Assisted-By: Claude Opus 5 (1M context) --- conversion/__init__.py | 1 + conversion/base.py | 48 ++- conversion/deepseek.py | 27 +- conversion/kimi_k3.py | 381 +++++++++++++++++++ gguf-py/gguf/constants.py | 77 +++- gguf-py/gguf/gguf_writer.py | 15 + gguf-py/gguf/tensor_mapping.py | 13 + src/llama-arch.cpp | 22 ++ src/llama-arch.h | 13 + src/llama-context.cpp | 1 + src/llama-graph.cpp | 15 + src/llama-graph.h | 1 + src/llama-hparams.h | 8 + src/llama-model.cpp | 3 + src/llama-model.h | 9 + src/models/kimi-k3.cpp | 645 +++++++++++++++++++++++++++++++++ src/models/models.h | 51 +++ 17 files changed, 1303 insertions(+), 27 deletions(-) create mode 100644 conversion/kimi_k3.py create mode 100644 src/models/kimi-k3.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index 695289b73ace..9b52fdb78f22 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -125,6 +125,7 @@ "JinaEmbeddingsV5Model": "bert", "KORMoForCausalLM": "qwen", "KimiK25ForConditionalGeneration": "deepseek", + "KimiK3ForConditionalGeneration": "kimi_k3", "KimiLinearForCausalLM": "kimi_linear", "KimiLinearModel": "kimi_linear", "KimiVLForConditionalGeneration": "deepseek", diff --git a/conversion/base.py b/conversion/base.py index 718d5394495e..6a9433c93da5 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -77,6 +77,49 @@ class ModelType(IntEnum): MMPROJ = 2 +def repack_mxfp4_blocks(packed: Tensor, scale: Tensor) -> np.ndarray: + """ + Repack 4-bit MX weights into ggml `block_mxfp4`. Lossless - this only moves + bits, it does not dequantize and requantize. + + Source (compressed-tensors "mxfp4-pack-quantized", and DeepSeek-V4's + equivalent weight/scale pair): + packed uint8 [rows, cols/2] two 4-bit codes per byte, element 2i in the + low nibble and 2i+1 in the high nibble + scale uint8 [rows, cols/32] one E8M0 biased exponent per 32-element group + + Destination, per 32-element group: one scale byte then 16 code bytes, where + byte j holds element j in the low nibble and element j+16 in the high nibble + (see dequantize_row_mxfp4 in ggml-quants.c). + + The 4-bit codes themselves need no remapping: both sides use sign in bit 3 + and a magnitude index into (0, .5, 1, 1.5, 2, 3, 4, 6), which is exactly + ggml's kvalues_mxfp4 order. ggml's kvalues are doubled and its scale is + halved (GGML_E8M0_TO_FP32_HALF), so the represented value is unchanged. + """ + p = packed.contiguous().view(torch.uint8) + s = scale.contiguous().view(torch.uint8) + + rows, packed_cols = p.shape + cols = packed_cols * 2 + if cols % 32 != 0: + raise ValueError(f"MXFP4 source row has {cols} values, expected a multiple of 32") + + n_blocks = cols // 32 + if tuple(s.shape) != (rows, n_blocks): + raise ValueError(f"MXFP4 scale shape {tuple(s.shape)} does not match {(rows, n_blocks)}") + + src = p.reshape(rows, n_blocks, 16) + lo = src & 0x0F # elements 0, 2, 4, ... + hi = (src >> 4) & 0x0F # elements 1, 3, 5, ... + + vals = torch.stack((lo, hi), dim=-1).reshape(rows, n_blocks, 32) + qs = vals[:, :, :16] | (vals[:, :, 16:] << 4) + + raw = torch.cat((s.unsqueeze(-1), qs.to(torch.uint8)), dim=-1) + return raw.reshape(rows, n_blocks * 17).cpu().numpy() + + class ModelBase: _model_classes: dict[ModelType, dict[str, type[ModelBase]]] = { ModelType.TEXT: {}, @@ -2661,7 +2704,10 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st # Step3-VL keeps text config under text_config but uses a custom top-level architecture. # For text conversion we route to a dedicated text-only class. # TODO: refactor this later to avoid adding exception here - if model_type == ModelType.TEXT and arch in ("StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration"): + # Kimi-K3's text_config reports "KimiLinearForCausalLM", which is the older + # Kimi-Linear-48B architecture and cannot load K3 (no attention residuals, + # latent MoE, situ, ...). Route on the top-level architecture instead. + if model_type == ModelType.TEXT and arch in ("StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"): return arch # if "architectures" is found in the sub-config, use that instead diff --git a/conversion/deepseek.py b/conversion/deepseek.py index 1846ca4010ec..e9d539a4ef61 100644 --- a/conversion/deepseek.py +++ b/conversion/deepseek.py @@ -12,7 +12,7 @@ if TYPE_CHECKING: from torch import Tensor -from .base import LazyTorchTensor, MmprojModel, ModelBase, TextModel, gguf, logger +from .base import LazyTorchTensor, MmprojModel, ModelBase, TextModel, gguf, logger, repack_mxfp4_blocks from .qwen import QwenModel @@ -709,30 +709,7 @@ def dequant_fp8_weight(weight: Tensor, scale: Tensor) -> Tensor: for name in tensors_to_remove: del self.model_tensors[name] - @staticmethod - def _pack_mxfp4_blocks(weight: Tensor, scale: Tensor) -> np.ndarray: - packed = weight.contiguous().view(torch.uint8) - scale_u8 = scale.contiguous().view(torch.uint8) - - out_features, packed_cols = packed.shape - logical_cols = packed_cols * 2 - if logical_cols % 32 != 0: - raise ValueError(f"MXFP4 source row has {logical_cols} values, expected a multiple of 32") - - n_blocks = logical_cols // 32 - if tuple(scale_u8.shape) != (out_features, n_blocks): - raise ValueError(f"MXFP4 scale shape {tuple(scale_u8.shape)} does not match {(out_features, n_blocks)}") - - src = packed.reshape(out_features, n_blocks, 16) - low = src & 0x0F - high = (src >> 4) & 0x0F - - # The safetensors bytes store adjacent values as low/high nibbles. - # ggml MXFP4 blocks store values 0..15 in low nibbles and 16..31 in high nibbles. - vals = torch.stack((low, high), dim=-1).reshape(out_features, n_blocks, 32) - qs = vals[:, :, :16] | (vals[:, :, 16:] << 4) - raw = torch.cat((scale_u8.unsqueeze(-1), qs.to(torch.uint8)), dim=-1) - return raw.reshape(out_features, n_blocks * 17).cpu().numpy() + _pack_mxfp4_blocks = staticmethod(repack_mxfp4_blocks) def _write_mxfp4_expert_tensor(self, bid: int, proj: str, tensor_key: gguf.MODEL_TENSOR) -> list[str]: n_experts = self.hparams["n_routed_experts"] diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py new file mode 100644 index 000000000000..58f7dbcac5c6 --- /dev/null +++ b/conversion/kimi_k3.py @@ -0,0 +1,381 @@ +from __future__ import annotations + +import re +from typing import Callable, Iterable, TYPE_CHECKING + +import numpy as np +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger, repack_mxfp4_blocks + +from .kimi_linear import KimiLinearModel + + +@ModelBase.register("KimiK3ForConditionalGeneration") +class KimiK3Model(TextModel): + """ + Kimi-K3 text model (KimiLinearForCausalLM under a `language_model.` prefix). + + Shares the hybrid MLA + KDA skeleton with Kimi-Linear-48B but is not + loadable by that converter: K3 adds cross-layer attention residuals, a + latent MoE, the situ activation, an MLA output gate and a full-rank KDA + gate, none of which exist in the older architecture. + + The vision tower and mm_projector are skipped - text only for now. + """ + + model_arch = gguf.MODEL_ARCH.KIMI_K3 + + _experts: list[dict[str, Tensor]] | None = None + + # `_res_norm.weight` and `_res_proj.weight` are only ever used as the + # elementwise product norm.weight * proj.weight (see _apply_attn_res in + # modeling_kimi_linear.py), so they are fused into a single [n_embd] vector + # at conversion time. They arrive as separate tensors, so buffer whichever + # comes first. + _res_parts: dict[str, Tensor] + + # HF suffix -> (gguf tensor, per-layer?) + _RES_FUSIONS = { + "self_attention_res": (gguf.MODEL_TENSOR.ATTN_RES_SCORE, True), + "mlp_res": (gguf.MODEL_TENSOR.FFN_RES_SCORE, True), + "output_attn_res": (gguf.MODEL_TENSOR.OUTPUT_RES_SCORE, False), + } + + # compressed-tensors MXFP4. The `language_model.` prefix is still present here: + # self.model_tensors is keyed by the raw checkpoint names, get_tensors() strips + # the prefix only on the way out. + _MXFP4_FORMAT = "mxfp4-pack-quantized" + _MXFP4_EXPERT_RE = re.compile( + r"^(?:language_model\.)?model\.layers\.(\d+)" + r"\.block_sparse_moe\.experts\.(\d+)\.(w[123])\.weight_packed$" + ) + _MXFP4_PROJ = { + "w1": gguf.MODEL_TENSOR.FFN_GATE_EXP, + "w2": gguf.MODEL_TENSOR.FFN_DOWN_EXP, + "w3": gguf.MODEL_TENSOR.FFN_UP_EXP, + } + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._res_parts = {} + + def set_vocab(self): + # K3 ships the same TikToken vocab as K2: its pre-tokenizer hashes to + # 81212dc7... which base.py already maps to "kimi-k2", so no new + # pre-tokenizer registration is needed. + KimiLinearModel.set_vocab(self) + + # ...but K2's converter ends by forcing eos to the tokenizer's own + # eos_id, which for K3 is 163585 = [EOS], the *document* terminator. + # K3's config and generation_config both say 163586 = <|end_of_msg|>, + # the chat turn terminator. Keeping [EOS] means generation never stops + # at the end of an assistant turn. Restore the configured value. + if (eos := self.hparams.get("eos_token_id")) is not None: + logger.info(f"restoring configured eos_token_id {eos} (kimi-linear forces the tokenizer's)") + self.gguf_writer.add_eos_token_id(eos) + + # + # compressed-tensors MXFP4 -> ggml MXFP4 + # + # The real checkpoint stores only the routed experts quantized (everything + # else is excluded by quantization_config["ignore"]), as a + # weight_packed/weight_scale pair per expert. Both sides are 4-bit E2M1 with + # a per-32 E8M0 scale, so this is a pure repack - see repack_mxfp4_blocks. + # + # Dequantizing instead would be catastrophic here: the routed experts are + # ~1.38 TB at 4 bits, so a bf16 round-trip would need ~5.5 TB of output. + # + + def _is_mxfp4_packed(self) -> bool: + quant_config = self.hparams.get("quantization_config") or {} + return (quant_config.get("quant_method") == "compressed-tensors" + and quant_config.get("format") == self._MXFP4_FORMAT) + + def dequant_model(self): + if not self._is_mxfp4_packed(): + return super().dequant_model() + + # Skipping base.py's dequant is only safe because the experts are the + # *only* quantized tensors. Verify that rather than assume it. + stray = [n for n in self.model_tensors + if n.endswith(".weight_packed") and not self._MXFP4_EXPERT_RE.match(n)] + if stray: + raise NotImplementedError( + f"{len(stray)} MXFP4 tensor(s) outside the routed experts, e.g. {stray[0]!r}; " + "only the routed experts have a repack path" + ) + + def _mxfp4_expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]): + """ + One stacked [n_expert, rows, cols] MXFP4 tensor, built lazily. + + Laziness is not an optimization here, it is the difference between + working and not: gguf_writer holds every added tensor until the final + write, so materializing this eagerly (as the DeepSeek-V4 and NVFP4 paths + do) would keep all ~1.38 TB of experts resident. Deferring it means only + the tensor currently being written is in memory, one expert at a time. + """ + # meta shapes, so this does not read any weights + rows, packed_cols = loaders[0][0]().shape + n_blocks = (packed_cols * 2) // 32 + byte_shape = (len(loaders), rows, n_blocks * 17) + + def load() -> np.ndarray: + out = np.empty(byte_shape, dtype=np.uint8) + for eid, (packed_fn, scale_fn) in enumerate(loaders): + out[eid] = repack_mxfp4_blocks( + LazyTorchTensor.to_eager(packed_fn()), + LazyTorchTensor.to_eager(scale_fn()), + ) + return out + + return gguf.LazyNumpyTensor( + meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(np.uint8, byte_shape), + func=load, + ) + + def _write_mxfp4_experts(self) -> None: + n_experts = self.hparams["num_experts"] + + # (bid, wid) -> {expert id: (packed name, scale name)} + groups: dict[tuple[int, str], dict[int, tuple[str, str]]] = {} + for name in self.model_tensors: + m = self._MXFP4_EXPERT_RE.match(name) + if m is None: + continue + bid, eid, wid = int(m.group(1)), int(m.group(2)), m.group(3) + scale_name = name.removesuffix("_packed") + "_scale" + if scale_name not in self.model_tensors: + raise KeyError(f"missing {scale_name} for {name}") + groups.setdefault((bid, wid), {})[eid] = (name, scale_name) + + consumed: list[str] = [] + for (bid, wid), experts in sorted(groups.items()): + missing = [e for e in range(n_experts) if e not in experts] + if missing: + raise KeyError( + f"layer {bid} {wid}: {len(missing)} of {n_experts} experts missing, " + f"first is {missing[0]}" + ) + if len(experts) != n_experts: + raise KeyError(f"layer {bid} {wid}: {len(experts)} experts, expected {n_experts}") + + loaders = [] + for eid in range(n_experts): + packed_name, scale_name = experts[eid] + loaders.append((self.model_tensors[packed_name], self.model_tensors[scale_name])) + consumed += [packed_name, scale_name] + + data = self._mxfp4_expert_tensor(loaders) + new_name = self.format_tensor_name(self._MXFP4_PROJ[wid], bid) + shape = gguf.quant_shape_from_byte_shape(data.shape, gguf.GGMLQuantizationType.MXFP4) + logger.info( + f"{new_name}: repacked {n_experts} experts to MXFP4, " + f"shape = {{{', '.join(str(n) for n in reversed(shape))}}}" + ) + self.gguf_writer.add_tensor(new_name, data, raw_dtype=gguf.GGMLQuantizationType.MXFP4) + + for name in consumed: + del self.model_tensors[name] + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + # Deliberately not a generator: base.py builds + # chain(generate_extra_tensors(), get_tensors()), so the tensors consumed + # here must be removed from model_tensors before get_tensors() starts. + if self._is_mxfp4_packed(): + self._write_mxfp4_experts() + return () + + def get_tensors(self) -> Iterable[tuple[str, Tensor]]: + for name, data in super().get_tensors(): + if name.startswith(("vision_tower.", "mm_projector.")): + continue # text only + if name.startswith("language_model."): + name = name[len("language_model."):] + yield name, data + + def set_gguf_parameters(self): + # MLA is served as MQA with a single large head, then decompressed + self.hparams["num_key_value_heads"] = 1 + + super().set_gguf_parameters() + self.gguf_writer.add_vocab_size(self.hparams["vocab_size"]) + + linear_attn_config = self.hparams["linear_attn_config"] + + # layer types: n_head_kv == 0 marks a KDA (recurrent) layer. + # KimiLinearConfig.is_kda_layer uses (layer_idx + 1) in kda_layers, so + # the lists are 1-indexed - an off-by-one here silently produces garbage. + full_attn_layers = linear_attn_config["full_attn_layers"] + n_kv_heads = [ + self.hparams["num_key_value_heads"] if (il + 1) in full_attn_layers else 0 + for il in range(self.hparams["num_hidden_layers"]) + ] + assert len(n_kv_heads) == self.hparams["num_hidden_layers"] + self.gguf_writer.add_head_count_kv(n_kv_heads) + + # --- KDA --- + self.gguf_writer.add_ssm_conv_kernel(linear_attn_config["short_conv_kernel_size"]) + self.gguf_writer.add_kda_head_dim(linear_attn_config["head_dim"]) + if (lb := linear_attn_config.get("gate_lower_bound")) is not None: + self.gguf_writer.add_kda_gate_lower_bound(lb) + + # --- MLA --- + if (q_lora_rank := self.hparams.get("q_lora_rank")) is not None: + self.gguf_writer.add_q_lora_rank(q_lora_rank) + kv_lora_rank = self.hparams["kv_lora_rank"] + self.gguf_writer.add_kv_lora_rank(kv_lora_rank) + + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + qk_rope_head_dim = self.hparams["qk_rope_head_dim"] + v_head_dim = self.hparams["v_head_dim"] + # K3 is nope-only; qk_rope_head_dim still sizes the un-absorbed part of K + assert self.hparams.get("mla_use_nope"), "K3 MLA is expected to be nope-only" + self.gguf_writer.add_rope_dimension_count(qk_rope_head_dim) + self.gguf_writer.add_key_length(kv_lora_rank + qk_rope_head_dim) + self.gguf_writer.add_key_length_mla(qk_nope_head_dim + qk_rope_head_dim) + self.gguf_writer.add_value_length_mla(v_head_dim) + + # --- MoE --- + self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"]) + self.gguf_writer.add_expert_shared_count(self.hparams["num_shared_experts"]) + self.gguf_writer.add_leading_dense_block_count(self.hparams["first_k_dense_replace"]) + self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(self.hparams["moe_renormalize"]) + assert self.hparams["moe_router_activation_func"] == "sigmoid" + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) + # latent MoE: routed experts live in a down-projected space + if (latent := self.hparams.get("routed_expert_hidden_size")) is not None: + self.gguf_writer.add_expert_latent_length(latent) + + # --- situ activation --- + assert self.hparams["hidden_act"] == "situ", \ + f"unexpected hidden_act {self.hparams['hidden_act']!r}" + self.gguf_writer.add_activation_situ_beta(self.hparams["activation_situ_beta"]) + self.gguf_writer.add_activation_situ_linear_beta(self.hparams["activation_situ_linear_beta"]) + + # --- cross-layer attention residuals --- + self.gguf_writer.add_attn_res_block_size(self.hparams["attn_res_block_size"]) + + def prepare_tensors(self): + super().prepare_tensors() + if self._experts is not None: + leftover = [k for d in self._experts for k in d.keys()] + if leftover: + raise ValueError(f"Unprocessed experts: {leftover}") + if self._res_parts: + raise ValueError(f"Unpaired attention-residual tensors: {sorted(self._res_parts)}") + if self._is_mxfp4_packed(): + # label the file for what it is; prepare_metadata runs after this + self._is_mxfp4 = True + self.ftype = gguf.LlamaFileType.MOSTLY_MXFP4_MOE + + def _try_fuse_res(self, data_torch: Tensor, name: str, bid: int | None): + """ + Pair _res_norm.weight with _res_proj.weight and emit their product. + + proj is [1, n_embd]; norm is [n_embd]. _apply_attn_res only ever uses + norm.weight * proj.weight.squeeze(0), so one vector is enough. + Returns None if this is not a res tensor, [] if buffered pending its pair. + """ + for prefix, (tensor_id, per_layer) in self._RES_FUSIONS.items(): + for kind in ("norm", "proj"): + if not name.endswith(f"{prefix}_{kind}.weight"): + continue + key = f"{prefix}.{bid}" + other = self._res_parts.pop(key, None) + if other is None: + self._res_parts[key] = (kind, data_torch) + return [] + other_kind, other_data = other + assert other_kind != kind, f"duplicate {kind} for {key}" + norm = data_torch if kind == "norm" else other_data + proj = data_torch if kind == "proj" else other_data + fused = norm.float().flatten() * proj.float().flatten() + # ".weight" suffix matches the convention map_tensor_name applies + new_name = (self.format_tensor_name(tensor_id, bid) if per_layer + else gguf.TENSOR_NAMES[tensor_id] + ".weight") + logger.info(f"fused {prefix}_norm * {prefix}_proj -> {new_name}") + return [(new_name, fused)] + return None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # --- cross-layer attention residuals: fuse norm * proj --- + fused = self._try_fuse_res(data_torch, name, bid) + if fused is not None: + yield from fused + return + + # --- KDA conv1d: HF [d_inner, 1, d_conv] -> ggml ne [d_conv, 1, d_inner, 1] --- + # GGUF reverses the numpy shape on write, so target numpy (1, d_inner, 1, d_conv). + # Both layouts have conv_step varying fastest, so this is a pure reshape. + if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")): + if data_torch.ndim == 3: # [d_inner, 1, d_conv] + d_inner, _, d_conv = data_torch.shape + elif data_torch.ndim == 2: # [d_inner, d_conv] + d_inner, d_conv = data_torch.shape + else: + raise ValueError(f"unexpected conv1d rank {data_torch.ndim} for {name}") + data_torch = data_torch.reshape(1, d_inner, 1, d_conv) + + # -exp(A_log) is folded here so the graph does not have to + if name.endswith(".A_log"): + data_torch = -torch.exp(data_torch) + + # dt_bias -> the name SSM_DT's mapping expects + if name.endswith(".dt_bias"): + name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias" + + # --- g_proj is two different tensors sharing one HF name --- + # KDA layers: full-rank gate, [d_inner, n_embd] (replaces g_a/g_b) + # MLA layers: output gate, [n_head*v_head_dim, n_embd] + # Name-based mapping cannot tell them apart, so resolve by layer type. + if name.endswith(".self_attn.g_proj.weight"): + assert bid is not None + is_kda = (bid + 1) not in self.hparams["linear_attn_config"]["full_attn_layers"] + tensor_id = gguf.MODEL_TENSOR.SSM_G if is_kda else gguf.MODEL_TENSOR.ATTN_GATE + yield self.format_tensor_name(tensor_id, bid), data_torch + return + + # --- routed experts: stack per-expert 2D weights into one 3D tensor --- + if ".block_sparse_moe.experts." in name: + n_experts = self.hparams["num_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) < n_experts * 3: + return + + # w1: gate, w2: down, w3: up + for wid, tensor_id in (("w1", gguf.MODEL_TENSOR.FFN_GATE_EXP), + ("w2", gguf.MODEL_TENSOR.FFN_DOWN_EXP), + ("w3", gguf.MODEL_TENSOR.FFN_UP_EXP)): + datas = [] + for xid in range(n_experts): + ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight" + datas.append(self._experts[bid].pop(ename)) + stacked = torch.stack(datas, dim=0) + yield from super().modify_tensors(stacked, self.format_tensor_name(tensor_id, bid), bid) + return + + # --- MLA absorption: split kv_b into k_b (transposed) and v_b --- + if name.endswith("kv_b_proj.weight"): + n_head_kv = self.hparams["num_key_value_heads"] + v_head_dim = self.hparams["v_head_dim"] + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + assert data_torch.shape[0] == n_head_kv * (v_head_dim + qk_nope_head_dim) + kv_b = data_torch.view(n_head_kv, v_head_dim + qk_nope_head_dim, data_torch.shape[-1]) + k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1) + k_b = k_b.transpose(1, 2) + yield from super().modify_tensors(k_b, name.replace("kv_b_proj", "k_b_proj"), bid) + yield from super().modify_tensors(v_b, name.replace("kv_b_proj", "v_b_proj"), bid) + return + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 98c4fa169143..814c3ef11af9 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -124,6 +124,7 @@ class LLM: EXPERT_WEIGHTS_NORM = "{arch}.expert_weights_norm" EXPERT_GATING_FUNC = "{arch}.expert_gating_func" EXPERT_GROUP_SCALE = "{arch}.expert_group_scale" + EXPERT_LATENT_LENGTH = "{arch}.expert_latent_length" EXPERTS_PER_GROUP = "{arch}.experts_per_group" MOE_EVERY_N_LAYERS = "{arch}.moe_every_n_layers" MOE_LATENT_SIZE = "{arch}.moe_latent_size" @@ -238,6 +239,13 @@ class Rope: SCALING_YARN_BETA_FAST = "{arch}.rope.scaling.yarn_beta_fast" SCALING_YARN_BETA_SLOW = "{arch}.rope.scaling.yarn_beta_slow" + class Activation: + SITU_BETA = "{arch}.activation.situ_beta" + SITU_LINEAR_BETA = "{arch}.activation.situ_linear_beta" + + class AttnRes: + BLOCK_SIZE = "{arch}.attn_res.block_size" + class Split: LLM_KV_SPLIT_NO = "split.no" LLM_KV_SPLIT_COUNT = "split.count" @@ -252,7 +260,8 @@ class SSM: DT_B_C_RMS = "{arch}.ssm.dt_b_c_rms" class KDA: - HEAD_DIM = "{arch}.kda.head_dim" + HEAD_DIM = "{arch}.kda.head_dim" + GATE_LOWER_BOUND = "{arch}.kda.gate_lower_bound" class WKV: HEAD_SIZE = "{arch}.wkv.head_size" @@ -579,6 +588,7 @@ class MODEL_ARCH(IntEnum): LLAMA_EMBED = auto() MAINCODER = auto() KIMI_LINEAR = auto() + KIMI_K3 = auto() TALKIE = auto() MELLUM = auto() NANBEIGE = auto() @@ -697,6 +707,13 @@ class MODEL_TENSOR(IntEnum): SSM_BETA = auto() # Kimi Linear qwen3.5 SSM_G_A = auto() # Kimi Linear SSM_G_B = auto() # Kimi Linear + SSM_G = auto() # Kimi K3 (full-rank KDA gate, replaces SSM_G_A/SSM_G_B) + ATTN_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, pre-attention) + FFN_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, pre-FFN) + OUTPUT_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, final) + FFN_ROUTED_DOWN = auto() # Kimi K3 (latent MoE: hidden -> latent) + FFN_ROUTED_UP = auto() # Kimi K3 (latent MoE: latent -> hidden) + FFN_ROUTED_NORM = auto() # Kimi K3 (latent MoE: norm on expert output) TIME_MIX_W0 = auto() TIME_MIX_W1 = auto() TIME_MIX_W2 = auto() @@ -1286,6 +1303,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.LLAMA_EMBED: "llama-embed", MODEL_ARCH.MAINCODER: "maincoder", MODEL_ARCH.KIMI_LINEAR: "kimi-linear", + MODEL_ARCH.KIMI_K3: "kimi-k3", MODEL_ARCH.TALKIE: "talkie", MODEL_ARCH.MELLUM: "mellum", MODEL_ARCH.NANBEIGE: "nanbeige", @@ -1402,6 +1420,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.SSM_BETA: "blk.{bid}.ssm_beta", # Kimi Linear qwen3.5 MODEL_TENSOR.SSM_G_A: "blk.{bid}.ssm_g_a", # Kimi Linear MODEL_TENSOR.SSM_G_B: "blk.{bid}.ssm_g_b", # Kimi Linear + MODEL_TENSOR.SSM_G: "blk.{bid}.ssm_g", # Kimi K3 + MODEL_TENSOR.ATTN_RES_SCORE: "blk.{bid}.attn_res_score", # Kimi K3 + MODEL_TENSOR.FFN_RES_SCORE: "blk.{bid}.ffn_res_score", # Kimi K3 + MODEL_TENSOR.OUTPUT_RES_SCORE: "output_res_score", # Kimi K3 + MODEL_TENSOR.FFN_ROUTED_DOWN: "blk.{bid}.ffn_routed_down", # Kimi K3 + MODEL_TENSOR.FFN_ROUTED_UP: "blk.{bid}.ffn_routed_up", # Kimi K3 + MODEL_TENSOR.FFN_ROUTED_NORM: "blk.{bid}.ffn_routed_norm", # Kimi K3 MODEL_TENSOR.TIME_MIX_W0: "blk.{bid}.time_mix_w0", MODEL_TENSOR.TIME_MIX_W1: "blk.{bid}.time_mix_w1", MODEL_TENSOR.TIME_MIX_W2: "blk.{bid}.time_mix_w2", @@ -4940,6 +4965,56 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_DOWN_SHEXP, MODEL_TENSOR.FFN_UP_SHEXP, ], + MODEL_ARCH.KIMI_K3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_RES_SCORE, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_RES_SCORE, + MODEL_TENSOR.FFN_RES_SCORE, + # MLA (full-attention layers) + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + # KDA (linear-attention layers) + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F_A, + MODEL_TENSOR.SSM_F_B, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_G, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + # FFN + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_ROUTED_DOWN, + MODEL_TENSOR.FFN_ROUTED_UP, + MODEL_TENSOR.FFN_ROUTED_NORM, + ], MODEL_ARCH.TALKIE: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 05f86396dc0a..0aee5168c023 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1103,6 +1103,21 @@ def add_ssm_group_count(self, value: int) -> None: def add_ssm_dt_b_c_rms(self, value: bool) -> None: self.add_bool(Keys.SSM.DT_B_C_RMS.format(arch=self.arch), value) + def add_kda_gate_lower_bound(self, value: float) -> None: + self.add_float32(Keys.KDA.GATE_LOWER_BOUND.format(arch=self.arch), value) + + def add_expert_latent_length(self, value: int) -> None: + self.add_uint32(Keys.LLM.EXPERT_LATENT_LENGTH.format(arch=self.arch), value) + + def add_activation_situ_beta(self, value: float) -> None: + self.add_float32(Keys.Activation.SITU_BETA.format(arch=self.arch), value) + + def add_activation_situ_linear_beta(self, value: float) -> None: + self.add_float32(Keys.Activation.SITU_LINEAR_BETA.format(arch=self.arch), value) + + def add_attn_res_block_size(self, value: int) -> None: + self.add_uint32(Keys.AttnRes.BLOCK_SIZE.format(arch=self.arch), value) + def add_kda_head_dim(self, value: int) -> None: self.add_uint32(Keys.KDA.HEAD_DIM.format(arch=self.arch), value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 79d270ab8fe9..61256e17df1a 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -910,6 +910,19 @@ class TensorNameMap: "model.layers.{bid}.linear_attn.in_proj_b", # qwen3.5 "model.layers.{bid}.self_attn.b_proj", # Kimi Linear ), + # Kimi K3 latent MoE: routed experts operate in a down-projected space + MODEL_TENSOR.FFN_ROUTED_DOWN: ( + "model.layers.{bid}.block_sparse_moe.routed_expert_down_proj", + ), + + MODEL_TENSOR.FFN_ROUTED_UP: ( + "model.layers.{bid}.block_sparse_moe.routed_expert_up_proj", + ), + + MODEL_TENSOR.FFN_ROUTED_NORM: ( + "model.layers.{bid}.block_sparse_moe.routed_expert_norm", + ), + MODEL_TENSOR.SSM_G_A: ( "model.layers.{bid}.self_attn.g_a_proj", ), diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 8ed9391d7c74..5e764047d1b9 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -143,6 +143,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_LLAMA_EMBED, "llama-embed" }, { LLM_ARCH_MAINCODER, "maincoder" }, { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, + { LLM_ARCH_KIMI_K3, "kimi-k3" }, { LLM_ARCH_TALKIE, "talkie" }, { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_NANBEIGE, "nanbeige" }, @@ -186,6 +187,9 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_FEATURES_LENGTH, "%s.features_length" }, { LLM_KV_BLOCK_COUNT, "%s.block_count" }, { LLM_KV_LEADING_DENSE_BLOCK_COUNT, "%s.leading_dense_block_count" }, + { LLM_KV_ATTN_RES_BLOCK_SIZE, "%s.attn_res.block_size" }, + { LLM_KV_ACTIVATION_SITU_BETA, "%s.activation.situ_beta" }, + { LLM_KV_ACTIVATION_SITU_LINEAR_BETA, "%s.activation.situ_linear_beta" }, { LLM_KV_FEED_FORWARD_LENGTH, "%s.feed_forward_length" }, { LLM_KV_EXPERT_FEED_FORWARD_LENGTH, "%s.expert_feed_forward_length" }, { LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, "%s.expert_shared_feed_forward_length" }, @@ -201,6 +205,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_EXPERT_GROUP_USED_COUNT, "%s.expert_group_used_count" }, { LLM_KV_EXPERT_WEIGHTS_SCALE, "%s.expert_weights_scale" }, { LLM_KV_EXPERT_WEIGHTS_NORM, "%s.expert_weights_norm" }, + { LLM_KV_EXPERT_LATENT_LENGTH, "%s.expert_latent_length" }, { LLM_KV_EXPERT_GATING_FUNC, "%s.expert_gating_func" }, { LLM_KV_EXPERT_GROUP_SCALE, "%s.expert_group_scale" }, { LLM_KV_EXPERTS_PER_GROUP, "%s.experts_per_group" }, @@ -312,6 +317,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_SSM_DT_B_C_RMS, "%s.ssm.dt_b_c_rms" }, { LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" }, + { LLM_KV_KDA_GATE_LOWER_BOUND, "%s.kda.gate_lower_bound" }, { LLM_KV_WKV_HEAD_SIZE, "%s.wkv.head_size" }, @@ -462,6 +468,13 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_SSM_F_B, "blk.%d.ssm_f_b" }, { LLM_TENSOR_SSM_BETA, "blk.%d.ssm_beta" }, { LLM_TENSOR_SSM_G_A, "blk.%d.ssm_g_a" }, + { LLM_TENSOR_SSM_G, "blk.%d.ssm_g" }, + { LLM_TENSOR_ATTN_RES_SCORE, "blk.%d.attn_res_score" }, + { LLM_TENSOR_FFN_RES_SCORE, "blk.%d.ffn_res_score" }, + { LLM_TENSOR_OUTPUT_RES_SCORE, "output_res_score" }, + { LLM_TENSOR_FFN_ROUTED_DOWN, "blk.%d.ffn_routed_down" }, + { LLM_TENSOR_FFN_ROUTED_UP, "blk.%d.ffn_routed_up" }, + { LLM_TENSOR_FFN_ROUTED_NORM, "blk.%d.ffn_routed_norm" }, { LLM_TENSOR_SSM_G_B, "blk.%d.ssm_g_b" }, { LLM_TENSOR_SSM_NORM, "blk.%d.ssm_norm" }, { LLM_TENSOR_ATTN_Q_A_NORM, "blk.%d.attn_q_a_norm" }, @@ -755,6 +768,13 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_SSM_F_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_BETA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SSM_G_A, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_SSM_G, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_RES_SCORE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_RES_SCORE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_OUTPUT_RES_SCORE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_ROUTED_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_ROUTED_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_ROUTED_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_SSM_G_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_TIME_MIX_LERP_X, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_TIME_MIX_LN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, @@ -975,6 +995,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_KIMI_K3: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: @@ -1035,6 +1056,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_KIMI_K3: case LLM_ARCH_QWEN3TTS: return false; default: diff --git a/src/llama-arch.h b/src/llama-arch.h index 18d9de186f75..d43d6fedef0f 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -145,6 +145,7 @@ enum llm_arch { LLM_ARCH_LLAMA_EMBED, LLM_ARCH_MAINCODER, LLM_ARCH_KIMI_LINEAR, + LLM_ARCH_KIMI_K3, LLM_ARCH_TALKIE, LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, @@ -191,6 +192,9 @@ enum llm_kv { LLM_KV_FEATURES_LENGTH, LLM_KV_BLOCK_COUNT, LLM_KV_LEADING_DENSE_BLOCK_COUNT, + LLM_KV_ATTN_RES_BLOCK_SIZE, + LLM_KV_ACTIVATION_SITU_BETA, + LLM_KV_ACTIVATION_SITU_LINEAR_BETA, LLM_KV_FEED_FORWARD_LENGTH, LLM_KV_EXPERT_FEED_FORWARD_LENGTH, LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, @@ -206,6 +210,7 @@ enum llm_kv { LLM_KV_EXPERT_GROUP_USED_COUNT, LLM_KV_EXPERT_WEIGHTS_SCALE, LLM_KV_EXPERT_WEIGHTS_NORM, + LLM_KV_EXPERT_LATENT_LENGTH, LLM_KV_EXPERT_GATING_FUNC, LLM_KV_EXPERT_GROUP_SCALE, LLM_KV_EXPERTS_PER_GROUP, @@ -317,6 +322,7 @@ enum llm_kv { LLM_KV_SSM_DT_B_C_RMS, LLM_KV_KDA_HEAD_DIM, + LLM_KV_KDA_GATE_LOWER_BOUND, LLM_KV_WKV_HEAD_SIZE, @@ -491,6 +497,13 @@ enum llm_tensor { LLM_TENSOR_SSM_BETA, // kimi: beta mixing coefficient and qwen3.5 LLM_TENSOR_SSM_G_A, // kimi: output gate projection A LLM_TENSOR_SSM_G_B, // kimi: output gate projection B + LLM_TENSOR_SSM_G, // kimi-k3: full-rank KDA gate + LLM_TENSOR_ATTN_RES_SCORE, // kimi-k3: fused res_norm*res_proj (pre-attn) + LLM_TENSOR_FFN_RES_SCORE, // kimi-k3: fused res_norm*res_proj (pre-ffn) + LLM_TENSOR_OUTPUT_RES_SCORE, // kimi-k3: fused res_norm*res_proj (final) + LLM_TENSOR_FFN_ROUTED_DOWN, // kimi-k3: latent MoE down + LLM_TENSOR_FFN_ROUTED_UP, // kimi-k3: latent MoE up + LLM_TENSOR_FFN_ROUTED_NORM, // kimi-k3: latent MoE norm LLM_TENSOR_TIME_MIX_W0, LLM_TENSOR_TIME_MIX_W1, LLM_TENSOR_TIME_MIX_W2, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0de3a68d1cb0..af0c5f37e1f1 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2295,6 +2295,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { uint32_t res; if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || + model.arch == LLM_ARCH_KIMI_K3 || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DEEPSEEK4 || diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 55d858024630..560233c9f353 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2174,6 +2174,21 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cur = ggml_silu(ctx0, cur); cb(cur, "ffn_moe_silu", il); } break; + case LLM_FFN_SITU: + { + // situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * lb*tanh(up/lb) + GGML_ASSERT(has_gate); + const float beta = hparams.situ_beta; + const float lb = hparams.situ_linear_beta; + + ggml_tensor * act = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, cur, 1.0f/beta)), beta); + act = ggml_mul(ctx0, act, ggml_sigmoid(ctx0, cur)); + if (lb > 0.0f) { + up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/lb)), lb); + } + cur = ggml_mul(ctx0, act, up); + cb(cur, "ffn_moe_situ", il); + } break; case LLM_FFN_GELU: if (has_gate) { cur = ggml_geglu_split(ctx0, cur, up); diff --git a/src/llama-graph.h b/src/llama-graph.h index 75bc0fe80dbc..af38e1cca519 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -59,6 +59,7 @@ enum llm_ffn_op_type : int { LLM_FFN_GEGLU, LLM_FFN_REGLU, LLM_FFN_SWIGLU_OAI_MOE, + LLM_FFN_SITU, // kimi-k3 (appended: do not renumber existing values) }; enum llm_ffn_gate_type { diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 57de808242bd..f6d9bf9c9494 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -4,6 +4,7 @@ #include #include +#include // bump if necessary #define LLAMA_MAX_LAYERS 512 @@ -167,6 +168,13 @@ struct llama_hparams { // for Kimi Linear KDA uint32_t n_embd_head_kda = 0; + // kimi-k3 + uint32_t n_expert_latent = 0; // routed_expert_hidden_size (0 = experts run at n_embd) + uint32_t attn_res_block_size = 0; // 0 = no cross-layer attention residuals + float kda_gate_lower_bound = -INFINITY; + float situ_beta = 1.0f; + float situ_linear_beta = 0.0f; // 0 = no linear-beta transform on the up branch + bool ssm_dt_b_c_rms = false; float f_clamp_kqv = 0.0f; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 0e27cb41713b..5bcf15d4827d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -320,6 +320,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_mimo2(params); case LLM_ARCH_KIMI_LINEAR: return new llama_model_kimi_linear(params); + case LLM_ARCH_KIMI_K3: + return new llama_model_kimi_k3(params); case LLM_ARCH_STEP35: return new llama_model_step35(params); default: @@ -2596,6 +2598,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_KIMI_K3: return LLAMA_ROPE_TYPE_NONE; // use what we call a normal RoPE, operating on pairs of consecutive head values diff --git a/src/llama-model.h b/src/llama-model.h index 341cb66fbafe..923862ca6d07 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -528,6 +528,14 @@ struct llama_layer { struct ggml_tensor * ssm_g_b = nullptr; struct ggml_tensor * ssm_o_norm = nullptr; + // kimi-k3 + struct ggml_tensor * ssm_g = nullptr; // full-rank KDA gate (replaces ssm_g_a/ssm_g_b) + struct ggml_tensor * attn_res_score = nullptr; // fused res_norm*res_proj, pre-attention + struct ggml_tensor * ffn_res_score = nullptr; // fused res_norm*res_proj, pre-FFN + struct ggml_tensor * ffn_routed_down = nullptr; // latent MoE: n_embd -> n_expert_latent + struct ggml_tensor * ffn_routed_up = nullptr; // latent MoE: n_expert_latent -> n_embd + struct ggml_tensor * ffn_routed_norm = nullptr; + // DSA (deepseek sparse attention) struct ggml_tensor * indexer_k_norm = nullptr; struct ggml_tensor * indexer_k_norm_b = nullptr; @@ -587,6 +595,7 @@ struct llama_model { struct ggml_tensor * tok_norm_b = nullptr; struct ggml_tensor * output_norm = nullptr; + struct ggml_tensor * output_res_score = nullptr; // kimi-k3: final cross-layer residual mix struct ggml_tensor * output_norm_b = nullptr; struct ggml_tensor * output = nullptr; struct ggml_tensor * output_b = nullptr; diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp new file mode 100644 index 000000000000..2043a6a83498 --- /dev/null +++ b/src/models/kimi-k3.cpp @@ -0,0 +1,645 @@ +#include "models.h" +#include "llama-memory-recurrent.h" + +// +// Kimi-K3 text model. +// +// Hybrid KDA (linear) + MLA (full) attention, as in Kimi-Linear-48B, plus five +// things that architecture does not have: +// +// 1. cross-layer residual attention (attn_res_block_size) +// 2. latent MoE (routed experts run at n_expert_latent) +// 3. situ activation (replaces SwiGLU everywhere) +// 4. MLA output gate (sigmoid gate before o_proj) +// 5. full-rank KDA gate (single ssm_g instead of ssm_g_a/ssm_g_b) +// + +void llama_model_kimi_k3::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q, false); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false); + + // n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-linear + for (uint32_t i = 0; i < hparams.n_layer(); ++i) { + hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0; + } + + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent, false); + + ml.get_key(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size, false); + ml.get_key(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta, false); + ml.get_key(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta, false); + + switch (hparams.n_layer()) { + case 93: type = LLM_TYPE_UNKNOWN; break; // Kimi-K3 + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_kimi_k3::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + if (hparams.attn_res_block_size > 0) { + output_res_score = create_tensor(tn(LLM_TENSOR_OUTPUT_RES_SCORE, "weight"), {n_embd}, 0); + } + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + if (hparams.attn_res_block_size > 0) { + layer.attn_res_score = create_tensor(tn(LLM_TENSOR_ATTN_RES_SCORE, "weight", i), {n_embd}, 0); + layer.ffn_res_score = create_tensor(tn(LLM_TENSOR_FFN_RES_SCORE, "weight", i), {n_embd}, 0); + } + + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = head_dim * n_head; + + if (hparams.is_recr(i)) { + // conv1d may be stored 4D [d_conv, 1, d_inner, 1] or 3D (quantization drops the trailing 1) + auto conv = [&](llm_tensor tid) { + ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED); + return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0); + }; + layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q); + layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K); + layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V); + + create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0); + + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, head_dim}, 0); + layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, 0); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); + + // K3's A_log is a plain 1-D [n_head] parameter (kimi-linear's is padded); + // accept the padded forms too so both layouts load. -exp() is folded at + // conversion time. Only the element count matters - the graph reshapes it. + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, TENSOR_NOT_REQUIRED); + if (!layer.ssm_a) { + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_head, 1, 1}, TENSOR_NOT_REQUIRED); + } + if (!layer.ssm_a) { + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_head}, 0); + } + + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0); + + // K3 uses a single full-rank gate instead of kimi-linear's g_a/g_b pair + layer.ssm_g = create_tensor(tn(LLM_TENSOR_SSM_G, "weight", i), {n_embd, d_inner}, 0); + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0); + } else { + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_embd_head_k = hparams.n_embd_head_k_mla(); + const int64_t n_embd_head_v = hparams.n_embd_head_v_mla(); + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_nope_head_dim = n_embd_head_k - qk_rope_head_dim; + + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, TENSOR_NOT_REQUIRED); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0); + + if (layer.attn_q_a_norm) { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, 0); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k}, 0); + } else { + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head * n_embd_head_k}, 0); + } + + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + qk_rope_head_dim}, 0); + layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i), + {kv_lora_rank, n_head * (qk_nope_head_dim + n_embd_head_v)}, + TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); + if (!layer.wkv_b) { + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {qk_nope_head_dim, kv_lora_rank, n_head}, 0); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, 0); + } + + // K3: sigmoid output gate applied to the attention output before o_proj + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head * n_embd_head_v}, TENSOR_NOT_REQUIRED); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, 0); + } + + if (i < (int) hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } else { + const int64_t n_ff_exp = hparams.n_ff_exp; + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + + // routed experts live in the latent space + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd_latent, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd_latent, n_ff_exp, n_expert}, 0); + + if (hparams.n_expert_latent > 0) { + layer.ffn_routed_down = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_DOWN, "weight", i), {n_embd, n_embd_latent}, 0); + layer.ffn_routed_up = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_UP, "weight", i), {n_embd_latent, n_embd}, 0); + layer.ffn_routed_norm = create_tensor(tn(LLM_TENSOR_FFN_ROUTED_NORM, "weight", i), {n_embd_latent}, TENSOR_NOT_REQUIRED); + } + + // shared experts stay at n_embd, width = moe_intermediate_size * n_expert_shared + const int64_t n_ff_shexp = n_ff_exp * (hparams.n_expert_shared > 0 ? hparams.n_expert_shared : 1); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, TENSOR_NOT_REQUIRED); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, TENSOR_NOT_REQUIRED); + } + } +} + +std::unique_ptr llama_model_kimi_k3::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +// +// situ activation: +// situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) +// linear_beta <= 0 disables the transform on the up branch. +// +static ggml_tensor * kimi_k3_situ(ggml_context * ctx0, ggml_tensor * gate, ggml_tensor * up, + float beta, float linear_beta) { + ggml_tensor * a = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, gate, 1.0f/beta)), beta); + a = ggml_mul(ctx0, a, ggml_sigmoid(ctx0, gate)); + + if (linear_beta > 0.0f) { + up = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, up, 1.0f/linear_beta)), linear_beta); + } + return ggml_mul(ctx0, a, up); +} + +// +// cross-layer residual attention +// + +void llama_model_kimi_k3::graph::res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens) { + ckpts.push_back(ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens)); +} + +ggml_tensor * llama_model_kimi_k3::graph::res_stack(int64_t n_embd, int64_t n_tokens) { + GGML_UNUSED(n_embd); + GGML_UNUSED(n_tokens); + if (stack_cache_n == (int) ckpts.size()) { + return stack_cache; // set unchanged since the last mix + } + ggml_tensor * acc = ckpts[0]; + for (size_t i = 1; i < ckpts.size(); ++i) { + acc = ggml_concat(ctx0, acc, ckpts[i], 1); + } + stack_cache = acc; + stack_cache_n = (int) ckpts.size(); + return acc; +} + +ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor * score_w, + int64_t n_embd, int64_t n_tokens, int il) { + const int n_ckpt = (int) ckpts.size(); + if (n_ckpt == 0) { + return cur; // layer 0: nothing banked yet + } + + const float eps = hparams.f_norm_rms_eps; + + ggml_tensor * src = res_stack(n_embd, n_tokens); // [n_embd, n_ckpt, n_tokens] + + // Scores for the banked checkpoints. One rms_norm covers all of them because + // ne0 is n_embd. NOTE: the scores use the *normalized* values, but the weighted + // sum below uses the *raw* ones - mirroring _apply_attn_res. + ggml_tensor * sc_src = ggml_rms_norm(ctx0, src, eps); + sc_src = ggml_mul(ctx0, sc_src, score_w); + sc_src = ggml_sum_rows(ctx0, sc_src); // [1, n_ckpt, n_tokens] + sc_src = ggml_reshape_2d(ctx0, sc_src, n_ckpt, n_tokens); + + // The current residual stream is scored separately and kept out of the stack, + // so the stack stays append-only. + ggml_tensor * sc_cur = ggml_rms_norm(ctx0, cur, eps); + sc_cur = ggml_mul(ctx0, sc_cur, score_w); + sc_cur = ggml_sum_rows(ctx0, sc_cur); // [1, n_tokens] + + ggml_tensor * scores = ggml_concat(ctx0, sc_src, sc_cur, 0); // [n_ckpt+1, n_tokens] + ggml_tensor * probs = ggml_soft_max(ctx0, scores); // over ne0 = n_ckpt+1 + cb(probs, "res_probs", il); + + // Split the convex combination: hc_pre reduces over ne1 for the stacked part, + // a plain broadcast-multiply handles the current stream. + ggml_tensor * p_src = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, n_ckpt, n_tokens, probs->nb[1], 0)); + ggml_tensor * p_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, 1, n_tokens, probs->nb[1], + probs->nb[0] * n_ckpt)); + + ggml_tensor * out = ggml_dsv4_hc_pre(ctx0, src, p_src); + out = ggml_add(ctx0, out, ggml_mul(ctx0, cur, p_cur)); + + return out; +} + +llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + cb(inpL, "inp_embd", -1); + + // K3 MLA is nope-only, so there is no position input + + auto * inp_kv = !hparams.is_mla() ? build_inp_mem_hybrid() : nullptr; + auto * inp_k = hparams.is_mla() ? build_inp_mem_hybrid_k() : nullptr; + auto * inp_rs = hparams.is_mla() ? inp_k->get_recr() : inp_kv->get_recr(); + auto * inp_attn_kv = !hparams.is_mla() ? inp_kv->get_attn() : nullptr; + auto * inp_attn_k = hparams.is_mla() ? inp_k->get_attn() : nullptr; + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const int64_t n_head_kda = hparams.n_head(); + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = n_head_kda * head_dim; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla(); + const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla(); + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_embd_head_qk_rope = hparams.n_rot(); + const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope; + const float kq_scale_mla = 1.0f / sqrtf((float) n_embd_head_k_mla); + + const uint32_t res_bs = hparams.attn_res_block_size; + const bool use_attn_res = res_bs > 0; + const int64_t n_embd_latent = hparams.n_expert_latent > 0 ? hparams.n_expert_latent : n_embd; + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + // `prefix_sum` is the residual stream. On checkpoint layers it is banked + // into res_stack and restarts from the attention output alone. + ggml_tensor * prefix_sum = inpL; + + cur = use_attn_res ? res_mix(prefix_sum, layer.attn_res_score, n_embd, n_tokens, il) + : prefix_sum; + + bool banked = false; + if (use_attn_res && (uint32_t) il % res_bs == 0) { + res_push(prefix_sum, n_embd, n_tokens); // banks the RAW layer input, not `cur` + banked = true; + } + + cur = build_norm(cur, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + ggml_build_forward_expand(gf, cur); + + if (hparams.is_recr(il)) { + cur = build_kda_layer(cur, layer, inp_rs, d_conv, head_dim, n_head_kda, + d_inner, n_seq_tokens, n_seqs, il); + } else { + cur = build_mla_layer(cur, layer, inp_attn_k, inp_attn_kv, + n_embd_head_k_mla, n_embd_head_v_mla, kv_lora_rank, + n_embd_head_qk_rope, n_embd_head_qk_nope, kq_scale_mla, il); + } + + prefix_sum = banked ? cur : ggml_add(ctx0, prefix_sum, cur); + cb(prefix_sum, "prefix_sum_attn", il); + + cur = use_attn_res ? res_mix(prefix_sum, layer.ffn_res_score, n_embd, n_tokens, il) + : prefix_sum; + + cur = build_norm(cur, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate, cur); + ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up, cur); + cur = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta); + cur = ggml_mul_mat(ctx0, layer.ffn_down, cur); + cb(cur, "ffn_out", il); + } else { + cur = build_latent_moe(cur, layer, n_embd_latent, il); + } + + prefix_sum = ggml_add(ctx0, prefix_sum, cur); + prefix_sum = build_cvec(prefix_sum, il); + cb(prefix_sum, "l_out", il); + + inpL = prefix_sum; + } + + cur = inpL; + + // final mix, then narrow to the output tokens + if (use_attn_res) { + cur = res_mix(cur, model.output_res_score, n_embd, n_tokens, -1); + } + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +// +// KDA layer +// + +// Causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use. +// Polled rather than blocking is not a concern here; this mirrors kimi-linear's helper. +static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0, + ggml_tensor * conv_states_all, ggml_tensor * conv_state_all, + int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w, + int64_t d_conv, int64_t head_dim, int64_t n_head, + int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head) { + const int64_t d_inner = head_dim * n_head; + const int64_t conv_state_size = (d_conv - 1) * d_inner; + const int64_t n_embd_r_total = 3 * conv_state_size; + + ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_state_all), + n_embd_r_total * ggml_element_size(conv_state_all), + qkv * conv_state_size * ggml_element_size(conv_state_all)); + + ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); + ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0); + + ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, last_conv_x, + ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), + n_embd_r_total * ggml_element_size(conv_states_all), + (kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + + ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); + ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight); + Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens); + Xcur = ggml_silu(ctx0, Xcur); + + return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs); +} + +ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer( + ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs, + int64_t d_conv, int64_t head_dim, int64_t n_head_kda, + int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) { + + const auto * mctx_cur = inp_rs->mctx; + const auto kv_head = mctx_cur->get_head(); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + + ggml_tensor * Qcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Kcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Vcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + cb(Qcur, "kda_q_conv", il); + cb(Kcur, "kda_k_conv", il); + cb(Vcur, "kda_v_conv", il); + + // The decay gate has two forms, selected by linear_attn_config.gate_lower_bound + // (fla/ops/kda/gate.py). `lower_bound` is NOT a clamp - when set it swaps the + // activation entirely: + // + // unset (kimi-linear): g = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias) + // set (K3, -5.0): g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias)) + // + // ssm_a holds -exp(A_log) (folded at conversion time), so exp(A_log) == -ssm_a. + ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); + ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a); + g1 = ggml_add(ctx0, g1, layer.ssm_dt_b); + + ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head_kda, 1); + + if (hparams.kda_gate_lower_bound > -INFINITY) { + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens); + g1 = ggml_mul(ctx0, g1, A); // -exp(A_log) * (...) + g1 = ggml_sigmoid(ctx0, ggml_scale(ctx0, g1, -1.0f)); + g1 = ggml_scale(ctx0, g1, hparams.kda_gate_lower_bound); + } else { + g1 = ggml_softplus(ctx0, g1); + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens); + g1 = ggml_mul(ctx0, g1, A); + } + cb(g1, "kda_g1", il); + + g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head_kda, n_seq_tokens, n_seqs); + + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); + beta = ggml_reshape_4d(ctx0, beta, 1, n_head_kda, n_seq_tokens, n_seqs); + beta = ggml_sigmoid(ctx0, beta); + cb(beta, "kda_beta", il); + + ggml_tensor * cur_3d = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs); + + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head_kda, n_seqs); + + const float eps = hparams.f_norm_rms_eps; + Qcur = ggml_l2_norm(ctx0, Qcur, eps); + Kcur = ggml_l2_norm(ctx0, Kcur, eps); + + auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il); + + ggml_tensor * output = ggml_cont(ctx0, attn_out.first); + cb(output, "kda_scan_out", il); + ggml_tensor * new_state = attn_out.second; + + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, new_state, + ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs, + kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all)))); + + // K3: single full-rank gate (kimi-linear factors this as g_b(g_a(x))) + ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur_3d, cur_3d->ne[0], n_seq_tokens * n_seqs); + ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g, cur_2d); + g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head_kda, n_seq_tokens * n_seqs); + + ggml_tensor * o = ggml_reshape_3d(ctx0, output, head_dim, n_head_kda, n_seq_tokens * n_seqs); + ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + cb(g2, "kda_g2", il); + cb(normed, "kda_normed", il); + ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, g2)); + + gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens); + cur = ggml_mul_mat(ctx0, layer.wo, gated); + cb(cur, "kda_out", il); + + return cur; +} + +// +// MLA layer (nope-only, with K3's sigmoid output gate) +// + +ggml_tensor * llama_model_kimi_k3::graph::build_mla_layer( + ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_attn_k * inp_attn_k, llm_graph_input_attn_kv * inp_attn_kv, + int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla, int64_t kv_lora_rank, + int64_t n_embd_head_qk_rope, int64_t n_embd_head_qk_nope, float kq_scale, int il) { + + ggml_tensor * inp_gate = cur; // the output gate reads the *normed* layer input + + ggml_tensor * Qcur; + if (layer.wq_a) { + Qcur = ggml_mul_mat(ctx0, layer.wq_a, cur); + Qcur = build_norm(Qcur, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + Qcur = ggml_mul_mat(ctx0, layer.wq_b, Qcur); + } else { + Qcur = ggml_mul_mat(ctx0, layer.wq, cur); + } + + ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + + ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0); + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens, + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), + ggml_row_size(kv_cmpr_pe->type, kv_lora_rank)); + + // no RoPE: mla_use_nope is asserted at conversion time + kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + + ggml_tensor * out; + if (layer.wk_b && layer.wv_b) { + ggml_tensor * q_nope = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_nope, n_head, n_tokens, + ggml_row_size(Qcur->type, n_embd_head_k_mla), + ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, 0); + ggml_tensor * q_pe = ggml_view_3d(ctx0, Qcur, n_embd_head_qk_rope, n_head, n_tokens, + ggml_row_size(Qcur->type, n_embd_head_k_mla), + ggml_row_size(Qcur->type, n_embd_head_k_mla) * n_head, + ggml_row_size(Qcur->type, n_embd_head_qk_nope)); + + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3); + + ggml_tensor * Q = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0); + ggml_tensor * kv_cmpr_3d = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); + ggml_tensor * K = ggml_concat(ctx0, kv_cmpr_3d, k_pe, 0); + ggml_tensor * V = kv_cmpr_3d; + + // wo == NULL: the output projection is applied after the gate below + out = build_attn(inp_attn_k, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, layer.wv_b, kq_scale, il); + } else { + ggml_tensor * Q = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k_mla, n_head, n_tokens); + ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_b, kv_cmpr); + const int64_t kv_per_head = n_embd_head_qk_nope + n_embd_head_v_mla; + + ggml_tensor * k_nope = ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens, + ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head), 0); + ggml_tensor * V = ggml_cont(ctx0, ggml_view_3d(ctx0, kv, n_embd_head_v_mla, n_head, n_tokens, + ggml_row_size(kv->type, kv_per_head), ggml_row_size(kv->type, kv_per_head * n_head), + ggml_row_size(kv->type, n_embd_head_qk_nope))); + + ggml_tensor * k_pe_t = ggml_new_tensor_3d(ctx0, k_pe->type, n_embd_head_qk_rope, n_head, n_tokens); + ggml_tensor * K = ggml_concat(ctx0, ggml_repeat(ctx0, k_pe, k_pe_t), k_nope, 0); + + out = build_attn(inp_attn_kv, nullptr, NULL, nullptr, Q, K, V, nullptr, nullptr, nullptr, kq_scale, il); + } + + // K3: attn_output *= sigmoid(g_proj(x)), then o_proj + if (layer.wqkv_gate) { + ggml_tensor * g = ggml_sigmoid(ctx0, ggml_mul_mat(ctx0, layer.wqkv_gate, inp_gate)); + out = ggml_mul(ctx0, out, g); + cb(out, "mla_gated", il); + } + + out = ggml_mul_mat(ctx0, layer.wo, out); + cb(out, "mla_out", il); + + return out; +} + +// +// latent MoE: down-project, run the routed experts in the latent space, norm, up-project; +// shared experts stay at n_embd and read the un-projected input. +// + +ggml_tensor * llama_model_kimi_k3::graph::build_latent_moe( + ggml_tensor * cur, const llama_layer & layer, int64_t n_embd_latent, int il) { + + ggml_tensor * identity = cur; + + ggml_tensor * routed_in = layer.ffn_routed_down + ? ggml_mul_mat(ctx0, layer.ffn_routed_down, cur) + : cur; + + // The router scores the FULL-WIDTH input while the experts consume the latent + // one, so the logits are computed here and handed to build_moe_ffn via + // `probs_in` (which substitutes for the gate_inp matmul). + ggml_tensor * logits = ggml_mul_mat(ctx0, layer.ffn_gate_inp, identity); + cb(logits, "ffn_moe_logits", il); + + ggml_tensor * moe_out = build_moe_ffn(routed_in, + nullptr, // gate_inp unused: logits supplied below + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + hparams.n_expert, + hparams.n_expert_used, + LLM_FFN_SITU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + logits); + cb(moe_out, "ffn_moe_out", il); + + if (layer.ffn_routed_norm) { + moe_out = build_norm(moe_out, layer.ffn_routed_norm, NULL, LLM_NORM_RMS, il); + } + if (layer.ffn_routed_up) { + moe_out = ggml_mul_mat(ctx0, layer.ffn_routed_up, moe_out); + } + GGML_UNUSED(n_embd_latent); + + if (layer.ffn_gate_shexp) { + ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate_shexp, identity); + ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up_shexp, identity); + ggml_tensor * sh = kimi_k3_situ(ctx0, g, u, hparams.situ_beta, hparams.situ_linear_beta); + sh = ggml_mul_mat(ctx0, layer.ffn_down_shexp, sh); + cb(sh, "ffn_shexp", il); + moe_out = ggml_add(ctx0, moe_out, sh); + } + + cb(moe_out, "ffn_out", il); + return moe_out; +} diff --git a/src/models/models.h b/src/models/models.h index ddb9ae2f1210..9bcf9f1d547c 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2272,6 +2272,57 @@ struct llama_model_mimo2 : public llama_model_base { }; +struct llama_model_kimi_k3 : public llama_model_base { + llama_model_kimi_k3(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_build_delta_net_base { + graph(const llama_model & model, const llm_graph_params & params); + + const llama_model & model; + + // Cross-layer residual attention (K3's `_apply_attn_res`). + // + // `src` holds the block-residual checkpoints, one [n_embd, n_token] slab + // per checkpoint, packed as [n_embd, n_ckpt, n_token] so that: + // - ggml_rms_norm reduces over ne0 = n_embd (scores all slabs at once) + // - ggml_dsv4_hc_pre reduces over ne1 = n_ckpt (the weighted sum) + // Both requirements are satisfied by the same layout, which is why this + // needs no transpose of the (large) stack. + // Each checkpoint is kept as its own [n_embd, 1, n_token] tensor and the + // [n_embd, n_ckpt, n_token] stack is materialised with ggml_concat only when + // the set changes (8 times per forward pass for the real model). A single + // preallocated buffer would be cheaper, but a bare ggml_new_tensor leaf has + // no backing buffer - ggml-alloc only allocates tensors that are op outputs. + std::vector ckpts; + ggml_tensor * stack_cache = nullptr; + int stack_cache_n = -1; + + void res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens); + ggml_tensor * res_stack(int64_t n_embd, int64_t n_tokens); + ggml_tensor * res_mix(ggml_tensor * cur, ggml_tensor * score_w, + int64_t n_embd, int64_t n_tokens, int il); + + ggml_tensor * build_kda_layer(ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_rs * inp_rs, + int64_t d_conv, int64_t head_dim, int64_t n_head_kda, + int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il); + + ggml_tensor * build_mla_layer(ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_attn_k * inp_attn_k, + llm_graph_input_attn_kv * inp_attn_kv, + int64_t n_embd_head_k_mla, int64_t n_embd_head_v_mla, + int64_t kv_lora_rank, int64_t n_embd_head_qk_rope, + int64_t n_embd_head_qk_nope, float kq_scale, int il); + + ggml_tensor * build_latent_moe(ggml_tensor * cur, const llama_layer & layer, + int64_t n_embd_latent, int il); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_kimi_linear : public llama_model_base { llama_model_kimi_linear(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; From d171e72cb6c6cec5d6233e43bea9c9d2b8a010e1 Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Mon, 27 Jul 2026 19:57:44 +0200 Subject: [PATCH 02/21] model: fix ty errors in the Kimi-K3 converter - `_res_parts` buffers (kind, tensor) pairs, not bare tensors - `get_tensors` must return an Iterator, matching ModelBase - LazyBase's `func` takes one argument, so pass the expert loaders through `args` instead of the closure - borrowing KimiLinearModel.set_vocab from an unrelated TextModel is deliberate and safe, but not expressible in the signature No behaviour change: the MXFP4 repack still dequantizes to the source weights with 0.0e+00 error and end-to-end logits are unchanged (8.386e-03 rel, corr 0.99996630). Assisted-By: Claude Opus 5 (1M context) --- conversion/kimi_k3.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py index 58f7dbcac5c6..277cf51db784 100644 --- a/conversion/kimi_k3.py +++ b/conversion/kimi_k3.py @@ -1,7 +1,7 @@ from __future__ import annotations import re -from typing import Callable, Iterable, TYPE_CHECKING +from typing import Callable, Iterable, Iterator, TYPE_CHECKING import numpy as np import torch @@ -35,8 +35,8 @@ class KimiK3Model(TextModel): # elementwise product norm.weight * proj.weight (see _apply_attn_res in # modeling_kimi_linear.py), so they are fused into a single [n_embd] vector # at conversion time. They arrive as separate tensors, so buffer whichever - # comes first. - _res_parts: dict[str, Tensor] + # comes first, tagged with which one it is. + _res_parts: dict[str, tuple[str, Tensor]] # HF suffix -> (gguf tensor, per-layer?) _RES_FUSIONS = { @@ -67,7 +67,11 @@ def set_vocab(self): # K3 ships the same TikToken vocab as K2: its pre-tokenizer hashes to # 81212dc7... which base.py already maps to "kimi-k2", so no new # pre-tokenizer registration is needed. - KimiLinearModel.set_vocab(self) + # + # Borrowed rather than inherited: K3 shares kimi-linear's vocab handling + # but none of its tensor layout. The method only touches TextModel + # members, so an unrelated TextModel is a valid receiver. + KimiLinearModel.set_vocab(self) # ty: ignore[invalid-argument-type] # ...but K2's converter ends by forcing eos to the tokenizer's own # eos_id, which for K3 is 163585 = [EOS], the *document* terminator. @@ -124,17 +128,21 @@ def _mxfp4_expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], Callabl n_blocks = (packed_cols * 2) // 32 byte_shape = (len(loaders), rows, n_blocks * 17) - def load() -> np.ndarray: + def load(fns: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]) -> np.ndarray: out = np.empty(byte_shape, dtype=np.uint8) - for eid, (packed_fn, scale_fn) in enumerate(loaders): + for eid, (packed_fn, scale_fn) in enumerate(fns): out[eid] = repack_mxfp4_blocks( LazyTorchTensor.to_eager(packed_fn()), LazyTorchTensor.to_eager(scale_fn()), ) return out + # loaders goes through args rather than the closure so that `func` matches + # LazyBase's single-argument shape; _recurse_apply passes plain callables + # through untouched. return gguf.LazyNumpyTensor( meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(np.uint8, byte_shape), + args=(loaders,), func=load, ) @@ -190,7 +198,7 @@ def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: self._write_mxfp4_experts() return () - def get_tensors(self) -> Iterable[tuple[str, Tensor]]: + def get_tensors(self) -> Iterator[tuple[str, Tensor]]: for name, data in super().get_tensors(): if name.startswith(("vision_tower.", "mm_projector.")): continue # text only From f9181dcb606d4ae0e95a9c7f0d2b5bd3d994c0ea Mon Sep 17 00:00:00 2001 From: "Piotr Wilkin (ilintar)" Date: Mon, 27 Jul 2026 21:07:57 +0200 Subject: [PATCH 03/21] Update conversion/kimi_k3.py Co-authored-by: Boris Dvorkin --- conversion/kimi_k3.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py index 277cf51db784..1064249e5205 100644 --- a/conversion/kimi_k3.py +++ b/conversion/kimi_k3.py @@ -332,7 +332,8 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter # -exp(A_log) is folded here so the graph does not have to if name.endswith(".A_log"): - data_torch = -torch.exp(data_torch) + n_head = self.hparams["num_attention_heads"] + data_torch = -torch.exp(data_torch.float()[:n_head]) # dt_bias -> the name SSM_DT's mapping expects if name.endswith(".dt_bias"): From 28a84b1d005e7617e750598c26ff88430c4fd218 Mon Sep 17 00:00:00 2001 From: "Piotr Wilkin (ilintar)" Date: Tue, 28 Jul 2026 00:05:05 +0200 Subject: [PATCH 04/21] Increase LLAMA_MAX_EXPERTS from 512 to 1024 --- src/llama-hparams.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llama-hparams.h b/src/llama-hparams.h index f6d9bf9c9494..66d45b411934 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -8,7 +8,7 @@ // bump if necessary #define LLAMA_MAX_LAYERS 512 -#define LLAMA_MAX_EXPERTS 512 // Qwen3 Next +#define LLAMA_MAX_EXPERTS 1024 // Kimi K3 enum llama_expert_gating_func_type { LLAMA_EXPERT_GATING_FUNC_TYPE_NONE = 0, From 50a54319d711170e751649f1cd2c2c5f80ad84a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Tue, 28 Jul 2026 12:33:10 +0200 Subject: [PATCH 05/21] tests : support for Kimi K3 in archs test --- tests/test-llama-archs.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index e900bdc0da9c..d77417d50463 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -105,6 +105,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_KIMI_LINEAR + || arch == LLM_ARCH_KIMI_K3 || arch == LLM_ARCH_MISTRAL4) { n_embd = 128; n_head = 1; @@ -145,7 +146,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_FULL_ATTENTION_INTERVAL, uint32_t(2)); if (arch == LLM_ARCH_PLAMO2 || arch == LLM_ARCH_JAMBA || arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE || - arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR) { + arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_KIMI_K3) { GGML_ASSERT(n_layer >= 2); std::vector n_head_per_layer; n_head_per_layer.reserve(n_layer); @@ -164,6 +165,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_KIMI_LINEAR + || arch == LLM_ARCH_KIMI_K3 || arch == LLM_ARCH_MISTRAL4) { ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(576)); ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512)); @@ -370,6 +372,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_PADDLEOCR: case LLM_ARCH_MIMO2: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_KIMI_K3: case LLM_ARCH_STEP35: case LLM_ARCH_MISTRAL4: case LLM_ARCH_MELLUM: @@ -436,7 +439,7 @@ static bool arch_supported(const llm_arch arch) { // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU - if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA) { + if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_KIMI_K3) { return false; } #endif // GGML_USE_WEBGPU From 03efbfe84fb20a0fdd0eb8b8cfcac17131146d7f Mon Sep 17 00:00:00 2001 From: Deepankar Singh Date: Tue, 28 Jul 2026 17:22:38 +0530 Subject: [PATCH 06/21] chat : add Kimi K3 chat format (reasoning, content, typed tool calls) K3's assistant output is an XTML-ish tagged format built by the template's open_tag/close_tag macros. Two properties break generic parsing: 1. The generation prompt ends with open_tag('think'), so the completion starts inside the think section with no opening marker in the output (thinking_forced_open). 2. Only <|open|>/<|close|>/<|sep|>/<|end_of_msg|> are special tokens; tag names ("think", "response", "message") are ordinary text tokens. Adds common_chat_params_init_kimi_k3 (PEG_NATIVE) with detection on the marker trio, reasoning extraction, response unwrapping, and tool-call parsing of the tools/call/argument tag structure with argument types taken from the tool schema. Includes the K3 chat template fixture and 9 test-chat cases derived from real generations of the full 2.8T model. Verified end-to-end against Kimi-K3-Q2_K (GrEarl/Kimi-K3-GGUF) on 8x B200: content, reasoning_content, streaming deltas, and tool_calls all correct; finish_reason stop/tool_calls as appropriate. Co-Authored-By: Claude Fable 5 --- common/chat.cpp | 178 ++++++++++++++++++ models/templates/Kimi-K3.jinja | 325 +++++++++++++++++++++++++++++++++ tests/test-chat.cpp | 105 +++++++++++ 3 files changed, 608 insertions(+) create mode 100644 models/templates/Kimi-K3.jinja diff --git a/common/chat.cpp b/common/chat.cpp index 01053ddde804..53805c8e67a6 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -2321,6 +2321,176 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha return data; } +// Kimi K3 - XTML-ish tagged format from the model's own template: +// open_tag(t, attrs) = <|open|>t k="v"...<|sep|> close_tag(t) = <|close|>t<|sep|> +// assistant := [think] [response] [tools] close_tag(message) <|end_of_msg|> +// Note the generation prompt already opens the think (or response) section, so +// the section opener is optional here - same situation as Kimi K2 Thinking. +static common_chat_params common_chat_params_init_kimi_k3(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + + const std::string SEP = "<|sep|>"; + const std::string MSG_START = "<|open|>message role=\"assistant\"<|sep|>"; + const std::string THINK_START = "<|open|>think<|sep|>"; + const std::string THINK_END = "<|close|>think<|sep|>"; + const std::string RESP_START = "<|open|>response<|sep|>"; + const std::string RESP_END = "<|close|>response<|sep|>"; + const std::string TOOLS_START = "<|open|>tools<|sep|>"; + const std::string TOOLS_END = "<|close|>tools<|sep|>"; + const std::string CALL_START = "<|open|>call tool=\""; + const std::string CALL_END = "<|close|>call<|sep|>"; + const std::string ARG_START = "<|open|>argument key=\""; + const std::string ARG_END = "<|close|>argument<|sep|>"; + const std::string MSG_END = "<|close|>message<|sep|>"; + const std::string EOM_TOKEN = "<|end_of_msg|>"; + + // The four markers are the only special tokens; tag names ("think", + // "response", "message") are ordinary tokens and must NOT be preserved, + // or ordinary prose containing those words would be mangled. + data.preserved_tokens = { + "<|open|>", + "<|close|>", + "<|sep|>", + "<|end_of_msg|>", + }; + + data.thinking_start_tag = THINK_START; + data.thinking_end_tags = { THINK_END }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; + + if (inputs.has_continuation()) { + const auto & msg = inputs.continue_msg; + + data.generation_prompt = MSG_START + THINK_START + msg.reasoning_content; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += THINK_END + RESP_START + msg.render_content(); + } + + data.prompt += data.generation_prompt; + } + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto end = p.end(); + + auto start = p.optional(p.literal(MSG_START)); + + // The think section is ALWAYS consumed, even when reasoning extraction + // is off: K3's generation prompt ends with open_tag('think'), so the + // opener is present on every request and would otherwise leak into + // content. With extraction off the thoughts fall into content, matching + // how the other reasoning models behave. + // Reasoning stops at its own closer, or at the response opener if the + // model skips the closer entirely (seen on short answers). + auto think_body = extract_reasoning ? p.reasoning(p.until_one_of({ THINK_END, RESP_START })) : + p.content(p.until_one_of({ THINK_END, RESP_START })); + + auto reasoning = p.optional(p.optional(p.literal(THINK_START)) + think_body + + p.optional(p.literal(THINK_END))); + + // Content runs to the response closer, or to whatever comes next if a + // truncated generation never emits one. + auto response = p.optional(p.literal(RESP_START)) + + p.content(p.until_one_of({ RESP_END, TOOLS_START, MSG_END })) + + p.optional(p.literal(RESP_END)); + + // The message closer is followed by the EOG token, which reaches the + // parser as text and must be consumed or the parse is left incomplete. + auto trailer = p.optional(p.literal(MSG_END)) + p.optional(p.literal(EOM_TOKEN)); + + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + return start + reasoning + response + trailer + end; + } + + auto tool_choices = p.choice(); + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + const json schema = function.contains("parameters") ? function.at("parameters") : json::object(); + + // Arguments arrive one tag per key, with the JSON type carried in a + // type="..." attribute. We take the type from the tool schema + // instead - it is authoritative, and it tells us whether the value + // should be parsed as JSON or kept as a literal string. + auto args = p.eps(); + if (schema.contains("properties") && !schema.at("properties").empty()) { + auto arg_choices = p.choice(); + for (const auto & prop : schema.at("properties").items()) { + const std::string & key = prop.key(); + + std::string type = "string"; + if (prop.value().is_object() && prop.value().contains("type") && + prop.value().at("type").is_string()) { + type = prop.value().at("type").get(); + } + + auto value = type == "string" ? p.tool_arg_string_value(p.until(ARG_END)) : + p.tool_arg_value(p.until(ARG_END)); + + // skip the trailing type="..." attribute: anything up to <|sep|> + arg_choices |= p.rule("kimi-k3-arg-" + name + "-" + key, + p.tool_arg(p.tool_arg_open(p.literal(ARG_START)) + + p.tool_arg_name(p.literal(key)) + p.literal("\"") + + p.until(SEP) + p.literal(SEP) + value + + p.tool_arg_close(p.literal(ARG_END)))); + } + args = p.zero_or_more(arg_choices); + } + + // skip the trailing index="N" attribute the same way + auto call = p.tool(p.tool_open(p.literal(CALL_START) + p.tool_name(p.literal(name)) + p.literal("\"") + + p.until(SEP) + p.literal(SEP)) + + p.tool_args(args) + p.tool_close(p.literal(CALL_END))); + + tool_choices |= p.rule("kimi-k3-tool-" + name, call); + }); + + // K3 emits every call inside one <|open|>tools<|sep|> section, then + // closes the message. The message closer is part of the trigger rule so + // that the lazy grammar still permits it once tool calls have started - + // otherwise constrained decoding rejects the model's own closing tag. + auto tools_section = + p.trigger_rule("kimi-k3-tool-call", p.literal(TOOLS_START) + p.one_or_more(tool_choices) + + p.literal(TOOLS_END) + p.optional(p.literal(MSG_END)) + + p.optional(p.literal(EOM_TOKEN))); + + auto tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? tools_section : + p.optional(tools_section); + + return start + reasoning + response + tools + trailer + end; + }); + + data.parser = parser.save(); + + if (include_grammar) { + data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED; + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + foreach_function(inputs.tools, [&](const json & tool) { + const auto & function = tool.at("function"); + if (function.contains("parameters")) { + auto schema = function.at("parameters"); + builder.resolve_refs(schema); + } + }); + parser.build_grammar(builder, data.grammar_lazy); + }); + + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_WORD, TOOLS_START }, + }; + } + + return data; +} + // Cohere2 MoE (a.k.a. "North Code") parser. // // The assistant turn is fully marker-wrapped: @@ -3287,6 +3457,14 @@ std::optional common_chat_try_specialized_template( return common_chat_params_init_kimi_k2(tmpl, params); } + // Kimi K3 - XTML-ish tagged format built from open_tag/close_tag macros. + // Detection: the <|open|>/<|close|>/<|sep|> marker trio is unique to K3. + if (src.find("<|open|>") != std::string::npos && src.find("<|close|>") != std::string::npos && + src.find("<|end_of_msg|>") != std::string::npos) { + LOG_DBG("Using specialized template: Kimi K3\n"); + return common_chat_params_init_kimi_k3(tmpl, params); + } + // Cohere2 MoE / North Code - marker-wrapped format with <|START_TEXT|> content and // <|START_ACTION|> JSON tool calls. <|START_TEXT|> is unique to this template (the older // Command-R templates use <|START_RESPONSE|>). diff --git a/models/templates/Kimi-K3.jinja b/models/templates/Kimi-K3.jinja new file mode 100644 index 000000000000..67e51a6af4b3 --- /dev/null +++ b/models/templates/Kimi-K3.jinja @@ -0,0 +1,325 @@ +{%- macro escape_attr(value) -%} +{{- value|string|replace('&', '&')|replace('"', '"') -}} +{%- endmacro -%} + +{%- macro open_tag(tag, attrs=[]) -%} +{{- '<|open|>' + tag -}} +{%- for attr in attrs -%} +{{- ' ' + attr[0] + '="' -}}{{- escape_attr(attr[1]) -}}{{- '"' -}} +{%- endfor -%} +{{- '<|sep|>' -}} +{%- endmacro -%} + +{%- macro close_tag(tag) -%} +{{- '<|close|>' + tag + '<|sep|>' -}} +{%- endmacro -%} + +{%- macro next_image(state) -%} +{%- if image_prompts is defined and image_prompts is not none -%} + {%- if state.image_index >= image_prompts|length -%} + {{- raise_exception('More image placeholders than image prompts.') -}} + {%- endif -%} + {{- image_prompts[state.image_index] -}} + {%- set state.image_index = state.image_index + 1 -%} +{%- else -%} + {{- '<|kimi_image_placeholder|>' -}} +{%- endif -%} +{%- endmacro -%} + +{%- macro render_text(text, state) -%} +{%- set text = text|string -%} +{%- if image_prompts is defined and image_prompts is not none and '<|kimi_image_placeholder|>' in text -%} + {%- set parts = text.split('<|kimi_image_placeholder|>') -%} + {%- for part in parts -%} + {{- part -}} + {%- if not loop.last -%}{{- next_image(state) -}}{%- endif -%} + {%- endfor -%} +{%- else -%} + {{- text -}} +{%- endif -%} +{%- endmacro -%} + +{%- macro render_content(content, state) -%} +{%- if content is string -%} + {{- render_text(content, state) -}} +{%- elif content is not none and content is defined -%} + {%- for part in content -%} + {%- if part.type in ['image', 'image_url'] -%} + {{- next_image(state) -}} + {%- else -%} + {{- render_text(part.text, state) -}} + {%- endif -%} + {%- endfor -%} +{%- endif -%} +{%- endmacro -%} + +{%- macro internal_system_message(message_type, body) -%} +{{- open_tag('message', [('role', 'system'), ('type', message_type)]) -}} +{{- body|trim -}} +{{- close_tag('message') -}} +{{- '<|end_of_msg|>' -}} +{%- endmacro -%} + +{%- macro json_sorted(value) -%} +{#- Minja は tojson(sort_keys=true) を実装していない。K3 参照実装は + deep_sort_dict() の後に compact JSON 化するため、dictsort と再帰マクロで + 同じバイト列を作る。配列の順序は保持し、mapping の各階層だけソートする。 -#} +{%- if value is mapping -%} +{{- '{' -}} +{%- for key, item in value|dictsort -%} +{%- if not loop.first -%}{{- ',' -}}{%- endif -%} +{{- key|tojson(ensure_ascii=false) -}}{{- ':' -}}{{- json_sorted(item) -}} +{%- endfor -%} +{{- '}' -}} +{%- elif value is string or value is number or value is boolean or value is none -%} +{{- value|tojson(ensure_ascii=false) -}} +{%- else -%} +{{- '[' -}} +{%- for item in value -%} +{%- if not loop.first -%}{{- ',' -}}{%- endif -%} +{{- json_sorted(item) -}} +{%- endfor -%} +{{- ']' -}} +{%- endif -%} +{%- endmacro -%} + +{%- macro render_tool_declare(tool_list, dynamic=false) -%} +{{- open_tag('message', [('role', 'system'), ('type', 'tool-declare')]) -}} +{%- if dynamic -%} +{{- '## New Tools Available\nThe system dynamically extends the toolset via lazy-loading.\nYou have access to all existing and extended tools.\nHere are the specs for the extended tools.\n\n```json\n' -}} +{%- else -%} +{{- '# Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n' -}} +{%- endif -%} +{{- json_sorted(tool_list) -}} +{{- '\n```' -}} +{{- close_tag('message') -}} +{{- '<|end_of_msg|>' -}} +{%- endmacro -%} + +{%- macro xtml_type(value) -%} +{%- if value is boolean -%}boolean +{%- elif value is none -%}null +{%- elif value is number -%}number +{%- elif value is string -%}string +{%- elif value is mapping -%}object +{%- else -%}array +{%- endif -%} +{%- endmacro -%} + +{%- macro xtml_value(value) -%} +{%- if value is string -%} +{{- value -}} +{%- else -%} +{{- value|tojson(ensure_ascii=false) -}} +{%- endif -%} +{%- endmacro -%} + +{%- macro render_assistant(message, state) -%} +{%- if thinking -%} + {%- set reasoning_content = message.get('reasoning_content') or message.get('reasoning') -%} + {{- open_tag('think') -}} + {%- if reasoning_content is not none and reasoning_content|string|trim -%} + {{- render_text(reasoning_content, state) -}} + {%- endif -%} + {{- close_tag('think') -}} +{%- endif -%} +{{- open_tag('response') -}} +{{- render_content(message.get('content'), state) -}} +{{- close_tag('response') -}} +{%- set tool_calls = message.get('tool_calls') -%} +{%- if tool_calls -%} + {{- open_tag('tools') -}} + {%- for tool_call in tool_calls -%} + {%- if tool_call is not mapping -%} + {{- raise_exception('Kimi K3 tool calls must be mappings.') -}} + {%- endif -%} + {%- set fn = tool_call.function if tool_call.function is defined and tool_call.function is mapping else tool_call -%} + {%- if fn.get('name') is none -%} + {{- raise_exception('Kimi K3 tool calls require a function name.') -}} + {%- endif -%} + {{- open_tag('call', [('tool', fn.name), ('index', loop.index)]) -}} + {%- set arguments = fn.get('arguments', {}) -%} + {%- set json_block = fn.get('_xtml_json_block') -%} + {%- if json_block is not none -%} + {{- open_tag('json', [('type', 'object')]) -}} + {{- render_text(json_block, state) -}} + {{- close_tag('json') -}} + {%- elif arguments is mapping -%} + {%- for key, value in arguments.items() -%} + {{- open_tag('argument', [('key', key), ('type', xtml_type(value))]) -}} + {{- render_text(xtml_value(value), state) -}} + {{- close_tag('argument') -}} + {%- endfor -%} + {%- elif arguments is string and arguments|trim -%} + {{- open_tag('json', [('type', 'object')]) -}} + {{- render_text(arguments, state) -}} + {{- close_tag('json') -}} + {%- elif arguments is not none and arguments is not string -%} + {{- raise_exception('Kimi K3 tool call arguments must be a mapping or a JSON object string.') -}} + {%- endif -%} + {{- close_tag('call') -}} + {%- endfor -%} + {{- close_tag('tools') -}} +{%- endif -%} +{%- endmacro -%} + +{%- macro render_tool_message(message, state, resolved_name=none) -%} +{%- set state.tool_index = state.tool_index + 1 -%} +{%- if resolved_name is not none -%} + {%- set tool_name = resolved_name -%} +{%- elif 'tool' in message -%} + {%- set tool_name = message.get('tool') -%} +{%- else -%} + {%- set tool_name = message.get('name') -%} +{%- endif -%} +{%- if tool_name is none and state.tool_calls is not none and state.tool_index <= state.tool_calls|length -%} + {%- set fallback_call = state.tool_calls[state.tool_index - 1] -%} + {%- set fallback_fn = fallback_call.function if fallback_call.function is defined and fallback_call.function is mapping else fallback_call -%} + {%- set tool_name = fallback_fn.name -%} +{%- endif -%} +{%- if tool_name is none -%} + {{- raise_exception('Kimi K3 tool messages need a resolvable tool name: carry `tool`/`name`, or match a preceding assistant tool_call by order.') -}} +{%- endif -%} +{{- open_tag('message', [('role', 'tool'), ('tool', tool_name), ('index', state.tool_index)]) -}} +{{- render_content(message.get('content'), state) -}} +{{- close_tag('message') -}} +{{- '<|end_of_msg|>' -}} +{%- endmacro -%} + +{%- if thinking is undefined -%} + {%- set thinking = true -%} +{%- endif -%} +{%- if thinking_effort is undefined -%} + {%- set thinking_effort = 'max' -%} +{%- endif -%} +{%- if thinking and thinking_effort is not none and thinking_effort not in ['low', 'high', 'max'] -%} + {{- raise_exception('Unsupported thinking_effort=' + thinking_effort|string + '; supported values are low, high, and max.') -}} +{%- endif -%} + +{%- set state = namespace(image_index=0, tool_calls=none, tool_index=0, response_schema=none) -%} + +{%- if tools is defined and tools -%} + {{- render_tool_declare(tools) -}} +{%- endif -%} + +{%- if thinking and thinking_effort in ['low', 'high', 'max'] -%} + {{- internal_system_message( + 'thinking-effort', + '`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`.\nNow the system is invoked with `thinking_effort=' + thinking_effort|string + '`.' + ) -}} +{%- endif -%} + +{%- for message in messages -%} + {%- if message is mapping -%} + {%- if 'role' not in message -%} + {{- raise_exception('Kimi K3 messages require a role.') -}} + {%- elif message.role == 'user' -%} + {%- set attrs = [('role', 'user')] -%} + {%- if message.get('name') -%}{%- set attrs = attrs + [('name', message.name)] -%}{%- endif -%} + {{- open_tag('message', attrs) -}} + {{- render_content(message.get('content'), state) -}} + {{- close_tag('message') -}} + {{- '<|end_of_msg|>' -}} + {%- elif message.role == 'system' and message.get('tools') -%} + {{- render_tool_declare(message.tools, dynamic=true) -}} + {%- elif message.role == 'system' -%} + {%- set attrs = [('role', 'system')] -%} + {%- if message.get('name') -%}{%- set attrs = attrs + [('name', message.name)] -%}{%- endif -%} + {{- open_tag('message', attrs) -}} + {{- render_content(message.get('content'), state) -}} + {{- close_tag('message') -}} + {{- '<|end_of_msg|>' -}} + {%- elif message.role == 'assistant' -%} + {%- set state.tool_calls = message.get('tool_calls') -%} + {%- set state.tool_index = 0 -%} + {%- set attrs = [('role', 'assistant')] -%} + {%- if message.get('name') -%}{%- set attrs = attrs + [('name', message.name)] -%}{%- endif -%} + {{- open_tag('message', attrs) -}} + {{- render_assistant(message, state) -}} + {{- close_tag('message') -}} + {{- '<|end_of_msg|>' -}} + {%- elif message.role == 'tool' and (loop.first or messages[loop.index0 - 1].role != 'tool') -%} + {%- set run = namespace(tool_messages=[], resolved_count=0) -%} + {%- for candidate in messages[loop.index0:] -%} + {%- if candidate is not mapping or candidate.role != 'tool' -%}{%- break -%}{%- endif -%} + {%- set run.tool_messages = run.tool_messages + [candidate] -%} + {%- set call_id = candidate.get('tool_call_id', candidate.get('id')) -%} + {%- set match = namespace(found=false) -%} + {%- if call_id is not none and state.tool_calls is not none -%} + {%- for tool_call in state.tool_calls -%} + {%- if not match.found and tool_call is mapping and tool_call.get('id') is not none and tool_call.get('id')|string == call_id|string -%} + {%- set match.found = true -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- if match.found -%}{%- set run.resolved_count = run.resolved_count + 1 -%}{%- endif -%} + {%- endfor -%} + {%- if run.tool_messages|length > 0 and run.resolved_count == run.tool_messages|length -%} + {%- set emitted = namespace(ids=[]) -%} + {%- for tool_call in state.tool_calls -%} + {%- if tool_call is mapping and tool_call.get('id') is not none and tool_call.get('id')|string not in emitted.ids -%} + {%- set emitted.ids = emitted.ids + [tool_call.get('id')|string] -%} + {%- set fn = tool_call.function if tool_call.function is defined and tool_call.function is mapping else tool_call -%} + {%- for tool_message in run.tool_messages -%} + {%- set result_id = tool_message.get('tool_call_id', tool_message.get('id')) -%} + {%- if result_id is not none and result_id|string == tool_call.get('id')|string -%} + {{- render_tool_message(tool_message, state, fn.get('name')) -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {%- for tool_message in run.tool_messages -%} + {{- render_tool_message(tool_message, state) -}} + {%- endfor -%} + {%- endif -%} + {%- endif -%} + {%- endif -%} +{%- endfor -%} + +{%- if tool_choice is defined and tool_choice == 'required' -%} + {{- internal_system_message('tool-choice', 'The system is invoked with `tool_choice=required`.\nYou MUST call tools in the next message.') -}} +{%- elif tool_choice is defined and tool_choice == 'none' -%} + {{- internal_system_message('tool-choice', 'The system is invoked with `tool_choice=none`.\nYou MUST NOT call any tools in the next message.') -}} +{%- endif -%} + +{%- if response_schema is defined -%} + {%- set state.response_schema = response_schema -%} +{%- elif response_format is defined and response_format is mapping and response_format.get('json_schema') is not none -%} + {%- set schema_wrapper = response_format.get('json_schema') -%} + {%- if schema_wrapper is mapping and 'schema' in schema_wrapper -%} + {%- set state.response_schema = schema_wrapper.get('schema') -%} + {%- elif schema_wrapper is mapping and 'json_schema' in schema_wrapper -%} + {%- set state.response_schema = schema_wrapper.get('json_schema') -%} + {%- else -%} + {%- set state.response_schema = schema_wrapper -%} + {%- endif -%} +{%- endif -%} + +{%- set response_format_type = none -%} +{%- if response_format is defined and response_format is mapping -%} + {%- set response_format_type = response_format.get('type') -%} +{%- elif response_format is defined -%} + {%- set response_format_type = response_format -%} +{%- endif -%} +{%- if response_format_type == 'json_object' -%} + {{- internal_system_message( + 'response-format', + 'The system is invoked with `response_format=json_object`.\nYour response must be raw JSON data without markdown code blocks (```json) or any additional formatting.' + ) -}} +{%- elif response_format_type == 'json_schema' -%} + {{- internal_system_message( + 'response-format', + 'The system is invoked with `response_format=json_schema`.\nYour response must be raw JSON data without markdown code blocks (```json) or any additional formatting.\nThe JSON data must match the following schema:\n```json\n' + json_sorted(state.response_schema) + '\n```' + ) -}} +{%- endif -%} + +{%- if add_generation_prompt -%} + {{- open_tag('message', [('role', 'assistant')]) -}} + {{- open_tag('think' if thinking else 'response') -}} +{%- endif -%} + +{%- if image_prompts is defined and image_prompts is not none and state.image_index != image_prompts|length -%} + {{- raise_exception('image prompt count ' + image_prompts|length|string + ' != consumed placeholder count ' + state.image_index|string) -}} +{%- endif -%} + diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index f54a58f9b67c..807869d048e3 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -4462,6 +4462,111 @@ static void test_template_output_peg_parsers(bool detailed_debug) { } } + // Kimi-K3 tests - custom parser + // Unique feature: XTML-ish tags built from <|open|>/<|close|>/<|sep|>, and a + // generation prompt that leaves the think section already open. + { + auto tst = peg_tester("models/templates/Kimi-K3.jinja", detailed_debug); + + // Content only. The response section is explicit even with no reasoning. + tst.test("<|open|>response<|sep|>Hello, world!\nWhat's up?<|close|>response<|sep|>" + "<|close|>message<|sep|>") + .expect(message_assist) + .run(); + + // Reasoning with NO opening tag - the generation prompt already opened + // it. This is the case that silently loses reasoning if unhandled. + tst.test("I'm thinking about this<|close|>think<|sep|>" + "<|open|>response<|sep|>Hello, world!\nWhat's up?<|close|>response<|sep|>" + "<|close|>message<|sep|>") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(simple_assist_msg("Hello, world!\nWhat's up?", "I'm thinking about this")) + .run(); + + // Prose that mentions the tag names must survive intact. + tst.test("<|open|>response<|sep|>Use the response tag, then message the handler." + "<|close|>response<|sep|><|close|>message<|sep|>") + .expect(simple_assist_msg("Use the response tag, then message the handler.")) + .run(); + + // Truncated mid-reasoning (hit the token budget): keep the reasoning. + tst.test("I was still thinking when the budget ran out") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect_reasoning("I was still thinking when the budget ran out") + .run(); + + // Single tool call, one argument. + tst.test("<|open|>response<|sep|><|close|>response<|sep|>" + "<|open|>tools<|sep|>" + "<|open|>call tool=\"special_function\" index=\"1\"<|sep|>" + "<|open|>argument key=\"arg1\" type=\"number\"<|sep|>1<|close|>argument<|sep|>" + "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>") + .tools({ special_function_tool }) + .expect_tool_calls({ + { "special_function", R"({"arg1":1})", "" }, + }) + .run(); + + // Tool call preceded by reasoning (no opening think tag) and content. + tst.test("I should call it<|close|>think<|sep|>" + "<|open|>response<|sep|>On it.<|close|>response<|sep|>" + "<|open|>tools<|sep|>" + "<|open|>call tool=\"special_function\" index=\"1\"<|sep|>" + "<|open|>argument key=\"arg1\" type=\"number\"<|sep|>1<|close|>argument<|sep|>" + "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .tools({ special_function_tool }) + .expect(simple_assist_msg("On it.", "I should call it", "special_function", + R"({"arg1":1})", "")) + .run(); + + // Multiple typed arguments: the type lives in an attribute, and the + // value must come back as a JSON number, not the string "2". + tst.test("<|open|>response<|sep|><|close|>response<|sep|>" + "<|open|>tools<|sep|>" + "<|open|>call tool=\"special_function_with_opt\" index=\"1\"<|sep|>" + "<|open|>argument key=\"arg1\" type=\"number\"<|sep|>1<|close|>argument<|sep|>" + "<|open|>argument key=\"arg2\" type=\"number\"<|sep|>2<|close|>argument<|sep|>" + "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>") + .tools({ special_function_tool_with_optional_param }) + .expect_tool_calls({ + { "special_function_with_opt", R"({"arg1":1,"arg2":2})", "" }, + }) + .run(); + + // Parallel tool calls in one <|open|>tools<|sep|> section. + tst.test("<|open|>response<|sep|><|close|>response<|sep|>" + "<|open|>tools<|sep|>" + "<|open|>call tool=\"special_function\" index=\"1\"<|sep|>" + "<|open|>argument key=\"arg1\" type=\"number\"<|sep|>1<|close|>argument<|sep|>" + "<|close|>call<|sep|>" + "<|open|>call tool=\"special_function_with_opt\" index=\"2\"<|sep|>" + "<|open|>argument key=\"arg1\" type=\"number\"<|sep|>1<|close|>argument<|sep|>" + "<|open|>argument key=\"arg2\" type=\"number\"<|sep|>2<|close|>argument<|sep|>" + "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>") + .parallel_tool_calls(true) + .tools({ special_function_tool, special_function_tool_with_optional_param }) + .expect_tool_calls({ + { "special_function", R"({"arg1":1})", "" }, + { "special_function_with_opt", R"({"arg1":1,"arg2":2})", "" }, + }) + .run(); + + // String-typed argument keeps its literal text (no JSON coercion). + tst.test("<|open|>response<|sep|><|close|>response<|sep|>" + "<|open|>tools<|sep|>" + "<|open|>call tool=\"python\" index=\"1\"<|sep|>" + "<|open|>argument key=\"code\" type=\"string\"<|sep|>print('hey')" + "<|close|>argument<|sep|>" + "<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|>") + .tools({ python_tool }) + .expect_tool_calls({ + // custom delimiter: the payload itself contains )" + { "python", R"JSON({"code":"print('hey')"})JSON", "" }, + }) + .run(); + } + // Kimi-K2-Thinking tests - custom parser // Unique feature: tool call ID embeds function name as functions.: { From 252486605e430cc9f483bf5f048e1d2c8adf6a11 Mon Sep 17 00:00:00 2001 From: Deepankar Singh Date: Tue, 28 Jul 2026 21:11:26 +0530 Subject: [PATCH 07/21] chat : add message_delimiters for Kimi K3 Per-role message-start markers for token-level span splitting. User and assistant messages carry only the role attribute, so their full opener (through <|sep|>) is used; system and tool messages continue with more attributes (type=/tool=/index=), so those delimiters stop after the role's closing quote. Verified against the K3 tiktoken vocabulary that the closing quote is always a standalone token across all attribute variants, so the token-level prefix match stays exact. Co-Authored-By: Claude Fable 5 --- common/chat.cpp | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/common/chat.cpp b/common/chat.cpp index 53805c8e67a6..6efbd06397e2 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -2363,6 +2363,19 @@ static common_chat_params common_chat_params_init_kimi_k3(const common_chat_temp data.thinking_start_tag = THINK_START; data.thinking_end_tags = { THINK_END }; + // Per-role message-start delimiters. User/assistant messages carry only the + // role attribute, so their full opener (through <|sep|>) is used. System and + // tool messages continue with more attributes (type=/tool=/index=), so those + // delimiters stop after the role's closing quote - verified against the K3 + // tokenizer that the quote is always its own token and never merges with the + // following attribute text, keeping the token-level prefix match exact. + data.message_delimiters = { + { COMMON_CHAT_ROLE_ASSISTANT, "<|open|>message role=\"assistant\"<|sep|>" }, + { COMMON_CHAT_ROLE_USER, "<|open|>message role=\"user\"<|sep|>" }, + { COMMON_CHAT_ROLE_TOOL, "<|open|>message role=\"tool\"" }, + { COMMON_CHAT_ROLE_SYSTEM, "<|open|>message role=\"system\"" }, + }; + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; From cf8b10bd85de697a6fdb8060e09838d56a6169f6 Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Fri, 31 Jul 2026 23:12:56 +0200 Subject: [PATCH 08/21] fix: apply nits from @ngxson and text fixes from @danielhanchen --- conversion/kimi_k3.py | 8 -------- src/llama-context.cpp | 5 ++++- src/llama-model.cpp | 1 + src/llama-model.h | 1 + src/models/kimi-k3.cpp | 37 ++++++++++++++----------------------- src/models/models.h | 10 ++++------ 6 files changed, 24 insertions(+), 38 deletions(-) diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py index 1064249e5205..f39be534a52c 100644 --- a/conversion/kimi_k3.py +++ b/conversion/kimi_k3.py @@ -85,14 +85,6 @@ def set_vocab(self): # # compressed-tensors MXFP4 -> ggml MXFP4 # - # The real checkpoint stores only the routed experts quantized (everything - # else is excluded by quantization_config["ignore"]), as a - # weight_packed/weight_scale pair per expert. Both sides are 4-bit E2M1 with - # a per-32 E8M0 scale, so this is a pure repack - see repack_mxfp4_blocks. - # - # Dequantizing instead would be catastrophic here: the routed experts are - # ~1.38 TB at 4 bits, so a bf16 round-trip would need ~5.5 TB of output. - # def _is_mxfp4_packed(self) -> bool: quant_config = self.hparams.get("quantization_config") or {} diff --git a/src/llama-context.cpp b/src/llama-context.cpp index af0c5f37e1f1..ff64022905f1 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2293,9 +2293,12 @@ void llama_context::output_reorder() { uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { uint32_t res; + if (model.arch == LLM_ARCH_KIMI_K3) { + // the n_tokens*40 budget below is exhausted at ubatch 3840 + res = std::max(n_tokens * 160, 64u * model.n_tensors()); + } if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || - model.arch == LLM_ARCH_KIMI_K3 || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DEEPSEEK4 || diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 5bcf15d4827d..fdae29bb83f7 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -844,6 +844,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_397B_A17B: return "397B.A17B"; case LLM_TYPE_685B_A37B: return "685B.A37B"; case LLM_TYPE_744B_A40B: return "744B.A40B"; + case LLM_TYPE_2_8T_A50B: return "2.8T.A50B"; case LLM_TYPE_E2B: return "E2B"; case LLM_TYPE_E4B: return "E4B"; default: return "?B"; diff --git a/src/llama-model.h b/src/llama-model.h index 923862ca6d07..0a6777378124 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -143,6 +143,7 @@ enum llm_type { LLM_TYPE_397B_A17B, // Qwen3.5 LLM_TYPE_685B_A37B, // DeepSeek V3.2 LLM_TYPE_744B_A40B, // GLM-5 + LLM_TYPE_2_8T_A50B, // Kimi-K3 LLM_TYPE_E2B, LLM_TYPE_E4B, }; diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp index 2043a6a83498..c8e6da608072 100644 --- a/src/models/kimi-k3.cpp +++ b/src/models/kimi-k3.cpp @@ -35,14 +35,14 @@ void llama_model_kimi_k3::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); - ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent, false); + ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent); - ml.get_key(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size, false); - ml.get_key(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta, false); - ml.get_key(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta, false); + ml.get_key(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size); + ml.get_key(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta); + ml.get_key(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta); switch (hparams.n_layer()) { - case 93: type = LLM_TYPE_UNKNOWN; break; // Kimi-K3 + case 93: type = LLM_TYPE_2_8T_A50B; break; // Kimi-K3 default: type = LLM_TYPE_UNKNOWN; } } @@ -93,16 +93,7 @@ void llama_model_kimi_k3::load_arch_tensors(llama_model_loader &) { layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); // K3's A_log is a plain 1-D [n_head] parameter (kimi-linear's is padded); - // accept the padded forms too so both layouts load. -exp() is folded at - // conversion time. Only the element count matters - the graph reshapes it. layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, TENSOR_NOT_REQUIRED); - if (!layer.ssm_a) { - layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_head, 1, 1}, TENSOR_NOT_REQUIRED); - } - if (!layer.ssm_a) { - layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_head}, 0); - } - layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0); // K3 uses a single full-rank gate instead of kimi-linear's g_a/g_b pair @@ -197,27 +188,27 @@ static ggml_tensor * kimi_k3_situ(ggml_context * ctx0, ggml_tensor * gate, ggml_ // void llama_model_kimi_k3::graph::res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens) { - ckpts.push_back(ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens)); + resi.push_back(ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens)); } ggml_tensor * llama_model_kimi_k3::graph::res_stack(int64_t n_embd, int64_t n_tokens) { GGML_UNUSED(n_embd); GGML_UNUSED(n_tokens); - if (stack_cache_n == (int) ckpts.size()) { - return stack_cache; // set unchanged since the last mix + if (resi_stack_n == (int) resi.size()) { + return resi_stack; // set unchanged since the last mix } - ggml_tensor * acc = ckpts[0]; - for (size_t i = 1; i < ckpts.size(); ++i) { - acc = ggml_concat(ctx0, acc, ckpts[i], 1); + ggml_tensor * acc = resi[0]; + for (size_t i = 1; i < resi.size(); ++i) { + acc = ggml_concat(ctx0, acc, resi[i], 1); } - stack_cache = acc; - stack_cache_n = (int) ckpts.size(); + resi_stack = acc; + resi_stack_n = (int) resi.size(); return acc; } ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor * score_w, int64_t n_embd, int64_t n_tokens, int il) { - const int n_ckpt = (int) ckpts.size(); + const int n_ckpt = (int) resi.size(); if (n_ckpt == 0) { return cur; // layer 0: nothing banked yet } diff --git a/src/models/models.h b/src/models/models.h index 9bcf9f1d547c..9a0b3001ece1 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2292,12 +2292,10 @@ struct llama_model_kimi_k3 : public llama_model_base { // needs no transpose of the (large) stack. // Each checkpoint is kept as its own [n_embd, 1, n_token] tensor and the // [n_embd, n_ckpt, n_token] stack is materialised with ggml_concat only when - // the set changes (8 times per forward pass for the real model). A single - // preallocated buffer would be cheaper, but a bare ggml_new_tensor leaf has - // no backing buffer - ggml-alloc only allocates tensors that are op outputs. - std::vector ckpts; - ggml_tensor * stack_cache = nullptr; - int stack_cache_n = -1; + // the set changes (8 times per forward pass for the real model). + std::vector resi; + ggml_tensor * resi_stack = nullptr; + int resi_stack_n = -1; void res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens); ggml_tensor * res_stack(int64_t n_embd, int64_t n_tokens); From 0302fa3c00fce7842234577b93201e9a71d01211 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Sat, 1 Aug 2026 11:43:27 +0200 Subject: [PATCH 09/21] tests : added missing hyperparameters and tensors for Kimi K3 in test-llama-archs --- src/llama-model-saver.cpp | 6 ++++++ tests/test-llama-archs.cpp | 4 ++++ 2 files changed, 10 insertions(+) diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index abca773a9a26..7de57a5f8084 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -213,6 +213,7 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); add_kv(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, true); add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + add_kv(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent); add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp); add_kv(LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH, hparams.n_ff_chexp); add_kv(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp); @@ -376,6 +377,10 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_XIELU_BETA, hparams.xielu_beta); add_kv(LLM_KV_XIELU_EPS, hparams.xielu_eps); + add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size); + add_kv(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta); + add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, hparams.situ_linear_beta); + // deprecated // add_kv(LLM_KV_TOKENIZER_PREFIX_ID, ???); // add_kv(LLM_KV_TOKENIZER_SUFFIX_ID, ???); @@ -403,6 +408,7 @@ void llama_model_saver::add_tensors_from_model() { add_tensor(model->output_norm_enc); add_tensor(model->output_s); add_tensor(model->output_in_s); + add_tensor(model->output_res_score); add_tensor(model->cls); add_tensor(model->cls_b); add_tensor(model->cls_out); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index d77417d50463..38f640984429 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -220,6 +220,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { if (moe) { ms.add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, n_ff); ms.add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, n_ff / 2); // distinct from n_ff so a saver key-clobber surfaces on reload + ms.add_kv(LLM_KV_EXPERT_LATENT_LENGTH, n_ff); ms.add_kv(LLM_KV_INTERLEAVE_MOE_LAYER_STEP, uint32_t(2)); ms.add_kv(LLM_KV_EXPERT_COUNT, uint32_t(2)); ms.add_kv(LLM_KV_EXPERT_USED_COUNT, uint32_t(1)); @@ -245,6 +246,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_KDA_HEAD_DIM, uint32_t(128)); ms.add_kv(LLM_KV_WKV_HEAD_SIZE, n_embd/n_head); ms.add_kv(LLM_KV_SHORTCONV_L_CACHE, uint32_t(3)); + ms.add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE, uint32_t(12)); + ms.add_kv(LLM_KV_ACTIVATION_SITU_BETA, 4.0f); + ms.add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, 25.0f); for (uint32_t il = 0; il < n_layer; il++) { ggml_tensor t; From c4954216aa922ed7145bd9911987780279710654 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Sat, 1 Aug 2026 13:32:22 +0200 Subject: [PATCH 10/21] chore : move overly verbose header file comments to Kimi K3 source file --- src/models/kimi-k3.cpp | 9 +++++++++ src/models/models.h | 10 ---------- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp index c8e6da608072..d1479791514d 100644 --- a/src/models/kimi-k3.cpp +++ b/src/models/kimi-k3.cpp @@ -215,6 +215,15 @@ ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor const float eps = hparams.f_norm_rms_eps; + // `src` holds the block-residual checkpoints, one [n_embd, n_token] slab + // per checkpoint, packed as [n_embd, n_ckpt, n_token] so that: + // - ggml_rms_norm reduces over ne0 = n_embd (scores all slabs at once) + // - ggml_dsv4_hc_pre reduces over ne1 = n_ckpt (the weighted sum) + // Both requirements are satisfied by the same layout, which is why this + // needs no transpose of the (large) stack. + // Each checkpoint is kept as its own [n_embd, 1, n_token] tensor and the + // [n_embd, n_ckpt, n_token] stack is materialised with ggml_concat only when + // the set changes (8 times per forward pass for the real model). ggml_tensor * src = res_stack(n_embd, n_tokens); // [n_embd, n_ckpt, n_tokens] // Scores for the banked checkpoints. One rms_norm covers all of them because diff --git a/src/models/models.h b/src/models/models.h index 9a0b3001ece1..f8edb65560d2 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2283,16 +2283,6 @@ struct llama_model_kimi_k3 : public llama_model_base { const llama_model & model; // Cross-layer residual attention (K3's `_apply_attn_res`). - // - // `src` holds the block-residual checkpoints, one [n_embd, n_token] slab - // per checkpoint, packed as [n_embd, n_ckpt, n_token] so that: - // - ggml_rms_norm reduces over ne0 = n_embd (scores all slabs at once) - // - ggml_dsv4_hc_pre reduces over ne1 = n_ckpt (the weighted sum) - // Both requirements are satisfied by the same layout, which is why this - // needs no transpose of the (large) stack. - // Each checkpoint is kept as its own [n_embd, 1, n_token] tensor and the - // [n_embd, n_ckpt, n_token] stack is materialised with ggml_concat only when - // the set changes (8 times per forward pass for the real model). std::vector resi; ggml_tensor * resi_stack = nullptr; int resi_stack_n = -1; From 4bb78d58e539263005dd050256a2898366693fec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Tue, 4 Aug 2026 19:11:38 +0200 Subject: [PATCH 11/21] tests : re-enabled KIMI_K3 in test-llama-archs for WebGPU backend --- tests/test-llama-archs.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 38f640984429..331d988db419 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -443,7 +443,7 @@ static bool arch_supported(const llm_arch arch) { // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU - if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_KIMI_K3) { + if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MINIMAX_M3 || arch == LLM_ARCH_KIMI_K3) { return false; } #endif // GGML_USE_WEBGPU From 0d5346be5de5ceb98a8158d49c2819c64b236955 Mon Sep 17 00:00:00 2001 From: Caleb DeLeeuw Date: Tue, 4 Aug 2026 12:22:03 -0700 Subject: [PATCH 12/21] model-saver : emit kda_gate_lower_bound for Kimi K3 Quick fix. The Kimi K3 loader reads kda_gate_lower_bound and gates a graph branch on it (it scales the KDA gate when the bound is above -INFINITY), but the model saver never wrote the key, so a save->load roundtrip silently dropped it back to the -INFINITY default and changed the model's output. The real K3 config sets gate_lower_bound = -5.0. I propose to emit it from the saver, and set it to -5.0 in the test-llama-archs K3 case so the roundtrip check exercises it (the roundtrip fails without the saver line). --- src/llama-model-saver.cpp | 1 + tests/test-llama-archs.cpp | 1 + 2 files changed, 2 insertions(+) diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 7de57a5f8084..2e2e2d3d28db 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -320,6 +320,7 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_SSM_DT_B_C_RMS, hparams.ssm_dt_b_c_rms); add_kv(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); add_kv(LLM_KV_WKV_HEAD_SIZE, hparams.wkv_head_size); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 331d988db419..a49cc827746f 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -249,6 +249,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTN_RES_BLOCK_SIZE, uint32_t(12)); ms.add_kv(LLM_KV_ACTIVATION_SITU_BETA, 4.0f); ms.add_kv(LLM_KV_ACTIVATION_SITU_LINEAR_BETA, 25.0f); + ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, -5.0f); for (uint32_t il = 0; il < n_layer; il++) { ggml_tensor t; From 3f362699542aba518c735211456e8fa14b6d0c39 Mon Sep 17 00:00:00 2001 From: "Piotr Wilkin (ilintar)" Date: Thu, 13 Aug 2026 23:27:19 +0200 Subject: [PATCH 13/21] Refactor conditional for model architecture check --- src/llama-context.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ff64022905f1..6eb792fe2be0 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2296,8 +2296,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { if (model.arch == LLM_ARCH_KIMI_K3) { // the n_tokens*40 budget below is exhausted at ubatch 3840 res = std::max(n_tokens * 160, 64u * model.n_tensors()); - } - if (model.arch == LLM_ARCH_QWEN3NEXT || + } else if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || From e88142750648be113de4f0c8c87f3feddbb484c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Fri, 14 Aug 2026 20:36:13 +0200 Subject: [PATCH 14/21] tests : re-enabled (again) KIMI_K3 and MINIMAX_M3 in test-llama-archs for WebGPU backend --- tests/test-llama-archs.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index a49cc827746f..d27a283fcd29 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -444,7 +444,7 @@ static bool arch_supported(const llm_arch arch) { // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU - if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_MINIMAX_M3 || arch == LLM_ARCH_KIMI_K3) { + if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA) { return false; } #endif // GGML_USE_WEBGPU From 3f99a416f263b468b00991c91fdc42d87377f2d2 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Sat, 15 Aug 2026 15:08:28 +0200 Subject: [PATCH 15/21] fix code comments --- common/chat.cpp | 55 +++++++++++--------------- conversion/base.py | 20 ++++------ conversion/kimi_k3.py | 71 +++++++++++++--------------------- models/templates/Kimi-K3.jinja | 5 +-- src/llama-graph.h | 2 +- src/models/kimi-k3.cpp | 60 +++++++++------------------- tests/test-chat.cpp | 8 ++-- 7 files changed, 81 insertions(+), 140 deletions(-) diff --git a/common/chat.cpp b/common/chat.cpp index 6efbd06397e2..8a7c5d3940ed 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -2321,11 +2321,11 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha return data; } -// Kimi K3 - XTML-ish tagged format from the model's own template: +// Kimi K3 - XTML tagged format, built by open_tag/close_tag macros: // open_tag(t, attrs) = <|open|>t k="v"...<|sep|> close_tag(t) = <|close|>t<|sep|> // assistant := [think] [response] [tools] close_tag(message) <|end_of_msg|> -// Note the generation prompt already opens the think (or response) section, so -// the section opener is optional here - same situation as Kimi K2 Thinking. +// the generation prompt already opens the think (or response) section, so the +// section opener is optional here - same as Kimi K2 Thinking static common_chat_params common_chat_params_init_kimi_k3(const common_chat_template & tmpl, const autoparser::generation_params & inputs) { common_chat_params data; @@ -2350,9 +2350,8 @@ static common_chat_params common_chat_params_init_kimi_k3(const common_chat_temp const std::string MSG_END = "<|close|>message<|sep|>"; const std::string EOM_TOKEN = "<|end_of_msg|>"; - // The four markers are the only special tokens; tag names ("think", - // "response", "message") are ordinary tokens and must NOT be preserved, - // or ordinary prose containing those words would be mangled. + // only the markers are special tokens. tag names ("think", "response", ...) are + // normal tokens and must not be preserved, or prose with those words is broken data.preserved_tokens = { "<|open|>", "<|close|>", @@ -2363,12 +2362,9 @@ static common_chat_params common_chat_params_init_kimi_k3(const common_chat_temp data.thinking_start_tag = THINK_START; data.thinking_end_tags = { THINK_END }; - // Per-role message-start delimiters. User/assistant messages carry only the - // role attribute, so their full opener (through <|sep|>) is used. System and - // tool messages continue with more attributes (type=/tool=/index=), so those - // delimiters stop after the role's closing quote - verified against the K3 - // tokenizer that the quote is always its own token and never merges with the - // following attribute text, keeping the token-level prefix match exact. + // per-role message-start delimiters. user/assistant messages only have the role + // attribute, so the full opener is used. system and tool messages have more + // attributes, so those delimiters stop after the closing quote of the role data.message_delimiters = { { COMMON_CHAT_ROLE_ASSISTANT, "<|open|>message role=\"assistant\"<|sep|>" }, { COMMON_CHAT_ROLE_USER, "<|open|>message role=\"user\"<|sep|>" }, @@ -2396,27 +2392,23 @@ static common_chat_params common_chat_params_init_kimi_k3(const common_chat_temp auto start = p.optional(p.literal(MSG_START)); - // The think section is ALWAYS consumed, even when reasoning extraction - // is off: K3's generation prompt ends with open_tag('think'), so the - // opener is present on every request and would otherwise leak into - // content. With extraction off the thoughts fall into content, matching - // how the other reasoning models behave. - // Reasoning stops at its own closer, or at the response opener if the - // model skips the closer entirely (seen on short answers). + // the think section is always consumed, even with reasoning extraction off: + // the generation prompt ends with open_tag('think'), so it is always present. + // reasoning stops at its own closer, or at the response opener if the model + // skips the closer auto think_body = extract_reasoning ? p.reasoning(p.until_one_of({ THINK_END, RESP_START })) : p.content(p.until_one_of({ THINK_END, RESP_START })); auto reasoning = p.optional(p.optional(p.literal(THINK_START)) + think_body + p.optional(p.literal(THINK_END))); - // Content runs to the response closer, or to whatever comes next if a - // truncated generation never emits one. + // content runs to the response closer, or to the next section if truncated auto response = p.optional(p.literal(RESP_START)) + p.content(p.until_one_of({ RESP_END, TOOLS_START, MSG_END })) + p.optional(p.literal(RESP_END)); - // The message closer is followed by the EOG token, which reaches the - // parser as text and must be consumed or the parse is left incomplete. + // the EOG token after the message closer reaches the parser as text, + // so it must be consumed or the parse stays incomplete auto trailer = p.optional(p.literal(MSG_END)) + p.optional(p.literal(EOM_TOKEN)); if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { @@ -2429,10 +2421,9 @@ static common_chat_params common_chat_params_init_kimi_k3(const common_chat_temp std::string name = function.at("name"); const json schema = function.contains("parameters") ? function.at("parameters") : json::object(); - // Arguments arrive one tag per key, with the JSON type carried in a - // type="..." attribute. We take the type from the tool schema - // instead - it is authoritative, and it tells us whether the value - // should be parsed as JSON or kept as a literal string. + // arguments come one tag per key, with the JSON type in a type="..." + // attribute. the type is taken from the tool schema instead, as it tells + // us if the value is JSON or a literal string auto args = p.eps(); if (schema.contains("properties") && !schema.at("properties").empty()) { auto arg_choices = p.choice(); @@ -2466,10 +2457,9 @@ static common_chat_params common_chat_params_init_kimi_k3(const common_chat_temp tool_choices |= p.rule("kimi-k3-tool-" + name, call); }); - // K3 emits every call inside one <|open|>tools<|sep|> section, then - // closes the message. The message closer is part of the trigger rule so - // that the lazy grammar still permits it once tool calls have started - - // otherwise constrained decoding rejects the model's own closing tag. + // all calls go inside one tools section, then the message is closed. the + // message closer is part of the trigger rule, or else the lazy grammar + // rejects it once tool calls have started auto tools_section = p.trigger_rule("kimi-k3-tool-call", p.literal(TOOLS_START) + p.one_or_more(tool_choices) + p.literal(TOOLS_END) + p.optional(p.literal(MSG_END)) + @@ -3470,8 +3460,7 @@ std::optional common_chat_try_specialized_template( return common_chat_params_init_kimi_k2(tmpl, params); } - // Kimi K3 - XTML-ish tagged format built from open_tag/close_tag macros. - // Detection: the <|open|>/<|close|>/<|sep|> marker trio is unique to K3. + // Kimi K3 - the <|open|>/<|close|>/<|end_of_msg|> markers are unique to it if (src.find("<|open|>") != std::string::npos && src.find("<|close|>") != std::string::npos && src.find("<|end_of_msg|>") != std::string::npos) { LOG_DBG("Using specialized template: Kimi K3\n"); diff --git a/conversion/base.py b/conversion/base.py index 6a9433c93da5..6150ccd70032 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -79,23 +79,17 @@ class ModelType(IntEnum): def repack_mxfp4_blocks(packed: Tensor, scale: Tensor) -> np.ndarray: """ - Repack 4-bit MX weights into ggml `block_mxfp4`. Lossless - this only moves - bits, it does not dequantize and requantize. + Repack 4-bit MX weights into ggml `block_mxfp4`. Lossless - only moves bits. - Source (compressed-tensors "mxfp4-pack-quantized", and DeepSeek-V4's - equivalent weight/scale pair): - packed uint8 [rows, cols/2] two 4-bit codes per byte, element 2i in the - low nibble and 2i+1 in the high nibble + Source (compressed-tensors "mxfp4-pack-quantized", also used by DeepSeek-V4): + packed uint8 [rows, cols/2] element 2i in the low nibble, 2i+1 in the high one scale uint8 [rows, cols/32] one E8M0 biased exponent per 32-element group - Destination, per 32-element group: one scale byte then 16 code bytes, where - byte j holds element j in the low nibble and element j+16 in the high nibble - (see dequantize_row_mxfp4 in ggml-quants.c). + Destination, per group: one scale byte then 16 code bytes, where byte j holds + element j in the low nibble and element j+16 in the high one. - The 4-bit codes themselves need no remapping: both sides use sign in bit 3 - and a magnitude index into (0, .5, 1, 1.5, 2, 3, 4, 6), which is exactly - ggml's kvalues_mxfp4 order. ggml's kvalues are doubled and its scale is - halved (GGML_E8M0_TO_FP32_HALF), so the represented value is unchanged. + The 4-bit codes need no remapping: both sides index into ggml's kvalues_mxfp4 + order. ggml doubles the kvalues and halves the scale, so the value is the same. """ p = packed.contiguous().view(torch.uint8) s = scale.contiguous().view(torch.uint8) diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py index f39be534a52c..3e2f6d65addc 100644 --- a/conversion/kimi_k3.py +++ b/conversion/kimi_k3.py @@ -19,10 +19,9 @@ class KimiK3Model(TextModel): """ Kimi-K3 text model (KimiLinearForCausalLM under a `language_model.` prefix). - Shares the hybrid MLA + KDA skeleton with Kimi-Linear-48B but is not - loadable by that converter: K3 adds cross-layer attention residuals, a - latent MoE, the situ activation, an MLA output gate and a full-rank KDA - gate, none of which exist in the older architecture. + Shares the hybrid MLA + KDA skeleton with kimi-linear, but that converter + cannot load it: K3 adds cross-layer attention residuals, a latent MoE, the + situ activation, an MLA output gate and a full-rank KDA gate. The vision tower and mm_projector are skipped - text only for now. """ @@ -31,11 +30,9 @@ class KimiK3Model(TextModel): _experts: list[dict[str, Tensor]] | None = None - # `_res_norm.weight` and `_res_proj.weight` are only ever used as the - # elementwise product norm.weight * proj.weight (see _apply_attn_res in - # modeling_kimi_linear.py), so they are fused into a single [n_embd] vector - # at conversion time. They arrive as separate tensors, so buffer whichever - # comes first, tagged with which one it is. + # `_res_norm.weight` and `_res_proj.weight` are only used as their + # elementwise product, so they are fused into one [n_embd] vector here. + # they arrive apart, so buffer the first one and tag it with its kind. _res_parts: dict[str, tuple[str, Tensor]] # HF suffix -> (gguf tensor, per-layer?) @@ -45,9 +42,8 @@ class KimiK3Model(TextModel): "output_attn_res": (gguf.MODEL_TENSOR.OUTPUT_RES_SCORE, False), } - # compressed-tensors MXFP4. The `language_model.` prefix is still present here: - # self.model_tensors is keyed by the raw checkpoint names, get_tensors() strips - # the prefix only on the way out. + # compressed-tensors MXFP4. the `language_model.` prefix is still there, as + # self.model_tensors is keyed by the raw checkpoint names _MXFP4_FORMAT = "mxfp4-pack-quantized" _MXFP4_EXPERT_RE = re.compile( r"^(?:language_model\.)?model\.layers\.(\d+)" @@ -64,20 +60,14 @@ def __init__(self, *args, **kwargs): self._res_parts = {} def set_vocab(self): - # K3 ships the same TikToken vocab as K2: its pre-tokenizer hashes to - # 81212dc7... which base.py already maps to "kimi-k2", so no new - # pre-tokenizer registration is needed. - # - # Borrowed rather than inherited: K3 shares kimi-linear's vocab handling - # but none of its tensor layout. The method only touches TextModel - # members, so an unrelated TextModel is a valid receiver. + # K3 has the same TikToken vocab as K2, so kimi-linear's vocab handling works. + # borrowed, not inherited: the method only touches TextModel members, and K3 + # shares none of kimi-linear's tensor layout. KimiLinearModel.set_vocab(self) # ty: ignore[invalid-argument-type] - # ...but K2's converter ends by forcing eos to the tokenizer's own - # eos_id, which for K3 is 163585 = [EOS], the *document* terminator. - # K3's config and generation_config both say 163586 = <|end_of_msg|>, - # the chat turn terminator. Keeping [EOS] means generation never stops - # at the end of an assistant turn. Restore the configured value. + # ...but that forces eos to the tokenizer's eos_id, which is [EOS], the + # document terminator. K3's config says <|end_of_msg|>, the turn terminator; + # with [EOS] the generation never stops at the end of a turn. if (eos := self.hparams.get("eos_token_id")) is not None: logger.info(f"restoring configured eos_token_id {eos} (kimi-linear forces the tokenizer's)") self.gguf_writer.add_eos_token_id(eos) @@ -95,8 +85,8 @@ def dequant_model(self): if not self._is_mxfp4_packed(): return super().dequant_model() - # Skipping base.py's dequant is only safe because the experts are the - # *only* quantized tensors. Verify that rather than assume it. + # skipping base.py's dequant is only safe if the experts are the only + # quantized tensors, so check it stray = [n for n in self.model_tensors if n.endswith(".weight_packed") and not self._MXFP4_EXPERT_RE.match(n)] if stray: @@ -109,11 +99,9 @@ def _mxfp4_expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], Callabl """ One stacked [n_expert, rows, cols] MXFP4 tensor, built lazily. - Laziness is not an optimization here, it is the difference between - working and not: gguf_writer holds every added tensor until the final - write, so materializing this eagerly (as the DeepSeek-V4 and NVFP4 paths - do) would keep all ~1.38 TB of experts resident. Deferring it means only - the tensor currently being written is in memory, one expert at a time. + gguf_writer holds every added tensor until the final write, so building + this eagerly (like the DeepSeek-V4 path does) keeps all ~1.38 TB of + experts in memory. lazy means only the tensor being written is resident. """ # meta shapes, so this does not read any weights rows, packed_cols = loaders[0][0]().shape @@ -129,9 +117,8 @@ def load(fns: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]) -> np.nda ) return out - # loaders goes through args rather than the closure so that `func` matches - # LazyBase's single-argument shape; _recurse_apply passes plain callables - # through untouched. + # loaders goes through args, not the closure, so that `func` matches + # LazyBase's single-argument shape return gguf.LazyNumpyTensor( meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(np.uint8, byte_shape), args=(loaders,), @@ -183,9 +170,8 @@ def _write_mxfp4_experts(self) -> None: del self.model_tensors[name] def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: - # Deliberately not a generator: base.py builds - # chain(generate_extra_tensors(), get_tensors()), so the tensors consumed - # here must be removed from model_tensors before get_tensors() starts. + # not a generator on purpose: base.py chains this with get_tensors(), so the + # tensors used here must be removed from model_tensors before that starts if self._is_mxfp4_packed(): self._write_mxfp4_experts() return () @@ -207,9 +193,8 @@ def set_gguf_parameters(self): linear_attn_config = self.hparams["linear_attn_config"] - # layer types: n_head_kv == 0 marks a KDA (recurrent) layer. - # KimiLinearConfig.is_kda_layer uses (layer_idx + 1) in kda_layers, so - # the lists are 1-indexed - an off-by-one here silently produces garbage. + # n_head_kv == 0 marks a KDA (recurrent) layer. the layer lists are 1-indexed, + # as KimiLinearConfig.is_kda_layer uses (layer_idx + 1) full_attn_layers = linear_attn_config["full_attn_layers"] n_kv_heads = [ self.hparams["num_key_value_heads"] if (il + 1) in full_attn_layers else 0 @@ -278,9 +263,7 @@ def _try_fuse_res(self, data_torch: Tensor, name: str, bid: int | None): """ Pair _res_norm.weight with _res_proj.weight and emit their product. - proj is [1, n_embd]; norm is [n_embd]. _apply_attn_res only ever uses - norm.weight * proj.weight.squeeze(0), so one vector is enough. - Returns None if this is not a res tensor, [] if buffered pending its pair. + Returns None if this is not a res tensor, [] if buffered until its pair. """ for prefix, (tensor_id, per_layer) in self._RES_FUSIONS.items(): for kind in ("norm", "proj"): @@ -312,7 +295,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter # --- KDA conv1d: HF [d_inner, 1, d_conv] -> ggml ne [d_conv, 1, d_inner, 1] --- # GGUF reverses the numpy shape on write, so target numpy (1, d_inner, 1, d_conv). - # Both layouts have conv_step varying fastest, so this is a pure reshape. + # conv_step varies fastest in both layouts, so this is a pure reshape. if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")): if data_torch.ndim == 3: # [d_inner, 1, d_conv] d_inner, _, d_conv = data_torch.shape diff --git a/models/templates/Kimi-K3.jinja b/models/templates/Kimi-K3.jinja index 67e51a6af4b3..48de47fc9029 100644 --- a/models/templates/Kimi-K3.jinja +++ b/models/templates/Kimi-K3.jinja @@ -61,9 +61,8 @@ {%- endmacro -%} {%- macro json_sorted(value) -%} -{#- Minja は tojson(sort_keys=true) を実装していない。K3 参照実装は - deep_sort_dict() の後に compact JSON 化するため、dictsort と再帰マクロで - 同じバイト列を作る。配列の順序は保持し、mapping の各階層だけソートする。 -#} +{#- tojson has no sort_keys, so sort each mapping level with dictsort to match the + reference implementation. Array order is kept as-is. -#} {%- if value is mapping -%} {{- '{' -}} {%- for key, item in value|dictsort -%} diff --git a/src/llama-graph.h b/src/llama-graph.h index af38e1cca519..94324c7457ed 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -59,7 +59,7 @@ enum llm_ffn_op_type : int { LLM_FFN_GEGLU, LLM_FFN_REGLU, LLM_FFN_SWIGLU_OAI_MOE, - LLM_FFN_SITU, // kimi-k3 (appended: do not renumber existing values) + LLM_FFN_SITU, // kimi-k3 }; enum llm_ffn_gate_type { diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp index d1479791514d..af341f1332cf 100644 --- a/src/models/kimi-k3.cpp +++ b/src/models/kimi-k3.cpp @@ -2,11 +2,8 @@ #include "llama-memory-recurrent.h" // -// Kimi-K3 text model. -// -// Hybrid KDA (linear) + MLA (full) attention, as in Kimi-Linear-48B, plus five -// things that architecture does not have: -// +// Kimi-K3 text model: hybrid KDA (linear) + MLA (full) attention, as in kimi-linear. +// Parts that kimi-linear does not have: // 1. cross-layer residual attention (attn_res_block_size) // 2. latent MoE (routed experts run at n_expert_latent) // 3. situ activation (replaces SwiGLU everywhere) @@ -92,7 +89,7 @@ void llama_model_kimi_k3::load_arch_tensors(llama_model_loader &) { layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, 0); layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); - // K3's A_log is a plain 1-D [n_head] parameter (kimi-linear's is padded); + // K3's A_log is a plain 1-D [n_head] tensor (kimi-linear's is padded) layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, TENSOR_NOT_REQUIRED); layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0); @@ -167,11 +164,8 @@ std::unique_ptr llama_model_kimi_k3::build_arch_graph(const l return std::make_unique(*this, params); } -// -// situ activation: -// situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) -// linear_beta <= 0 disables the transform on the up branch. -// +// situ(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * linear_beta*tanh(up/linear_beta) +// linear_beta <= 0 disables the transform on the up branch static ggml_tensor * kimi_k3_situ(ggml_context * ctx0, ggml_tensor * gate, ggml_tensor * up, float beta, float linear_beta) { ggml_tensor * a = ggml_scale(ctx0, ggml_tanh(ctx0, ggml_scale(ctx0, gate, 1.0f/beta)), beta); @@ -215,27 +209,18 @@ ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor const float eps = hparams.f_norm_rms_eps; - // `src` holds the block-residual checkpoints, one [n_embd, n_token] slab - // per checkpoint, packed as [n_embd, n_ckpt, n_token] so that: - // - ggml_rms_norm reduces over ne0 = n_embd (scores all slabs at once) - // - ggml_dsv4_hc_pre reduces over ne1 = n_ckpt (the weighted sum) - // Both requirements are satisfied by the same layout, which is why this - // needs no transpose of the (large) stack. - // Each checkpoint is kept as its own [n_embd, 1, n_token] tensor and the - // [n_embd, n_ckpt, n_token] stack is materialised with ggml_concat only when - // the set changes (8 times per forward pass for the real model). + // the checkpoint stack is packed as [n_embd, n_ckpt, n_tokens]: rms_norm reduces over + // ne0 = n_embd and dsv4_hc_pre reduces over ne1 = n_ckpt, so no transpose is needed ggml_tensor * src = res_stack(n_embd, n_tokens); // [n_embd, n_ckpt, n_tokens] - // Scores for the banked checkpoints. One rms_norm covers all of them because - // ne0 is n_embd. NOTE: the scores use the *normalized* values, but the weighted - // sum below uses the *raw* ones - mirroring _apply_attn_res. + // one rms_norm scores all checkpoints at once + // note: the scores use the normalized values, but the sum below uses the raw ones ggml_tensor * sc_src = ggml_rms_norm(ctx0, src, eps); sc_src = ggml_mul(ctx0, sc_src, score_w); sc_src = ggml_sum_rows(ctx0, sc_src); // [1, n_ckpt, n_tokens] sc_src = ggml_reshape_2d(ctx0, sc_src, n_ckpt, n_tokens); - // The current residual stream is scored separately and kept out of the stack, - // so the stack stays append-only. + // the current residual stream is scored apart, so the stack stays append-only ggml_tensor * sc_cur = ggml_rms_norm(ctx0, cur, eps); sc_cur = ggml_mul(ctx0, sc_cur, score_w); sc_cur = ggml_sum_rows(ctx0, sc_cur); // [1, n_tokens] @@ -244,8 +229,7 @@ ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor ggml_tensor * probs = ggml_soft_max(ctx0, scores); // over ne0 = n_ckpt+1 cb(probs, "res_probs", il); - // Split the convex combination: hc_pre reduces over ne1 for the stacked part, - // a plain broadcast-multiply handles the current stream. + // split the sum: hc_pre handles the stack, a broadcast-multiply the current stream ggml_tensor * p_src = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, n_ckpt, n_tokens, probs->nb[1], 0)); ggml_tensor * p_cur = ggml_cont(ctx0, ggml_view_2d(ctx0, probs, 1, n_tokens, probs->nb[1], probs->nb[0] * n_ckpt)); @@ -300,8 +284,8 @@ llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_par for (int il = 0; il < n_layer; ++il) { const auto & layer = model.layers[il]; - // `prefix_sum` is the residual stream. On checkpoint layers it is banked - // into res_stack and restarts from the attention output alone. + // the residual stream, banked on checkpoint layers and then restarted + // from the attention output alone ggml_tensor * prefix_sum = inpL; cur = use_attn_res ? res_mix(prefix_sum, layer.attn_res_score, n_embd, n_tokens, il) @@ -377,8 +361,7 @@ llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_par // KDA layer // -// Causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use. -// Polled rather than blocking is not a concern here; this mirrors kimi-linear's helper. +// causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all, int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w, @@ -432,14 +415,10 @@ ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer( cb(Kcur, "kda_k_conv", il); cb(Vcur, "kda_v_conv", il); - // The decay gate has two forms, selected by linear_attn_config.gate_lower_bound - // (fla/ops/kda/gate.py). `lower_bound` is NOT a clamp - when set it swaps the - // activation entirely: - // + // gate_lower_bound is not a clamp - when set, it swaps the decay gate activation: // unset (kimi-linear): g = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias) // set (K3, -5.0): g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias)) - // - // ssm_a holds -exp(A_log) (folded at conversion time), so exp(A_log) == -ssm_a. + // ssm_a holds -exp(A_log) (folded at conversion time), so exp(A_log) == -ssm_a ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a); g1 = ggml_add(ctx0, g1, layer.ssm_dt_b); @@ -602,14 +581,13 @@ ggml_tensor * llama_model_kimi_k3::graph::build_latent_moe( ? ggml_mul_mat(ctx0, layer.ffn_routed_down, cur) : cur; - // The router scores the FULL-WIDTH input while the experts consume the latent - // one, so the logits are computed here and handed to build_moe_ffn via - // `probs_in` (which substitutes for the gate_inp matmul). + // the router scores the full-width input while the experts take the latent one, + // so the logits are computed here and passed to build_moe_ffn ggml_tensor * logits = ggml_mul_mat(ctx0, layer.ffn_gate_inp, identity); cb(logits, "ffn_moe_logits", il); ggml_tensor * moe_out = build_moe_ffn(routed_in, - nullptr, // gate_inp unused: logits supplied below + nullptr, // gate_inp unused: the logits above are passed instead layer.ffn_up_exps, layer.ffn_gate_exps, layer.ffn_down_exps, diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 807869d048e3..3b3f73a1346e 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -4463,7 +4463,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) { } // Kimi-K3 tests - custom parser - // Unique feature: XTML-ish tags built from <|open|>/<|close|>/<|sep|>, and a + // Unique feature: XTML tags built from <|open|>/<|close|>/<|sep|>, and a // generation prompt that leaves the think section already open. { auto tst = peg_tester("models/templates/Kimi-K3.jinja", detailed_debug); @@ -4474,8 +4474,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .expect(message_assist) .run(); - // Reasoning with NO opening tag - the generation prompt already opened - // it. This is the case that silently loses reasoning if unhandled. + // Reasoning with no opening tag - the generation prompt already opened it tst.test("I'm thinking about this<|close|>think<|sep|>" "<|open|>response<|sep|>Hello, world!\nWhat's up?<|close|>response<|sep|>" "<|close|>message<|sep|>") @@ -4520,8 +4519,7 @@ static void test_template_output_peg_parsers(bool detailed_debug) { R"({"arg1":1})", "")) .run(); - // Multiple typed arguments: the type lives in an attribute, and the - // value must come back as a JSON number, not the string "2". + // Multiple typed arguments: values must come back as JSON numbers, not strings tst.test("<|open|>response<|sep|><|close|>response<|sep|>" "<|open|>tools<|sep|>" "<|open|>call tool=\"special_function_with_opt\" index=\"1\"<|sep|>" From 493ad3b897a144704fb1802b159f39d8e4948e45 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Sat, 15 Aug 2026 15:11:27 +0200 Subject: [PATCH 16/21] add template on conversion --- conversion/kimi_k3.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py index 3e2f6d65addc..e9ecfd902954 100644 --- a/conversion/kimi_k3.py +++ b/conversion/kimi_k3.py @@ -1,6 +1,7 @@ from __future__ import annotations import re +from pathlib import Path from typing import Callable, Iterable, Iterator, TYPE_CHECKING import numpy as np @@ -72,6 +73,13 @@ def set_vocab(self): logger.info(f"restoring configured eos_token_id {eos} (kimi-linear forces the tokenizer's)") self.gguf_writer.add_eos_token_id(eos) + # K3 renders chats in python (encoding_k3.py) and ships no jinja template, + # so add the bundled one when the model has none + if gguf.SpecialVocab(self.dir_model, load_merges=False).chat_template is None: + template_path = Path(__file__).parent.parent / "models" / "templates" / "Kimi-K3.jinja" + logger.info(f"gguf: model has no chat template, using {template_path.name}") + self.gguf_writer.add_chat_template(template_path.read_text(encoding="utf-8")) + # # compressed-tensors MXFP4 -> ggml MXFP4 # From e1039c233fae3dc830693736bf4f8ae51bc90298 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Sat, 15 Aug 2026 15:14:15 +0200 Subject: [PATCH 17/21] move repack_mxfp4_blocks to model base --- conversion/base.py | 74 +++++++++++++++++++++--------------------- conversion/deepseek.py | 6 ++-- conversion/kimi_k3.py | 4 +-- 3 files changed, 41 insertions(+), 43 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index 6150ccd70032..4f4fbca347aa 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -77,43 +77,6 @@ class ModelType(IntEnum): MMPROJ = 2 -def repack_mxfp4_blocks(packed: Tensor, scale: Tensor) -> np.ndarray: - """ - Repack 4-bit MX weights into ggml `block_mxfp4`. Lossless - only moves bits. - - Source (compressed-tensors "mxfp4-pack-quantized", also used by DeepSeek-V4): - packed uint8 [rows, cols/2] element 2i in the low nibble, 2i+1 in the high one - scale uint8 [rows, cols/32] one E8M0 biased exponent per 32-element group - - Destination, per group: one scale byte then 16 code bytes, where byte j holds - element j in the low nibble and element j+16 in the high one. - - The 4-bit codes need no remapping: both sides index into ggml's kvalues_mxfp4 - order. ggml doubles the kvalues and halves the scale, so the value is the same. - """ - p = packed.contiguous().view(torch.uint8) - s = scale.contiguous().view(torch.uint8) - - rows, packed_cols = p.shape - cols = packed_cols * 2 - if cols % 32 != 0: - raise ValueError(f"MXFP4 source row has {cols} values, expected a multiple of 32") - - n_blocks = cols // 32 - if tuple(s.shape) != (rows, n_blocks): - raise ValueError(f"MXFP4 scale shape {tuple(s.shape)} does not match {(rows, n_blocks)}") - - src = p.reshape(rows, n_blocks, 16) - lo = src & 0x0F # elements 0, 2, 4, ... - hi = (src >> 4) & 0x0F # elements 1, 3, 5, ... - - vals = torch.stack((lo, hi), dim=-1).reshape(rows, n_blocks, 32) - qs = vals[:, :, :16] | (vals[:, :, 16:] << 4) - - raw = torch.cat((s.unsqueeze(-1), qs.to(torch.uint8)), dim=-1) - return raw.reshape(rows, n_blocks * 17).cpu().numpy() - - class ModelBase: _model_classes: dict[ModelType, dict[str, type[ModelBase]]] = { ModelType.TEXT: {}, @@ -695,6 +658,43 @@ def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: return () + @staticmethod + def repack_mxfp4_blocks(packed: Tensor, scale: Tensor) -> np.ndarray: + """ + Repack 4-bit MX weights into ggml `block_mxfp4`. Lossless - only moves bits. + + Source (compressed-tensors "mxfp4-pack-quantized", also used by DeepSeek-V4): + packed uint8 [rows, cols/2] element 2i in the low nibble, 2i+1 in the high one + scale uint8 [rows, cols/32] one E8M0 biased exponent per 32-element group + + Destination, per group: one scale byte then 16 code bytes, where byte j holds + element j in the low nibble and element j+16 in the high one. + + The 4-bit codes need no remapping: both sides index into ggml's kvalues_mxfp4 + order. ggml doubles the kvalues and halves the scale, so the value is the same. + """ + p = packed.contiguous().view(torch.uint8) + s = scale.contiguous().view(torch.uint8) + + rows, packed_cols = p.shape + cols = packed_cols * 2 + if cols % 32 != 0: + raise ValueError(f"MXFP4 source row has {cols} values, expected a multiple of 32") + + n_blocks = cols // 32 + if tuple(s.shape) != (rows, n_blocks): + raise ValueError(f"MXFP4 scale shape {tuple(s.shape)} does not match {(rows, n_blocks)}") + + src = p.reshape(rows, n_blocks, 16) + lo = src & 0x0F # elements 0, 2, 4, ... + hi = (src >> 4) & 0x0F # elements 1, 3, 5, ... + + vals = torch.stack((lo, hi), dim=-1).reshape(rows, n_blocks, 32) + qs = vals[:, :, :16] | (vals[:, :, 16:] << 4) + + raw = torch.cat((s.unsqueeze(-1), qs.to(torch.uint8)), dim=-1) + return raw.reshape(rows, n_blocks * 17).cpu().numpy() + @staticmethod def _nvfp4_pack(weight: Tensor, scale: Tensor) -> tuple[np.ndarray, list[int]]: """Repack NVFP4 ModelOpt tensors into ggml super-block layout. diff --git a/conversion/deepseek.py b/conversion/deepseek.py index e9d539a4ef61..dbfa96c2ee22 100644 --- a/conversion/deepseek.py +++ b/conversion/deepseek.py @@ -12,7 +12,7 @@ if TYPE_CHECKING: from torch import Tensor -from .base import LazyTorchTensor, MmprojModel, ModelBase, TextModel, gguf, logger, repack_mxfp4_blocks +from .base import LazyTorchTensor, MmprojModel, ModelBase, TextModel, gguf, logger from .qwen import QwenModel @@ -709,8 +709,6 @@ def dequant_fp8_weight(weight: Tensor, scale: Tensor) -> Tensor: for name in tensors_to_remove: del self.model_tensors[name] - _pack_mxfp4_blocks = staticmethod(repack_mxfp4_blocks) - def _write_mxfp4_expert_tensor(self, bid: int, proj: str, tensor_key: gguf.MODEL_TENSOR) -> list[str]: n_experts = self.hparams["n_routed_experts"] data: np.ndarray | None = None @@ -724,7 +722,7 @@ def _write_mxfp4_expert_tensor(self, bid: int, proj: str, tensor_key: gguf.MODEL weight = LazyTorchTensor.to_eager(self.model_tensors[weight_name]()) scale = LazyTorchTensor.to_eager(self.model_tensors[scale_name]()) - packed = self._pack_mxfp4_blocks(weight, scale) + packed = self.repack_mxfp4_blocks(weight, scale) if data is None: data = np.empty((n_experts, *packed.shape), dtype=packed.dtype) data[eid] = packed diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py index e9ecfd902954..3adfbb0a0474 100644 --- a/conversion/kimi_k3.py +++ b/conversion/kimi_k3.py @@ -10,7 +10,7 @@ if TYPE_CHECKING: from torch import Tensor -from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger, repack_mxfp4_blocks +from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger from .kimi_linear import KimiLinearModel @@ -119,7 +119,7 @@ def _mxfp4_expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], Callabl def load(fns: list[tuple[Callable[[], Tensor], Callable[[], Tensor]]]) -> np.ndarray: out = np.empty(byte_shape, dtype=np.uint8) for eid, (packed_fn, scale_fn) in enumerate(fns): - out[eid] = repack_mxfp4_blocks( + out[eid] = self.repack_mxfp4_blocks( LazyTorchTensor.to_eager(packed_fn()), LazyTorchTensor.to_eager(scale_fn()), ) From 7b902958b99cc5f3cec6278af5debc978cca44c3 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Sat, 15 Aug 2026 15:44:06 +0200 Subject: [PATCH 18/21] nits --- src/models/kimi-k3.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp index af341f1332cf..10e47e4b5a4d 100644 --- a/src/models/kimi-k3.cpp +++ b/src/models/kimi-k3.cpp @@ -32,7 +32,7 @@ void llama_model_kimi_k3::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); - ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent); + ml.get_key(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent, false); ml.get_key(LLM_KV_ATTN_RES_BLOCK_SIZE, hparams.attn_res_block_size); ml.get_key(LLM_KV_ACTIVATION_SITU_BETA, hparams.situ_beta); @@ -90,7 +90,7 @@ void llama_model_kimi_k3::load_arch_tensors(llama_model_loader &) { layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); // K3's A_log is a plain 1-D [n_head] tensor (kimi-linear's is padded) - layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, TENSOR_NOT_REQUIRED); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, 0); layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0); // K3 uses a single full-rank gate instead of kimi-linear's g_a/g_b pair From 29f5f9c1e0e6475be873c067edce955e870f48b3 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Sat, 15 Aug 2026 15:56:46 +0200 Subject: [PATCH 19/21] add_value_length --- conversion/kimi_k3.py | 2 ++ src/models/kimi-k3.cpp | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/conversion/kimi_k3.py b/conversion/kimi_k3.py index 3adfbb0a0474..a5005e15bd62 100644 --- a/conversion/kimi_k3.py +++ b/conversion/kimi_k3.py @@ -229,7 +229,9 @@ def set_gguf_parameters(self): # K3 is nope-only; qk_rope_head_dim still sizes the un-absorbed part of K assert self.hparams.get("mla_use_nope"), "K3 MLA is expected to be nope-only" self.gguf_writer.add_rope_dimension_count(qk_rope_head_dim) + # MLA is served as MQA, so the cache holds the compressed latent self.gguf_writer.add_key_length(kv_lora_rank + qk_rope_head_dim) + self.gguf_writer.add_value_length(kv_lora_rank) self.gguf_writer.add_key_length_mla(qk_nope_head_dim + qk_rope_head_dim) self.gguf_writer.add_value_length_mla(v_head_dim) diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp index 10e47e4b5a4d..33a41225caa0 100644 --- a/src/models/kimi-k3.cpp +++ b/src/models/kimi-k3.cpp @@ -21,6 +21,10 @@ void llama_model_kimi_k3::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false); + // MLA is served as MQA, so the cache holds the compressed latent. set it here too, + // as GGUFs converted before value_length was written fall back to n_embd/n_head + hparams.n_embd_head_v_full = hparams.n_lora_kv; + // n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-linear for (uint32_t i = 0; i < hparams.n_layer(); ++i) { hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0; From d8980d69301d109d0d5c705ff304fc2bc8ee822b Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Sat, 15 Aug 2026 16:11:27 +0200 Subject: [PATCH 20/21] optimize res_stack construction --- src/models/kimi-k3.cpp | 41 ++++++++++++++--------------------------- src/models/models.h | 7 ++----- 2 files changed, 16 insertions(+), 32 deletions(-) diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp index 33a41225caa0..d952d72cdf13 100644 --- a/src/models/kimi-k3.cpp +++ b/src/models/kimi-k3.cpp @@ -21,8 +21,8 @@ void llama_model_kimi_k3::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false); - // MLA is served as MQA, so the cache holds the compressed latent. set it here too, - // as GGUFs converted before value_length was written fall back to n_embd/n_head + // the MLA cache holds the compressed latent + // set it here too, as older GGUFs have no value_length key hparams.n_embd_head_v_full = hparams.n_lora_kv; // n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-linear @@ -185,37 +185,24 @@ static ggml_tensor * kimi_k3_situ(ggml_context * ctx0, ggml_tensor * gate, ggml_ // cross-layer residual attention // +// layout is [n_embd, n_ckpt, n_tokens]: rms_norm reduces over ne0, dsv4_hc_pre over ne1 +// append the new checkpoint, do not re-fold the whole chain void llama_model_kimi_k3::graph::res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens) { - resi.push_back(ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens)); -} + ggml_tensor * ckpt = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens); -ggml_tensor * llama_model_kimi_k3::graph::res_stack(int64_t n_embd, int64_t n_tokens) { - GGML_UNUSED(n_embd); - GGML_UNUSED(n_tokens); - if (resi_stack_n == (int) resi.size()) { - return resi_stack; // set unchanged since the last mix - } - ggml_tensor * acc = resi[0]; - for (size_t i = 1; i < resi.size(); ++i) { - acc = ggml_concat(ctx0, acc, resi[i], 1); - } - resi_stack = acc; - resi_stack_n = (int) resi.size(); - return acc; + resi_stack = resi_stack ? ggml_concat(ctx0, resi_stack, ckpt, 1) : ckpt; } ggml_tensor * llama_model_kimi_k3::graph::res_mix(ggml_tensor * cur, ggml_tensor * score_w, - int64_t n_embd, int64_t n_tokens, int il) { - const int n_ckpt = (int) resi.size(); - if (n_ckpt == 0) { + int64_t n_tokens, int il) { + if (!resi_stack) { return cur; // layer 0: nothing banked yet } - const float eps = hparams.f_norm_rms_eps; + const int n_ckpt = (int) resi_stack->ne[1]; + const float eps = hparams.f_norm_rms_eps; - // the checkpoint stack is packed as [n_embd, n_ckpt, n_tokens]: rms_norm reduces over - // ne0 = n_embd and dsv4_hc_pre reduces over ne1 = n_ckpt, so no transpose is needed - ggml_tensor * src = res_stack(n_embd, n_tokens); // [n_embd, n_ckpt, n_tokens] + ggml_tensor * src = resi_stack; // [n_embd, n_ckpt, n_tokens] // one rms_norm scores all checkpoints at once // note: the scores use the normalized values, but the sum below uses the raw ones @@ -292,7 +279,7 @@ llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_par // from the attention output alone ggml_tensor * prefix_sum = inpL; - cur = use_attn_res ? res_mix(prefix_sum, layer.attn_res_score, n_embd, n_tokens, il) + cur = use_attn_res ? res_mix(prefix_sum, layer.attn_res_score, n_tokens, il) : prefix_sum; bool banked = false; @@ -317,7 +304,7 @@ llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_par prefix_sum = banked ? cur : ggml_add(ctx0, prefix_sum, cur); cb(prefix_sum, "prefix_sum_attn", il); - cur = use_attn_res ? res_mix(prefix_sum, layer.ffn_res_score, n_embd, n_tokens, il) + cur = use_attn_res ? res_mix(prefix_sum, layer.ffn_res_score, n_tokens, il) : prefix_sum; cur = build_norm(cur, layer.ffn_norm, NULL, LLM_NORM_RMS, il); @@ -344,7 +331,7 @@ llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_par // final mix, then narrow to the output tokens if (use_attn_res) { - cur = res_mix(cur, model.output_res_score, n_embd, n_tokens, -1); + cur = res_mix(cur, model.output_res_score, n_tokens, -1); } if (inp_out_ids) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); diff --git a/src/models/models.h b/src/models/models.h index f8edb65560d2..df217e939e8a 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2283,14 +2283,11 @@ struct llama_model_kimi_k3 : public llama_model_base { const llama_model & model; // Cross-layer residual attention (K3's `_apply_attn_res`). - std::vector resi; - ggml_tensor * resi_stack = nullptr; - int resi_stack_n = -1; + ggml_tensor * resi_stack = nullptr; void res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens); - ggml_tensor * res_stack(int64_t n_embd, int64_t n_tokens); ggml_tensor * res_mix(ggml_tensor * cur, ggml_tensor * score_w, - int64_t n_embd, int64_t n_tokens, int il); + int64_t n_tokens, int il); ggml_tensor * build_kda_layer(ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs, From 914db7f972ccc341d636ecd2036a74ff6a0cae9b Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Sat, 15 Aug 2026 16:26:59 +0200 Subject: [PATCH 21/21] nits --- src/llama-graph.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 560233c9f353..1896758c5da5 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1835,6 +1835,8 @@ ggml_tensor * llm_graph_context::build_ffn( cur = ggml_reglu(ctx0, cur); cb(cur, "ffn_reglu", il); } break; + case LLM_FFN_SITU: + GGML_ABORT("not yet supported"); default: GGML_ABORT("fatal error"); }