Compare commits

..
4 Commits
Author SHA1 Message Date
ece963f41b ui: mask API Key field in settings and error splash to stop browser a… (#26562)
* ui: mask API Key field in settings and error splash to stop browser autofill

* ui: set autocomplete=new-password on private fields

The password input type makes browsers offer to save the API key
in the password manager and autofill saved site credentials into
the field. The new-password autocomplete value disables both.

---------

Co-authored-by: Pascal <admin@serveurperso.com>
2026-08-15 22:52:55 +02:00
Gautam0507andGitHub 0d9ceae1e3 ui: read structuredContent from MCP tool result when content is empty (#26691) 2026-08-15 20:00:18 +02:00
ad1de39e07 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>
2026-08-15 17:11:05 +02:00
Xuan-Son NguyenandGitHub 22b8e310b9 server: re-design yield_to_queue thread model (#27133)
* run common_speculative_process in worker

* swap worker <--> main thread design
2026-08-15 16:48:40 +02:00
32 changed files with 2006 additions and 82 deletions
+180
View File
@@ -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|>).
+1
View File
@@ -125,6 +125,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"JinaEmbeddingsV5Model": "bert",
"KORMoForCausalLM": "qwen",
"KimiK25ForConditionalGeneration": "deepseek",
"KimiK3ForConditionalGeneration": "kimi_k3",
"KimiLinearForCausalLM": "kimi_linear",
"KimiLinearModel": "kimi_linear",
"KimiVLForConditionalGeneration": "deepseek",
+41 -1
View File
@@ -658,6 +658,43 @@ class ModelBase:
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.
@@ -2661,7 +2698,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
+1 -26
View File
@@ -709,31 +709,6 @@ class DeepseekV4Model(TextModel):
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()
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
@@ -747,7 +722,7 @@ class DeepseekV4Model(TextModel):
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
+375
View File
@@ -0,0 +1,375 @@
from __future__ import annotations
import re
from pathlib import Path
from typing import Callable, Iterable, Iterator, TYPE_CHECKING
import numpy as np
import torch
if TYPE_CHECKING:
from torch import Tensor
from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger
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, 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.
"""
model_arch = gguf.MODEL_ARCH.KIMI_K3
_experts: list[dict[str, Tensor]] | None = None
# `<x>_res_norm.weight` and `<x>_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?)
_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 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+)"
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 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 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)
# 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
#
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 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:
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.
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
n_blocks = (packed_cols * 2) // 32
byte_shape = (len(loaders), rows, n_blocks * 17)
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] = self.repack_mxfp4_blocks(
LazyTorchTensor.to_eager(packed_fn()),
LazyTorchTensor.to_eager(scale_fn()),
)
return out
# 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,),
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]]:
# 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 ()
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
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"]
# 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
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)
# 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)
# --- 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 <x>_res_norm.weight with <x>_res_proj.weight and emit their product.
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"):
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).
# 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
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"):
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"):
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)
+76 -1
View File
@@ -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,
+15
View File
@@ -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)
+13
View File
@@ -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",
),
+324
View File
@@ -0,0 +1,324 @@
{%- macro escape_attr(value) -%}
{{- value|string|replace('&', '&amp;')|replace('"', '&quot;') -}}
{%- 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) -%}
{#- 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 -%}
{%- 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 -%}
+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:
+13
View File
@@ -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,
@@ -192,6 +193,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,
@@ -207,6 +211,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,
@@ -318,6 +323,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,
@@ -492,6 +498,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,
+4 -1
View File
@@ -2293,7 +2293,10 @@ 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_QWEN3NEXT ||
if (model.arch == LLM_ARCH_KIMI_K3) {
// the n_tokens*40 budget below is exhausted at ubatch 3840
res = std::max<uint32_t>(n_tokens * 160, 64u * model.n_tensors());
} else if (model.arch == LLM_ARCH_QWEN3NEXT ||
model.arch == LLM_ARCH_KIMI_LINEAR ||
model.arch == LLM_ARCH_QWEN35 ||
model.arch == LLM_ARCH_QWEN35MOE ||
+17
View File
@@ -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");
}
@@ -2174,6 +2176,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);
+1
View File
@@ -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
};
enum llm_ffn_gate_type {
+9 -1
View File
@@ -4,10 +4,11 @@
#include <array>
#include <cassert>
#include <cmath>
// 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,
@@ -170,6 +171,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;
+7
View File
@@ -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);
@@ -319,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);
@@ -376,6 +378,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 +409,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);
+4
View File
@@ -322,6 +322,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:
@@ -845,6 +847,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";
@@ -2602,6 +2605,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
+10
View File
@@ -144,6 +144,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,
};
@@ -530,6 +531,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;
@@ -589,6 +598,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;
+614
View File
@@ -0,0 +1,614 @@
#include "models.h"
#include "llama-memory-recurrent.h"
//
// 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)
// 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);
// 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
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);
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_2_8T_A50B; 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] tensor (kimi-linear's is padded)
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
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<llm_graph_context> llama_model_kimi_k3::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
// 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
//
// 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) {
ggml_tensor * ckpt = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);
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_tokens, int il) {
if (!resi_stack) {
return cur; // layer 0: nothing banked yet
}
const int n_ckpt = (int) resi_stack->ne[1];
const float eps = hparams.f_norm_rms_eps;
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
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 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]
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 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));
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];
// 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_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_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_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
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);
// 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
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 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: the logits above are passed instead
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;
}
+36
View File
@@ -2285,6 +2285,42 @@ 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`).
ggml_tensor * resi_stack = nullptr;
void res_push(ggml_tensor * cur, int64_t n_embd, int64_t n_tokens);
ggml_tensor * res_mix(ggml_tensor * cur, ggml_tensor * score_w,
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<llm_graph_context> 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;
+103
View File
@@ -4462,6 +4462,109 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
}
}
// Kimi-K3 tests - custom parser
// 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);
// 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
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: 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|>"
"<|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.<name>:<counter>
{
+9 -1
View File
@@ -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<uint32_t> 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));
@@ -218,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));
@@ -244,6 +247,10 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
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_RESIDUAL_SCALE, 3.5565588200778455f);
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;
@@ -372,6 +379,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:
+17 -8
View File
@@ -2947,8 +2947,10 @@ private:
});
// generate the actual drafts (if any)
{
common_speculative_draft(spec.get());
if (!drafting.empty()) {
queue_tasks.yield_to_queue([&]() {
common_speculative_draft(spec.get());
});
}
// make checkpoints if needed
@@ -3578,8 +3580,8 @@ private:
has_output |= batch.tokens[i].output;
}
// decode on the worker thread, so we can still handle metrics tasks while waiting
// note: the sync is done here too, so that the wait also happens off the main thread
// yield to the queue, so we can still handle metrics tasks while decoding
// note: the sync is done here too, so that the wait is also covered by the yield
int ret = 0;
queue_tasks.yield_to_queue([&]() {
ret = llama_decode(ctx_tgt, batch_view);
@@ -3644,11 +3646,18 @@ private:
// TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL]
// for now, always re-evaluate for simplicity
// ref: https://github.com/ggml-org/llama.cpp/pull/22728#issuecomment-4400925384
if (!common_speculative_process(spec.get(), batch_view)) {
SRV_ERR("%s", "failed to process speculative batch\n");
if (spec) {
bool ok = true;
queue_tasks.yield_to_queue([&]() {
ok = common_speculative_process(spec.get(), batch_view);
});
// TODO: handle error
throw std::runtime_error("failed to process speculative batch");
if (!ok) {
SRV_ERR("%s", "failed to process speculative batch\n");
// TODO: handle error
throw std::runtime_error("failed to process speculative batch");
}
}
// handle `n_cmpl > 1` tasks - when the main prompt is processed, activate all child tasks too
+54 -32
View File
@@ -150,31 +150,46 @@ bool server_queue::process_new_tasks(bool is_yielding) {
void server_queue::worker_loop() {
while (true) {
std::function<void()> work;
{
std::unique_lock<std::mutex> lock(mutex_tasks);
// wait on busy instead of yielding - busy stays set even when the yield already ended
worker.cv.wait(lock, [&]{
return worker.stop || worker.work != nullptr;
return worker.stop || worker.busy;
});
if (worker.stop) {
return;
}
work = std::move(worker.work);
worker.work = nullptr;
}
// note: do not hold any lock here, work() may post new tasks
std::exception_ptr exception;
try {
work();
} catch (...) {
exception = std::current_exception();
// process tasks while the yield is active
while (true) {
bool terminated = false;
try {
// note: do not hold any lock here, the callback may post new tasks
terminated = process_new_tasks(true);
} catch (...) {
std::unique_lock<std::mutex> lock(mutex_tasks);
worker.exception = std::current_exception();
break;
}
std::unique_lock<std::mutex> lock(mutex_tasks);
if (terminated || worker.stop || !worker.yielding) {
break;
}
if (!queue_tasks.empty()) {
continue; // a new task arrived in the meantime
}
condition_tasks.wait(lock, [&]{
return worker.stop || !running || !worker.yielding || !queue_tasks.empty();
});
}
// signal completion to yield_to_queue()
std::unique_lock<std::mutex> lock(mutex_tasks);
worker.exception = std::move(exception);
worker.busy = false;
// signal to yield_to_queue() that no more tasks will be processed
{
std::unique_lock<std::mutex> lock(mutex_tasks);
worker.busy = false;
}
condition_tasks.notify_all();
}
}
@@ -188,6 +203,7 @@ void server_queue::worker_stop() {
worker.stop = true;
}
worker.cv.notify_one();
condition_tasks.notify_all();
worker.thread.join();
}
@@ -199,29 +215,29 @@ void server_queue::yield_to_queue(std::function<void()> && work) {
{
std::unique_lock<std::mutex> lock(mutex_tasks);
GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested");
worker.busy = true;
worker.work = std::move(work);
worker.busy = true;
worker.yielding = true;
}
worker.cv.notify_one();
while (true) {
// note: on terminate this is a no-op, but we still wait for the work to finish
process_new_tasks(true);
std::unique_lock<std::mutex> lock(mutex_tasks);
// declined tasks are moved to queue_tasks_unhandled, so a non-empty queue always has something new
condition_tasks.wait(lock, [&]{
return !worker.busy || (running && !queue_tasks.empty());
});
if (!worker.busy) {
break;
}
// run the work on the current thread, so that all ggml compute stays on the same thread
std::exception_ptr exception;
try {
work();
} catch (...) {
exception = std::current_exception();
}
std::exception_ptr exception;
{
std::unique_lock<std::mutex> lock(mutex_tasks);
// the yield is over, wait for the worker to finish its current task
worker.yielding = false;
condition_tasks.notify_all();
condition_tasks.wait(lock, [&]{
return !worker.busy;
});
// put the declined tasks back, keeping their order
while (!queue_tasks_unhandled.empty()) {
queue_tasks.push_front(std::move(queue_tasks_unhandled.back()));
@@ -231,8 +247,12 @@ void server_queue::yield_to_queue(std::function<void()> && work) {
// make sure to avoid idle timeout here
time_last_task = ggml_time_ms();
// the worker is idle now, take the exception it may have left behind
std::swap(exception, worker.exception);
// an exception from work() takes precedence over the one from the worker
if (!exception) {
std::swap(exception, worker.exception);
} else {
worker.exception = nullptr;
}
}
QUE_DBG("%s", "done yielding to queue\n");
@@ -249,7 +269,9 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
// spawn the worker thread used by yield_to_queue()
GGML_ASSERT(!worker.thread.joinable() && "start_loop() is already running");
worker.stop = false;
worker.stop = false;
worker.busy = false;
worker.yielding = false;
worker.thread = std::thread([this]() { worker_loop(); });
constexpr auto max_wait_time = std::chrono::seconds(1);
+7 -6
View File
@@ -33,11 +33,11 @@ private:
// used by yield_to_queue, all fields are guarded by mutex_tasks
struct worker_t {
std::thread thread;
std::condition_variable cv; // the worker sleeps on this until there is work
std::function<void()> work; // pending work, picked up by the thread
std::exception_ptr exception; // exception thrown by work(), if any
bool stop = false;
bool busy = false;
std::condition_variable cv; // the worker sleeps on this until a yield starts
std::exception_ptr exception; // exception thrown while processing tasks, if any
bool stop = false;
bool busy = false; // set by yield_to_queue(), cleared by the worker once it is done processing tasks
bool yielding = false; // work() is still running on the start_loop() thread
};
worker_t worker;
@@ -93,7 +93,7 @@ public:
*/
void start_loop(int64_t idle_sleep_ms = -1);
// run work() on a separate thread, while the current thread calls process_new_tasks
// while waiting for work() to finish, run process_new_tasks on the worker thread
// returns once work() is done (may throw exceptions)
// must be called from start_loop() thread (ideally inside callback_update_slots)
// use case: return metrics while encode/decode is running
@@ -116,6 +116,7 @@ public:
// the second argument tells whether the queue is currently yielding (see yield_to_queue)
// only then may the callback return false to decline the task, and it must leave it
// untouched, so that it can be put back in the queue later
// note: while yielding, the callback runs on worker thread, not main thread
void on_new_task(std::function<bool(server_task &&, bool)> callback) {
callback_new_task = std::move(callback);
}
@@ -158,6 +158,8 @@
<div class="relative">
<Input
id="api-key-input"
type="password"
autocomplete="new-password"
placeholder="Enter your API key..."
bind:value={apiKeyInput}
onkeydown={handleApiKeyKeydown}
@@ -81,7 +81,8 @@
<div class="relative w-full">
<Input
id={field.key}
type={field.isPositiveInteger ? 'number' : 'text'}
type={field.isPrivate ? 'password' : field.isPositiveInteger ? 'number' : 'text'}
autocomplete={field.isPrivate ? 'new-password' : undefined}
{...field.isPositiveInteger
? {
min: String(field.min ?? 1),
@@ -324,6 +324,7 @@ const SETTINGS_REGISTRY: Record<string, SettingsSectionEntry> = {
{
defaultValue: '',
help: `Set the API Key if you are using <code> ${CLI_FLAGS.API_KEY} </code> option for the server.`,
isPrivate: true,
key: SETTINGS_KEYS.API_KEY,
label: 'API Key',
section: SETTINGS_SECTION_SLUGS.GENERAL,
@@ -713,6 +714,7 @@ export const SETTINGS_CHAT_SECTIONS: SettingsSection[] = [
help: s.help,
isExperimental: s.isExperimental,
isPositiveInteger: s.isPositiveInteger,
isPrivate: s.isPrivate,
key: s.key,
label: s.label,
max: s.max,
+15 -3
View File
@@ -18,7 +18,8 @@ import {
DEFAULT_CLIENT_VERSION,
DEFAULT_IMAGE_MIME_TYPE,
DEFAULT_MCP_CONFIG,
HEADERS
HEADERS,
NEWLINE
} from '$lib/constants';
import {
MCPConnectionPhase,
@@ -70,6 +71,7 @@ interface ToolResultContentItem {
interface ToolCallResult {
content?: ToolResultContentItem[];
structuredContent?: Record<string, unknown>;
isError?: boolean;
_meta?: Record<string, unknown>;
}
@@ -1012,10 +1014,20 @@ export class MCPService {
if (!Array.isArray(content)) return '';
return content
const formatted = content
.map((item) => this.formatSingleContent(item))
.filter(Boolean)
.join('\n');
.join(NEWLINE);
if (formatted !== '') {
return formatted;
}
if (result.structuredContent && typeof result.structuredContent === 'object') {
return JSON.stringify(result.structuredContent);
}
return '';
}
private static formatSingleContent(content: ToolResultContentItem): string {
+2
View File
@@ -31,6 +31,7 @@ export interface SettingsEntry {
radioOptions?: Array<{ value: string; label: string; key: string; isExperimental?: boolean }>;
isExperimental?: boolean;
isPositiveInteger?: boolean;
isPrivate?: boolean;
placeholder?: string;
min?: number;
max?: number;
@@ -55,6 +56,7 @@ export interface SettingsFieldConfig {
type: SettingsFieldType;
isExperimental?: boolean;
isPositiveInteger?: boolean;
isPrivate?: boolean;
placeholder?: string;
min?: number;
max?: number;
+18 -1
View File
@@ -2,7 +2,7 @@ import { Client } from '@modelcontextprotocol/sdk/client';
import { CORS_PROXY } from '$lib/constants';
import { MCPConnectionPhase, MCPTransportType } from '$lib/enums';
import { MCPService } from '$lib/services/mcp.service';
import type { MCPConnectionLog, MCPServerConfig } from '$lib/types';
import type { MCPConnection, MCPConnectionLog, MCPServerConfig } from '$lib/types';
import { afterEach, describe, expect, it, vi } from 'vitest';
type DiagnosticFetchFactory = (
@@ -329,4 +329,21 @@ describe('MCPService', () => {
)
).toHaveLength(0);
});
it('falls back to structuredContent when content array is empty', async () => {
const connection = {
client: {
callTool: vi.fn().mockResolvedValue({
content: [],
structuredContent: { accounts: [{ id: 1 }], total: 1 }
})
},
requestTimeoutMs: 9000,
serverName: 'test-server'
} as unknown as MCPConnection;
const result = await MCPService.callTool(connection, { arguments: {}, name: 'tool' });
expect(result.isError).toBe(false);
expect(result.content).toBe('{"accounts":[{"id":1}],"total":1}');
});
});
@@ -0,0 +1,12 @@
import { SETTINGS_CHAT_SECTIONS, SETTINGS_KEYS } from '$lib/constants';
import { describe, expect, it } from 'vitest';
describe('checkApiKeyField', () => {
it('should have isPrivate set to true', () => {
const fields = SETTINGS_CHAT_SECTIONS.flatMap((section) => section.fields);
const apiKeyField = fields.find((field) => field?.key === SETTINGS_KEYS.API_KEY);
expect(apiKeyField).toBeDefined();
expect(apiKeyField?.isPrivate).toBe(true);
});
});