🔨 feat(kimi-k3): поддержка Kimi K3 — конверсия MXFP4→BF16 и инференс-граф

Архитектура kimi-k3 (KimiK3ForConditionalGeneration, text_config kimi_linear):
- деквант compressed-tensors mxfp4-pack-quantized (E2M1 + E8M0 scale, group 32)
- KDA safe gate: g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias)),
  full-rank выходной гейт g_proj, срез A_log[:n_head]
- gated MLA (NoPE): attn * sigmoid(g_proj(x)) перед o_proj, q_lora 1536
- Stable LatentMoE: роутер на полном hidden, эксперты в латенте 3584,
  RMSNorm после взвешенной суммы, up-проекция обратно
- AttnRes: банк снапшотов residual-потока каждые attn_res_block_size слоёв,
  softmax-смеси перед attention/FFN и на выходе модели
- SiTU-GLU: beta*tanh(g/beta)*sigmoid(g) * lbeta*tanh(up/lbeta) (LLM_FFN_SITU)

Проверено на синтетическом мини-K3: convert → BF16 GGUF → Q8_0 → генерация.
Деквант MXFP4 сверен с реальными тензорами moonshotai/Kimi-K3 (shard 2).

Author: worthant

conversion/__init__.py

@@ -121,6 +121,7 @@
     "JinaEmbeddingsV5Model": "bert",
     "KORMoForCausalLM": "qwen",
     "KimiK25ForConditionalGeneration": "deepseek",
+    "KimiK3ForConditionalGeneration": "kimi_k3",
     "KimiLinearForCausalLM": "kimi_linear",
     "KimiLinearModel": "kimi_linear",
     "KimiVLForConditionalGeneration": "deepseek",

conversion/base.py

@@ -386,6 +386,29 @@ def dequant_gptq(g_idx: Tensor, qweight: Tensor, qzeros: Tensor, scales: Tensor)
 
                 return (scales[g_idx].float() * (weight - zeros[g_idx]).float()).T
 
+            def dequant_mxfp4_packed(w: Tensor, scale: Tensor, group_size: int) -> Tensor:
+                # compressed-tensors "mxfp4-pack-quantized":
+                # w: uint8 [..., K/2], two FP4 (E2M1) values per byte, low nibble = even index
+                # scale: uint8 [..., K/group_size], E8M0 exponent, scale = 2^(x - 127)
+                assert w.dtype == torch.uint8
+                assert scale.dtype == torch.uint8
+
+                kvalues = torch.tensor(
+                    [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
+                    dtype=torch.float32,
+                )
+                if self.lazy:
+                    kvalues = LazyTorchTensor.from_eager(kvalues)
+
+                lo = (w & 0x0F).to(torch.long)
+                hi = (w >> 4).to(torch.long)
+                # interleave: [..., K/2, 2] -> [..., K]
+                vals = torch.stack((kvalues[lo], kvalues[hi]), dim=-1).reshape(*w.shape[:-1], w.shape[-1] * 2)
+
+                exp = torch.ldexp(torch.ones_like(scale, dtype=torch.float32), scale.to(torch.int32) - 127)
+                vals = vals.reshape(*vals.shape[:-1], -1, group_size) * exp.unsqueeze(-1)
+                return vals.reshape(*w.shape[:-1], w.shape[-1] * 2)
+
             def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: Tensor | None, num_bits: int, group_size: int):
                 assert w.dtype == torch.int32
                 shape = tuple(shape_tensor.tolist())
@@ -531,6 +554,23 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T
                             tensors_to_remove += [base_name + n for n in ("_packed", "_shape", "_scale")]
                             if (base_name + "_zero_point") in self.model_tensors:
                                 tensors_to_remove.append(base_name + "_zero_point")
