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)
2026-08-15 17:11:05 +02:00
committed by GitHub
co-authored by Boris Dvorkin Stanisław Szymczyk Deepankar Singh Claude Fable 5 Caleb DeLeeuw Xuan Son Nguyen
parent 22b8e310b9
commit ad1de39e07
22 changed files with 1875 additions and 31 deletions
+22
View File
@@ -144,6 +144,7 @@ static const std::map<llm_arch, const char *> 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" },
@@ -187,6 +188,9 @@ static const std::map<llm_kv, const char *> 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" },
@@ -202,6 +206,7 @@ static const std::map<llm_kv, const char *> 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" },
@@ -313,6 +318,7 @@ static const std::map<llm_kv, const char *> 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" },
@@ -463,6 +469,13 @@ static const std::map<llm_tensor, const char *> 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" },
@@ -756,6 +769,13 @@ static const std::map<llm_tensor, llm_tensor_info> 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}},
@@ -976,6 +996,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:
@@ -1040,6 +1061,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: