mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-08-16 09:45:30 +02:00
Compare commits
10 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 10bf611e53 | |||
| ece963f41b | |||
| 0d9ceae1e3 | |||
| ad1de39e07 | |||
| 22b8e310b9 | |||
| adb55e5148 | |||
| 77140d247c | |||
| 5f754ea0e2 | |||
| 27df9199d1 | |||
| 9b0a2ce859 |
+180
@@ -2325,6 +2325,179 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
|
||||
return data;
|
||||
}
|
||||
|
||||
// Kimi K3 - XTML tagged format, built by open_tag/close_tag macros:
|
||||
// open_tag(t, attrs) = <|open|>t k="v"...<|sep|> close_tag(t) = <|close|>t<|sep|>
|
||||
// assistant := [think] [response] [tools] close_tag(message) <|end_of_msg|>
|
||||
// the generation prompt already opens the think (or response) section, so the
|
||||
// section opener is optional here - same as Kimi K2 Thinking
|
||||
static common_chat_params common_chat_params_init_kimi_k3(const common_chat_template & tmpl,
|
||||
const autoparser::generation_params & inputs) {
|
||||
common_chat_params data;
|
||||
|
||||
data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs);
|
||||
data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs);
|
||||
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
|
||||
data.supports_thinking = true;
|
||||
|
||||
const std::string SEP = "<|sep|>";
|
||||
const std::string MSG_START = "<|open|>message role=\"assistant\"<|sep|>";
|
||||
const std::string THINK_START = "<|open|>think<|sep|>";
|
||||
const std::string THINK_END = "<|close|>think<|sep|>";
|
||||
const std::string RESP_START = "<|open|>response<|sep|>";
|
||||
const std::string RESP_END = "<|close|>response<|sep|>";
|
||||
const std::string TOOLS_START = "<|open|>tools<|sep|>";
|
||||
const std::string TOOLS_END = "<|close|>tools<|sep|>";
|
||||
const std::string CALL_START = "<|open|>call tool=\"";
|
||||
const std::string CALL_END = "<|close|>call<|sep|>";
|
||||
const std::string ARG_START = "<|open|>argument key=\"";
|
||||
const std::string ARG_END = "<|close|>argument<|sep|>";
|
||||
const std::string MSG_END = "<|close|>message<|sep|>";
|
||||
const std::string EOM_TOKEN = "<|end_of_msg|>";
|
||||
|
||||
// only the markers are special tokens. tag names ("think", "response", ...) are
|
||||
// normal tokens and must not be preserved, or prose with those words is broken
|
||||
data.preserved_tokens = {
|
||||
"<|open|>",
|
||||
"<|close|>",
|
||||
"<|sep|>",
|
||||
"<|end_of_msg|>",
|
||||
};
|
||||
|
||||
data.thinking_start_tag = THINK_START;
|
||||
data.thinking_end_tags = { THINK_END };
|
||||
|
||||
// per-role message-start delimiters. user/assistant messages only have the role
|
||||
// attribute, so the full opener is used. system and tool messages have more
|
||||
// attributes, so those delimiters stop after the closing quote of the role
|
||||
data.message_delimiters = {
|
||||
{ COMMON_CHAT_ROLE_ASSISTANT, "<|open|>message role=\"assistant\"<|sep|>" },
|
||||
{ COMMON_CHAT_ROLE_USER, "<|open|>message role=\"user\"<|sep|>" },
|
||||
{ COMMON_CHAT_ROLE_TOOL, "<|open|>message role=\"tool\"" },
|
||||
{ COMMON_CHAT_ROLE_SYSTEM, "<|open|>message role=\"system\"" },
|
||||
};
|
||||
|
||||
auto has_tools = inputs.tools.is_array() && !inputs.tools.empty();
|
||||
auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE;
|
||||
auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE;
|
||||
|
||||
if (inputs.has_continuation()) {
|
||||
const auto & msg = inputs.continue_msg;
|
||||
|
||||
data.generation_prompt = MSG_START + THINK_START + msg.reasoning_content;
|
||||
if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) {
|
||||
data.generation_prompt += THINK_END + RESP_START + msg.render_content();
|
||||
}
|
||||
|
||||
data.prompt += data.generation_prompt;
|
||||
}
|
||||
|
||||
auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
|
||||
auto end = p.end();
|
||||
|
||||
auto start = p.optional(p.literal(MSG_START));
|
||||
|
||||
// the think section is always consumed, even with reasoning extraction off:
|
||||
// the generation prompt ends with open_tag('think'), so it is always present.
|
||||
// reasoning stops at its own closer, or at the response opener if the model
|
||||
// skips the closer
|
||||
auto think_body = extract_reasoning ? p.reasoning(p.until_one_of({ THINK_END, RESP_START })) :
|
||||
p.content(p.until_one_of({ THINK_END, RESP_START }));
|
||||
|
||||
auto reasoning = p.optional(p.optional(p.literal(THINK_START)) + think_body +
|
||||
p.optional(p.literal(THINK_END)));
|
||||
|
||||
// content runs to the response closer, or to the next section if truncated
|
||||
auto response = p.optional(p.literal(RESP_START)) +
|
||||
p.content(p.until_one_of({ RESP_END, TOOLS_START, MSG_END })) +
|
||||
p.optional(p.literal(RESP_END));
|
||||
|
||||
// the EOG token after the message closer reaches the parser as text,
|
||||
// so it must be consumed or the parse stays incomplete
|
||||
auto trailer = p.optional(p.literal(MSG_END)) + p.optional(p.literal(EOM_TOKEN));
|
||||
|
||||
if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) {
|
||||
return start + reasoning + response + trailer + end;
|
||||
}
|
||||
|
||||
auto tool_choices = p.choice();
|
||||
foreach_function(inputs.tools, [&](const json & tool) {
|
||||
const auto & function = tool.at("function");
|
||||
std::string name = function.at("name");
|
||||
const json schema = function.contains("parameters") ? function.at("parameters") : json::object();
|
||||
|
||||
// arguments come one tag per key, with the JSON type in a type="..."
|
||||
// attribute. the type is taken from the tool schema instead, as it tells
|
||||
// us if the value is JSON or a literal string
|
||||
auto args = p.eps();
|
||||
if (schema.contains("properties") && !schema.at("properties").empty()) {
|
||||
auto arg_choices = p.choice();
|
||||
for (const auto & prop : schema.at("properties").items()) {
|
||||
const std::string & key = prop.key();
|
||||
|
||||
std::string type = "string";
|
||||
if (prop.value().is_object() && prop.value().contains("type") &&
|
||||
prop.value().at("type").is_string()) {
|
||||
type = prop.value().at("type").get<std::string>();
|
||||
}
|
||||
|
||||
auto value = type == "string" ? p.tool_arg_string_value(p.until(ARG_END)) :
|
||||
p.tool_arg_value(p.until(ARG_END));
|
||||
|
||||
// skip the trailing type="..." attribute: anything up to <|sep|>
|
||||
arg_choices |= p.rule("kimi-k3-arg-" + name + "-" + key,
|
||||
p.tool_arg(p.tool_arg_open(p.literal(ARG_START)) +
|
||||
p.tool_arg_name(p.literal(key)) + p.literal("\"") +
|
||||
p.until(SEP) + p.literal(SEP) + value +
|
||||
p.tool_arg_close(p.literal(ARG_END))));
|
||||
}
|
||||
args = p.zero_or_more(arg_choices);
|
||||
}
|
||||
|
||||
// skip the trailing index="N" attribute the same way
|
||||
auto call = p.tool(p.tool_open(p.literal(CALL_START) + p.tool_name(p.literal(name)) + p.literal("\"") +
|
||||
p.until(SEP) + p.literal(SEP)) +
|
||||
p.tool_args(args) + p.tool_close(p.literal(CALL_END)));
|
||||
|
||||
tool_choices |= p.rule("kimi-k3-tool-" + name, call);
|
||||
});
|
||||
|
||||
// all calls go inside one tools section, then the message is closed. the
|
||||
// message closer is part of the trigger rule, or else the lazy grammar
|
||||
// rejects it once tool calls have started
|
||||
auto tools_section =
|
||||
p.trigger_rule("kimi-k3-tool-call", p.literal(TOOLS_START) + p.one_or_more(tool_choices) +
|
||||
p.literal(TOOLS_END) + p.optional(p.literal(MSG_END)) +
|
||||
p.optional(p.literal(EOM_TOKEN)));
|
||||
|
||||
auto tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? tools_section :
|
||||
p.optional(tools_section);
|
||||
|
||||
return start + reasoning + response + tools + trailer + end;
|
||||
});
|
||||
|
||||
data.parser = parser.save();
|
||||
|
||||
if (include_grammar) {
|
||||
data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED;
|
||||
data.grammar = build_grammar([&](const common_grammar_builder & builder) {
|
||||
foreach_function(inputs.tools, [&](const json & tool) {
|
||||
const auto & function = tool.at("function");
|
||||
if (function.contains("parameters")) {
|
||||
auto schema = function.at("parameters");
|
||||
builder.resolve_refs(schema);
|
||||
}
|
||||
});
|
||||
parser.build_grammar(builder, data.grammar_lazy);
|
||||
});
|
||||
|
||||
data.grammar_triggers = {
|
||||
{ COMMON_GRAMMAR_TRIGGER_TYPE_WORD, TOOLS_START },
|
||||
};
|
||||
}
|
||||
|
||||
return data;
|
||||
}
|
||||
|
||||
// Cohere2 MoE (a.k.a. "North Code") parser.
|
||||
//
|
||||
// The assistant turn is fully marker-wrapped:
|
||||
@@ -3293,6 +3466,13 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
|
||||
return common_chat_params_init_kimi_k2(tmpl, params);
|
||||
}
|
||||
|
||||
// Kimi K3 - the <|open|>/<|close|>/<|end_of_msg|> markers are unique to it
|
||||
if (src.find("<|open|>") != std::string::npos && src.find("<|close|>") != std::string::npos &&
|
||||
src.find("<|end_of_msg|>") != std::string::npos) {
|
||||
LOG_DBG("Using specialized template: Kimi K3\n");
|
||||
return common_chat_params_init_kimi_k3(tmpl, params);
|
||||
}
|
||||
|
||||
// Cohere2 MoE / North Code - marker-wrapped format with <|START_TEXT|> content and
|
||||
// <|START_ACTION|> JSON tool calls. <|START_TEXT|> is unique to this template (the older
|
||||
// Command-R templates use <|START_RESPONSE|>).
|
||||
|
||||
@@ -102,7 +102,8 @@ bool common_imatrix_load(const std::string & fname, common_imatrix & imatrix) {
|
||||
const int64_t chunk_count_key = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_COUNT);
|
||||
const int64_t chunk_size_key = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_SIZE);
|
||||
|
||||
if (datasets_key != -1 && gguf_get_arr_type(ctx_gguf, datasets_key) == GGUF_TYPE_STRING) {
|
||||
if (datasets_key != -1 && gguf_get_kv_type(ctx_gguf, datasets_key) == GGUF_TYPE_ARRAY &&
|
||||
gguf_get_arr_type(ctx_gguf, datasets_key) == GGUF_TYPE_STRING) {
|
||||
const int64_t n = gguf_get_arr_n(ctx_gguf, datasets_key);
|
||||
imatrix.datasets.reserve(imatrix.datasets.size() + n);
|
||||
for (int64_t i = 0; i < n; ++i) {
|
||||
@@ -143,6 +144,13 @@ bool common_imatrix_load(const std::string & fname, common_imatrix & imatrix) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (in_sum2->type != GGML_TYPE_F32 || counts->type != GGML_TYPE_F32) {
|
||||
LOG_ERR("%s: sums and counts for %s must be F32\n", __func__, name.c_str());
|
||||
gguf_free(ctx_gguf);
|
||||
ggml_free(ctx);
|
||||
return false;
|
||||
}
|
||||
|
||||
auto & e = imatrix.entries[name];
|
||||
|
||||
const int64_t nval = ggml_nelements(in_sum2);
|
||||
|
||||
+33
-3
@@ -365,8 +365,25 @@ struct local_model {
|
||||
std::string name;
|
||||
std::string path;
|
||||
std::string path_mmproj;
|
||||
std::string path_draft;
|
||||
};
|
||||
|
||||
// TODO @ngxson: handle "eagle3-" when it's supported by common_speculative_types_from_gguf()
|
||||
static const char * draft_prefixes[] = { "mtp-", "dspark-", "dflash-" };
|
||||
|
||||
static bool is_mmproj_file(const std::string & fname) {
|
||||
return fname.find("mmproj") != std::string::npos;
|
||||
}
|
||||
|
||||
static bool is_draft_file(const std::string & fname) {
|
||||
for (const auto & prefix : draft_prefixes) {
|
||||
if (fname.rfind(prefix, 0) == 0) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
common_presets common_preset_context::load_from_models_dir(const std::string & models_dir) const {
|
||||
if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) {
|
||||
throw std::runtime_error(string_format("error: '%s' does not exist or is not a directory\n", models_dir.c_str()));
|
||||
@@ -378,10 +395,15 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
|
||||
common_file_info model_file;
|
||||
common_file_info first_shard_file;
|
||||
common_file_info mmproj_file;
|
||||
common_file_info draft_file;
|
||||
for (const auto & file : files) {
|
||||
if (string_ends_with(file.name, ".gguf")) {
|
||||
if (file.name.find("mmproj") != std::string::npos) {
|
||||
if (is_mmproj_file(file.name)) {
|
||||
mmproj_file = file;
|
||||
} else if (is_draft_file(file.name)) {
|
||||
if (draft_file.path.empty()) {
|
||||
draft_file = file; // first sidecar found wins
|
||||
}
|
||||
} else if (file.name.find("-00001-of-") != std::string::npos) {
|
||||
first_shard_file = file;
|
||||
} else {
|
||||
@@ -393,7 +415,8 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
|
||||
local_model model{
|
||||
/* name */ name,
|
||||
/* path */ first_shard_file.path.empty() ? model_file.path : first_shard_file.path,
|
||||
/* path_mmproj */ mmproj_file.path // can be empty
|
||||
/* path_mmproj */ mmproj_file.path, // can be empty
|
||||
/* path_draft */ draft_file.path // can be empty
|
||||
};
|
||||
if (!model.path.empty()) {
|
||||
models.push_back(model);
|
||||
@@ -405,13 +428,17 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
|
||||
if (file.is_dir) {
|
||||
scan_subdir(file.path, file.name);
|
||||
} else if (string_ends_with(file.name, ".gguf")) {
|
||||
if (is_mmproj_file(file.name) || is_draft_file(file.name)) {
|
||||
continue; // companion file, cannot be loaded as a model on its own
|
||||
}
|
||||
// single file model
|
||||
std::string name = file.name;
|
||||
string_replace_all(name, ".gguf", "");
|
||||
local_model model{
|
||||
/* name */ name,
|
||||
/* path */ file.path,
|
||||
/* path_mmproj */ ""
|
||||
/* path_mmproj */ "",
|
||||
/* path_draft */ ""
|
||||
};
|
||||
models.push_back(model);
|
||||
}
|
||||
@@ -426,6 +453,9 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m
|
||||
if (!model.path_mmproj.empty()) {
|
||||
preset.set_option(*this, "LLAMA_ARG_MMPROJ", model.path_mmproj);
|
||||
}
|
||||
if (!model.path_draft.empty()) {
|
||||
preset.set_option(*this, "LLAMA_ARG_SPEC_DRAFT_MODEL", model.path_draft);
|
||||
}
|
||||
out[preset.name] = preset;
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -549,20 +549,34 @@ static void load_vocab(const char * filename, const Config * config, struct my_l
|
||||
|
||||
const int token_idx = gguf_find_key(ctx, KV_TOKENIZER_LIST);
|
||||
GGML_ASSERT(token_idx >= 0);
|
||||
|
||||
const int score_idx = gguf_find_key(ctx, KV_TOKENIZER_SCORES);
|
||||
GGML_ASSERT(score_idx >= 0);
|
||||
const float * scores = (const float * ) gguf_get_arr_data(ctx, score_idx);
|
||||
|
||||
const int toktype_idx = gguf_find_key(ctx, KV_TOKENIZER_TOKEN_TYPE);
|
||||
GGML_ASSERT(toktype_idx >= 0);
|
||||
const int * toktypes = (const int * ) gguf_get_arr_data(ctx, toktype_idx);
|
||||
if (gguf_get_kv_type(ctx, token_idx) != GGUF_TYPE_ARRAY ||
|
||||
gguf_get_arr_type(ctx, token_idx) != GGUF_TYPE_STRING) {
|
||||
die_fmt("invalid gguf type for %s", KV_TOKENIZER_LIST);
|
||||
}
|
||||
|
||||
const uint32_t n_vocab = gguf_get_arr_n(ctx, token_idx);
|
||||
if (n_vocab != static_cast<uint32_t>(config->vocab_size)) {
|
||||
die_fmt("vocab size mismatch: (gguf) %u != (llama2c) %d", n_vocab, config->vocab_size);
|
||||
}
|
||||
|
||||
const int score_idx = gguf_find_key(ctx, KV_TOKENIZER_SCORES);
|
||||
GGML_ASSERT(score_idx >= 0);
|
||||
if (gguf_get_kv_type(ctx, score_idx) != GGUF_TYPE_ARRAY ||
|
||||
gguf_get_arr_type(ctx, score_idx) != GGUF_TYPE_FLOAT32 ||
|
||||
gguf_get_arr_n(ctx, score_idx) < n_vocab) {
|
||||
die_fmt("invalid gguf type or size for %s", KV_TOKENIZER_SCORES);
|
||||
}
|
||||
const float * scores = (const float * ) gguf_get_arr_data(ctx, score_idx);
|
||||
|
||||
const int toktype_idx = gguf_find_key(ctx, KV_TOKENIZER_TOKEN_TYPE);
|
||||
GGML_ASSERT(toktype_idx >= 0);
|
||||
if (gguf_get_kv_type(ctx, toktype_idx) != GGUF_TYPE_ARRAY ||
|
||||
gguf_get_arr_type(ctx, toktype_idx) != GGUF_TYPE_INT32 ||
|
||||
gguf_get_arr_n(ctx, toktype_idx) < n_vocab) {
|
||||
die_fmt("invalid gguf type or size for %s", KV_TOKENIZER_TOKEN_TYPE);
|
||||
}
|
||||
const int * toktypes = (const int * ) gguf_get_arr_data(ctx, toktype_idx);
|
||||
|
||||
vocab->id_to_token.resize(n_vocab);
|
||||
|
||||
for (uint32_t i = 0; i < n_vocab; i++) {
|
||||
|
||||
@@ -3962,7 +3962,10 @@ static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vec
|
||||
}
|
||||
|
||||
// Needs to be kept up to date on shader changes
|
||||
const uint32_t bank_conflict_offset = device->coopmat_support ? 8 : 1;
|
||||
// Needs to stay aligned with ggml_vk_mul_mm_spec.
|
||||
const bool intel_shmem_stride_pad_zero = device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support &&
|
||||
device->driver_id == vk::DriverId::eIntelProprietaryWindows;
|
||||
const uint32_t bank_conflict_offset = intel_shmem_stride_pad_zero ? 0 : (device->coopmat_support ? 8 : 1);
|
||||
const uint32_t type_size = device->fp16 ? sizeof(ggml_fp16_t) : sizeof(float);
|
||||
const uint32_t warps = warptile[0] / warptile[10];
|
||||
|
||||
@@ -4579,8 +4582,13 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
|
||||
}
|
||||
#endif
|
||||
|
||||
auto const &ggml_vk_mul_mm_spec = [](std::vector<uint32_t> spec, bool aligned) {
|
||||
spec.push_back(aligned ? 1u : 0u);
|
||||
auto const &ggml_vk_mul_mm_spec = [&device](std::vector<uint32_t> spec, bool aligned) {
|
||||
spec.push_back(aligned ? 1u : 0u); // constantID=11: ALIGNED
|
||||
if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support &&
|
||||
device->driver_id == vk::DriverId::eIntelProprietaryWindows) {
|
||||
spec.push_back(0u); // constantID=12: SHMEM_STRIDE_PAD = 0
|
||||
spec.push_back(1u); // constantID=13: APPLY_SLM_A_RESHAPE = true
|
||||
}
|
||||
return spec;
|
||||
};
|
||||
|
||||
|
||||
@@ -119,10 +119,13 @@ layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working wit
|
||||
#endif
|
||||
|
||||
#ifdef COOPMAT
|
||||
#define SHMEM_STRIDE (BK / 2 + 4)
|
||||
layout(constant_id = 12) const uint SHMEM_STRIDE_PAD = 4;
|
||||
layout(constant_id = 13) const bool APPLY_SLM_A_RESHAPE = false;
|
||||
#else
|
||||
#define SHMEM_STRIDE (BK / 2 + 1)
|
||||
const uint SHMEM_STRIDE_PAD = 1;
|
||||
const bool APPLY_SLM_A_RESHAPE = false;
|
||||
#endif
|
||||
#define SHMEM_STRIDE (BK / 2 + SHMEM_STRIDE_PAD)
|
||||
|
||||
shared FLOAT_TYPEV2 buf_a[BM * SHMEM_STRIDE];
|
||||
shared FLOAT_TYPEV2 buf_b[BN * SHMEM_STRIDE];
|
||||
@@ -302,7 +305,7 @@ void main() {
|
||||
[[unroll]] for (uint i = 0; i < BK; i += TK) {
|
||||
[[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
|
||||
// Load from shared into cache
|
||||
coopMatLoad(cache_a, buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor);
|
||||
coopMatLoad(cache_a, buf_a, a_shmem_index(warp_r * WM + cm_row * TM, i / 2), a_shmem_stride(), gl_CooperativeMatrixLayoutRowMajor);
|
||||
|
||||
[[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
|
||||
coopMatLoad(cache_b, buf_b, (warp_c * WN + cm_col * TN) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutColumnMajor);
|
||||
|
||||
@@ -1,60 +1,76 @@
|
||||
// k_pair is the K coordinate measured in FLOAT_TYPEV2 elements.
|
||||
uint a_shmem_index(uint m, uint k_pair) {
|
||||
if (APPLY_SLM_A_RESHAPE) {
|
||||
const uint tile_width = TK / 2;
|
||||
return (k_pair / tile_width) * BM * tile_width
|
||||
+ m * tile_width
|
||||
+ k_pair % tile_width;
|
||||
}
|
||||
return m * SHMEM_STRIDE + k_pair;
|
||||
}
|
||||
|
||||
uint a_shmem_stride() {
|
||||
return APPLY_SLM_A_RESHAPE ? TK / 2 : SHMEM_STRIDE;
|
||||
}
|
||||
|
||||
void store_a(uint m, uint k_pair, FLOAT_TYPEV2 value) {
|
||||
buf_a[a_shmem_index(m, k_pair)] = value;
|
||||
}
|
||||
|
||||
void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uint idx_m, const uint block, const uint end_k) {
|
||||
#if defined(DATA_A_F32) || defined(DATA_A_F16)
|
||||
#if LOAD_VEC_A == 8
|
||||
if (ALIGNED != 0) {
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
FLOAT_TYPEV8 aa = FLOAT_TYPEV8(data_a[idx]);
|
||||
buf_a[buf_idx ] = aa[0].xy;
|
||||
buf_a[buf_idx + 1] = aa[0].zw;
|
||||
buf_a[buf_idx + 2] = aa[1].xy;
|
||||
buf_a[buf_idx + 3] = aa[1].zw;
|
||||
store_a(col, k_pair, aa[0].xy);
|
||||
store_a(col, k_pair + 1, aa[0].zw);
|
||||
store_a(col, k_pair + 2, aa[1].xy);
|
||||
store_a(col, k_pair + 3, aa[1].zw);
|
||||
return;
|
||||
}
|
||||
#elif LOAD_VEC_A == 4
|
||||
if (ALIGNED != 0) {
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
FLOAT_TYPEV4 aa = FLOAT_TYPEV4(data_a[idx]);
|
||||
buf_a[buf_idx ] = aa.xy;
|
||||
buf_a[buf_idx + 1] = aa.zw;
|
||||
store_a(col, k_pair, aa.xy);
|
||||
store_a(col, k_pair + 1, aa.zw);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
const uint idx = pos_a + col * p.stride_a + row * 2;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row;
|
||||
if (idx_m < p.M && block + row * 2 + 1 < end_k) {
|
||||
buf_a[buf_idx] = FLOAT_TYPEV2(data_a_scalar[idx],
|
||||
data_a_scalar[idx + 1]);
|
||||
store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx],
|
||||
data_a_scalar[idx + 1]));
|
||||
} else if (idx_m < p.M && block + row * 2 < end_k) {
|
||||
buf_a[buf_idx] = FLOAT_TYPEV2(data_a_scalar[idx], 0.0f);
|
||||
store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx], 0.0f));
|
||||
} else {
|
||||
buf_a[buf_idx] = FLOAT_TYPEV2(0.0f);
|
||||
store_a(col, row, FLOAT_TYPEV2(0.0f));
|
||||
}
|
||||
#elif defined(DATA_A_BF16)
|
||||
#if LOAD_VEC_A == 4
|
||||
if (ALIGNED != 0) {
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
FLOAT_TYPEV4 aa = FLOAT_TYPEV4(TO_FLOAT_TYPE(data_a[idx]));
|
||||
buf_a[buf_idx ] = aa.xy;
|
||||
buf_a[buf_idx + 1] = aa.zw;
|
||||
store_a(col, k_pair, aa.xy);
|
||||
store_a(col, k_pair + 1, aa.zw);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
const uint idx = pos_a + col * p.stride_a + row * 2;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row;
|
||||
if (idx_m < p.M && block + row * 2 + 1 < end_k) {
|
||||
buf_a[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]),
|
||||
TO_FLOAT_TYPE(data_a_scalar[idx + 1]));
|
||||
store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]),
|
||||
TO_FLOAT_TYPE(data_a_scalar[idx + 1])));
|
||||
} else if (idx_m < p.M && block + row * 2 < end_k) {
|
||||
buf_a[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), 0.0f);
|
||||
store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), 0.0f));
|
||||
} else {
|
||||
buf_a[buf_idx] = FLOAT_TYPEV2(0.0f);
|
||||
store_a(col, row, FLOAT_TYPEV2(0.0f));
|
||||
}
|
||||
#elif defined(DATA_A_Q4_0)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
|
||||
|
||||
const uint ib = idx / 4;
|
||||
const uint iqs = idx & 0x03;
|
||||
@@ -64,13 +80,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const vec4 v0 = (vec4(unpack8(vui & 0x0F0F0F0F)) - 8.0f) * d;
|
||||
const vec4 v1 = (vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) - 8.0f) * d;
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xy);
|
||||
buf_a[buf_idx + 1] = FLOAT_TYPEV2(v0.zw);
|
||||
buf_a[buf_idx + 8] = FLOAT_TYPEV2(v1.xy);
|
||||
buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.zw);
|
||||
const uint k_pair = row * LOAD_VEC_A / 4;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2(v0.xy));
|
||||
store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw));
|
||||
store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy));
|
||||
store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw));
|
||||
#elif defined(DATA_A_Q4_1)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
|
||||
|
||||
const uint ib = idx / 4;
|
||||
const uint iqs = idx & 0x03;
|
||||
@@ -80,13 +96,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const vec4 v0 = vec4(unpack8(vui & 0x0F0F0F0F)) * dm.x + dm.y;
|
||||
const vec4 v1 = vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) * dm.x + dm.y;
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xy);
|
||||
buf_a[buf_idx + 1 ] = FLOAT_TYPEV2(v0.zw);
|
||||
buf_a[buf_idx + 8 ] = FLOAT_TYPEV2(v1.xy);
|
||||
buf_a[buf_idx + 9 ] = FLOAT_TYPEV2(v1.zw);
|
||||
const uint k_pair = row * LOAD_VEC_A / 4;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2(v0.xy));
|
||||
store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw));
|
||||
store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy));
|
||||
store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw));
|
||||
#elif defined(DATA_A_Q5_0)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
|
||||
|
||||
const uint ib = idx / 8;
|
||||
const uint iqs = idx & 0x07;
|
||||
@@ -98,12 +114,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
|
||||
const uint vui = uint(data_a_packed16[ib].qs[iqs]);
|
||||
const vec4 v = (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f) * d;
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(v.xz);
|
||||
buf_a[buf_idx + 8] = FLOAT_TYPEV2(v.yw);
|
||||
store_a(col, row, FLOAT_TYPEV2(v.xz));
|
||||
store_a(col, row + 8, FLOAT_TYPEV2(v.yw));
|
||||
#elif defined(DATA_A_Q5_1)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
|
||||
|
||||
const uint ib = idx / 4;
|
||||
const uint iqs = idx & 0x03;
|
||||
@@ -119,13 +133,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const vec4 v0 = vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, ((vui >> 12) & 0xF) | qh1.y) * dm.x + dm.y;
|
||||
const vec4 v1 = vec4(((vui >> 16) & 0xF) | qh2.x, ((vui >> 20) & 0xF) | qh2.y, ((vui >> 24) & 0xF) | qh3.x, ((vui >> 28) & 0xF) | qh3.y) * dm.x + dm.y;
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xz);
|
||||
buf_a[buf_idx + 1] = FLOAT_TYPEV2(v1.xz);
|
||||
buf_a[buf_idx + 8] = FLOAT_TYPEV2(v0.yw);
|
||||
buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.yw);
|
||||
const uint k_pair = row * LOAD_VEC_A / 4;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2(v0.xz));
|
||||
store_a(col, k_pair + 1, FLOAT_TYPEV2(v1.xz));
|
||||
store_a(col, k_pair + 8, FLOAT_TYPEV2(v0.yw));
|
||||
store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.yw));
|
||||
#elif defined(DATA_A_Q8_0)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 8;
|
||||
const uint iqs = idx & 0x07;
|
||||
@@ -135,11 +149,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const i8vec2 v1 = unpack8(int32_t(data_a_packed16[ib].qs[2*iqs + 1])).xy;
|
||||
const vec4 v = vec4(v0.x, v0.y, v1.x, v1.y) * d;
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(v.xy);
|
||||
buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw);
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2(v.xy));
|
||||
store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw));
|
||||
#elif defined(DATA_A_Q1_0)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 16;
|
||||
const uint iqs = idx & 0xfu;
|
||||
@@ -147,13 +161,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const float d = float(data_a[ib].d);
|
||||
const uint bits = uint(data_a[ib].qs[iqs]);
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2((bits & 0x01u) != 0u ? d : -d, (bits & 0x02u) != 0u ? d : -d);
|
||||
buf_a[buf_idx + 1] = FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d);
|
||||
buf_a[buf_idx + 2] = FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d);
|
||||
buf_a[buf_idx + 3] = FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d);
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2((bits & 0x01u) != 0u ? d : -d, (bits & 0x02u) != 0u ? d : -d));
|
||||
store_a(col, k_pair + 1, FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d));
|
||||
store_a(col, k_pair + 2, FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d));
|
||||
store_a(col, k_pair + 3, FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d));
|
||||
#elif defined(DATA_A_Q2_0)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 16;
|
||||
const uint iqs = idx & 0xfu;
|
||||
@@ -161,11 +175,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib].d);
|
||||
const uint bits = uint(data_a[ib].qs[iqs]);
|
||||
|
||||
buf_a[buf_idx ] = d * (FLOAT_TYPEV2(bits & 3u, (bits >> 2u) & 3u) - FLOAT_TYPEV2(1.0f));
|
||||
buf_a[buf_idx + 1] = d * (FLOAT_TYPEV2((bits >> 4u) & 3u, bits >> 6u) - FLOAT_TYPEV2(1.0f));
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, d * (FLOAT_TYPEV2(bits & 3u, (bits >> 2u) & 3u) - FLOAT_TYPEV2(1.0f)));
|
||||
store_a(col, k_pair + 1, d * (FLOAT_TYPEV2((bits >> 4u) & 3u, bits >> 6u) - FLOAT_TYPEV2(1.0f)));
|
||||
#elif defined(DATA_A_Q2_K)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 64; // 4 values per idx
|
||||
const uint iqs = (idx % 64) * 2; // 0,2,4..126
|
||||
@@ -180,11 +194,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
|
||||
const vec4 v = dm.x * float(scales & 0xF) * qs - dm.y * float(scales >> 4);
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(v.xy);
|
||||
buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw);
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2(v.xy));
|
||||
store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw));
|
||||
#elif defined(DATA_A_TQ2_0)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 128; // 2 values per idx
|
||||
const uint iqs = (idx % 128) * 2; // elem 0,2,4..254
|
||||
@@ -197,10 +211,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
|
||||
const vec2 v = d * (vec2((qs >> shift) & 3) - 1.0);
|
||||
|
||||
buf_a[buf_idx] = FLOAT_TYPEV2(v.xy);
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2(v.xy));
|
||||
#elif defined(DATA_A_Q3_K)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 128; // 2 values per idx
|
||||
const uint iqs = idx % 128; // 0..127
|
||||
@@ -220,11 +234,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const vec2 qs = vec2(unpack8((uint(data_a_packed16[ib].qs[qsi / 2]) >> qsshift) & 0x0303).xy);
|
||||
const vec2 hm = vec2(unpack8(((uint(data_a_packed16[ib].hmask[hmi / 2]) >> (4 * n + halfsplit)) & 0x0101 ^ 0x0101) << 2).xy);
|
||||
|
||||
buf_a[buf_idx] = FLOAT_TYPEV2(dl * (qs.x - hm.x),
|
||||
dl * (qs.y - hm.y));
|
||||
store_a(col, row * LOAD_VEC_A / 2, FLOAT_TYPEV2(dl * (qs.x - hm.x),
|
||||
dl * (qs.y - hm.y)));
|
||||
#elif defined(DATA_A_Q4_K)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 64; // 4 values per idx
|
||||
const uint iqs = (idx % 64) * 2; // 0,2,4..126
|
||||
@@ -256,11 +269,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
|
||||
const vec4 q = vec4(unpack8((data_a_packed32[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F));
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m));
|
||||
buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m));
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m)));
|
||||
store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)));
|
||||
#elif defined(DATA_A_Q5_K)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 64; // 4 values per idx
|
||||
const uint iqs = (idx % 64) * 2; // 0,2,4..126
|
||||
@@ -295,11 +308,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const uint qh = ((data_a_packed32[ib].qh[qhi / 4] >> (iqs / 16)) & 0x01010101) << 4;
|
||||
const vec4 q = vec4(unpack8(qs | qh));
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m));
|
||||
buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m));
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m)));
|
||||
store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)));
|
||||
#elif defined(DATA_A_Q6_K)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 128; // 2 values per idx
|
||||
const uint iqs = idx % 128; // 0..127
|
||||
@@ -318,10 +331,9 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const uint qh = (uint(data_a_packed16[ib].qh[qhi]) >> qhshift) & 0x0303;
|
||||
const vec2 q = (vec2(unpack8(ql | (qh << 4)).xy) - 32) * dscale;
|
||||
|
||||
buf_a[buf_idx] = FLOAT_TYPEV2(q.x, q.y);
|
||||
store_a(col, row * LOAD_VEC_A / 2, FLOAT_TYPEV2(q.x, q.y));
|
||||
#elif defined(DATA_A_IQ1_S)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 32; // 8 values per idx
|
||||
const uint ib32 = (idx % 32) / 4; // 0..7
|
||||
@@ -334,13 +346,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA;
|
||||
const int16_t grid = int16_t(iq1s_grid[qs | (bitfieldExtract(qh, 3 * int(ib8 & 3), 3) << 8)]);
|
||||
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
[[unroll]] for (int k = 0; k < 4; ++k) {
|
||||
buf_a[buf_idx + k] = FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta),
|
||||
dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta));
|
||||
store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta),
|
||||
dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta)));
|
||||
}
|
||||
#elif defined(DATA_A_IQ1_M)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 32; // 8 values per idx
|
||||
const uint ib8 = idx % 32;
|
||||
@@ -356,13 +368,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA;
|
||||
const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]);
|
||||
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
[[unroll]] for (int k = 0; k < 4; ++k) {
|
||||
buf_a[buf_idx + k] = FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta),
|
||||
dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta));
|
||||
store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta),
|
||||
dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta)));
|
||||
}
|
||||
#elif defined(DATA_A_IQ2_XXS)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 32; // 8 values per idx
|
||||
const uint ib32 = (idx % 32) / 4; // 0..7
|
||||
@@ -383,17 +395,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const vec4 grid0 = vec4(unpack8(grid.x));
|
||||
const vec4 grid1 = vec4(unpack8(grid.y));
|
||||
|
||||
buf_a[buf_idx ] = db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x,
|
||||
(sign & 2) != 0 ? -grid0.y : grid0.y);
|
||||
buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z,
|
||||
(sign & 8) != 0 ? -grid0.w : grid0.w);
|
||||
buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x,
|
||||
(sign & 32) != 0 ? -grid1.y : grid1.y);
|
||||
buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z,
|
||||
(sign & 128) != 0 ? -grid1.w : grid1.w);
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x,
|
||||
(sign & 2) != 0 ? -grid0.y : grid0.y));
|
||||
store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z,
|
||||
(sign & 8) != 0 ? -grid0.w : grid0.w));
|
||||
store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x,
|
||||
(sign & 32) != 0 ? -grid1.y : grid1.y));
|
||||
store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z,
|
||||
(sign & 128) != 0 ? -grid1.w : grid1.w));
|
||||
#elif defined(DATA_A_IQ2_XS)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 32; // 8 values per idx
|
||||
const uint ib32 = (idx % 32) / 4; // 0..7
|
||||
@@ -409,17 +421,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const vec4 grid0 = vec4(unpack8(grid.x));
|
||||
const vec4 grid1 = vec4(unpack8(grid.y));
|
||||
|
||||
buf_a[buf_idx ] = db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x,
|
||||
(sign & 2) != 0 ? -grid0.y : grid0.y);
|
||||
buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z,
|
||||
(sign & 8) != 0 ? -grid0.w : grid0.w);
|
||||
buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x,
|
||||
(sign & 32) != 0 ? -grid1.y : grid1.y);
|
||||
buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z,
|
||||
(sign & 128) != 0 ? -grid1.w : grid1.w);
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x,
|
||||
(sign & 2) != 0 ? -grid0.y : grid0.y));
|
||||
store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z,
|
||||
(sign & 8) != 0 ? -grid0.w : grid0.w));
|
||||
store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x,
|
||||
(sign & 32) != 0 ? -grid1.y : grid1.y));
|
||||
store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z,
|
||||
(sign & 128) != 0 ? -grid1.w : grid1.w));
|
||||
#elif defined(DATA_A_IQ2_S)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 32; // 8 values per idx
|
||||
const uint ib8 = idx % 32; // 0..31
|
||||
@@ -437,17 +449,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const vec4 grid0 = vec4(unpack8(grid.x));
|
||||
const vec4 grid1 = vec4(unpack8(grid.y));
|
||||
|
||||
buf_a[buf_idx ] = db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x,
|
||||
(sign & 2) != 0 ? -grid0.y : grid0.y);
|
||||
buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z,
|
||||
(sign & 8) != 0 ? -grid0.w : grid0.w);
|
||||
buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x,
|
||||
(sign & 32) != 0 ? -grid1.y : grid1.y);
|
||||
buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z,
|
||||
(sign & 128) != 0 ? -grid1.w : grid1.w);
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x,
|
||||
(sign & 2) != 0 ? -grid0.y : grid0.y));
|
||||
store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z,
|
||||
(sign & 8) != 0 ? -grid0.w : grid0.w));
|
||||
store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x,
|
||||
(sign & 32) != 0 ? -grid1.y : grid1.y));
|
||||
store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z,
|
||||
(sign & 128) != 0 ? -grid1.w : grid1.w));
|
||||
#elif defined(DATA_A_IQ3_XXS)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 64; // 4 values per idx
|
||||
const uint iqs = idx % 64; // 0..63
|
||||
@@ -465,13 +477,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const uint grid = iq3xxs_grid[qs];
|
||||
const vec4 v = db * vec4(unpack8(grid));
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x,
|
||||
(sign & 2) != 0 ? -v.y : v.y);
|
||||
buf_a[buf_idx + 1] = FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z,
|
||||
(sign & 8) != 0 ? -v.w : v.w);
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x,
|
||||
(sign & 2) != 0 ? -v.y : v.y));
|
||||
store_a(col, k_pair + 1, FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z,
|
||||
(sign & 8) != 0 ? -v.w : v.w));
|
||||
#elif defined(DATA_A_IQ3_S)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 64; // 4 values per idx
|
||||
const uint iqs = idx % 64; // 0..63
|
||||
@@ -487,13 +499,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const uint32_t grid = iq3s_grid[qs | ((qh << (8 - (iqs % 8))) & 256)];
|
||||
const vec4 v = db * vec4(unpack8(grid));
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x,
|
||||
(sign & 2) != 0 ? -v.y : v.y);
|
||||
buf_a[buf_idx + 1] = FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z,
|
||||
(sign & 8) != 0 ? -v.w : v.w);
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x,
|
||||
(sign & 2) != 0 ? -v.y : v.y));
|
||||
store_a(col, k_pair + 1, FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z,
|
||||
(sign & 8) != 0 ? -v.w : v.w));
|
||||
#elif defined(DATA_A_IQ4_XS)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2;
|
||||
|
||||
const uint ib = idx / 64; // 4 values per idx
|
||||
const uint ib32 = (idx % 64) / 8; // 0..7
|
||||
@@ -507,11 +519,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const float d = float(data_a[ib].d);
|
||||
const vec4 v = d * float(int(sl | (sh << 4)) - 32) * vec4(kvalues_iq4nl[qs.x], kvalues_iq4nl[qs.y], kvalues_iq4nl[qs.z], kvalues_iq4nl[qs.w]);
|
||||
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(v.xy);
|
||||
buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw);
|
||||
const uint k_pair = row * LOAD_VEC_A / 2;
|
||||
store_a(col, k_pair, FLOAT_TYPEV2(v.xy));
|
||||
store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw));
|
||||
#elif defined(DATA_A_IQ4_NL)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
|
||||
|
||||
const uint ib = idx / 8;
|
||||
const uint iqs = idx & 0x07;
|
||||
@@ -519,13 +531,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
const FLOAT_TYPE d = FLOAT_TYPE(data_a_packed16[ib].d);
|
||||
const uint vui = uint(data_a_packed16[ib].qs[iqs]);
|
||||
|
||||
buf_a[buf_idx ] = d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF],
|
||||
kvalues_iq4nl[bitfieldExtract(vui, 8, 4)]);
|
||||
buf_a[buf_idx + 8] = d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)],
|
||||
kvalues_iq4nl[vui >> 12]);
|
||||
const uint k_pair = row * LOAD_VEC_A / 4;
|
||||
store_a(col, k_pair, d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF],
|
||||
kvalues_iq4nl[bitfieldExtract(vui, 8, 4)]));
|
||||
store_a(col, k_pair + 8, d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)],
|
||||
kvalues_iq4nl[vui >> 12]));
|
||||
#elif defined(DATA_A_MXFP4)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4;
|
||||
|
||||
const uint ib = idx / 8;
|
||||
const uint iqs = (idx & 0x07) * 2;
|
||||
@@ -536,38 +548,37 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin
|
||||
#ifdef USE_OCP_FP4
|
||||
const float d = e8m0_to_fp32(data_a[ib].e);
|
||||
const u8vec2 packed = u8vec2(vui, vui2);
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * FLOAT_TYPE(d);
|
||||
buf_a[buf_idx + 8] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * FLOAT_TYPE(d);
|
||||
store_a(col, row, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * FLOAT_TYPE(d));
|
||||
store_a(col, row + 8, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * FLOAT_TYPE(d));
|
||||
#else
|
||||
const float d = e8m0_to_fp32(data_a[ib].e) * 0.5;
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d,
|
||||
kvalues_mxfp4[vui2 & 0xF] * d);
|
||||
buf_a[buf_idx + 8] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d,
|
||||
kvalues_mxfp4[vui2 >> 4] * d);
|
||||
store_a(col, row, FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d,
|
||||
kvalues_mxfp4[vui2 & 0xF] * d));
|
||||
store_a(col, row + 8, FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d,
|
||||
kvalues_mxfp4[vui2 >> 4] * d));
|
||||
#endif
|
||||
#elif defined(DATA_A_NVFP4)
|
||||
const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row;
|
||||
// lo and hi nibbles are 8 elements apart, which doesn't quite line up with
|
||||
// how the thread mapping and buf_idx calculation works for other types.
|
||||
const uint buf_idx = col * SHMEM_STRIDE + (row & 3) + (row & ~3) * 2;
|
||||
|
||||
const uint ib = idx / 16u;
|
||||
const uint sub = (idx & 0xC) >> 2;
|
||||
const uint iqs = (idx & 0xF) * 2;
|
||||
const uint vui = uint(data_a[ib].qs[iqs]);
|
||||
const uint vui2 = uint(data_a[ib].qs[iqs+1]);
|
||||
|
||||
// lo and hi nibbles are 8 elements apart, which doesn't quite line up with
|
||||
// how the thread mapping and buf_idx calculation works for other types.
|
||||
const uint eff_row = (row & 3) + (row & ~3) * 2;
|
||||
#ifdef USE_OCP_FP4
|
||||
const FLOAT_TYPE d = FLOAT_TYPE(ue4m3_from_bits(data_a[ib].d[sub]));
|
||||
const u8vec2 packed = u8vec2(vui, vui2);
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * d;
|
||||
buf_a[buf_idx + 4] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * d;
|
||||
store_a(col, eff_row, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * d);
|
||||
store_a(col, eff_row + 4, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * d);
|
||||
#else
|
||||
const float d = ue4m3_to_fp32(data_a[ib].d[sub]) * 0.5;
|
||||
buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d,
|
||||
kvalues_mxfp4[vui2 & 0xF] * d);
|
||||
buf_a[buf_idx + 4] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d,
|
||||
kvalues_mxfp4[vui2 >> 4] * d);
|
||||
store_a(col, eff_row, FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d,
|
||||
kvalues_mxfp4[vui2 & 0xF] * d));
|
||||
store_a(col, eff_row + 4, FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d,
|
||||
kvalues_mxfp4[vui2 >> 4] * d));
|
||||
#endif
|
||||
#endif
|
||||
}
|
||||
|
||||
@@ -124,6 +124,7 @@ class Keys:
|
||||
EXPERT_WEIGHTS_NORM = "{arch}.expert_weights_norm"
|
||||
EXPERT_GATING_FUNC = "{arch}.expert_gating_func"
|
||||
EXPERT_GROUP_SCALE = "{arch}.expert_group_scale"
|
||||
EXPERT_LATENT_LENGTH = "{arch}.expert_latent_length"
|
||||
EXPERTS_PER_GROUP = "{arch}.experts_per_group"
|
||||
MOE_EVERY_N_LAYERS = "{arch}.moe_every_n_layers"
|
||||
MOE_LATENT_SIZE = "{arch}.moe_latent_size"
|
||||
@@ -238,6 +239,13 @@ class Keys:
|
||||
SCALING_YARN_BETA_FAST = "{arch}.rope.scaling.yarn_beta_fast"
|
||||
SCALING_YARN_BETA_SLOW = "{arch}.rope.scaling.yarn_beta_slow"
|
||||
|
||||
class Activation:
|
||||
SITU_BETA = "{arch}.activation.situ_beta"
|
||||
SITU_LINEAR_BETA = "{arch}.activation.situ_linear_beta"
|
||||
|
||||
class AttnRes:
|
||||
BLOCK_SIZE = "{arch}.attn_res.block_size"
|
||||
|
||||
class Split:
|
||||
LLM_KV_SPLIT_NO = "split.no"
|
||||
LLM_KV_SPLIT_COUNT = "split.count"
|
||||
@@ -252,7 +260,8 @@ class Keys:
|
||||
DT_B_C_RMS = "{arch}.ssm.dt_b_c_rms"
|
||||
|
||||
class KDA:
|
||||
HEAD_DIM = "{arch}.kda.head_dim"
|
||||
HEAD_DIM = "{arch}.kda.head_dim"
|
||||
GATE_LOWER_BOUND = "{arch}.kda.gate_lower_bound"
|
||||
|
||||
class WKV:
|
||||
HEAD_SIZE = "{arch}.wkv.head_size"
|
||||
@@ -580,6 +589,7 @@ class MODEL_ARCH(IntEnum):
|
||||
LLAMA_EMBED = auto()
|
||||
MAINCODER = auto()
|
||||
KIMI_LINEAR = auto()
|
||||
KIMI_K3 = auto()
|
||||
TALKIE = auto()
|
||||
MELLUM = auto()
|
||||
NANBEIGE = auto()
|
||||
@@ -698,6 +708,13 @@ class MODEL_TENSOR(IntEnum):
|
||||
SSM_BETA = auto() # Kimi Linear qwen3.5
|
||||
SSM_G_A = auto() # Kimi Linear
|
||||
SSM_G_B = auto() # Kimi Linear
|
||||
SSM_G = auto() # Kimi K3 (full-rank KDA gate, replaces SSM_G_A/SSM_G_B)
|
||||
ATTN_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, pre-attention)
|
||||
FFN_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, pre-FFN)
|
||||
OUTPUT_RES_SCORE = auto() # Kimi K3 (fused res_norm * res_proj, final)
|
||||
FFN_ROUTED_DOWN = auto() # Kimi K3 (latent MoE: hidden -> latent)
|
||||
FFN_ROUTED_UP = auto() # Kimi K3 (latent MoE: latent -> hidden)
|
||||
FFN_ROUTED_NORM = auto() # Kimi K3 (latent MoE: norm on expert output)
|
||||
TIME_MIX_W0 = auto()
|
||||
TIME_MIX_W1 = auto()
|
||||
TIME_MIX_W2 = auto()
|
||||
@@ -1288,6 +1305,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
|
||||
MODEL_ARCH.LLAMA_EMBED: "llama-embed",
|
||||
MODEL_ARCH.MAINCODER: "maincoder",
|
||||
MODEL_ARCH.KIMI_LINEAR: "kimi-linear",
|
||||
MODEL_ARCH.KIMI_K3: "kimi-k3",
|
||||
MODEL_ARCH.TALKIE: "talkie",
|
||||
MODEL_ARCH.MELLUM: "mellum",
|
||||
MODEL_ARCH.NANBEIGE: "nanbeige",
|
||||
@@ -1404,6 +1422,13 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
|
||||
MODEL_TENSOR.SSM_BETA: "blk.{bid}.ssm_beta", # Kimi Linear qwen3.5
|
||||
MODEL_TENSOR.SSM_G_A: "blk.{bid}.ssm_g_a", # Kimi Linear
|
||||
MODEL_TENSOR.SSM_G_B: "blk.{bid}.ssm_g_b", # Kimi Linear
|
||||
MODEL_TENSOR.SSM_G: "blk.{bid}.ssm_g", # Kimi K3
|
||||
MODEL_TENSOR.ATTN_RES_SCORE: "blk.{bid}.attn_res_score", # Kimi K3
|
||||
MODEL_TENSOR.FFN_RES_SCORE: "blk.{bid}.ffn_res_score", # Kimi K3
|
||||
MODEL_TENSOR.OUTPUT_RES_SCORE: "output_res_score", # Kimi K3
|
||||
MODEL_TENSOR.FFN_ROUTED_DOWN: "blk.{bid}.ffn_routed_down", # Kimi K3
|
||||
MODEL_TENSOR.FFN_ROUTED_UP: "blk.{bid}.ffn_routed_up", # Kimi K3
|
||||
MODEL_TENSOR.FFN_ROUTED_NORM: "blk.{bid}.ffn_routed_norm", # Kimi K3
|
||||
MODEL_TENSOR.TIME_MIX_W0: "blk.{bid}.time_mix_w0",
|
||||
MODEL_TENSOR.TIME_MIX_W1: "blk.{bid}.time_mix_w1",
|
||||
MODEL_TENSOR.TIME_MIX_W2: "blk.{bid}.time_mix_w2",
|
||||
@@ -4960,6 +4985,56 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||
MODEL_TENSOR.FFN_DOWN_SHEXP,
|
||||
MODEL_TENSOR.FFN_UP_SHEXP,
|
||||
],
|
||||
MODEL_ARCH.KIMI_K3: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT_NORM,
|
||||
MODEL_TENSOR.OUTPUT,
|
||||
MODEL_TENSOR.OUTPUT_RES_SCORE,
|
||||
MODEL_TENSOR.ATTN_NORM,
|
||||
MODEL_TENSOR.ATTN_RES_SCORE,
|
||||
MODEL_TENSOR.FFN_RES_SCORE,
|
||||
# MLA (full-attention layers)
|
||||
MODEL_TENSOR.ATTN_Q,
|
||||
MODEL_TENSOR.ATTN_K,
|
||||
MODEL_TENSOR.ATTN_V,
|
||||
MODEL_TENSOR.ATTN_OUT,
|
||||
MODEL_TENSOR.ATTN_GATE,
|
||||
MODEL_TENSOR.ATTN_Q_A,
|
||||
MODEL_TENSOR.ATTN_Q_B,
|
||||
MODEL_TENSOR.ATTN_KV_A_MQA,
|
||||
MODEL_TENSOR.ATTN_KV_B,
|
||||
MODEL_TENSOR.ATTN_K_B,
|
||||
MODEL_TENSOR.ATTN_V_B,
|
||||
MODEL_TENSOR.ATTN_Q_A_NORM,
|
||||
MODEL_TENSOR.ATTN_KV_A_NORM,
|
||||
# KDA (linear-attention layers)
|
||||
MODEL_TENSOR.SSM_CONV1D_Q,
|
||||
MODEL_TENSOR.SSM_CONV1D_K,
|
||||
MODEL_TENSOR.SSM_CONV1D_V,
|
||||
MODEL_TENSOR.SSM_F_A,
|
||||
MODEL_TENSOR.SSM_F_B,
|
||||
MODEL_TENSOR.SSM_BETA,
|
||||
MODEL_TENSOR.SSM_A,
|
||||
MODEL_TENSOR.SSM_G,
|
||||
MODEL_TENSOR.SSM_DT,
|
||||
MODEL_TENSOR.SSM_NORM,
|
||||
# FFN
|
||||
MODEL_TENSOR.FFN_NORM,
|
||||
MODEL_TENSOR.FFN_GATE,
|
||||
MODEL_TENSOR.FFN_DOWN,
|
||||
MODEL_TENSOR.FFN_UP,
|
||||
MODEL_TENSOR.FFN_GATE_INP,
|
||||
MODEL_TENSOR.FFN_EXP_PROBS_B,
|
||||
MODEL_TENSOR.FFN_GATE_EXP,
|
||||
MODEL_TENSOR.FFN_DOWN_EXP,
|
||||
MODEL_TENSOR.FFN_UP_EXP,
|
||||
MODEL_TENSOR.FFN_GATE_SHEXP,
|
||||
MODEL_TENSOR.FFN_DOWN_SHEXP,
|
||||
MODEL_TENSOR.FFN_UP_SHEXP,
|
||||
MODEL_TENSOR.FFN_ROUTED_DOWN,
|
||||
MODEL_TENSOR.FFN_ROUTED_UP,
|
||||
MODEL_TENSOR.FFN_ROUTED_NORM,
|
||||
],
|
||||
MODEL_ARCH.TALKIE: [
|
||||
MODEL_TENSOR.TOKEN_EMBD,
|
||||
MODEL_TENSOR.OUTPUT,
|
||||
|
||||
@@ -1103,6 +1103,21 @@ class GGUFWriter:
|
||||
def add_ssm_dt_b_c_rms(self, value: bool) -> None:
|
||||
self.add_bool(Keys.SSM.DT_B_C_RMS.format(arch=self.arch), value)
|
||||
|
||||
def add_kda_gate_lower_bound(self, value: float) -> None:
|
||||
self.add_float32(Keys.KDA.GATE_LOWER_BOUND.format(arch=self.arch), value)
|
||||
|
||||
def add_expert_latent_length(self, value: int) -> None:
|
||||
self.add_uint32(Keys.LLM.EXPERT_LATENT_LENGTH.format(arch=self.arch), value)
|
||||
|
||||
def add_activation_situ_beta(self, value: float) -> None:
|
||||
self.add_float32(Keys.Activation.SITU_BETA.format(arch=self.arch), value)
|
||||
|
||||
def add_activation_situ_linear_beta(self, value: float) -> None:
|
||||
self.add_float32(Keys.Activation.SITU_LINEAR_BETA.format(arch=self.arch), value)
|
||||
|
||||
def add_attn_res_block_size(self, value: int) -> None:
|
||||
self.add_uint32(Keys.AttnRes.BLOCK_SIZE.format(arch=self.arch), value)
|
||||
|
||||
def add_kda_head_dim(self, value: int) -> None:
|
||||
self.add_uint32(Keys.KDA.HEAD_DIM.format(arch=self.arch), value)
|
||||
|
||||
|
||||
@@ -912,6 +912,19 @@ class TensorNameMap:
|
||||
"model.layers.{bid}.linear_attn.in_proj_b", # qwen3.5
|
||||
"model.layers.{bid}.self_attn.b_proj", # Kimi Linear
|
||||
),
|
||||
# Kimi K3 latent MoE: routed experts operate in a down-projected space
|
||||
MODEL_TENSOR.FFN_ROUTED_DOWN: (
|
||||
"model.layers.{bid}.block_sparse_moe.routed_expert_down_proj",
|
||||
),
|
||||
|
||||
MODEL_TENSOR.FFN_ROUTED_UP: (
|
||||
"model.layers.{bid}.block_sparse_moe.routed_expert_up_proj",
|
||||
),
|
||||
|
||||
MODEL_TENSOR.FFN_ROUTED_NORM: (
|
||||
"model.layers.{bid}.block_sparse_moe.routed_expert_norm",
|
||||
),
|
||||
|
||||
MODEL_TENSOR.SSM_G_A: (
|
||||
"model.layers.{bid}.self_attn.g_a_proj",
|
||||
),
|
||||
|
||||
@@ -0,0 +1,324 @@
|
||||
{%- macro escape_attr(value) -%}
|
||||
{{- value|string|replace('&', '&')|replace('"', '"') -}}
|
||||
{%- 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 -%}
|
||||
|
||||
@@ -5,7 +5,7 @@ import os
|
||||
import sys
|
||||
import subprocess
|
||||
|
||||
HTTPLIB_VERSION = "refs/tags/v0.53.0"
|
||||
HTTPLIB_VERSION = "refs/tags/v0.53.1"
|
||||
|
||||
vendor = {
|
||||
"https://github.com/nlohmann/json/releases/latest/download/json.hpp": "vendor/nlohmann/json.hpp",
|
||||
|
||||
@@ -46,6 +46,8 @@ Mandatory on every review; any finding here is **blocking**. Rule of thumb: GGUF
|
||||
|
||||
- **Sizes/counts from tensor dims:** validate before allocating. Products like `ne[i]*nb[i]`/nbytes can overflow on crafted dims into an undersized alloc then heap overflow. Overflow checks must run BEFORE the arithmetic they guard - padding/alignment macros wrap to 0 near `SIZE_MAX`, so a guard after the pad passes.
|
||||
- **GGUF strings/arrays:** cap declared lengths and element counts before using them to size a loop or buffer; validate element type and length before casting an array to a pointer or reading fixed indices (`[i+1]`, `[0..2]`).
|
||||
- **Element-type confusion:** casting `gguf_get_arr_data()` or `tensor->data` to `float *`/`int32_t *` needs an element-type check first (`gguf_get_kv_type() == GGUF_TYPE_ARRAY` then `gguf_get_arr_type()`; `type == GGML_TYPE_F32` for tensors). A `UINT8` array or `I8` tensor passes every length check, then gets read 4 bytes per element - a nearby length check is not a type check.
|
||||
- **Loaders:** `GGML_ASSERT` on a file-derived value aborts the process; throw instead where the caller already catches (vocab, model loader, clip).
|
||||
- **File-supplied counts indexing fixed arrays:** bound any count (e.g. layer/block count into a `LLAMA_MAX_*` array) before indexing; watch checks that only fire when an optional key is present.
|
||||
- **Declared vs actual array length:** check the declared length of a GGUF array against the count actually read, not just against a buffer size.
|
||||
- **Bounds comparisons:** flag narrowing casts (`size_t`->`int32_t`) and signed/unsigned mixing that can bypass a length check and copy past a buffer.
|
||||
|
||||
@@ -396,8 +396,11 @@ static void llama_adapter_lora_init_impl(llama_model & model, const char * path_
|
||||
llama_file gguf_file(path_lora, "rb");
|
||||
std::vector<uint8_t> read_buf;
|
||||
auto set_tensor = [&](ggml_tensor * orig, ggml_tensor * dev) {
|
||||
size_t offs = gguf_get_data_offset(ctx_gguf.get()) + gguf_get_tensor_offset(ctx_gguf.get(), gguf_find_tensor(ctx_gguf.get(), orig->name));
|
||||
size_t size = ggml_nbytes(orig);
|
||||
const size_t offs = gguf_get_data_offset(ctx_gguf.get()) + gguf_get_tensor_offset(ctx_gguf.get(), gguf_find_tensor(ctx_gguf.get(), orig->name));
|
||||
const size_t size = ggml_nbytes(orig);
|
||||
if (offs + size < offs || offs + size > gguf_file.size()) {
|
||||
throw std::runtime_error(format("LoRA tensor '%s' data is not within the file bounds, file is corrupted or incomplete", orig->name));
|
||||
}
|
||||
read_buf.resize(size);
|
||||
gguf_file.seek(offs, SEEK_SET);
|
||||
gguf_file.read_raw(read_buf.data(), size);
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 ||
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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
@@ -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;
|
||||
|
||||
@@ -316,15 +316,19 @@ namespace GGUFMeta {
|
||||
struct GGUFMeta::ArrayInfo arr_info =
|
||||
GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(ctx, kid);
|
||||
|
||||
bool type_ok = false;
|
||||
switch (arr_info.gt) {
|
||||
case GGUF_TYPE_UINT32:
|
||||
case GGUF_TYPE_INT32: GGML_ASSERT((std::is_same<T, int32_t>::value) ||
|
||||
(std::is_same<T, uint32_t>::value)); break;
|
||||
case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same<T, float>::value)); break;
|
||||
case GGUF_TYPE_STRING: GGML_ASSERT((std::is_same<T, std::string>::value)); break;
|
||||
case GGUF_TYPE_INT32: type_ok = (std::is_same<T, int32_t>::value) ||
|
||||
(std::is_same<T, uint32_t>::value); break;
|
||||
case GGUF_TYPE_FLOAT32: type_ok = (std::is_same<T, float>::value); break;
|
||||
case GGUF_TYPE_STRING: type_ok = (std::is_same<T, std::string>::value); break;
|
||||
default:
|
||||
throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str()));
|
||||
}
|
||||
if (!type_ok) {
|
||||
throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt)));
|
||||
}
|
||||
|
||||
if constexpr (std::is_same<T, std::string>::value) {
|
||||
const size_t n_items = gguf_get_arr_n(ctx, kid);
|
||||
@@ -357,16 +361,20 @@ namespace GGUFMeta {
|
||||
struct GGUFMeta::ArrayInfo arr_info =
|
||||
GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(ctx, kid);
|
||||
|
||||
bool type_ok = false;
|
||||
switch (arr_info.gt) {
|
||||
case GGUF_TYPE_BOOL:
|
||||
case GGUF_TYPE_UINT32:
|
||||
case GGUF_TYPE_INT32: GGML_ASSERT((std::is_same<T, int32_t>::value) ||
|
||||
(std::is_same<T, uint32_t>::value)); break;
|
||||
case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same<T, float>::value)); break;
|
||||
case GGUF_TYPE_STRING: GGML_ASSERT((std::is_same<T, std::string>::value)); break;
|
||||
case GGUF_TYPE_INT32: type_ok = (std::is_same<T, int32_t>::value) ||
|
||||
(std::is_same<T, uint32_t>::value); break;
|
||||
case GGUF_TYPE_FLOAT32: type_ok = (std::is_same<T, float>::value); break;
|
||||
case GGUF_TYPE_STRING: type_ok = (std::is_same<T, std::string>::value); break;
|
||||
default:
|
||||
throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str()));
|
||||
}
|
||||
if (!type_ok) {
|
||||
throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt)));
|
||||
}
|
||||
|
||||
if (arr_info.length > N_MAX) {
|
||||
throw std::runtime_error(format("array length %u for key %s exceeds max %u", (uint32_t) arr_info.length, key.c_str(), (uint32_t) N_MAX));
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
+31
-1
@@ -1989,6 +1989,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|
||||
// Kimi-K2 doesn't need merges, skip
|
||||
LLAMA_LOG_INFO("%s: Kimi-K2 tokenizer detected, skipping BPE merges\n", __func__);
|
||||
} else {
|
||||
if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY ||
|
||||
gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) {
|
||||
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
|
||||
}
|
||||
const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
|
||||
for (int i = 0; i < n_merges; i++) {
|
||||
const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
|
||||
@@ -2028,8 +2032,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|
||||
|
||||
const int precompiled_charsmap_keyidx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str());
|
||||
if (precompiled_charsmap_keyidx != -1) {
|
||||
if (gguf_get_kv_type(ctx, precompiled_charsmap_keyidx) != GGUF_TYPE_ARRAY) {
|
||||
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str()));
|
||||
}
|
||||
const gguf_type pc_type = gguf_get_arr_type(ctx, precompiled_charsmap_keyidx);
|
||||
GGML_ASSERT(pc_type == GGUF_TYPE_INT8 || pc_type == GGUF_TYPE_UINT8);
|
||||
if (pc_type != GGUF_TYPE_INT8 && pc_type != GGUF_TYPE_UINT8) {
|
||||
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str()));
|
||||
}
|
||||
|
||||
const size_t n_precompiled_charsmap = gguf_get_arr_n(ctx, precompiled_charsmap_keyidx);
|
||||
const char * pc = (const char *) gguf_get_arr_data(ctx, precompiled_charsmap_keyidx);
|
||||
@@ -2081,6 +2090,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|
||||
throw std::runtime_error("cannot find tokenizer merges in model file\n");
|
||||
}
|
||||
{
|
||||
if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY ||
|
||||
gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) {
|
||||
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str()));
|
||||
}
|
||||
const int n_merges = gguf_get_arr_n(ctx, merges_keyidx);
|
||||
for (int i = 0; i < n_merges; i++) {
|
||||
const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i);
|
||||
@@ -2407,11 +2420,20 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|
||||
throw std::runtime_error("cannot find tokenizer vocab in model file\n");
|
||||
}
|
||||
|
||||
if (gguf_get_kv_type(ctx, token_idx) != GGUF_TYPE_ARRAY ||
|
||||
gguf_get_arr_type(ctx, token_idx) != GGUF_TYPE_STRING) {
|
||||
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_LIST).c_str()));
|
||||
}
|
||||
|
||||
const uint32_t n_tokens = gguf_get_arr_n(ctx, token_idx);
|
||||
|
||||
const float * scores = nullptr;
|
||||
const int score_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SCORES).c_str());
|
||||
if (score_idx != -1) {
|
||||
if (gguf_get_kv_type(ctx, score_idx) != GGUF_TYPE_ARRAY ||
|
||||
gguf_get_arr_type(ctx, score_idx) != GGUF_TYPE_FLOAT32) {
|
||||
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SCORES).c_str()));
|
||||
}
|
||||
const uint32_t n_scores = gguf_get_arr_n(ctx, score_idx);
|
||||
if (n_scores < n_tokens) {
|
||||
throw std::runtime_error("Index out of array bounds for scores (" + std::to_string(n_scores) + " < " + std::to_string(n_tokens) + ")\n");
|
||||
@@ -2422,6 +2444,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|
||||
const int * toktypes = nullptr;
|
||||
const int toktype_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str());
|
||||
if (toktype_idx != -1) {
|
||||
if (gguf_get_kv_type(ctx, toktype_idx) != GGUF_TYPE_ARRAY ||
|
||||
gguf_get_arr_type(ctx, toktype_idx) != GGUF_TYPE_INT32) {
|
||||
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str()));
|
||||
}
|
||||
const uint32_t n_toktypes = gguf_get_arr_n(ctx, toktype_idx);
|
||||
if (n_toktypes < n_tokens) {
|
||||
throw std::runtime_error("Index out of array bounds for toktypes (" + std::to_string(n_toktypes) + " < " + std::to_string(n_tokens) + ")\n");
|
||||
@@ -2584,6 +2610,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|
||||
{
|
||||
const int suppress_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str());
|
||||
if (suppress_idx != -1) {
|
||||
if (gguf_get_kv_type(ctx, suppress_idx) != GGUF_TYPE_ARRAY ||
|
||||
gguf_get_arr_type(ctx, suppress_idx) != GGUF_TYPE_INT32) {
|
||||
throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str()));
|
||||
}
|
||||
const int n = gguf_get_arr_n(ctx, suppress_idx);
|
||||
const int32_t * data = (const int32_t *) gguf_get_arr_data(ctx, suppress_idx);
|
||||
// drop out-of-range ids
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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>
|
||||
{
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -3746,6 +3746,9 @@ struct clip_model_loader {
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (gguf_get_kv_type(ctx_gguf.get(), i) != GGUF_TYPE_ARRAY) {
|
||||
throw std::runtime_error(string_format("%s: key '%s' is not an array\n", __func__, key.c_str()));
|
||||
}
|
||||
const auto type = gguf_get_arr_type(ctx_gguf.get(), i);
|
||||
if (type != GGUF_TYPE_FLOAT32) {
|
||||
throw std::runtime_error(string_format("%s: array '%s' has type %d, expected %d (GGUF_TYPE_FLOAT32)\n", __func__, key.c_str(), type, GGUF_TYPE_FLOAT32));
|
||||
@@ -3780,6 +3783,9 @@ struct clip_model_loader {
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (gguf_get_kv_type(ctx_gguf.get(), i) != GGUF_TYPE_ARRAY) {
|
||||
throw std::runtime_error(string_format("%s: key '%s' is not an array\n", __func__, key.c_str()));
|
||||
}
|
||||
const auto type = gguf_get_arr_type(ctx_gguf.get(), i);
|
||||
if (type != GGUF_TYPE_INT32) {
|
||||
throw std::runtime_error(string_format("%s: array '%s' has type %d, expected %d (GGUF_TYPE_INT32)\n", __func__, key.c_str(), type, GGUF_TYPE_INT32));
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Vendored
+2
@@ -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;
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
Vendored
+1
-1
@@ -41,7 +41,7 @@ if (LLAMA_BUILD_BORINGSSL)
|
||||
set(FIPS OFF CACHE BOOL "Enable FIPS (BoringSSL)")
|
||||
|
||||
set(BORINGSSL_GIT "https://boringssl.googlesource.com/boringssl" CACHE STRING "BoringSSL git repository")
|
||||
set(BORINGSSL_VERSION "0.20260803.0" CACHE STRING "BoringSSL version")
|
||||
set(BORINGSSL_VERSION "0.20260813.0" CACHE STRING "BoringSSL version")
|
||||
|
||||
message(STATUS "Fetching BoringSSL version ${BORINGSSL_VERSION}")
|
||||
|
||||
|
||||
Vendored
+27
-11
@@ -7146,6 +7146,10 @@ PathParamsMatcher::PathParamsMatcher(const std::string &pattern)
|
||||
bool PathParamsMatcher::match(Request &request) const {
|
||||
request.matches = std::smatch();
|
||||
request.path_params.clear();
|
||||
|
||||
// A pattern without parameters is just a literal path to compare against
|
||||
if (param_names_.empty()) { return request.path == pattern(); }
|
||||
|
||||
request.path_params.reserve(param_names_.size());
|
||||
|
||||
// One past the position at which the path matched the pattern last time
|
||||
@@ -7188,6 +7192,11 @@ bool PathParamsMatcher::match(Request &request) const {
|
||||
|
||||
bool RegexMatcher::match(Request &request) const {
|
||||
request.path_params.clear();
|
||||
// See CPPHTTPLIB_REGEX_ROUTE_PATH_MAX_LENGTH: an overlong path is treated as
|
||||
// a non-match rather than risking a stack overflow in std::regex_match.
|
||||
if (request.path.length() > CPPHTTPLIB_REGEX_ROUTE_PATH_MAX_LENGTH) {
|
||||
return false;
|
||||
}
|
||||
return std::regex_match(request.path, request.matches, regex_);
|
||||
}
|
||||
|
||||
@@ -7613,11 +7622,21 @@ Server::~Server() = default;
|
||||
|
||||
std::unique_ptr<detail::MatcherBase>
|
||||
Server::make_matcher(const std::string &pattern) {
|
||||
// Path params take precedence, so "/users/:id/(.*)" keeps being matched as
|
||||
// a path params pattern
|
||||
if (pattern.find("/:") != std::string::npos) {
|
||||
return detail::make_unique<detail::PathParamsMatcher>(pattern);
|
||||
} else {
|
||||
return detail::make_unique<detail::RegexMatcher>(pattern);
|
||||
}
|
||||
|
||||
// A pattern with no regex metacharacter only has to be compared literally,
|
||||
// which is what PathParamsMatcher already does when it captures no
|
||||
// parameter, so std::regex is only worth building for the patterns that
|
||||
// actually need it
|
||||
if (pattern.find_first_of(".^$|()[]{}*+?\\") == std::string::npos) {
|
||||
return detail::make_unique<detail::PathParamsMatcher>(pattern);
|
||||
}
|
||||
|
||||
return detail::make_unique<detail::RegexMatcher>(pattern);
|
||||
}
|
||||
|
||||
Server &Server::Get(const std::string &pattern, Handler handler) {
|
||||
@@ -8259,15 +8278,12 @@ bool Server::read_content_core(
|
||||
}
|
||||
}
|
||||
if (has_data) {
|
||||
auto result =
|
||||
detail::read_content_without_length(strm, payload_max_length_, out);
|
||||
if (result == detail::ReadContentResult::PayloadTooLarge) {
|
||||
res.status = StatusCode::PayloadTooLarge_413;
|
||||
return false;
|
||||
} else if (result != detail::ReadContentResult::Success) {
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
// Route through the same decompressing reader used by the
|
||||
// length-framed and chunked paths below, so payload_max_length_ is
|
||||
// enforced on the decompressed size here too instead of only on the
|
||||
// compressed wire bytes.
|
||||
return detail::read_content(strm, req, payload_max_length_, res.status,
|
||||
nullptr, out, true);
|
||||
}
|
||||
}
|
||||
return true;
|
||||
|
||||
Vendored
+23
-2
@@ -8,8 +8,8 @@
|
||||
#ifndef CPPHTTPLIB_HTTPLIB_H
|
||||
#define CPPHTTPLIB_HTTPLIB_H
|
||||
|
||||
#define CPPHTTPLIB_VERSION "0.53.0"
|
||||
#define CPPHTTPLIB_VERSION_NUM "0x003500"
|
||||
#define CPPHTTPLIB_VERSION "0.53.1"
|
||||
#define CPPHTTPLIB_VERSION_NUM "0x003501"
|
||||
|
||||
#ifdef _WIN32
|
||||
#if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00
|
||||
@@ -138,6 +138,18 @@
|
||||
#define CPPHTTPLIB_RANGE_MAX_COUNT 1024
|
||||
#endif
|
||||
|
||||
// std::regex_match's backtracking implementation (most acutely on libstdc++)
|
||||
// recurses roughly once per matched character for quantified patterns such
|
||||
// as "(.*)", so a long enough path can exhaust the calling thread's stack; on
|
||||
// a default ~8MB thread stack that has been observed to take on the order of
|
||||
// a couple thousand characters for a simple pattern. 256 leaves a wide safety
|
||||
// margin below that (well under the 8192-byte request URI limit) while still
|
||||
// fitting any realistic route segment; raise it if a route legitimately needs
|
||||
// longer paths. Regex routes are never applied to paths longer than this.
|
||||
#ifndef CPPHTTPLIB_REGEX_ROUTE_PATH_MAX_LENGTH
|
||||
#define CPPHTTPLIB_REGEX_ROUTE_PATH_MAX_LENGTH 256
|
||||
#endif
|
||||
|
||||
#ifndef CPPHTTPLIB_TCP_NODELAY
|
||||
#define CPPHTTPLIB_TCP_NODELAY false
|
||||
#endif
|
||||
@@ -839,6 +851,15 @@ inline bool parse_url(const std::string &url, UrlComponents &uc) {
|
||||
}
|
||||
|
||||
pos = close + 1;
|
||||
|
||||
// The IPv6 literal is the whole host, so ']' must be followed by a port,
|
||||
// path, query or fragment delimiter (or the end of input). Otherwise the
|
||||
// trailing bytes would be folded into the path while the connection
|
||||
// still targets the bracketed address.
|
||||
if (pos < url.size()) {
|
||||
auto c = url[pos];
|
||||
if (c != ':' && c != '/' && c != '?' && c != '#') { return false; }
|
||||
}
|
||||
} else {
|
||||
auto end = url.find_first_of(":/?#", pos);
|
||||
if (end == std::string::npos) { end = url.size(); }
|
||||
|
||||
Reference in New Issue
Block a user