+                elif quant_format == "mxfp4-pack-quantized":
+                    assert weight_config.get("strategy") == "group"
+                    assert weight_config.get("type") == "float"
+                    assert weight_config.get("num_bits") == 4
+                    group_size = weight_config.get("group_size")
+                    assert isinstance(group_size, int)
+                    for name in self.model_tensors.keys():
+                        if name.endswith(".weight_packed"):
+                            base_name = name.removesuffix("_packed")
+                            w = self.model_tensors[name]
+                            scale = self.model_tensors[base_name + "_scale"]
+                            new_tensors[base_name] = (
+                                lambda w=w, scale=scale: dequant_mxfp4_packed(w(), scale(), group_size)
+                            )
+                            tensors_to_remove += [base_name + n for n in ("_packed", "_scale")]
+                            if (base_name + "_shape") in self.model_tensors:
+                                tensors_to_remove.append(base_name + "_shape")
                 elif nvfp4_compressed_tensors:
                     # Don't error from compressed-tensors, we'll handle them in _generate_nvfp4_tensors
                     pass
@@ -2633,7 +2673,7 @@ 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"):
+    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

conversion/kimi_k3.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from typing import Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, gguf
+from .kimi_linear import KimiLinearModel
+
+
+@ModelBase.register("KimiK3ForConditionalGeneration")
+class KimiK3Model(KimiLinearModel):
+    """Kimi K3: hybrid KDA + gated-MLA (NoPE) with Attention Residuals and Stable LatentMoE.
+
+    Text config is `kimi_linear` with K3 extensions:
+      - SiTU-GLU activation (soft-capped SiLU) in dense MLP, shared and routed experts
+      - AttnRes: residual-stream snapshot bank every `attn_res_block_size` layers,
+        softmax mixtures before attention, before MLP and at model output
+      - Stable LatentMoE: routed experts run in a `routed_expert_hidden_size` latent
+        space (down proj -> experts -> weighted sum -> RMSNorm -> up proj)
+      - KDA safe gate: g_log = gate_lower_bound * sigmoid(exp(A_log) * (g_raw + dt_bias))
+        with a full-rank output gate g_proj instead of the low-rank g_a/g_b pair
+      - MLA output gate: attn = attn * sigmoid(g_proj(x)) before o_proj
+    Routed expert weights are MXFP4 (compressed-tensors), dequantized in ModelBase.
+    """
+    model_arch = gguf.MODEL_ARCH.KIMI_K3
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+
+        # Stable LatentMoE
+        self.gguf_writer.add_moe_latent_size(self.hparams["routed_expert_hidden_size"])
+
+        # SiTU-GLU activation parameters
+        self.gguf_writer.add_situ_beta(self.hparams["activation_situ_beta"])
+        self.gguf_writer.add_situ_linear_beta(self.hparams["activation_situ_linear_beta"])
+
+        # Attention residuals
+        self.gguf_writer.add_attn_res_block_size(self.hparams["attn_res_block_size"])
+
+        # KDA safe gate lower bound
+        self.gguf_writer.add_kda_gate_lower_bound(self.hparams["linear_attn_config"]["gate_lower_bound"])
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # text-only conversion: vision tensors are handled by the mmproj path
+        if name.startswith(("vision_tower.", "mm_projector.")):
+            return
+
+        name = name.removeprefix("language_model.")
+
+        # K3 checkpoints store A_log as [head_dim] (128) but only the first
+        # num_heads (96) entries are used. The safe-gate formula is
+        #   g_log = gate_lower_bound * sigmoid(exp(A_log) * (g_raw + dt_bias))
+        # so we store exp(A_log) directly (unlike Kimi-Linear's -exp(A_log)).
+        if name.endswith(".A_log"):
+            n_head = self.hparams["num_attention_heads"]
+            data_torch = torch.exp(data_torch.float()[:n_head])
+            # skip KimiLinearModel's -exp(A_log) handling
+            yield from super(KimiLinearModel, self).modify_tensors(data_torch, name, bid)
+            return
+
+        # res projections are stored as [1, n_embd]: flatten to [n_embd]
+        if name.endswith(("_res_proj.weight", "_res_norm.weight")):
+            data_torch = data_torch.reshape(-1)
+
+        yield from super().modify_tensors(data_torch, name, bid)

gguf-py/gguf/constants.py

