model: add Kimi-K3 text model (#26185)
* 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) <noreply@anthropic.com> * 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) <noreply@anthropic.com> * Update conversion/kimi_k3.py Co-authored-by: Boris Dvorkin <b_dvorkin@niuitmo.ru> * Increase LLAMA_MAX_EXPERTS from 512 to 1024 * tests : support for Kimi K3 in archs test * 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 <noreply@anthropic.com> * 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 <noreply@anthropic.com> * fix: apply nits from @ngxson and text fixes from @danielhanchen * tests : added missing hyperparameters and tensors for Kimi K3 in test-llama-archs * chore : move overly verbose header file comments to Kimi K3 source file * tests : re-enabled KIMI_K3 in test-llama-archs for WebGPU backend * 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). * Refactor conditional for model architecture check * tests : re-enabled (again) KIMI_K3 and MINIMAX_M3 in test-llama-archs for WebGPU backend * fix code comments * add template on conversion * move repack_mxfp4_blocks to model base * nits * add_value_length * optimize res_stack construction * nits --------- Co-authored-by: Boris Dvorkin <b_dvorkin@niuitmo.ru> Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com> Co-authored-by: Deepankar Singh <singh.deepankar39@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Caleb DeLeeuw <caleb.deleeuw@gmail.com> Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
This commit is contained in:
Piotr Wilkin (ilintar)
co-authored by
Boris Dvorkin
Stanisław Szymczyk
Deepankar Singh
Claude Fable 5
Caleb DeLeeuw
Xuan Son Nguyen
Boris Dvorkin
Stanisław Szymczyk
Deepankar Singh
Claude Fable 5
Caleb DeLeeuw
Xuan Son Nguyen
parent
22b8e310b9
commit
ad1de39e07
@@ -124,6 +124,7 @@ class Keys:
|
||||
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 Keys:
|
||||
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 Keys:
|
||||
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"
|
||||
@@ -580,6 +589,7 @@ class MODEL_ARCH(IntEnum):
|
||||
LLAMA_EMBED = auto()
|
||||
MAINCODER = auto()
|
||||
KIMI_LINEAR = auto()
|
||||
KIMI_K3 = auto()
|
||||
TALKIE = auto()
|
||||
MELLUM = auto()
|
||||
NANBEIGE = auto()
|
||||
@@ -698,6 +708,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()
|
||||
@@ -1288,6 +1305,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
|
||||
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",
|
||||
@@ -1404,6 +1422,13 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
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",
|
||||
@@ -4960,6 +4985,56 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
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,
|
||||
|
||||
@@ -1103,6 +1103,21 @@ class GGUFWriter:
|
||||
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)
|
||||
|
||||
|
||||
@@ -912,6 +912,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",
|
||||
),
|
||||
|
||||
Reference in New Issue
Block a user