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:
co-authored by
Boris Dvorkin
Stanisław Szymczyk
Deepankar Singh
Claude Fable 5
Caleb DeLeeuw
Xuan Son Nguyen
parent
22b8e310b9
commit
ad1de39e07
+180
@@ -2325,6 +2325,179 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
|
||||
return data;
|
||||
}
|
||||
|
||||
// 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|>
|
||||
// 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;
|
||||
|
||||
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|>";
|
||||
|
||||
// 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|>",
|
||||
"<|sep|>",
|
||||
"<|end_of_msg|>",
|
||||
};
|
||||
|
||||
data.thinking_start_tag = THINK_START;
|
||||
data.thinking_end_tags = { THINK_END };
|
||||
|
||||
// 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|>" },
|
||||
{ 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;
|
||||
|
||||
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 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 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 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) {
|
||||
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 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();
|
||||
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<std::string>();
|
||||
}
|
||||
|
||||
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);
|
||||
});
|
||||
|
||||
// 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)) +
|
||||
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:
|
||||
@@ -3293,6 +3466,13 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
|
||||
return common_chat_params_init_kimi_k2(tmpl, params);
|
||||
}
|
||||
|
||||
// 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");
|
||||
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|>).
|
||||
|
||||
Reference in New Issue
Block a user