@@ -126,6 +126,9 @@ class LLM:
         EXPERTS_PER_GROUP                 = "{arch}.experts_per_group"
         MOE_EVERY_N_LAYERS                = "{arch}.moe_every_n_layers"
         MOE_LATENT_SIZE                   = "{arch}.moe_latent_size"
+        SITU_BETA                         = "{arch}.situ_beta"
+        SITU_LINEAR_BETA                  = "{arch}.situ_linear_beta"
+        ATTN_RES_BLOCK_SIZE               = "{arch}.attn_res_block_size"
         NEXTN_PREDICT_LAYERS              = "{arch}.nextn_predict_layers"
         NUM_DEEPSTACK_LAYERS              = "{arch}.n_deepstack_layers"
         DEEPSTACK_MAPPING                 = "{arch}.deepstack_mapping"
@@ -243,7 +246,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"
@@ -545,6 +549,7 @@ class MODEL_ARCH(IntEnum):
     LLAMA_EMBED      = auto()
     MAINCODER        = auto()
     KIMI_LINEAR      = auto()
+    KIMI_K3          = auto()
     TALKIE           = auto()
     MELLUM           = auto()
     NANBEIGE         = auto()
@@ -620,6 +625,13 @@ class MODEL_TENSOR(IntEnum):
     FFN_GATE_TID2EID     = auto()
     MOE_LATENT_DOWN      = auto() # nemotron 3 super
     MOE_LATENT_UP        = auto() # nemotron 3 super
+    MOE_LATENT_NORM      = auto() # kimi k3
+    ATTN_RES_NORM        = auto() # kimi k3
+    ATTN_RES_PROJ        = auto() # kimi k3
+    FFN_RES_NORM         = auto() # kimi k3
+    FFN_RES_PROJ         = auto() # kimi k3
+    OUTPUT_RES_NORM      = auto() # kimi k3
+    OUTPUT_RES_PROJ      = auto() # kimi k3
     ATTN_Q_NORM          = auto()
     ATTN_K_NORM          = auto()
     LAYER_OUT_NORM       = auto()
@@ -1135,6 +1147,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",
@@ -1210,6 +1223,13 @@ class MODEL_TENSOR(IntEnum):
     MODEL_TENSOR.FFN_GATE_TID2EID:          "blk.{bid}.ffn_gate_tid2eid",
     MODEL_TENSOR.MOE_LATENT_DOWN:           "blk.{bid}.ffn_latent_down",      # nemotron 3 super
     MODEL_TENSOR.MOE_LATENT_UP:             "blk.{bid}.ffn_latent_up",        # nemotron 3 super
+    MODEL_TENSOR.MOE_LATENT_NORM:           "blk.{bid}.ffn_latent_norm",      # kimi k3
+    MODEL_TENSOR.ATTN_RES_NORM:             "blk.{bid}.attn_res_norm",        # kimi k3
+    MODEL_TENSOR.ATTN_RES_PROJ:             "blk.{bid}.attn_res_proj",        # kimi k3
+    MODEL_TENSOR.FFN_RES_NORM:              "blk.{bid}.ffn_res_norm",         # kimi k3
+    MODEL_TENSOR.FFN_RES_PROJ:              "blk.{bid}.ffn_res_proj",         # kimi k3
+    MODEL_TENSOR.OUTPUT_RES_NORM:           "output_res_norm",                # kimi k3
+    MODEL_TENSOR.OUTPUT_RES_PROJ:           "output_res_proj",                # kimi k3
     MODEL_TENSOR.LAYER_OUT_NORM:            "blk.{bid}.layer_output_norm",
     MODEL_TENSOR.LAYER_OUT_SCALE:           "blk.{bid}.layer_output_scale",
     MODEL_TENSOR.PER_LAYER_TOKEN_EMBD:      "per_layer_token_embd",           # gemma3n
@@ -4438,6 +4458,55 @@ class MODEL_TENSOR(IntEnum):
         MODEL_TENSOR.FFN_DOWN,
         MODEL_TENSOR.FFN_UP,
     ],
+    MODEL_ARCH.KIMI_K3: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.OUTPUT_RES_NORM,
+        MODEL_TENSOR.OUTPUT_RES_PROJ,
+        MODEL_TENSOR.ATTN_NORM,
+        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,
+        MODEL_TENSOR.ATTN_RES_NORM,
+        MODEL_TENSOR.ATTN_RES_PROJ,
+        MODEL_TENSOR.FFN_RES_NORM,
+        MODEL_TENSOR.FFN_RES_PROJ,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.MOE_LATENT_DOWN,
+        MODEL_TENSOR.MOE_LATENT_NORM,
+        MODEL_TENSOR.MOE_LATENT_UP,
+        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_DT,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+    ],
     MODEL_ARCH.KIMI_LINEAR: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,

gguf-py/gguf/gguf_writer.py

@@ -881,6 +881,15 @@ def add_moe_every_n_layers(self, value: int) -> None:
     def add_moe_latent_size(self, value: int) -> None:
         self.add_uint32(Keys.LLM.MOE_LATENT_SIZE.format(arch=self.arch), value)
 
+    def add_situ_beta(self, value: float) -> None:
+        self.add_float32(Keys.LLM.SITU_BETA.format(arch=self.arch), value)
+
+    def add_situ_linear_beta(self, value: float) -> None:
+        self.add_float32(Keys.LLM.SITU_LINEAR_BETA.format(arch=self.arch), value)
+
+    def add_attn_res_block_size(self, value: int) -> None:
+        self.add_uint32(Keys.LLM.ATTN_RES_BLOCK_SIZE.format(arch=self.arch), value)
+
     def add_nextn_predict_layers(self, count: int) -> None:
         self.add_uint32(Keys.LLM.NEXTN_PREDICT_LAYERS.format(arch=self.arch), count)
 
@@ -1084,6 +1093,9 @@ def add_ssm_dt_b_c_rms(self, value: bool) -> None:
     def add_kda_head_dim(self, value: int) -> None:
         self.add_uint32(Keys.KDA.HEAD_DIM.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_tokenizer_model(self, model: str) -> None:
         self.add_string(Keys.Tokenizer.MODEL, model)
 

gguf-py/gguf/tensor_mapping.py

@@ -118,6 +118,14 @@ class TensorNameMap:
             "model.norm",                              # cogvlm
         ),
 
+        MODEL_TENSOR.OUTPUT_RES_NORM: (
+            "model.output_attn_res_norm",              # kimi k3
+        ),
+
+        MODEL_TENSOR.OUTPUT_RES_PROJ: (
+            "model.output_attn_res_proj",              # kimi k3
+        ),
+
         # Rope frequencies
         MODEL_TENSOR.ROPE_FREQS: (
             "rope.freqs",  # llama-pth
@@ -611,10 +619,32 @@ class TensorNameMap:
 
         MODEL_TENSOR.MOE_LATENT_DOWN: (
             "backbone.layers.{bid}.mixer.fc1_latent_proj",                 # nemotron 3 super
+            "model.layers.{bid}.block_sparse_moe.routed_expert_down_proj", # kimi k3
         ),
 
         MODEL_TENSOR.MOE_LATENT_UP: (
             "backbone.layers.{bid}.mixer.fc2_latent_proj",                 # nemotron 3 super
+            "model.layers.{bid}.block_sparse_moe.routed_expert_up_proj",   # kimi k3
+        ),
+
+        MODEL_TENSOR.MOE_LATENT_NORM: (
+            "model.layers.{bid}.block_sparse_moe.routed_expert_norm",      # kimi k3
+        ),
+
+        MODEL_TENSOR.ATTN_RES_NORM: (
+            "model.layers.{bid}.self_attention_res_norm",                  # kimi k3
+        ),
+
+        MODEL_TENSOR.ATTN_RES_PROJ: (
+            "model.layers.{bid}.self_attention_res_proj",                  # kimi k3
+        ),
+
+        MODEL_TENSOR.FFN_RES_NORM: (
+            "model.layers.{bid}.mlp_res_norm",                             # kimi k3
+        ),
+
+        MODEL_TENSOR.FFN_RES_PROJ: (
+            "model.layers.{bid}.mlp_res_proj",                             # kimi k3
         ),
 
         # Feed-forward down