Merge commit '033df86b69' into concedo_experimental

# Conflicts:
#	docs/backend/SYCL.md
#	docs/development/parsing.md
#	examples/speculative-simple/speculative-simple.cpp
#	ggml/src/ggml-hexagon/ggml-hexagon.cpp
#	ggml/src/ggml-hexagon/htp/get-rows-ops.c
#	ggml/src/ggml-hexagon/htp/get-rows-ops.h
#	ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h
#	ggml/src/ggml-hexagon/htp/htp-ops.h
#	ggml/src/ggml-hexagon/htp/hvx-erf.h
#	ggml/src/ggml-hexagon/htp/main.c
#	ggml/src/ggml-hexagon/htp/matmul-ops.c
#	ggml/src/ggml-hexagon/htp/unary-ops.c
#	ggml/src/ggml-hexagon/htp/unary-ops.h
#	ggml/src/ggml-sycl/CMakeLists.txt
#	ggml/src/ggml-sycl/common.hpp
#	ggml/src/ggml-sycl/element_wise.cpp
#	ggml/src/ggml-sycl/element_wise.hpp
#	ggml/src/ggml-sycl/fusion.cpp
#	ggml/src/ggml-sycl/fwht.cpp
#	ggml/src/ggml-sycl/ggml-sycl.cpp
#	ggml/src/ggml-sycl/im2col.hpp
#	ggml/src/ggml-sycl/pad.hpp
#	ggml/src/ggml-sycl/rope.hpp
#	ggml/src/ggml-sycl/upscale.hpp
#	scripts/snapdragon/run.py
#	scripts/sync_vendor.py
#	src/CMakeLists.txt
#	tests/test-backend-ops.cpp
#	tests/test-backend-sampler.cpp
#	tests/test-chat.cpp
#	tests/test-llama-archs.cpp
#	tools/imatrix/imatrix.cpp
#	tools/server/README.md
This commit is contained in:
Concedo 2026-10-10 13:19:02 +08:00
commit 0890a5ee77
114 changed files with 3800 additions and 763 deletions

View file

@ -681,7 +681,7 @@ ggml-vulkan-shaders-noext.o: ggml/src/ggml-vulkan-shaders-noext.cpp ggml/include
$(CXX) $(CXXFLAGS) $(VKGEN_NOEXT_FORCE) $(VULKAN_FLAGS) -c $< -o $@
# intermediate objects
llama.o: src/llama.cpp ggml/include/ggml.h ggml/include/ggml-alloc.h ggml/include/ggml-backend.h ggml/include/ggml-cuda.h ggml/include/ggml-metal.h include/llama.h otherarch/llama-util.h src/llama-chat.cpp src/llama-mmap.cpp src/llama-context.cpp src/llama-adapter.cpp src/llama-arch.cpp src/llama-batch.cpp src/llama-vocab.cpp src/llama-grammar.cpp src/llama-sampler.cpp src/llama-kv-cache.cpp src/llama-kv-cache-dsa.cpp src/llama-kv-cache-dsv4.cpp src/llama-kv-cache-iswa.cpp src/llama-kv-cache-msa.cpp src/llama-memory-hybrid.cpp src/llama-memory-hybrid-iswa.cpp src/llama-memory-recurrent.cpp src/llama-model-loader.cpp src/llama-model-saver.cpp src/llama-quant.cpp src/llama-hparams.cpp src/llama-graph.cpp src/llama-io.cpp src/llama-memory.cpp common/fit.cpp ggml/include/ggml.h ggml/include/ggml-cpu.h ggml/include/ggml-cuda.h include/llama.h otherarch/llama-util.h
llama.o: src/llama.cpp ggml/include/ggml.h ggml/include/ggml-alloc.h ggml/include/ggml-backend.h ggml/include/ggml-cuda.h ggml/include/ggml-metal.h include/llama.h otherarch/llama-util.h src/llama-chat.cpp src/llama-mmap.cpp src/llama-context.cpp src/llama-adapter.cpp src/llama-arch.cpp src/llama-batch.cpp src/llama-vocab.cpp src/llama-grammar.cpp src/llama-sampler.cpp src/llama-kv-cache.cpp src/llama-kv-cache-dsa.cpp src/llama-kv-cache-dsv4.cpp src/llama-kv-cache-iswa.cpp src/llama-kv-cache-msa.cpp src/llama-memory-hybrid.cpp src/llama-memory-hybrid-iswa.cpp src/llama-memory-recurrent.cpp src/llama-model-loader.cpp src/llama-model-saver.cpp src/llama-quant.cpp src/llama-hparams.cpp src/llama-graph.cpp src/llama-io.cpp src/llama-memory.cpp src/llama-moe-cache.cpp common/fit.cpp ggml/include/ggml.h ggml/include/ggml-cpu.h ggml/include/ggml-cuda.h include/llama.h otherarch/llama-util.h
$(CXX) $(CXXFLAGS) -c $< -o $@
llama-model.o: src/llama-model.cpp src/llama-model.h src/models/models.h ggml/include/ggml.h include/llama.h
$(CXX) $(CXXFLAGS) -c $< -o $@

View file

@ -2777,6 +2777,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.tensor_buft_overrides);
}
).set_env("LLAMA_ARG_N_CPU_MOE"));
add_opt(common_arg(
{"--moe-cache-mib"}, "N",
"GPU cache size in MiB for the MoE experts kept in the CPU (default: 0, disabled)",
[](common_params & params, int value) {
if (value < 0) {
throw std::invalid_argument("invalid value");
}
params.moe_cache_size = (size_t) value*1024*1024;
}
).set_env("LLAMA_ARG_MOE_CACHE_MIB"));
add_opt(common_arg(
{"-ncffn", "--n-cpu-ffn"}, "N",
"keep the dense FFN weights of the first N layers in the CPU\n"

View file

@ -291,7 +291,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context
common_peg_parser tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & func = tool.at("function");
std::string name = func.at("name");
const auto schema = common_chat_tool_parameters(func);
@ -308,7 +308,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context
}
have_call_id = true;
}
auto args_parser = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema));
auto args_parser = p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema));
if (!arguments.start.empty()) {
args_parser = p.literal(arguments.start) + args_parser;
}
@ -318,7 +318,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_json(parser_build_context
auto atomic_peek = !arguments.start.empty() ? std::optional(p.peek(p.literal(arguments.start))) : std::nullopt;
auto func_parser = build_func_parser(p, name, call_id_section, have_call_id, args_parser, atomic_peek);
tool_choice |= p.rule("tool-" + name, func_parser);
tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
});
auto require_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;
@ -364,14 +364,14 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
common_peg_parser tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & func = tool.at("function");
std::string name = func.at("name");
// Build parser for each argument, separating required and optional
std::vector<common_peg_parser> required_parsers;
std::vector<common_peg_parser> optional_parsers;
foreach_parameter(func, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
foreach_parameter(func, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
auto arg =
p.tool_arg(p.tool_arg_open(arguments.name_prefix + p.tool_arg_name(p.literal(param.name)) +
arguments.name_suffix) +
@ -380,10 +380,10 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
p.ac(p.tool_arg_string_value(until_suffix) +
p.tool_arg_close(p.literal(arguments.value_suffix)), arguments.value_suffix) :
(p.tool_arg_json_value(p.schema(
p.json(), "tool-" + name + "-arg-" + param.name + "-schema", doc, *param.schema)) +
p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema", doc, *param.schema)) +
p.tool_arg_close(p.literal(arguments.value_suffix)))));
auto named_arg = p.rule("tool-" + name + "-arg-" + param.name, arg);
auto named_arg = p.rule("tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index), arg);
if (param.required) {
required_parsers.push_back(named_arg);
} else {
@ -434,7 +434,7 @@ common_peg_parser analyze_tools::build_tool_parser_tag_tagged(parser_build_conte
auto atomic_peek = (!arguments.name_prefix.empty() && !required_parsers.empty()) ?
std::optional(p.peek(p.literal(arguments.name_prefix))) : std::nullopt;
auto func_parser = build_func_parser(p, name, call_id_section, have_call_id, args_seq, atomic_peek);
tool_choice |= p.rule("tool-" + name, func_parser);
tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
});
auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;

View file

@ -483,7 +483,9 @@ common_peg_parser common_chat_peg_builder::standard_constructed_tools(
// Build tool choices for tagged format
auto tool_choices = choice();
for (const auto & tool_def : tools) {
for (size_t i = 0; i < tools.size(); i++) {
const auto & tool_def = tools[i];
if (!tool_def.contains("function")) {
continue;
}
@ -513,7 +515,7 @@ common_peg_parser common_chat_peg_builder::standard_constructed_tools(
auto tool_parser = tool(tool_open(literal(func_opener) + tool_name(literal(name)) + literal(func_name_suffix)) +
space() + tool_args(args) + space() + tool_close(literal(func_closer)));
tool_choices |= rule("tool-" + name, tool_parser);
tool_choices |= rule("tool-" + std::to_string(i), tool_parser);
}
// Build the section with markers
@ -560,7 +562,8 @@ common_peg_parser common_chat_peg_builder::python_style_tool_calls(
auto tool_choices = choice();
for (const auto & tool_def : tools) {
for (size_t i = 0; i < tools.size(); i++) {
const auto & tool_def = tools[i];
if (!tool_def.contains("function")) {
continue;
}
@ -607,7 +610,7 @@ common_peg_parser common_chat_peg_builder::python_style_tool_calls(
space() + tool_args(args) + space() + tool_close(literal(")"))
);
tool_choices |= rule("tool-" + name, tool_parser);
tool_choices |= rule("tool-" + std::to_string(i), tool_parser);
}
if (parallel_tool_calls) {
@ -635,7 +638,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key(
auto tool_choices = choice();
for (const auto & tool_def : tools) {
for (size_t i = 0; i < tools.size(); i++) {
const auto & tool_def = tools[i];
if (!tool_def.contains("function")) {
continue;
}
@ -668,10 +672,10 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key(
// Arguments — either wrapped in args_key or parsed directly
common_peg_parser args_parser = eps();
if (args_key.empty()) {
args_parser = tool_args(schema(json(), "tool-" + name + "-schema", params));
args_parser = tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
} else {
args_parser = literal("\"" + effective_args_key + "\"") + space() + literal(":") + space() +
tool_args(schema(json(), "tool-" + name + "-schema", params));
tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
}
inner_fields.push_back(args_parser);
@ -698,7 +702,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key(
space() + tool_close(literal("}"))
);
tool_choices |= rule("tool-" + name, tool_parser);
tool_choices |= rule("tool-" + std::to_string(i), tool_parser);
}
return tool_choices;
@ -721,7 +725,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys(
std::string nested_name_field = !name_spec.first.empty() ? name_spec.second : effective_name_key;
std::string nested_args_field = !args_spec.first.empty() ? args_spec.second : effective_args_key;
for (const auto & tool_def : tools) {
for (size_t i = 0; i < tools.size(); i++) {
const auto & tool_def = tools[i];
if (!tool_def.contains("function")) {
continue;
}
@ -732,7 +737,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys(
auto nested_name = literal("\"" + nested_name_field + "\"") + space() + literal(":") + space() +
atomic(literal("\"") + tool_name(literal(name)) + literal("\""));
auto nested_args = literal("\"" + nested_args_field + "\"") + space() + literal(":") + space() +
tool_args(schema(json(), "tool-" + name + "-schema", params));
tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
auto nested_object = literal("{") + space() +
nested_name + space() + literal(",") + space() +
@ -770,7 +775,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys(
auto nested_field = literal("\"" + nested_prefix + "\"") + space() + literal(":") + space() + nested_object;
tool_parser_body = tool_parser_body + nested_field + space() + tool_close(literal("}"));
tool_choices |= rule("tool-" + name, tool(tool_parser_body));
tool_choices |= rule("tool-" + std::to_string(i), tool(tool_parser_body));
}
return tool_choices;
@ -790,7 +795,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys(
auto name_key_parser = literal("\"" + effective_name_key + "\"");
auto args_key_parser = literal("\"" + effective_args_key + "\"");
for (const auto & tool_def : tools) {
for (size_t i = 0; i < tools.size(); i++) {
const auto & tool_def = tools[i];
if (!tool_def.contains("function")) {
continue;
}
@ -801,7 +807,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys(
auto tool_name_ = name_key_parser + space() + literal(":") + space() +
atomic(literal("\"") + tool_name(literal(name)) + literal("\""));
auto tool_args_ = args_key_parser + space() + literal(":") + space() +
tool_args(schema(json(), "tool-" + name + "-schema", params));
tool_args(schema(json(), "tool-" + std::to_string(i) + "-schema", params));
// Build ID parsers if keys are provided
common_peg_parser id_parser = eps();
@ -861,7 +867,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys(
}
ordered_body = ordered_body + space() + tool_close(literal("}"));
tool_choices |= rule("tool-" + name, tool(ordered_body));
tool_choices |= rule("tool-" + std::to_string(i), tool(ordered_body));
}
return tool_choices;

View file

@ -31,6 +31,7 @@
#include "parsers/ministral3.cpp"
#include "parsers/muse-glimmer.cpp"
#include "parsers/qwen3-coder.cpp"
#include "parsers/translate-gemma.cpp"
#include "jinja/lexer.cpp"
#include "jinja/parser.cpp"
@ -1257,6 +1258,13 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
return common_chat_params_init_minicpm5(tmpl, params);
}
// TranslateGemma - user content must follow a custom schema with language codes
if (src.find("[source_lang_code]") != std::string::npos &&
src.find("[target_lang_code]") != std::string::npos) {
LOG_DBG("Using specialized template: TranslateGemma\n");
return common_chat_params_init_translate_gemma(tmpl, params);
}
// Qwen3-Coder XML tool calls, also used by Nemotron Nano 3, Qwen3.5 and StepFun-3.5-Flash
if (src.find("<tool_call>") != std::string::npos &&
src.find("<function=") != std::string::npos &&
@ -1486,14 +1494,70 @@ common_chat_params common_chat_templates_apply(const struct common_chat_template
common_chat_templates_apply_legacy(tmpls, inputs);
}
common_chat_msg common_chat_parse(const std::string & input,
void common_chat_input::append(const std::string & piece, llama_token token) {
if (piece.empty()) {
return;
}
tokens.push_back(token);
tokens.resize(tokens.size() + piece.size() - 1, LLAMA_TOKEN_NULL);
text += piece;
}
void common_chat_input::append(const common_chat_input & chunk) {
tokens.insert(tokens.end(), chunk.tokens.begin(), chunk.tokens.end());
text += chunk.text;
}
void common_chat_input::truncate(size_t pos) {
if (pos < text.size()) {
text.erase(pos);
tokens.resize(pos);
}
}
common_chat_input common_chat_input::substr(size_t pos, size_t n) const {
common_chat_input out;
out.text = text.substr(pos, n);
out.tokens.assign(tokens.begin() + pos, tokens.begin() + pos + out.size());
return out;
}
void common_chat_input::prepend(const std::string & prefix) {
tokens.insert(tokens.begin(), prefix.size(), LLAMA_TOKEN_NULL);
text = prefix + text;
}
void common_chat_input::prepend(const common_chat_input & prefix) {
tokens.insert(tokens.begin(), prefix.tokens.begin(), prefix.tokens.end());
text = prefix.text + text;
}
common_chat_input common_chat_input_tokenize(const llama_vocab * vocab, const std::string & text) {
common_chat_input input;
auto tokens = common_tokenize(vocab, text, false, true);
for (size_t i = 0; i < tokens.size(); i++) {
std::string piece = common_token_to_piece(vocab, tokens[i], true);
if (i == 0 && std::isspace(piece[0]) && !std::isspace(text[0])) {
// Some tokenizers will add a space before the first special token, need to exclude
continue;
}
input.append(piece, tokens[i]);
}
if (input.text != text) {
// the pieces do not give back the same text, keep the text without tokens
return common_chat_input(text);
}
return input;
}
common_chat_msg common_chat_parse(const common_chat_input & input,
bool is_partial,
const common_chat_parser_params & params) {
return common_chat_peg_parse(params.parser, input, is_partial, params);
}
common_chat_msg common_chat_peg_parse(const common_peg_arena & src_parser,
const std::string & input,
const common_chat_input & input,
bool is_partial,
const common_chat_parser_params & params) {
const common_peg_arena & parser = src_parser.empty() ?
@ -1504,18 +1568,17 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena & src_pars
LOG_DBG("No parser definition detected, assuming pure content parser.");
}
const std::string effective_input = params.generation_prompt.empty()
? input
: params.generation_prompt + input;
common_chat_input effective_input = input;
effective_input.prepend(params.generation_prompt);
//LOG_DBG("Parsing PEG input with format %s: %s\n", common_chat_format_name(params.format), effective_input.c_str());
//LOG_DBG("Parsing PEG input with format %s: %s\n", common_chat_format_name(params.format), effective_input.text.c_str());
common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_LENIENT;
if (params.debug) {
flags |= COMMON_PEG_PARSE_FLAG_DEBUG;
}
common_peg_parse_context ctx(effective_input, flags);
common_peg_parse_context ctx(std::move(effective_input.text), std::move(effective_input.tokens), flags);
auto result = parser.parse(ctx);
if (result.fail()) {
@ -1541,8 +1604,8 @@ common_chat_msg common_chat_peg_parse(const common_peg_arena & src_pars
}
return msg;
}
LOG_WRN("%s: unparsed %s output: %s\n", __func__, common_chat_format_name(params.format), effective_input.substr(result.end).c_str());
LOG_DBG("%s: full %s output triggering error:\n=== BEGIN ===\n%s\n=== END ===\n", __func__, common_chat_format_name(params.format), effective_input.c_str());
LOG_WRN("%s: unparsed %s output: %s\n", __func__, common_chat_format_name(params.format), ctx.input.substr(result.end).c_str());
LOG_DBG("%s: full %s output triggering error:\n=== BEGIN ===\n%s\n=== END ===\n", __func__, common_chat_format_name(params.format), ctx.input.c_str());
throw std::runtime_error(std::string("The model produced output that does not match the expected ") + common_chat_format_name(params.format) + " format");
}

View file

@ -282,6 +282,31 @@ struct common_chat_params {
common_chat_msg_delimiters message_delimiters;
};
struct common_chat_input {
std::string text;
std::vector<llama_token> tokens;
common_chat_input() = default;
// plain text, with no tokens
explicit common_chat_input(std::string text) : text(std::move(text)), tokens(this->text.size(), LLAMA_TOKEN_NULL) {}
size_t size() const { return text.size(); }
bool empty() const { return text.empty(); }
void append(const std::string & piece, llama_token token);
void append(const common_chat_input & chunk);
void prepend(const std::string & prefix);
void prepend(const common_chat_input & prefix);
void truncate(size_t pos);
common_chat_input substr(size_t pos, size_t n = std::string::npos) const;
};
common_chat_input common_chat_input_tokenize(const llama_vocab * vocab, const std::string & text);
// per-message parsing syntax
// should be derived from common_chat_params
struct common_chat_parser_params {
@ -289,7 +314,7 @@ struct common_chat_parser_params {
common_reasoning_format reasoning_format = COMMON_REASONING_FORMAT_NONE; // TODO: refactor this to "bool parse_reasoning"
// Whether reasoning_content should be inlined in the content (e.g. for reasoning_format=deepseek in stream mode)
bool reasoning_in_content = false;
std::string generation_prompt;
common_chat_input generation_prompt;
bool parse_tool_calls = true;
bool is_continuation = false;
bool echo = false; // Include assistant prefilled msg in output
@ -298,7 +323,7 @@ struct common_chat_parser_params {
common_chat_parser_params() = default;
common_chat_parser_params(const common_chat_params & chat_params) {
format = chat_params.format;
generation_prompt = chat_params.generation_prompt;
generation_prompt = common_chat_input(chat_params.generation_prompt);
}
};
@ -337,8 +362,8 @@ std::string common_chat_format_example(const struct common_chat_templates *
const std::map<std::string, std::string> & chat_template_kwargs);
const char * common_chat_format_name(common_chat_format format);
common_chat_msg common_chat_parse(const std::string & input, bool is_partial, const common_chat_parser_params & params);
common_chat_msg common_chat_peg_parse(const common_peg_arena & src_parser, const std::string & input, bool is_partial, const common_chat_parser_params & params);
common_chat_msg common_chat_parse(const common_chat_input & input, bool is_partial, const common_chat_parser_params & params);
common_chat_msg common_chat_peg_parse(const common_peg_arena & src_parser, const common_chat_input & input, bool is_partial, const common_chat_parser_params & params);
// used by arg and server
const char * common_reasoning_format_name(common_reasoning_format format);

View file

@ -1177,6 +1177,8 @@ static const std::map<common_decision_type, std::string> COMMON_DECISION_TYPE_NA
{ COMMON_DECISION_TYPE_LAYA, "laya" },
{ COMMON_DECISION_TYPE_CLEF, "clef" },
{ COMMON_DECISION_TYPE_PPLX_DECIDER, "pplx-decider" },
{ COMMON_DECISION_TYPE_LFM2_D1, "lfm2-d1" },
{ COMMON_DECISION_TYPE_LFM2_D1_OMNI, "lfm2-d1-omni" },
};
static common_decision_type common_decision_type_from_string(const std::string & str) {
@ -1290,7 +1292,8 @@ common_init_result::common_init_result(common_params & params, bool model_only)
// these decision models return a score for each token via the embeddings output
// TODO: maybe improve this in the future
const auto decision_type = common_get_decision_type(model);
if (decision_type == COMMON_DECISION_TYPE_LAYA || decision_type == COMMON_DECISION_TYPE_KEV || decision_type == COMMON_DECISION_TYPE_CLEF) {
if (decision_type == COMMON_DECISION_TYPE_LAYA || decision_type == COMMON_DECISION_TYPE_KEV || decision_type == COMMON_DECISION_TYPE_CLEF ||
decision_type == COMMON_DECISION_TYPE_LFM2_D1_OMNI) {
params.embedding = true;
params.pooling_type = LLAMA_POOLING_TYPE_NONE;
@ -1729,6 +1732,8 @@ struct llama_context_params common_context_params_to_llama(const common_params &
cparams.type_k = params.cache_type_k;
cparams.type_v = params.cache_type_v;
cparams.moe_cache_size = params.moe_cache_size;
return cparams;
}
@ -2390,40 +2395,36 @@ void common_prompt_checkpoint::update_dft(
}
}
void common_prompt_checkpoint::load_tgt(
bool common_prompt_checkpoint::load_tgt(
llama_context * ctx,
llama_seq_id seq_id,
llama_state_seq_flags flags) const {
if (ctx == nullptr) {
return;
return true;
}
if (data_tgt.empty()) {
return;
return true;
}
const size_t n = llama_state_seq_set_data_ext(ctx, data_tgt.data(), data_tgt.size(), seq_id, flags);
if (n != data_tgt.size()) {
GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", data_tgt.size(), n);
}
return n == data_tgt.size();
}
void common_prompt_checkpoint::load_dft(
bool common_prompt_checkpoint::load_dft(
llama_context * ctx,
llama_seq_id seq_id,
llama_state_seq_flags flags) const {
if (ctx == nullptr) {
return;
return true;
}
if (data_dft.empty()) {
return;
return true;
}
const size_t n = llama_state_seq_set_data_ext(ctx, data_dft.data(), data_dft.size(), seq_id, flags);
if (n != data_dft.size()) {
GGML_ABORT("checkpoint size mismatch: expected %zu, got %zu\n", data_dft.size(), n);
}
return n == data_dft.size();
}
void common_prompt_checkpoint::clear_tgt() {

View file

@ -594,6 +594,8 @@ struct common_params {
ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K
ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V
size_t moe_cache_size = 0; // GPU cache size in bytes for the MoE experts kept in the CPU
common_conversation_mode conversation_mode = COMMON_CONVERSATION_MODE_AUTO;
// multimodal models (see tools/mtmd)
@ -962,6 +964,8 @@ enum common_decision_type {
COMMON_DECISION_TYPE_LAYA, // score of one marker token per option, read from the embeddings output
COMMON_DECISION_TYPE_CLEF, // all questions in one prompt, score of option i read from the embeddings output at row i
COMMON_DECISION_TYPE_PPLX_DECIDER, // same as openjev, label codes of 1 or 2 letters
COMMON_DECISION_TYPE_LFM2_D1, // same as openjev, the labels depend on the question type
COMMON_DECISION_TYPE_LFM2_D1_OMNI, // same as laya, other prompt layout
COMMON_DECISION_TYPE_UNKNOWN, // a decision model of a type that is not supported
};
@ -1293,12 +1297,13 @@ struct common_prompt_checkpoint {
llama_seq_id seq_id,
llama_state_seq_flags flags);
void load_tgt(
// return false if the state could not be restored
bool load_tgt(
llama_context * ctx,
llama_seq_id seq_id,
llama_state_seq_flags flags) const;
void load_dft(
bool load_dft(
llama_context * ctx,
llama_seq_id seq_id,
llama_state_seq_flags flags) const;

View file

@ -152,13 +152,13 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
// build tool call section first since we might need it in reasoning
auto tool_choice = p.choice();
if (has_tool_calls) {
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
std::vector<common_peg_parser> required_parsers;
std::vector<common_peg_parser> optional_parsers;
foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
bool is_string = param.schema->may_be_string();
auto arg = p.tool_arg(
@ -166,11 +166,11 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
p.literal("\" string=\"" + std::string(is_string ? "true" : "false") + "\">")) +
(is_string ?
p.tool_arg_string_value(p.until(PARAM_END)) :
p.tool_arg_json_value(p.schema(p.json(), "tool-" + name + "-arg-" + param.name + "-schema",
p.tool_arg_json_value(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema",
doc, *param.schema))) +
p.tool_arg_close(p.literal(PARAM_END)));
auto named_arg = p.rule("tool-" + name + "-arg-" + param.name, arg);
auto named_arg = p.rule("tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index), arg);
if (param.required) {
required_parsers.push_back(named_arg);
} else {
@ -199,7 +199,7 @@ common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_templ
p.tool_name(p.literal(name)) + p.literal("\">\n")) +
invoke_body + p.space() + p.tool_close(p.literal(INVOKE_END)));
tool_choice |= p.rule("tool-" + name, func_parser);
tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
});
}

View file

@ -42,7 +42,7 @@ common_chat_params common_chat_params_init_functionary_v3_2(const common_chat_te
// Build tool call parsers for each available function
auto tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
const auto schema = common_chat_tool_parameters(function);
@ -50,10 +50,10 @@ common_chat_params common_chat_params_init_functionary_v3_2(const common_chat_te
// Tool format: >>>function_name\n{json_args}
auto tool_parser = p.tool(
p.tool_open(p.tool_name(p.literal(name)) + p.literal("\n")) +
p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema))
p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema))
);
tool_choice |= p.rule("tool-" + name, tool_parser);
tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
});
auto content_only = content_until_end;

View file

@ -254,13 +254,13 @@ common_chat_params common_chat_params_init_gemma4(const common_chat_template &
auto tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
// TODO @aldehir : need to extend json-schema-to-grammar to produce more than JSON rules
// const auto & params = function.at("parameters");
tool_choice |= p.rule("tool-" + name, p.tool(p.sequence({
tool_choice |= p.rule("tool-" + std::to_string(tool_index), p.tool(p.sequence({
p.tool_open(p.tool_name(p.literal(name)) + p.peek(p.literal("{"))),
p.tool_args(p.ref("gemma4-dict")),
})));

View file

@ -30,17 +30,18 @@ common_chat_params common_chat_params_init_gigachat_v3(
if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
// Build a choice of all available tools
auto tool_choice = p.choice();
for (const auto & tool : inputs.tools) {
for (size_t i = 0; i < inputs.tools.size(); i++) {
const auto & tool = inputs.tools[i];
const auto & function = tool.at("function");
std::string name = function.at("name");
const auto schema = common_chat_tool_parameters(function);
auto tool_name = p.json_member("name", "\"" + p.tool_name(p.literal(name)) + "\"");
auto tool_args = p.json_member("arguments", p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)));
auto tool_args = p.json_member("arguments", p.tool_args(p.schema(p.json(), "tool-" + std::to_string(i) + "-schema", schema)));
auto tool_open = p.tool_open(p.literal("{") << tool_name);
tool_choice |= p.rule("tool-" + name, tool_open << "," << tool_args << "}");
tool_choice |= p.rule("tool-" + std::to_string(i), tool_open << "," << tool_args << "}");
}
// Define the tool call structure

View file

@ -106,14 +106,14 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
auto tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
const auto params = common_chat_tool_parameters(function);
auto func_name = p.literal(" to=functions.") + p.tool_name(p.literal(name));
auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type);
auto args = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", params));
auto args = p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", params));
// recipient in role header
// <|start|>assistant to=functions.NAME<|channel|>(commentary|analysis)[constraint]<|message|>ARGS
@ -123,7 +123,7 @@ common_chat_params common_chat_params_init_gpt_oss(const common_chat_template &
// <|channel|>(commentary|analysis) to=functions.NAME[constraint]<|message|>ARGS
auto tool_in_channel = p.tool(p.tool_open(channel + func_name + constraint + p.literal("<|message|>")) + args);
tool_choice |= p.rule("tool-" + name, tool_in_role | tool_in_channel);
tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_in_role | tool_in_channel);
});
auto tool_call = p.trigger_rule("tool-call", tool_choice);

View file

@ -110,14 +110,14 @@ common_chat_params common_chat_params_init_k2_horizon(const common_chat_template
// The models leave out <ifm|arg_type> even when asked for xml_typed
auto arg_type = call_format == "xml_typed" ? p.optional(ARG_TYPE + p.until(ARG_TYPE_END) + ARG_TYPE_END + p.space()) : p.eps();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
std::vector<common_peg_parser> required_args;
std::vector<common_peg_parser> optional_args;
foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
auto rule_name = "tool-" + name + "-arg-" + param.name;
foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
auto rule_name = "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index);
auto types = param.schema->value_types();
auto arg_value = arg_string;
if (!types.has(common_chat_schema::TYPE_STRING)) {
@ -149,12 +149,12 @@ common_chat_params common_chat_params_init_k2_horizon(const common_chat_template
(param.required ? required_args : optional_args).push_back(p.rule(rule_name, arg));
});
auto args = p.permute("tool-" + name + "-args", required_args);
auto args = p.permute("tool-" + std::to_string(tool_index) + "-args", required_args);
if (!optional_args.empty()) {
args = args + p.zero_or_more(p.choice(optional_args));
}
tool_choice |= p.rule("tool-" + name, p.tool(
tool_choice |= p.rule("tool-" + std::to_string(tool_index), p.tool(
p.tool_open(CALL_START + p.tool_name(p.literal(name)) + "\n") + p.tool_args(args) << p.tool_close(p.literal(CALL_END))));
});
}

View file

@ -79,7 +79,7 @@ common_chat_params common_chat_params_init_kimi_k2(const common_chat_template &
// The ID format is: functions.<name>:<index>
// We need to match: functions.<name>:<digits>
auto tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
const auto schema = common_chat_tool_parameters(function);
@ -89,11 +89,11 @@ common_chat_params common_chat_params_init_kimi_k2(const common_chat_template &
auto tool_id = p.tool_id(p.literal("functions.") + p.tool_name(p.literal(name)) + p.literal(":") + p.chars("[0-9]", 1, -1));
auto tool_parser = p.tool(
p.tool_open(tool_id + p.literal(ARGS_BEGIN)) +
p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)) +
p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema)) +
p.tool_close(p.optional((p.literal(CALL_END))))
);
tool_choice |= p.rule("tool-" + name, tool_parser);
tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
});
// Tool calls section: <|tool_calls_section_begin|> tool_calls <|tool_calls_section_end|>

View file

@ -95,7 +95,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
}
auto tool_choices = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
const json schema = common_chat_tool_parameters(function);
@ -106,6 +106,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
auto args = p.eps();
if (schema.contains("properties") && !schema.at("properties").empty()) {
auto arg_choices = p.choice();
size_t param_index = 0;
for (const auto & prop : schema.at("properties").items()) {
const std::string & key = prop.key();
@ -119,7 +120,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
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,
arg_choices |= p.rule("kimi-k3-arg-" + std::to_string(tool_index) + "-" + std::to_string(param_index++),
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 +
@ -133,7 +134,7 @@ common_chat_params common_chat_params_init_kimi_k3(const common_chat_template &
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);
tool_choices |= p.rule("kimi-k3-tool-" + std::to_string(tool_index), call);
});
// all calls go inside one tools section, then the message is closed. the

View file

@ -118,7 +118,7 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
auto arg_string = p.rule("ling3-arg-string",
p.tool_arg_string_value(p.until(ARG_VAL_END)) + arg_close);
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
@ -127,8 +127,8 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
// each argument may be preceded by whitespace: the model emits
// newlines between arguments, the template history does not
foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
auto rule_name = "ling3-arg-" + name + "-" + param.name;
foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
auto rule_name = "ling3-arg-" + std::to_string(tool_index) + "-" + std::to_string(param_index);
auto types = param.schema->value_types();
@ -159,7 +159,7 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
// required arguments in any order (as Qwen3-Coder does), then
// optional ones in any order and number
auto args = p.permute("ling3-" + name + "-args", required_args);
auto args = p.permute("ling3-" + std::to_string(tool_index) + "-args", required_args);
if (!optional_args.empty()) {
args = args + p.zero_or_more(p.choice(optional_args));
}
@ -169,7 +169,7 @@ common_chat_params common_chat_params_init_ling3(const common_chat_template &
p.tool_args(args) +
p.tool_close(p.optional(p.space()) + p.literal(CALL_END)));
tool_choices |= p.rule("ling3-tool-" + name, call);
tool_choices |= p.rule("ling3-tool-" + std::to_string(tool_index), call);
});
auto calls = inputs.parallel_tool_calls ?

View file

@ -109,13 +109,13 @@ common_chat_params common_chat_params_init_llm_jp_harmony(const common_chat_temp
if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
auto tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
const auto params = common_chat_tool_parameters(function);
auto func_name = p.literal(" to=functions.") + p.tool_name(p.literal(name));
auto args = p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", params));
auto args = p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", params));
// recipient in role header
// <|start|>assistant to=functions.NAME<|channel|>(commentary|analysis)[constraint]<|message|>ARGS
@ -125,7 +125,7 @@ common_chat_params common_chat_params_init_llm_jp_harmony(const common_chat_temp
// <|channel|>(commentary|analysis) to=functions.NAME[constraint]<|message|>ARGS
auto tool_in_channel = p.tool(p.tool_open(channel + func_name + constraint + message) + args);
tool_choice |= p.rule("tool-" + name, tool_in_role | tool_in_channel);
tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_in_role | tool_in_channel);
});
// parallel calls are separated by <|end|>; inside the trigger rule so the lazy grammar covers all of them

View file

@ -68,18 +68,18 @@ common_chat_params common_chat_params_init_minicpm5(const common_chat_template &
});
auto tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
const std::string name = function.at("name");
std::vector<common_peg_parser> arg_rules;
foreach_parameter(function, [&](const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
auto value_parser = p.eps();
if (prop.schema->may_be_string()) {
value_parser = string_value;
} else {
value_parser = p.tool_arg_json_value(
p.schema(p.json(), "tool-" + name + "-arg-" + prop.name + "-schema", doc, *prop.schema)
p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema", doc, *prop.schema)
) + p.tool_arg_close(p.literal("</param>"));
}
@ -99,7 +99,7 @@ common_chat_params common_chat_params_init_minicpm5(const common_chat_template &
<< p.tool_args(args)
<< p.tool_close(p.literal("</function>")));
tool_choice |= p.rule("tool-" + name, tool_parser);
tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
});
auto max_calls = inputs.parallel_tool_calls ? -1 : 1;

View file

@ -85,7 +85,7 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
}
auto tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
auto params = common_chat_tool_parameters(function);
@ -154,8 +154,9 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
members_of = [&](const common_chat_schema_object & object, const std::string & rule_prefix) -> common_peg_parser {
std::vector<common_peg_parser> required_elements;
std::vector<common_peg_parser> optional_elements;
for (const auto & prop : object.properties) {
auto element = element_of(prop.name, *prop.schema, rule_prefix + "-" + prop.name);
for (size_t i = 0; i < object.properties.size(); i++) {
const auto & prop = object.properties[i];
auto element = element_of(prop.name, *prop.schema, rule_prefix + "-" + std::to_string(i));
(prop.required ? required_elements : optional_elements).push_back(element);
}
@ -180,7 +181,7 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
common_peg_parser invoke_body = p.eps();
if (doc->root->kind() == common_chat_schema::KIND_OBJECT) {
invoke_body = members_of(static_cast<const common_chat_schema_object &>(*doc->root), "tool-" + name + "-arg");
invoke_body = members_of(static_cast<const common_chat_schema_object &>(*doc->root), "tool-" + std::to_string(tool_index) + "-arg");
}
auto func_parser = p.tool(
@ -189,7 +190,7 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template
p.space() + invoke_body + p.space() +
p.tool_close(p.literal(INVOKE_END)));
tool_choice |= p.rule("tool-" + name, func_parser);
tool_choice |= p.rule("tool-" + std::to_string(tool_index), func_parser);
});
auto require_tools = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED;

View file

@ -86,14 +86,14 @@ common_chat_params common_chat_params_init_ministral_3(const common_chat_templat
// Tool call parser
if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) {
auto tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
const auto schema = common_chat_tool_parameters(function);
tool_choice |=
p.rule("tool-" + name, p.tool_open(p.tool_name(p.literal(name)) + "[ARGS]") +
p.tool_args(p.schema(p.json(), "tool-" + name + "-schema", schema)));
p.rule("tool-" + std::to_string(tool_index), p.tool_open(p.tool_name(p.literal(name)) + "[ARGS]") +
p.tool_args(p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-schema", schema)));
});
auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0;

View file

@ -81,18 +81,18 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa
"</atem:parameter>");
auto tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
const std::string name = function.at("name");
std::vector<common_peg_parser> arg_rules;
foreach_parameter(function, [&](const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) {
auto value_parser = p.eps();
if (prop.schema->may_be_string()) {
value_parser = string_value;
} else {
value_parser = p.tool_arg_json_value(
p.schema(p.json(), "tool-" + name + "-arg-" + prop.name + "-schema", doc, *prop.schema))
p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema", doc, *prop.schema))
+ p.tool_arg_close(p.literal("</atem:parameter>"));
}
@ -113,7 +113,7 @@ common_chat_params common_chat_params_init_muse_glimmer(const common_chat_templa
<< p.tool_args(args)
<< p.tool_close(p.literal("</atem:invoke>") + p.space() + p.literal("</atem:function_calls>")));
tool_choice |= p.rule("tool-" + name, tool_parser);
tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser);
});
auto tool_calls = inputs.parallel_tool_calls

View file

@ -2,24 +2,25 @@
#include "log.h"
void foreach_function(const json & tools, const std::function<void(const json &)> & fn) {
for (const auto & tool : tools) {
void foreach_function(const json & tools, const std::function<void(size_t, const json &)> & fn) {
for (size_t i = 0; i < tools.size(); i++) {
const auto & tool = tools[i];
if (!tool.contains("type") || tool.at("type") != "function" || !tool.contains("function")) {
LOG_INF("Skipping tool without function: %s", tool.dump(2).c_str());
continue;
}
fn(tool);
fn(i, tool);
}
}
void foreach_parameter(const json & function, const std::function<void(const common_chat_schema_property &, const common_chat_schema_document_ptr &)> & fn) {
void foreach_parameter(const json & function, const std::function<void(size_t, const common_chat_schema_property &, const common_chat_schema_document_ptr &)> & fn) {
auto params = common_chat_tool_parameters(function);
auto doc = std::make_shared<const common_chat_schema_document>(common_chat_schema_from_json(params));
const auto * object = dynamic_cast<const common_chat_schema_object *>(doc->root.get());
if (!object) {
return;
}
for (const auto & prop : object->properties) {
fn(prop, doc);
for (size_t i = 0; i < object->properties.size(); i++) {
fn(i, object->properties[i], doc);
}
}

View file

@ -17,11 +17,11 @@
using json = common_json;
// iterate over the function tools of an OpenAI-style tools array
void foreach_function(const json & tools, const std::function<void(const json &)> & fn);
// iterate over the function tools of an OpenAI-style tools array, passing each tool with its index in the array
void foreach_function(const json & tools, const std::function<void(size_t, const json &)> & fn);
// iterate over the parameters of a function tool, with the document that owns them
void foreach_parameter(const json & function, const std::function<void(const common_chat_schema_property &, const common_chat_schema_document_ptr &)> & fn);
// iterate over the parameters of a function tool, passing each parameter with its index and the document that owns it
void foreach_parameter(const json & function, const std::function<void(size_t, const common_chat_schema_property &, const common_chat_schema_document_ptr &)> & fn);
// render a template; the override arguments let a parser feed in messages, tools or context it has rewritten
std::string common_chat_template_direct_apply_impl(
@ -81,3 +81,5 @@ common_chat_params common_chat_params_init_ministral_3(const common_chat_templat
common_chat_params common_chat_params_init_muse_glimmer(const common_chat_template & tmpl, const autoparser::generation_params & inputs);
common_chat_params common_chat_params_init_qwen3_coder(const common_chat_template & tmpl, const autoparser::generation_params & inputs);
common_chat_params common_chat_params_init_translate_gemma(const common_chat_template & tmpl, const autoparser::generation_params & inputs);

View file

@ -65,7 +65,7 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
// Match complete <function=name> opener for Qwen3-Coder models that occasionally omit the
// starting <tool_call>. The model may hallucinate a tool name, but it is preferable over
// constraining on <function which may occur in valid content generation, e.g. #include <functional>
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t, const json & tool) {
const std::string name = tool.at("function").at("name");
tool_call_starts.push_back("<function=" + name + ">");
});
@ -93,15 +93,15 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
p.ac(p.tool_arg_string_value(p.until("\n</parameter>\n")) + arg_close, "\n</parameter>\n"));
auto tool_choice = p.choice();
foreach_function(inputs.tools, [&](const json & tool) {
foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) {
const auto & function = tool.at("function");
std::string name = function.at("name");
std::vector<common_peg_parser> required_args;
std::vector<common_peg_parser> optional_args;
foreach_parameter(function, [&](const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
auto rule_name = "tool-" + name + "-arg-" + param.name;
foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) {
auto rule_name = "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index);
auto arg_open = p.tool_arg_open("<parameter=" + p.tool_arg_name(p.literal(param.name)) + ">\n");
@ -141,7 +141,7 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
// Accept required arguments in any order, as Qwen does not always adhere to the
// order provided.
auto args = p.permute("tool-" + name + "-args", required_args);
auto args = p.permute("tool-" + std::to_string(tool_index) + "-args", required_args);
if (!optional_args.empty()) {
args = args + p.zero_or_more(p.choice(optional_args));
}
@ -150,7 +150,7 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat
p.tool_args(args) +
p.tool_close(p.literal("</function>\n")));
tool_choice |= p.rule("tool-" + name, func);
tool_choice |= p.rule("tool-" + std::to_string(tool_index), func);
});
auto min_calls = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED ? 1 : 0;

View file

@ -20,4 +20,5 @@ set(LLAMA_CHAT_PARSERS_SOURCES
${CMAKE_CURRENT_LIST_DIR}/ministral3.cpp
${CMAKE_CURRENT_LIST_DIR}/muse-glimmer.cpp
${CMAKE_CURRENT_LIST_DIR}/qwen3-coder.cpp
${CMAKE_CURRENT_LIST_DIR}/translate-gemma.cpp
)

View file

@ -0,0 +1,63 @@
#include "parsers.h"
#include "log.h"
// TranslateGemma does not support tools or reasoning, it only needs user messages in its own content schema
common_chat_params common_chat_params_init_translate_gemma(
const common_chat_template & tmpl,
const autoparser::generation_params & inputs) {
common_chat_params data;
// default to chat_template_kwargs, or en-GB if not specified
std::string src_lang = inputs.extra_context.value("source_lang_code", "en-GB");
std::string tgt_lang = inputs.extra_context.value("target_lang_code", "en-GB");
for (const char * key : { "source_lang_code", "target_lang_code" }) {
if (!inputs.extra_context.contains(key)) {
LOG_WRN("TranslateGemma: %s not set in chat_template_kwargs, defaulting to en-GB\n", key);
}
}
json messages = inputs.messages;
for (auto & message : messages) {
if (message.value("role", "") != "user") {
continue;
}
std::string text;
const auto & content = message.contains("content") ? message.at("content") : json();
if (content.is_string()) {
text = content.get<std::string>();
} else if (content.is_array()) {
for (const auto & part : content) {
if (!text.empty()) {
text += "\n";
}
text += part.value("text", "");
}
}
message["content"] = json::array({
json{
{"type", "text"},
{"text", text},
{"source_lang_code", src_lang},
{"target_lang_code", tgt_lang},
}
});
}
data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs, messages);
data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs, messages);
data.format = COMMON_CHAT_FORMAT_PEG_NATIVE;
data.supports_thinking = false;
if (inputs.has_continuation()) {
data.generation_prompt = "<start_of_turn>model\n" + inputs.continue_msg.render_content();
data.prompt += data.generation_prompt;
}
auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) {
return p.literal(data.generation_prompt) << p.content(p.rest());
});
data.parser = parser.save();
return data;
}

View file

@ -2,6 +2,7 @@
#include "json-schema.h"
#include "json.h"
#include "llama.h"
#include <memory>
#include <set>
@ -182,7 +183,8 @@ inline common_peg_parse_flags operator~(common_peg_parse_flags a) {
}
struct common_peg_parse_context {
std::string input;
std::string input; // [h, e, l, l, o, _, w, o, r, l, d]
std::vector<llama_token> tokens; // [id, -1, -1, -1, -1, id, -1, -1, -1, -1, -1]
common_peg_parse_flags flags;
common_peg_ast_arena ast;
@ -194,6 +196,11 @@ struct common_peg_parse_context {
common_peg_parse_context(const std::string & input, common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE)
: input(input), flags(flags), parse_depth(0) {}
common_peg_parse_context(std::string input, std::vector<llama_token> tokens, common_peg_parse_flags flags = COMMON_PEG_PARSE_FLAG_NONE)
: input(std::move(input)), tokens(std::move(tokens)), flags(flags), parse_depth(0) {
GGML_ASSERT(this->tokens.empty() || this->tokens.size() == this->input.size());
}
bool is_lenient() const { return flags & COMMON_PEG_PARSE_FLAG_LENIENT; }
bool is_debug() const { return flags & COMMON_PEG_PARSE_FLAG_DEBUG; }
};

View file

@ -399,8 +399,11 @@ struct common_sampler * common_sampler_init(
// only if user explicitly included adaptive-p sampler
samplers.push_back(llama_sampler_init_adaptive_p(params.adaptive_target, params.adaptive_decay, params.seed));
} else {
// default: sample from distribution
samplers.push_back(llama_sampler_init_dist(params.seed));
// Keep distribution sampling when callers request probabilities.
const bool greedy = params.n_probs == 0 && !params.samplers.empty() &&
((params.samplers.back() == COMMON_SAMPLER_TYPE_TEMPERATURE && params.temp == 0.0f && params.dynatemp_range == 0.0f) ||
(params.samplers.back() == COMMON_SAMPLER_TYPE_TOP_K && params.top_k == 1));
samplers.push_back(greedy ? llama_sampler_init_greedy() : llama_sampler_init_dist(params.seed));
}
} else if (params.mirostat == 1) {
samplers.push_back(llama_sampler_init_temp(params.temp));

View file

@ -2561,6 +2561,9 @@ common_params common_base_params_to_speculative(const common_params & params) {
result.n_outputs_max = params.n_parallel;
result.n_outputs_max_per_seq = 1;
// the MoE cache is only used by the target context
result.moe_cache_size = 0;
// dflash/dspark decode the whole noise block in a single pass and sample every block position on the backend
// TODO: refactor such properties to be announced by the speculative types
// something like `struct common_speculative_type_props common_speculative_type_get_props(...);`

View file

@ -161,6 +161,8 @@ TEXT_MODEL_MAP: dict[str, str] = {
"Lfm2BidirectionalForMaskedLM": "lfm2",
"Lfm2BidirectionalModel": "lfm2",
"Lfm2ForCausalLM": "lfm2",
"D1Model": "lfm2",
"D1OmniModel": "lfm2",
"Lfm2Model": "lfm2",
"Lfm2MoeForCausalLM": "lfm2",
"Llama4ForCausalLM": "llama",
@ -253,6 +255,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"Qwen3_5ForConditionalGeneration": "qwen",
"Qwen3_5MoeForCausalLM": "qwen",
"Qwen3_5MoeForConditionalGeneration": "qwen",
"Qwen3_5TextModel": "qwen",
"Qwen4ExpForCausalLM": "qwen4exp",
"Qwen4ExpForConditionalGeneration": "qwen4exp",
"RND1": "qwen",
@ -303,6 +306,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
"AudioFlamingo3ForConditionalGeneration": "ultravox",
"ClefModel": "clef",
"CogVLMForCausalLM": "cogvlm",
"Cohere2VisionForConditionalGeneration": "command_r",
"PplxDeciderModel": "pplx_decider",
"DeepseekOCR2ForCausalLM": "deepseek",
"DeepseekOCRForCausalLM": "deepseek",
@ -333,6 +337,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
"KimiK25ForConditionalGeneration": "kimivl",
"KimiVLForConditionalGeneration": "kimivl",
"Lfm2AudioForConditionalGeneration": "lfm2",
"D1OmniModel": "lfm2",
"Lfm2VlForConditionalGeneration": "lfm2",
"LightOnOCRForConditionalGeneration": "lighton_ocr",
"Llama4ForConditionalGeneration": "llama4",

View file

@ -2336,12 +2336,26 @@ class TextModel(ModelBase):
else:
raise NotImplementedError("Only MEAN, CLS, and LAST pooling types supported")
self.gguf_writer.add_pooling_type(pooling_type)
else:
embedding_config_path = self.dir_model / "embedding_config.json"
if embedding_config_path.is_file():
with open(embedding_config_path, encoding="utf-8") as f:
embedding_config = json.load(f)
pooling = embedding_config.get("pooling")
if pooling == "last_token":
self.gguf_writer.add_pooling_type(gguf.PoolingType.LAST)
elif pooling is not None:
raise NotImplementedError(f"unsupported embedding_config.json pooling {pooling!r}")
# pooling before a classification head (e.g. ModernBertForSequenceClassification)
if (classifier_pooling := self.hparams.get("classifier_pooling")) is not None:
if classifier_pooling not in ("cls", "mean"):
raise NotImplementedError(f"Unsupported classifier_pooling: {classifier_pooling}")
self.gguf_writer.add_classifier_pooling_type(mode_mapping[classifier_pooling])
if (classifier_activation := self.hparams.get("classifier_activation")) is not None:
if classifier_activation not in ("gelu", "silu", "tanh"):
raise NotImplementedError(f"Unsupported classifier_activation: {classifier_activation}")
self.gguf_writer.add_classifier_activation(classifier_activation)
def _set_vocab_glmedge(self):
from transformers import AutoTokenizer
@ -2499,7 +2513,11 @@ class TextModel(ModelBase):
if template is not None:
self.gguf_writer.add_chat_template(template)
def _set_vocab_plamo(self):
def _set_vocab_plamo(
self,
eot_token: str,
normal_tokens: Iterable[str] = (),
):
# PLaMo models use a custom tokenizer with a .jsonl file
tokenizer_jsonl_path = self.dir_model / "tokenizer.jsonl"
tokenizer_config_path = self.dir_model / "tokenizer_config.json"
@ -2523,27 +2541,30 @@ class TextModel(ModelBase):
tokens = []
scores = []
toktypes = []
normal_tokens = set(normal_tokens)
with open(tokenizer_jsonl_path, "r", encoding="utf-8") as f:
for line_num, line in enumerate(f):
if line.strip():
token_data = json.loads(line)
# Format: [token, score, type, ?, ?, ?, ?]
token = token_data[0].encode("utf-8")
token_str = token_data[0]
token = token_str.encode("utf-8")
score = float(token_data[1])
token_type_str = token_data[2] if len(token_data) > 2 else "NORMAL"
tokens.append(token)
scores.append(score)
if token_type_str == "UNKNOWN":
if token_str in normal_tokens:
toktypes.append(gguf.TokenType.NORMAL)
elif token_type_str == "UNKNOWN":
toktypes.append(gguf.TokenType.UNKNOWN)
elif token_type_str == "CONTROL":
toktypes.append(gguf.TokenType.CONTROL)
elif token_type_str == "BYTE":
toktypes.append(gguf.TokenType.BYTE)
else:
token_str = token_data[0]
if token_str.startswith("<|plamo:") and token_str.endswith("|>"):
toktypes.append(gguf.TokenType.CONTROL)
else:
@ -2580,8 +2601,7 @@ class TextModel(ModelBase):
token_id = tokens.index(tokenizer_config["unk_token"].encode("utf-8"))
self.gguf_writer.add_unk_token_id(token_id)
# Add <|plamo:op|> as EOT to ensure appropriate end of generation
self.gguf_writer.add_eot_token_id(4)
self.gguf_writer.add_eot_token_id(tokens.index(eot_token.encode("utf-8")))
self.gguf_writer.add_add_space_prefix(False)

View file

@ -1,14 +1,14 @@
from __future__ import annotations
import re
from typing import Iterable, TYPE_CHECKING
from typing import Callable, Iterable, TYPE_CHECKING
import torch
if TYPE_CHECKING:
from torch import Tensor
from .base import ModelBase, TextModel, gguf, logger
from .base import MmprojModel, ModelBase, TextModel, gguf, logger
@ModelBase.register("CohereForCausalLM")
@ -180,3 +180,28 @@ class Cohere2MoeModel(TextModel):
experts = [k for d in self._experts for k in d.keys()]
if len(experts) > 0:
raise ValueError(f"Unprocessed experts: {experts}")
@ModelBase.register("Cohere2VisionForConditionalGeneration")
# [TAG_HF_EXAMPLE_GATED] CohereLabs/command-a-vision-07-2025 is gated
@ModelBase.example("CohereLabs/command-a-plus-05-2026-bf16")
class Cohere2VisionModel(MmprojModel):
def set_gguf_parameters(self):
super().set_gguf_parameters()
self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.COHERE2V)
self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams["layer_norm_eps"])
self.gguf_writer.add_vision_projector_scale_factor(self.global_config["downsample_factor"])
self.gguf_writer.add_vision_preproc_max_tiles(self.preprocessor_config["max_patches"])
self.gguf_writer.add_vision_use_gelu(True)
def tensor_force_quant(self, name, new_name, bid, n_dims):
if ".embeddings." in name:
return gguf.GGMLQuantizationType.F32
return super().tensor_force_quant(name, new_name, bid, n_dims)
@classmethod
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
name, gen = item
if not name.startswith(("model.vision_tower.", "model.multi_modal_projector.")):
return None
return super().filter_tensors((name, gen))

View file

@ -1,5 +1,8 @@
from __future__ import annotations
import json
from pathlib import Path
from typing import Any, Callable, Iterable, TYPE_CHECKING
import torch
@ -7,7 +10,7 @@ import torch
if TYPE_CHECKING:
from torch import Tensor
from .base import MmprojModel, ModelBase, TextModel, gguf
from .base import MmprojModel, ModelBase, TextModel, gguf, jinja_str_or_json, logger
from .gemma import ConformerAudioModel
@ -65,6 +68,68 @@ class LFM2Model(TextModel):
yield from super().modify_tensors(data_torch, name, bid)
def _is_d1_checkpoint(dir_model: Path) -> bool:
if not (dir_model / "config.json").is_file():
return False
with open(dir_model / "config.json", encoding="utf-8") as f:
return json.load(f).get("auto_map", {}).get("AutoModel", "").endswith(".D1Model")
@ModelBase.register_hparams_loader(_is_d1_checkpoint)
def _load_d1_hparams(dir_model: Path) -> dict[str, Any]:
logger.info("gguf: detected d1 checkpoint")
hparams = ModelBase.load_hparams(dir_model, False, guess=False)
# the mmproj stays LFM2-VL
hparams["text_config"]["architectures"] = ["D1Model"]
return hparams
@ModelBase.register("D1Model")
@ModelBase.example("LiquidAI/d1-3b")
class D1Model(LFM2Model):
model_arch = gguf.MODEL_ARCH.LFM2
def set_vocab(self):
super().set_vocab()
self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}])
@staticmethod
def _systemone_template() -> str:
# follows prompt.py of the model repo
description = jinja_str_or_json("o.description")
choice = (
"{{ '\\n\\nOptions:\\n' }}"
"{% for o in options %}{{ o.label }} {% if o.description %}" + description + "{% else %}{{ o.key | replace('_', ' ') }}{% endif %}"
"{% if not loop.last %}{{ '\\n' }}{% endif %}{% endfor %}"
"{{ '\\n\\nReply with the option code only.' }}"
)
# with criteria, a missing description is written as None
noul = (
"{% set ns = namespace(criteria=false) %}{% for o in options %}{% if o.description is not none %}{% set ns.criteria = true %}{% endif %}{% endfor %}"
"{% if ns.criteria %}"
"{% for o in options %}{{ '\\nYes: ' if o.key == 'true' else '\\nNo: ' }}"
"{% if o.description is none %}None{% else %}" + description + "{% endif %}{% endfor %}{% endif %}"
"{{ '\\n\\nReply with yes or no only.' }}"
)
score = (
"{{ '\\n\\n' }}{% for o in options %}{{ o.key }} " + description + "{{ '\\n' }}{% endfor %}"
"{{ '\\nReply with a single digit 0-' }}{{ options | length - 1 }}{{ ' only.' }}"
)
return (
"<|startoftext|><|im_start|>user\n"
"{% for image in images %}{{ image }}{% endfor %}"
"{% if state is not none %}{% if state is string %}{{ state }}{% else %}{{ state | tojson(indent=2) }}{% endif %}"
"{{ '\\n\\n\\nQUESTION:\\n' }}{% endif %}"
+ jinja_str_or_json("instructions")
+ "{% if type == 'choice' %}" + choice + "{% elif type == 'noul' %}" + noul + "{% else %}" + score + "{% endif %}"
"{{ '<|im_end|>\\n<|im_start|>assistant\\n' }}"
)
def set_gguf_parameters(self):
super().set_gguf_parameters()
self.gguf_writer.add_decision_type(gguf.DecisionType.LFM2_D1)
@ModelBase.register("Lfm2Model", "Lfm2BidirectionalModel", "Lfm2BidirectionalForMaskedLM")
@ModelBase.example("LiquidAI/LFM2.5-ColBERT-350M", "LiquidAI/LFM2.5-Embedding-350M", "LiquidAI/LFM2.5-Encoder-350M", "LiquidAI/LFM2.5-Encoder-230M")
class LFM2ColBertModel(LFM2Model):
@ -96,6 +161,121 @@ class LFM2ColBertModel(LFM2Model):
yield f"{self.dense_tensor_name}.weight", tensor.clone()
def _is_d1_omni_checkpoint(dir_model: Path) -> bool:
if not (dir_model / "config.json").is_file():
return False
with open(dir_model / "config.json", encoding="utf-8") as f:
return json.load(f).get("model_type") == "d1_omni"
@ModelBase.register_hparams_loader(_is_d1_omni_checkpoint)
def _load_d1_omni_hparams(dir_model: Path) -> dict[str, Any]:
logger.info("gguf: detected d1-omni checkpoint")
hparams = ModelBase.load_hparams(dir_model, False, guess=False)
text = hparams["text_config"]
n_layer, n_layer_head = text["num_hidden_layers"], hparams["head_layers"]
# the trunk uses the LFM2 FFN sizing, the head blocks are appended with a plain 4x MLP
n_ff = int(text["block_ffn_dim_multiplier"] * int(2 * text["intermediate_size"] / 3))
n_ff = text["block_multiple_of"] * ((n_ff + text["block_multiple_of"] - 1) // text["block_multiple_of"])
text["num_hidden_layers"] = n_layer + n_layer_head
text["intermediate_size"] = [n_ff] * n_layer + [4 * text["hidden_size"]] * n_layer_head
text["block_auto_adjust_ff_dim"] = False
return hparams
@ModelBase.register("D1OmniModel")
@ModelBase.example("LiquidAI/d1-omni-600M")
class D1OmniModel(LFM2Model):
model_arch = gguf.MODEL_ARCH.LFM2
# the server cuts the text to these lengths, see server-decision.cpp
_MAX_LENGTH = 16384
_IMAGE_TEXT_LENGTH = 896
_AUDIO_TEXT_LENGTH = 15360
def set_vocab(self):
super().set_vocab()
# the systemone template writes the BOS, after the media
self.gguf_writer.remove_key(gguf.Keys.Tokenizer.ADD_BOS)
self.gguf_writer.add_add_bos_token(False)
self.gguf_writer.add_token_type_count(3) # choice, score, noul
self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}])
@staticmethod
def _systemone_template() -> str:
# follows prompt.py of the model repo, the server cuts each marked piece to its token budget
# the media (images, or an audio clip if audio is true) come first
description = jinja_str_or_json("o.description")
has_description = "o.description is not none and o.description != ''"
yes_no = "{{ 'yes' if o.key == 'true' else 'no' }}"
option_code = "{% if loop.index0 < 10 %}00{% elif loop.index0 < 100 %}0{% endif %}{{ loop.index0 }}"
option = (
"{% if type == 'choice' and audio %}option_" + option_code + ": "
"{% if " + has_description + " %}" + description + "{% else %}{{ o.key }}{% endif %}"
"{% elif type == 'choice' %}{{ o.key }}{% if " + has_description + " %}: " + description + "{% endif %}"
"{% elif type == 'score' %}level {{ o.key }}: " + description
+ "{% elif audio %}{{ o.key }}: " + yes_no
+ "{% else %}{{ o.key }}: {% if " + has_description + " %}" + description
+ "{% elif images and not ns.criteria %}" + yes_no
+ "{% elif o.key == 'true' %}yes, the statement holds"
"{% else %}no, the statement does not hold{% endif %}{% endif %}"
)
state = "{% if state is string %}{{ state }}{% elif state is not none %}{{ state | tojson }}{% elif audio %}{}{% endif %}"
return (
"{% set ns = namespace(criteria=false) %}"
"{% for o in options %}{% if o.description is not none %}{% set ns.criteria = true %}{% endif %}{% endfor %}"
"{% for image in images %}{{ image }}{% endfor %}{{ sep }}"
"<|startoftext|><|reserved_7|>{{ sep }}{{ mark_state }}" + state
+ "{{ sep }}{{ mark_question }}<|reserved_8|>" + jinja_str_or_json("instructions")
+ "{% for o in options %}{{ sep }}<|reserved_9|><|mask|>{{ sep }}{{ mark_option }} " + option
+ "{{ sep }}<|reserved_10|>{% endfor %}{{ sep }}<|reserved_11|>"
)
def set_gguf_parameters(self):
lengths = (self.hparams["max_length"], self.hparams["image_text_length"], self.hparams["audio_text_length"])
if lengths != (self._MAX_LENGTH, self._IMAGE_TEXT_LENGTH, self._AUDIO_TEXT_LENGTH):
raise ValueError(f"unexpected text lengths: {lengths}")
n_head, n_layer_head = self.hparams["num_attention_heads"], self.hparams["head_layers"]
self.hparams["num_key_value_heads"] = [
self.hparams["num_key_value_heads"] if t != "conv" else 0 for t in self.hparams["layer_types"]
] + [n_head] * n_layer_head
# the head needs per-layer sizes, LFM2Model writes a single feed forward length
TextModel.set_gguf_parameters(self)
self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])
self.gguf_writer.add_shortconv_l_cache(self.hparams["conv_L_cache"])
self.gguf_writer.add_layer_norm_eps(1e-5) # nn.LayerNorm of the head
self.gguf_writer.add_causal_attention(False)
self.gguf_writer.add_decision_type(gguf.DecisionType.LFM2_D1_OMNI)
self.gguf_writer.add_decision_block_count(n_layer_head)
# "choice:3-5" -> "choice.3_5", "choice:11+" -> "choice.11"
for name, value in self.hparams["temperatures"].items():
self.gguf_writer.add_decision_temperature(name.replace(":", ".").replace("-", "_").rstrip("+"), value)
@classmethod
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
name, gen = item
if name.startswith(("vision.", "audio.")):
return None
name = name.replace("encoder.", "model.", 1) if name.startswith("encoder.") else name
name = name.replace("head.head.layers.", "head.layers.").replace("in_proj_", "in_proj.")
name = name.removeprefix("head.") if name.startswith(("head.type_emb", "head.scorer")) else name
return super().filter_tensors((name, gen))
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if name.startswith("head.layers.") and bid is not None:
# the head blocks come after the trunk blocks
suffix = name.split(".", 3)[3]
bid += self.block_count - self.hparams["head_layers"]
name = f"head.layers.{bid}.{suffix}"
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("Lfm2MoeForCausalLM")
@ModelBase.example("LiquidAI/LFM2-8B-A1B")
class LFM2MoeModel(TextModel):
@ -188,6 +368,12 @@ class LFM2VLModel(MmprojModel):
# python notation, e.g. for vision_feature_layer == -1, we pick last layer -> vision_feature_layers_to_drop = 0
vision_feature_layers_to_drop = -(self.global_config.get("vision_feature_layer", -1) + 1)
self.gguf_writer.add_vision_block_count(self.find_vparam(self.n_block_keys) - vision_feature_layers_to_drop)
# PIL resample enum
if (resample := self.preprocessor_config.get("resample")) is not None:
resize_algo = {1: "lanczos", 2: "bilinear", 3: "bicubic"}.get(resample)
if resize_algo is None:
raise ValueError(f"unsupported resample: {resample}")
self.gguf_writer.add_vision_image_resize_algo(resize_algo)
@classmethod
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
@ -205,6 +391,58 @@ class LFM2VLModel(MmprojModel):
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("D1OmniModel")
@ModelBase.example("LiquidAI/d1-omni-600M")
class D1OmniMmprojModel(ConformerAudioModel):
has_vision_encoder = True
has_audio_encoder = True
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
assert self.hparams_vision is not None and self.hparams_audio is not None
# dynamic resolution, as LFM2VLModel
self.hparams_vision["image_size"] = 256
# the images are normalized to [-1, 1] (vision.py of the model repo)
self.preprocessor_config = {**self.preprocessor_config, "image_mean": [0.5] * 3, "image_std": [0.5] * 3}
self.hparams_audio["hidden_size"] = self.hparams_audio["d_model"]
self.hparams_audio["intermediate_size"] = self.hparams_audio["d_model"] * self.hparams_audio["ff_expansion_factor"]
self.hparams_audio["num_attention_heads"] = self.hparams_audio["n_heads"]
def set_gguf_parameters(self):
super().set_gguf_parameters()
self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.D1OMNI_V)
self.gguf_writer.add_vision_attention_layernorm_eps(self.find_vparam(["layer_norm_eps"]))
self.gguf_writer.add_vision_projector_scale_factor(self.global_config.get("downsample_factor", 2))
self.gguf_writer.add_vision_use_gelu(True)
assert self.hparams_audio is not None
self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.D1OMNI_A)
self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["feat_in"])
self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
@classmethod
def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
name, gen = item
if name.startswith(("encoder.", "head.")):
return None
name = name.replace("vision.tower.", "vision_tower.").replace("vision.projector.", "multi_modal_projector.")
name = name.replace("audio.encoder.", "conformer.")
# the residual block continues the adapter: norm, linear, gelu, linear, then norm, down, up
for old, new in (("adapter.norm", 0), ("adapter.linear_1", 1), ("adapter.linear_2", 3),
("residual.ln", 4), ("residual.down", 5), ("residual.up", 6)):
name = name.replace(f"audio.{old}.", f"audio_adapter.model.{new}.")
return super().filter_tensors((name, gen))
def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
if "patch_embedding.weight" in name:
data_torch = data_torch.view(data_torch.shape[0], 16, 16, 3).permute(0, 3, 1, 2)
yield from super().modify_tensors(data_torch, name, bid)
@ModelBase.register("Lfm2AudioForConditionalGeneration")
@ModelBase.example("LiquidAI/LFM2.5-Audio-1.5B", "LiquidAI/LFM2-Audio-1.5B")
class LFM2AudioModel(ConformerAudioModel):

View file

@ -64,7 +64,7 @@ class Plamo2Model(TextModel):
model_arch = gguf.MODEL_ARCH.PLAMO2
def set_vocab(self):
self._set_vocab_plamo()
self._set_vocab_plamo(eot_token="<|plamo:op|>")
def set_gguf_parameters(self):
hparams = self.hparams
@ -170,7 +170,10 @@ class Plamo3Model(TextModel):
})
def set_vocab(self):
self._set_vocab_plamo()
self._set_vocab_plamo(
eot_token="<|plamo:tag|>",
normal_tokens=("<|plamo:begin_", "<|plamo:end_", ":plamo|>"),
)
tokenizer_config_path = self.dir_model / "tokenizer_config.json"
tokenizer_config = {}

View file

@ -650,11 +650,24 @@ class _Qwen35MRopeMixin:
self.gguf_writer.add_rope_dimension_sections(self._QWEN35_DEFAULT_MROPE_SECTION)
@ModelBase.register("Qwen3_5ForConditionalGeneration", "Qwen3_5ForCausalLM")
@ModelBase.register("Qwen3_5ForConditionalGeneration", "Qwen3_5ForCausalLM", "Qwen3_5TextModel")
@ModelBase.example("Qwen/Qwen3.5-9B")
class Qwen3_5TextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
model_arch = gguf.MODEL_ARCH.QWEN35
def __init__(self, dir_model, *args, **kwargs):
# Inner TextModel does not own mtp.*. Set no_mtp before mixin bumps block_count.
hparams = kwargs.pop("hparams", None)
if hparams is None:
hparams = ModelBase.load_hparams(dir_model, self.is_mistral_format)
if get_model_architecture(hparams, ModelType.TEXT) == "Qwen3_5TextModel":
self.no_mtp = True
super().__init__(dir_model, *args, hparams=hparams, **kwargs)
def set_gguf_parameters(self):
super().set_gguf_parameters()
self._try_set_pooling_type()
def _is_openjev_checkpoint(dir_model: Path) -> bool:
return (dir_model / "helper" / "shim.py").is_file() and (dir_model / "config.json").is_file()

View file

@ -218,8 +218,10 @@ class Qwen3TTSSpeakerEncoderModel(MmprojModel):
if hparams is None:
hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
hparams["text_config"] = {"hidden_size": hparams["talker_config"]["hidden_size"]}
# ECAPA-TDNN has a fixed 4-stage backbone, but MmprojModel.__init__ needs a n_block_keys
hparams["speaker_encoder_config"]["n_layers"] = 4
# ECAPA-TDNN has a fixed 4-stage backbone, but MmprojModel.__init__ needs a n_block_keys.
# The CustomVoice variant ships no speaker encoder, so its config lacks this key entirely.
if "speaker_encoder_config" in hparams:
hparams["speaker_encoder_config"]["n_layers"] = 4
super().__init__(dir_model, *args, hparams=hparams, **kwargs)
self._wav_config_cache = None

View file

@ -2,6 +2,9 @@
#include "convert.cuh"
#include "fwht.cuh"
// wide FWHT blocks use one row per thread block with this many threads
#define GGML_CUDA_FWHT_BLOCK_NT 256
template <int N, typename T>
__launch_bounds__(4*ggml_cuda_get_physical_warp_size(), 1)
__global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, const float scale) {
@ -59,6 +62,87 @@ __global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, cons
}
}
// Wide blocks: one row per thread block instead of per warp, so each thread keeps N/NT
// values rather than N/32. Stages below the warp width still shuffle, those up to the
// block width go through shared memory, and the rest stay in registers.
template <int N, int NT, typename T>
__launch_bounds__(NT, 1)
__global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows, const float scale) {
constexpr int warp_size = ggml_cuda_get_physical_warp_size();
constexpr int NE = N / NT;
static_assert(NE >= 1 && N % NT == 0 && NT % warp_size == 0, "bad FWHT block shape");
__shared__ float s[N];
const int64_t r = blockIdx.x;
if (r >= n_rows) {
return;
}
src += r * N;
dst += r * N;
const int tid = threadIdx.x;
const int lane = tid % warp_size;
ggml_cuda_pdl_sync();
float reg[NE];
#pragma unroll
for (int i = 0; i < NE; ++i) {
reg[i] = ggml_cuda_cast<float>(src[i * NT + tid]) * scale;
}
// stages within a warp: partner differs in the lane bits
#pragma unroll
for (int h = 1; h < warp_size; h *= 2) {
#pragma unroll
for (int j = 0; j < NE; j++) {
const float val = reg[j];
const float val2 = __shfl_xor_sync(0xFFFFFFFF, val, h, warp_size);
reg[j] = (lane & h) == 0 ? val + val2 : val2 - val;
}
}
// stages across warps: partner differs in the thread-index bits above the lane
#pragma unroll
for (int h = warp_size; h < NT; h *= 2) {
#pragma unroll
for (int j = 0; j < NE; j++) {
s[j * NT + tid] = reg[j];
}
__syncthreads();
#pragma unroll
for (int j = 0; j < NE; j++) {
const float val = reg[j];
const float val2 = s[j * NT + (tid ^ h)];
reg[j] = (tid & h) == 0 ? val + val2 : val2 - val;
}
__syncthreads();
}
// stages above the block width: partner is another register of the same thread
#pragma unroll
for (int h = NT; h < N; h *= 2) {
const int step = h / NT;
#pragma unroll
for (int j = 0; j < NE; j += 2 * step) {
#pragma unroll
for (int k = 0; k < step; k++) {
const float x = reg[j + k];
const float y = reg[j + k + step];
reg[j + k] = x + y;
reg[j + k + step] = x - y;
}
}
}
#pragma unroll
for (int i = 0; i < NE; ++i) {
dst[i * NT + tid] = reg[i];
}
}
template <typename T>
static bool ggml_cuda_op_fwht_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) {
const int n = src->ne[0];
@ -94,7 +178,34 @@ static bool ggml_cuda_op_fwht_impl(ggml_backend_cuda_context & ctx, const ggml_t
ggml_cuda_kernel_launch(fwht_cuda<512, T>, launch_params, src_d, dst_d, rows, scale);
return true;
default:
return false;
break;
}
// wide blocks: one row per thread block
{
constexpr int nt = GGML_CUDA_FWHT_BLOCK_NT;
dim3 grid_dims_w(rows, 1, 1);
dim3 block_dims_w(nt, 1, 1);
const ggml_cuda_kernel_launch_params launch_params_w =
ggml_cuda_kernel_launch_params(grid_dims_w, block_dims_w, 0, stream);
switch (n) {
case 1024:
ggml_cuda_kernel_launch(fwht_cuda_block<1024, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
return true;
case 2048:
ggml_cuda_kernel_launch(fwht_cuda_block<2048, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
return true;
case 4096:
ggml_cuda_kernel_launch(fwht_cuda_block<4096, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
return true;
case 8192:
ggml_cuda_kernel_launch(fwht_cuda_block<8192, nt, T>, launch_params_w, src_d, dst_d, rows, scale);
return true;
default:
return false;
}
}
}

View file

@ -1,7 +1,9 @@
#include "gated_delta_net.cuh"
#include "ggml-cuda/common.cuh"
template <int S_v, bool KDA, bool keep_rs_t>
constexpr int gdn_cols_per_warp = 4;
template <int S_v, bool KDA, bool keep_rs_t, int cols_per_warp = gdn_cols_per_warp>
__global__ void __launch_bounds__((ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v) * 4, 2)
gated_delta_net_cuda(const float * q,
const float * k,
@ -30,9 +32,19 @@ gated_delta_net_cuda(const float * q,
int K) {
const uint32_t h_idx = blockIdx.x;
const uint32_t sequence = blockIdx.y;
// each warp owns one column, using warp-level primitives to reduce across rows
const int lane = threadIdx.x;
const int col = blockIdx.z * blockDim.y + threadIdx.y;
constexpr int warp_size = ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v;
static_assert(S_v % warp_size == 0, "S_v must be a multiple of warp_size");
// the warp is split into cols_per_warp segments of lanes_per_col lanes; each segment owns
// one state column and reduces within itself
constexpr int lanes_per_col = warp_size / cols_per_warp;
constexpr int rows_per_lane = S_v / lanes_per_col;
static_assert(S_v % lanes_per_col == 0, "S_v must be a multiple of lanes_per_col");
const int lane = threadIdx.x;
const int col_in_warp = lane / lanes_per_col; // column slot within the warp
const int lane_in_col = lane - col_in_warp * lanes_per_col; // lane within the column's reduction segment
const int col = (blockIdx.z * blockDim.y + threadIdx.y) * cols_per_warp + col_in_warp;
const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic);
const uint32_t iq3 = fastdiv(sequence, rq3_magic);
@ -47,16 +59,13 @@ gated_delta_net_cuda(const float * q,
curr_state += state_in_offset + col * S_v;
attn_data += (sequence * n_tokens * H + h_idx) * S_v;
constexpr int warp_size = ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v;
static_assert(S_v % warp_size == 0, "S_v must be a multiple of warp_size");
constexpr int rows_per_lane = (S_v + warp_size - 1) / warp_size;
float s_shard[rows_per_lane];
// state is stored transposed: M[col][i] = S[i][col], row col is contiguous
ggml_cuda_pdl_sync();
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
const int i = r * warp_size + lane;
const int i = r * lanes_per_col + lane_in_col;
s_shard[r] = curr_state[i];
}
@ -76,7 +85,7 @@ gated_delta_net_cuda(const float * q,
float q_reg[rows_per_lane];
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
const int i = r * warp_size + lane;
const int i = r * lanes_per_col + lane_in_col;
k_reg[r] = k_t[i];
q_reg[r] = q_t[i];
}
@ -90,7 +99,7 @@ gated_delta_net_cuda(const float * q,
for (int r = 0; r < rows_per_lane; r++) {
kv_shard += s_shard[r] * k_reg[r];
}
float kv_col = warp_reduce_sum<warp_size>(kv_shard);
float kv_col = warp_reduce_sum<lanes_per_col>(kv_shard);
// delta[col] = (v[col] - g * kv[col]) * beta
float delta_col = (v_t[col] - g_val * kv_col) * beta_val;
@ -104,9 +113,9 @@ gated_delta_net_cuda(const float * q,
attn_partial += s_shard[r] * q_reg[r];
}
float attn_col = warp_reduce_sum<warp_size>(attn_partial);
float attn_col = warp_reduce_sum<lanes_per_col>(attn_partial);
if (lane == 0) {
if (lane_in_col == 0) {
attn_data[col] = attn_col * scale;
}
} else {
@ -114,11 +123,11 @@ gated_delta_net_cuda(const float * q,
float kv_shard = 0.0f;
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
const int i = r * warp_size + lane;
const int i = r * lanes_per_col + lane_in_col;
kv_shard += expf(g_t[i]) * s_shard[r] * k_reg[r];
}
float kv_col = warp_reduce_sum<warp_size>(kv_shard);
float kv_col = warp_reduce_sum<lanes_per_col>(kv_shard);
// delta[col] = (v[col] - kv[col]) * beta
float delta_col = (v_t[col] - kv_col) * beta_val;
@ -128,14 +137,14 @@ gated_delta_net_cuda(const float * q,
float attn_partial = 0.0f;
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
const int i = r * warp_size + lane;
const int i = r * lanes_per_col + lane_in_col;
s_shard[r] = expf(g_t[i]) * s_shard[r] + k_reg[r] * delta_col;
attn_partial += s_shard[r] * q_reg[r];
}
float attn_col = warp_reduce_sum<warp_size>(attn_partial);
float attn_col = warp_reduce_sum<lanes_per_col>(attn_partial);
if (lane == 0) {
if (lane_in_col == 0) {
attn_data[col] = attn_col * scale;
}
}
@ -150,7 +159,7 @@ gated_delta_net_cuda(const float * q,
float * curr_state = state + target_slot * state_slot_stride;
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
const int i = r * warp_size + lane;
const int i = r * lanes_per_col + lane_in_col;
curr_state[col * S_v + i] = s_shard[r];
}
}
@ -160,7 +169,7 @@ gated_delta_net_cuda(const float * q,
if constexpr (!keep_rs_t) {
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
const int i = r * warp_size + lane;
const int i = r * lanes_per_col + lane_in_col;
state[col * S_v + i] = s_shard[r];
}
}
@ -179,8 +188,16 @@ static void launch_gated_delta_net(
float scale, int64_t state_slot_stride, int K, cudaStream_t stream) {
//TODO: Add chunked kernel for even faster pre-fill
const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
const int num_warps = 4;
dim3 grid_dims(H, n_seqs, (S_v + num_warps - 1) / num_warps);
// four columns per warp (see the kernel); shrink the CTA when the wider CTA would leave
// SMs without a CTA, so small head counts keep the device filled
const int nsm = ggml_cuda_info().devices[ggml_cuda_get_device()].nsm;
const int cols_per_warp = gdn_cols_per_warp;
int num_warps = 4;
while (num_warps > 1 && H*n_seqs*(S_v / (cols_per_warp * num_warps)) < nsm) {
num_warps /= 2;
}
// one CTA covers cols_per_warp*num_warps columns (see the kernel)
dim3 grid_dims(H, n_seqs, (S_v + cols_per_warp * num_warps - 1) / (cols_per_warp * num_warps));
dim3 block_dims(warp_size <= S_v ? warp_size : S_v, num_warps, 1);
const uint3 neqk1_magic = init_fastdiv_values(neqk1);

View file

@ -5329,9 +5329,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
case GGML_UNARY_OP_CEIL:
case GGML_UNARY_OP_ROUND:
case GGML_UNARY_OP_TRUNC:
// TODO: should become:
//return ggml_is_contiguous_rows(op->src[0]);
return ggml_is_contiguous(op->src[0]);
if (op->src[0]->type == GGML_TYPE_BF16 && ggml_get_unary_op(op) == GGML_UNARY_OP_XIELU) {
return false;
}
return op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_BF16;
default:
return false;
}
@ -5673,7 +5674,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
return max_bias == 0.0f;
}
case GGML_OP_ROLL:
if(op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0])) {
if(op->src[0]->type == GGML_TYPE_F32) {
return true;
}
return false;

View file

@ -239,6 +239,14 @@ static __global__ void lightning_indexer_kernel_wmma(
// tokens scored per block by the tile kernel
#define LIGHTNING_INDEXER_TILE_TOKENS 8
// heads whose queries the tile kernel stages per pass, MUSA arch 21 caps static shared memory
// at 28 KB and the queries of four heads do not fit there next to the key tile
#if defined(GGML_USE_MUSA) && defined(__MUSA_ARCH__) && __MUSA_ARCH__ < 220
#define LIGHTNING_INDEXER_TILE_HEADS_PER_PASS 2
#else
#define LIGHTNING_INDEXER_TILE_HEADS_PER_PASS 4
#endif
// TODO there is one ugly assumption used in this kernel - that WARP_SIZE is equal to 32
// thanks to that one warp operating on float4 processes whole indexer K/Q vectors
// 32 * 4 = 128 (N_EMBD)
@ -406,9 +414,11 @@ static __global__ void lightning_indexer_kernel_tile(
constexpr int KEY_LANES = THREADS_PER_BLOCK / TOKENS_PER_BLOCK;
constexpr int KEYS_PER_THREAD = K_VECS_PER_BLOCK / KEY_LANES;
constexpr int N_EMBD_H2 = N_EMBD / 2;
constexpr int HEADS_PER_PASS = N_HEAD < LIGHTNING_INDEXER_TILE_HEADS_PER_PASS ? N_HEAD : LIGHTNING_INDEXER_TILE_HEADS_PER_PASS;
static_assert(THREADS_PER_BLOCK % TOKENS_PER_BLOCK == 0, "threads must cover the token tile");
static_assert(K_VECS_PER_BLOCK % KEY_LANES == 0, "key lanes must cover the key tile");
static_assert(N_HEAD % HEADS_PER_PASS == 0, "head passes must cover the heads");
const int tid = threadIdx.y * WARP_SIZE + threadIdx.x;
const int start_kv = blockIdx.x * K_VECS_PER_BLOCK;
@ -417,7 +427,7 @@ static __global__ void lightning_indexer_kernel_tile(
// the row padding keeps the keys of consecutive threads in distinct banks
__shared__ half2 k_shared[K_VECS_PER_BLOCK][N_EMBD_H2 + 1];
__shared__ float2 q_shared[N_HEAD][TOKENS_PER_BLOCK][N_EMBD_H2];
__shared__ float2 q_shared[HEADS_PER_PASS][TOKENS_PER_BLOCK][N_EMBD_H2];
__shared__ float w_shared[N_HEAD][TOKENS_PER_BLOCK];
// phase 1 - stage the key tile four elements at a time, rows past n_kv are zero
@ -451,22 +461,7 @@ static __global__ void lightning_indexer_kernel_tile(
k_shared[r][2*c4 + 1] = hi;
}
// phase 2 - stage the queries and weights of every head, tokens past n_batch are zero
#pragma unroll
for (int i = tid; i < N_HEAD * TOKENS_PER_BLOCK * (N_EMBD / 4); i += THREADS_PER_BLOCK) {
const int h = i / (TOKENS_PER_BLOCK * (N_EMBD / 4));
const int r = i / (N_EMBD / 4) % TOKENS_PER_BLOCK;
const int c4 = i % (N_EMBD / 4);
float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
if (start_batch + r < n_batch) {
v = *(const float4 *) ((const char *) Q + h*nbq1 + (start_batch + r)*nbq2 + i_stream*nbq3 + c4*sizeof(float4));
}
q_shared[h][r][2*c4 + 0] = make_float2(v.x, v.y);
q_shared[h][r][2*c4 + 1] = make_float2(v.z, v.w);
}
// phase 2 - stage the weights of every head, tokens past n_batch are zero
if (tid < N_HEAD * TOKENS_PER_BLOCK) {
const int h = tid / TOKENS_PER_BLOCK;
@ -475,33 +470,60 @@ static __global__ void lightning_indexer_kernel_tile(
((const float *) ((const char *) W + (start_batch + r)*nbw1 + i_stream*nbw3))[h] : 0.0f;
}
__syncthreads();
// phase 3 - float products of the widened keys for every head, ReLU, weight
const int kl = tid % KEY_LANES;
const int tl = tid / KEY_LANES;
float qk[N_HEAD][KEYS_PER_THREAD] = { { 0.0f } };
#pragma unroll 8
for (int c = 0; c < N_EMBD_H2; ++c) {
float2 k_val[KEYS_PER_THREAD];
#pragma unroll
for (int j = 0; j < KEYS_PER_THREAD; ++j) {
k_val[j] = __half22float2(k_shared[kl + j*KEY_LANES][c]);
for (int h0 = 0; h0 < N_HEAD; h0 += HEADS_PER_PASS) {
// the previous pass is fully consumed before its queries are replaced
if (h0 > 0) {
__syncthreads();
}
// phase 3 - stage the queries of the heads of this pass, tokens past n_batch are zero
#pragma unroll
for (int h = 0; h < N_HEAD; ++h) {
const float2 q_val = q_shared[h][tl][c];
for (int i = tid; i < HEADS_PER_PASS * TOKENS_PER_BLOCK * (N_EMBD / 4); i += THREADS_PER_BLOCK) {
const int h = i / (TOKENS_PER_BLOCK * (N_EMBD / 4));
const int r = i / (N_EMBD / 4) % TOKENS_PER_BLOCK;
const int c4 = i % (N_EMBD / 4);
float4 v = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
if (start_batch + r < n_batch) {
v = *(const float4 *) ((const char *) Q + (h0 + h)*nbq1 + (start_batch + r)*nbq2 + i_stream*nbq3 + c4*sizeof(float4));
}
q_shared[h][r][2*c4 + 0] = make_float2(v.x, v.y);
q_shared[h][r][2*c4 + 1] = make_float2(v.z, v.w);
}
__syncthreads();
// phase 4 - float products of the widened keys for the heads of this pass
#pragma unroll 8
for (int c = 0; c < N_EMBD_H2; ++c) {
float2 k_val[KEYS_PER_THREAD];
#pragma unroll
for (int j = 0; j < KEYS_PER_THREAD; ++j) {
qk[h][j] = fmaf(k_val[j].x, q_val.x, qk[h][j]);
qk[h][j] = fmaf(k_val[j].y, q_val.y, qk[h][j]);
k_val[j] = __half22float2(k_shared[kl + j*KEY_LANES][c]);
}
#pragma unroll
for (int h = 0; h < HEADS_PER_PASS; ++h) {
const float2 q_val = q_shared[h][tl][c];
#pragma unroll
for (int j = 0; j < KEYS_PER_THREAD; ++j) {
qk[h0 + h][j] = fmaf(k_val[j].x, q_val.x, qk[h0 + h][j]);
qk[h0 + h][j] = fmaf(k_val[j].y, q_val.y, qk[h0 + h][j]);
}
}
}
}
// phase 5 - ReLU, weight, add the mask and write, consecutive threads write consecutive keys
float score[KEYS_PER_THREAD] = { 0.0f };
#pragma unroll
@ -512,8 +534,6 @@ static __global__ void lightning_indexer_kernel_tile(
}
}
// phase 4 - add the mask and write, consecutive threads write consecutive keys
const int i_batch = start_batch + tl;
if (i_batch >= n_batch) {
return;
@ -677,8 +697,6 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor *
LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_vec, 128, 32, k, GGML_TYPE_F32)
GGML_ABORT("fatal error");
}
#ifndef GGML_USE_MUSA
// MUSA archs 21 and 22 cap static shared memory at 28 KB, below what the tile kernel stages
} else if (n_embd == 128 && n_head == 4 && n_batch >= LIGHTNING_INDEXER_TILE_TOKENS) {
// too few heads for a wmma tile, the tile kernel shares the keys across the tokens
constexpr int WARPS_PER_BLOCK = 8;
@ -698,9 +716,8 @@ void ggml_cuda_lightning_indexer(ggml_backend_cuda_context & ctx, ggml_tensor *
LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_BF16)
LIGHTNING_INDEXER_CASE(lightning_indexer_kernel_tile, 128, 4, k, GGML_TYPE_F32)
GGML_ABORT("fatal error");
#endif // GGML_USE_MUSA
} else if (n_embd == 128 && n_head == 4) {
// a batch smaller than a token tile, or MUSA, use vector kernel
// a batch smaller than a token tile, use vector kernel
constexpr int K_VECS_PER_WARP = 8;
constexpr int WARPS_PER_BLOCK = 8;
constexpr int K_VECS_PER_BLOCK = K_VECS_PER_WARP * WARPS_PER_BLOCK;

View file

@ -3,38 +3,46 @@
template <int block_size>
static __global__ void norm_f32(
const float * x, float * dst, const int ncols, const int64_t stride_row, const int64_t stride_channel,
const int64_t stride_sample, const float eps) {
const int nrows = gridDim.x;
const int nchannels = gridDim.y;
const float * x, float * dst, const int ncols, const int nchannels, const int nsamples,
const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps) {
const int nrows = gridDim.x;
const int row = blockIdx.x;
const int tid = threadIdx.x;
const int row = blockIdx.x;
const int channel = blockIdx.y;
const int sample = blockIdx.z;
const int tid = threadIdx.x;
x += sample*stride_sample + channel*stride_channel + row*stride_row;
dst += ((sample*nchannels + channel)*nrows + row)*ncols;
float2 mean_var = make_float2(0.0f, 0.0f);
extern __shared__ float2 s_sum2[];
ggml_cuda_pdl_sync();
for (int col = tid; col < ncols; col += block_size) {
const float xi = x[col];
mean_var.x += xi;
mean_var.y += xi * xi;
}
// sum up partial sums
extern __shared__ float2 s_sum2[];
mean_var = block_reduce<block_reduce_method::SUM, block_size>(mean_var, s_sum2);
// grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples
for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) {
for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) {
const float * xc = x + sample*stride_sample + channel*stride_channel + row*stride_row;
float * dstc = dst + ((sample*nchannels + channel)*nrows + row)*ncols;
const float mean = mean_var.x / ncols;
const float var = mean_var.y / ncols - mean * mean;
const float inv_std = rsqrtf(var + eps);
float2 mean_var = make_float2(0.0f, 0.0f);
for (int col = tid; col < ncols; col += block_size) {
dst[col] = (x[col] - mean) * inv_std;
for (int col = tid; col < ncols; col += block_size) {
const float xi = xc[col];
mean_var.x += xi;
mean_var.y += xi * xi;
}
// sum up partial sums
mean_var = block_reduce<block_reduce_method::SUM, block_size>(mean_var, s_sum2);
const float mean = mean_var.x / ncols;
const float var = mean_var.y / ncols - mean * mean;
const float inv_std = rsqrtf(var + eps);
for (int col = tid; col < ncols; col += block_size) {
dstc[col] = (xc[col] - mean) * inv_std;
}
if constexpr (block_size > WARP_SIZE) {
// sync is needed as we reuse s_sum2 across block_reduce invocations, see #26385
__syncthreads();
}
}
}
}
@ -77,6 +85,8 @@ template <int block_size, bool do_multiply = false, bool do_add = false, bool do
static __global__ void rms_norm_f32(const float * x,
float * dst,
const int ncols,
const int nchannels,
const int nsamples,
const int64_t stride_row,
const int64_t stride_channel,
const int64_t stride_sample,
@ -99,61 +109,71 @@ static __global__ void rms_norm_f32(const float * x,
const uint3 add_nsamples_packed = make_uint3(0, 0, 0),
const float scale_out = 1.0f) {
ggml_cuda_pdl_lc();
const int nrows = gridDim.x;
const int nchannels = gridDim.y;
const int row = blockIdx.x;
const int channel = blockIdx.y;
const int sample = blockIdx.z;
const int tid = threadIdx.x;
const int nrows = gridDim.x;
const int row = blockIdx.x;
const int tid = threadIdx.x;
static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying");
static_assert(!do_scale || !do_multiply, "fusing scale is not supported with multiplying");
x += sample*stride_sample + channel*stride_channel + row*stride_row;
dst += ((sample*nchannels + channel)*nrows + row)*ncols;
if constexpr (do_multiply) {
const uint32_t mul_row = fastmodulo(row, mul_nrows_packed);
const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed);
const uint32_t mul_sample = fastmodulo(sample, mul_nsamples_packed);
mul += mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row;
}
if constexpr (do_add) {
const int add_row = fastmodulo(row, add_nrows_packed);
const int add_channel = fastmodulo(channel, add_nchannels_packed);
const int add_sample = fastmodulo(sample, add_nsamples_packed);
add += add_sample * add_stride_sample + add_channel * add_stride_channel + add_row * add_stride_row;
}
float tmp = 0.0f; // partial sum for thread in warp
extern __shared__ float s_sum[];
ggml_cuda_pdl_sync();
for (int col = tid; col < ncols; col += block_size) {
const float xi = x[col];
tmp += xi * xi;
}
// sum up partial sums
extern __shared__ float s_sum[];
tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
// grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples
for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) {
for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) {
const float * xc = x + sample*stride_sample + channel*stride_channel + row*stride_row;
float * dstc = dst + ((sample*nchannels + channel)*nrows + row)*ncols;
const float mean = tmp / ncols;
const float scale = rsqrtf(mean + eps);
[[maybe_unused]] const float * mulc = nullptr;
if constexpr (do_multiply) {
const uint32_t mul_row = fastmodulo(row, mul_nrows_packed);
const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed);
const uint32_t mul_sample = fastmodulo(sample, mul_nsamples_packed);
mulc = mul + mul_sample * mul_stride_sample + mul_channel * mul_stride_channel + mul_row * mul_stride_row;
}
for (int col = tid; col < ncols; col += block_size) {
if constexpr (do_multiply && do_add) {
const int mul_col = fastmodulo(col, mul_ncols_packed);
const int add_col = fastmodulo(col, add_ncols_packed);
dst[col] = scale * x[col] * mul[mul_col] + add[add_col];
} else if constexpr (do_multiply) {
const int mul_col = fastmodulo(col, mul_ncols_packed);
dst[col] = scale * x[col] * mul[mul_col];
} else if constexpr (do_scale) {
dst[col] = scale_out * (scale * x[col]);
} else {
dst[col] = scale * x[col];
[[maybe_unused]] const float * addc = nullptr;
if constexpr (do_add) {
const int add_row = fastmodulo(row, add_nrows_packed);
const int add_channel = fastmodulo(channel, add_nchannels_packed);
const int add_sample = fastmodulo(sample, add_nsamples_packed);
addc = add + add_sample * add_stride_sample + add_channel * add_stride_channel + add_row * add_stride_row;
}
float tmp = 0.0f; // partial sum for thread in warp
for (int col = tid; col < ncols; col += block_size) {
const float xi = xc[col];
tmp += xi * xi;
}
// sum up partial sums
tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
const float mean = tmp / ncols;
const float scale = rsqrtf(mean + eps);
for (int col = tid; col < ncols; col += block_size) {
if constexpr (do_multiply && do_add) {
const int mul_col = fastmodulo(col, mul_ncols_packed);
const int add_col = fastmodulo(col, add_ncols_packed);
dstc[col] = scale * xc[col] * mulc[mul_col] + addc[add_col];
} else if constexpr (do_multiply) {
const int mul_col = fastmodulo(col, mul_ncols_packed);
dstc[col] = scale * xc[col] * mulc[mul_col];
} else if constexpr (do_scale) {
dstc[col] = scale_out * (scale * xc[col]);
} else {
dstc[col] = scale * xc[col];
}
}
if constexpr (block_size > WARP_SIZE) {
// sync is needed as we reuse s_sum across block_reduce invocations, see #26385
__syncthreads();
}
}
}
}
@ -247,50 +267,57 @@ static __global__ void rms_norm_back_f32(
template <int block_size>
static __global__ void l2_norm_f32(
const float * x, float * dst, const int ncols, const int64_t stride_row, const int64_t stride_channel,
const int64_t stride_sample, const float eps) {
const int nrows = gridDim.x;
const int nchannels = gridDim.y;
const float * x, float * dst, const int ncols, const int nchannels, const int nsamples,
const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps) {
const int nrows = gridDim.x;
const int row = blockIdx.x;
const int tid = threadIdx.x;
const int row = blockIdx.x;
const int channel = blockIdx.y;
const int sample = blockIdx.z;
const int tid = threadIdx.x;
x += sample*stride_sample + channel*stride_channel + row*stride_row;
dst += ((sample*nchannels + channel)*nrows + row)*ncols;
float tmp = 0.0f; // partial sum for thread in warp
extern __shared__ float s_sum[];
ggml_cuda_pdl_sync();
for (int col = tid; col < ncols; col += block_size) {
const float xi = x[col];
tmp += xi * xi;
}
// sum up partial sums
extern __shared__ float s_sum[];
tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
ggml_cuda_pdl_lc();
// grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples
for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) {
for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) {
const float * xc = x + sample*stride_sample + channel*stride_channel + row*stride_row;
float * dstc = dst + ((sample*nchannels + channel)*nrows + row)*ncols;
// from https://pytorch.org/docs/stable/generated/torch.nn.functional.normalize.html
const float scale = rsqrtf(fmaxf(tmp, eps * eps));
float tmp = 0.0f; // partial sum for thread in warp
for (int col = tid; col < ncols; col += block_size) {
dst[col] = scale * x[col];
for (int col = tid; col < ncols; col += block_size) {
const float xi = xc[col];
tmp += xi * xi;
}
// sum up partial sums
tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
// from https://pytorch.org/docs/stable/generated/torch.nn.functional.normalize.html
const float scale = rsqrtf(fmaxf(tmp, eps * eps));
for (int col = tid; col < ncols; col += block_size) {
dstc[col] = scale * xc[col];
}
if constexpr (block_size > WARP_SIZE) {
// sync is needed as we reuse s_sum across block_reduce invocations, see #26385
__syncthreads();
}
}
}
}
static void norm_f32_cuda(
const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples,
const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) {
const dim3 blocks_num(nrows, nchannels, nsamples);
const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
if (ncols < 1024) {
const dim3 block_dims(WARP_SIZE, 1, 1);
norm_f32<WARP_SIZE><<<blocks_num, block_dims, 0, stream>>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps);
norm_f32<WARP_SIZE><<<blocks_num, block_dims, 0, stream>>>(x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps);
} else {
const dim3 block_dims(1024, 1, 1);
norm_f32<1024><<<blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float2): 0, stream>>>(x, dst, ncols, stride_row, stride_channel, stride_sample, eps);
norm_f32<1024><<<blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float2): 0, stream>>>(x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps);
}
}
@ -310,19 +337,19 @@ static void rms_norm_f32_cuda(
const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples,
const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream,
const float scale_out = 1.0f) {
const dim3 blocks_num(nrows, nchannels, nsamples);
const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
if (ncols < 1024) {
const dim3 block_dims(256, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
ggml_cuda_kernel_launch(rms_norm_f32<256, false, false, do_scale>, launch_params,
x, dst, ncols, stride_row, stride_channel, stride_sample, eps,
x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
// underlying cudaLaunchKernelEx does not support default params
nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0),
nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), scale_out);
} else {
const dim3 block_dims(1024, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
ggml_cuda_kernel_launch(rms_norm_f32<1024, false, false, do_scale>, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps,
ggml_cuda_kernel_launch(rms_norm_f32<1024, false, false, do_scale>, launch_params, x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps,
// underlying cudaLaunchKernelEx does not support default params
nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0),
nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), scale_out);
@ -356,7 +383,7 @@ static void rms_norm_mul_f32_cuda(const float * x,
const uint32_t add_nsamples,
const float eps,
cudaStream_t stream) {
const dim3 blocks_num(nrows, nchannels, nsamples);
const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
if (mul == nullptr) {
rms_norm_f32_cuda(x, dst, ncols, nrows, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, stream);
return;
@ -370,7 +397,7 @@ static void rms_norm_mul_f32_cuda(const float * x,
const dim3 block_dims(256, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
ggml_cuda_kernel_launch(rms_norm_f32<256, true>, launch_params,
x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
// underlying cudaLaunchKernelEx does not support default params
nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), 1.0f);
@ -378,7 +405,7 @@ static void rms_norm_mul_f32_cuda(const float * x,
const dim3 block_dims(1024, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
ggml_cuda_kernel_launch(rms_norm_f32<1024, true>, launch_params,
x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
// underlying cudaLaunchKernelEx does not support default params
nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), 1.0f);
@ -397,7 +424,7 @@ static void rms_norm_mul_f32_cuda(const float * x,
const dim3 block_dims(256, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims,block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
ggml_cuda_kernel_launch(rms_norm_f32<256, true, true>, launch_params,
x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add,
add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed,
add_nchannels_packed, add_nsamples_packed, 1.0f);
@ -405,7 +432,7 @@ static void rms_norm_mul_f32_cuda(const float * x,
const dim3 block_dims(1024, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
ggml_cuda_kernel_launch(rms_norm_f32<1024, true, true>, launch_params,
x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel,
mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add,
add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed,
add_nchannels_packed, add_nsamples_packed, 1.0f);
@ -426,15 +453,15 @@ static void rms_norm_back_f32_cuda(const float * grad, const float * xf, float *
static void l2_norm_f32_cuda(
const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples,
const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) {
const dim3 blocks_num(nrows, nchannels, nsamples);
const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
if (ncols < 1024) {
const dim3 block_dims(WARP_SIZE, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, 0, stream};
ggml_cuda_kernel_launch(l2_norm_f32<WARP_SIZE>, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps);
ggml_cuda_kernel_launch(l2_norm_f32<WARP_SIZE>, launch_params, x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps);
} else {
const dim3 block_dims(1024, 1, 1);
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream};
ggml_cuda_kernel_launch(l2_norm_f32<1024>, launch_params, x, dst, ncols, stride_row, stride_channel, stride_sample, eps);
ggml_cuda_kernel_launch(l2_norm_f32<1024>, launch_params, x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps);
}
}

View file

@ -50,6 +50,65 @@ static __global__ void pool2d_nchw_kernel(
o_ptr[cur_oh * ow + cur_ow] = res;
}
template <typename Ti, typename To>
static __global__ void pool2d_nchw_kernel_warp(
const int ih, const int iw, const int oh, const int ow,
const int kh, const int kw, const int sh, const int sw,
const int ph, const int pw, const int parallel_elements,
const Ti * __restrict__ src, To * __restrict__ dst, const enum ggml_op_pool op) {
const int warp_id = (threadIdx.x + blockIdx.x * blockDim.x) / WARP_SIZE;
const int lane = threadIdx.x % WARP_SIZE;
if (warp_id >= parallel_elements) {
return;
}
const int I_HW = ih * iw;
const int O_HW = oh * ow;
const int nc = warp_id / O_HW;
const int cur_oh = warp_id % O_HW / ow;
const int cur_ow = warp_id % O_HW % ow;
const Ti* i_ptr = src + nc * I_HW;
const int start_h = cur_oh * sh - ph;
const int bh = max(0, start_h);
const int eh = min(ih, start_h + kh);
const int start_w = cur_ow * sw - pw;
const int bw = max(0, start_w);
const int ew = min(iw, start_w + kw);
const int win_w = ew - bw;
const int win_elems = (eh - bh) * win_w;
const To scale = 1. / (kh * kw);
To res;
switch (op) {
case GGML_OP_POOL_AVG: res = 0; break;
case GGML_OP_POOL_MAX: res = -FLT_MAX; break;
default: res = 0; assert(false);
}
for (int t = lane; t < win_elems; t += WARP_SIZE) {
const int i = bh + t / win_w;
const int j = bw + t % win_w;
const Ti cur = i_ptr[i * iw + j];
switch (op) {
case GGML_OP_POOL_AVG: res += cur * scale; break;
case GGML_OP_POOL_MAX: res = max(res, (To)cur); break;
default: break;
}
}
#pragma unroll
for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) {
const To other = __shfl_xor_sync(0xFFFFFFFF, res, offset, WARP_SIZE);
res = (op == GGML_OP_POOL_MAX) ? max(res, other) : res + other;
}
if (lane == 0) {
dst[nc * O_HW + cur_oh * ow + cur_ow] = res;
}
}
static void pool2d_nchw_kernel_f32_f32_cuda(
const int ih, const int iw, const int oh, const int ow,
const int kh, const int kw, const int sh, const int sw,
@ -57,6 +116,13 @@ static void pool2d_nchw_kernel_f32_f32_cuda(
const float * src, float * dst, const enum ggml_op_pool op,
cudaStream_t stream) {
if (kh * kw >= CUDA_POOL2D_WARP_KERNEL_MIN_WINDOW) {
const int warps_per_block = CUDA_POOL2D_BLOCK_SIZE / WARP_SIZE;
const int num_blocks = (parallel_elements + warps_per_block - 1) / warps_per_block;
pool2d_nchw_kernel_warp<<<num_blocks, CUDA_POOL2D_BLOCK_SIZE, 0, stream>>>(ih, iw, oh, ow, kh, kw, sh, sw, ph, pw, parallel_elements, src, dst, op);
return;
}
const int num_blocks = (parallel_elements + CUDA_POOL2D_BLOCK_SIZE - 1) / CUDA_POOL2D_BLOCK_SIZE;
dim3 block_nums(num_blocks);
pool2d_nchw_kernel<<<block_nums, CUDA_POOL2D_BLOCK_SIZE, 0, stream>>>(ih, iw, oh, ow, kh, kw, sh, sw, ph, pw, parallel_elements, src, dst, op);

View file

@ -1,5 +1,6 @@
#include "common.cuh"
#define CUDA_POOL2D_BLOCK_SIZE 256
#define CUDA_POOL2D_WARP_KERNEL_MIN_WINDOW 32
void ggml_cuda_op_pool2d(ggml_backend_cuda_context & ctx, ggml_tensor * dst);

View file

@ -17,6 +17,10 @@ static __global__ void roll_f32_cuda(const float * __restrict__ src,
const int64_t ne01,
const int64_t ne02,
const int64_t ne03,
const int64_t nb00,
const int64_t nb01,
const int64_t nb02,
const int64_t nb03,
const int s0,
const int s1,
const int s2,
@ -39,7 +43,7 @@ static __global__ void roll_f32_cuda(const float * __restrict__ src,
const int64_t d3 = wrap_index(i3 - s3, ne03);
dst[i3 * (ne00 * ne01 * ne02) + i2 * (ne01 * ne00) + i1 * ne00 + i0] =
src[d3 * (ne00 * ne01 * ne02) + d2 * (ne01 * ne00) + d1 * ne00 + d0];
src[(d3 * nb03 + d2 * nb02 + d1 * nb01 + d0 * nb00) / sizeof(float)];
}
void ggml_cuda_op_roll(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
@ -63,5 +67,5 @@ void ggml_cuda_op_roll(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
int64_t num_blocks = (sz + CUDA_ROLL_BLOCK_SIZE - 1) / CUDA_ROLL_BLOCK_SIZE;
roll_f32_cuda<<<num_blocks, CUDA_ROLL_BLOCK_SIZE, 0, stream>>>(
src0_d, dst_d, ne00, ne01, ne02, ne03, s0, s1, s2, s3);
src0_d, dst_d, ne00, ne01, ne02, ne03, nb00, nb01, nb02, nb03, s0, s1, s2, s3);
}

View file

@ -709,7 +709,7 @@ void ggml_cuda_op_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * rope
// one block per row: block_reduce gives the norm scale, then each thread applies mul and rope to the elements it owns
template <int block_size, bool has_ff, typename D>
static __global__ void rms_norm_mul_rope_f32(
const float * x, D * dst, const int ncols,
const float * x, D * dst, const int ncols, const int nchannels, const int nsamples,
const int64_t s01, const int64_t s02, const int64_t s03,
const int64_t s1, const int64_t s2, const int64_t s3,
const float eps,
@ -724,66 +724,76 @@ static __global__ void rms_norm_mul_rope_f32(
const int64_t * row_indices, const int set_rows_stride,
const bool is_neox) {
ggml_cuda_pdl_lc();
const int row = blockIdx.x;
const int channel = blockIdx.y;
const int sample = blockIdx.z;
const int tid = threadIdx.x;
x += sample*s03 + channel*s02 + row*s01;
const uint32_t mul_row = fastmodulo(row, mul_nrows_packed);
const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed);
const uint32_t mul_sample = fastmodulo(sample, mul_nsamples_packed);
mul += mul_sample*mul_s03 + mul_channel*mul_s02 + mul_row*mul_s01;
float tmp = 0.0f;
ggml_cuda_pdl_sync();
for (int col = tid; col < ncols; col += block_size) {
const float xi = x[col];
tmp += xi * xi;
}
const int row = blockIdx.x;
const int tid = threadIdx.x;
extern __shared__ float s_sum[];
tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
const float scale = rsqrtf(tmp/ncols + eps);
ggml_cuda_pdl_sync();
int64_t idst = sample*s3 + channel*s2 + row*s1;
if (set_rows_stride != 0) {
idst = row*s1 + row_indices[channel]*set_rows_stride;
}
dst += idst;
// grid.y and grid.z are clamped to the CUDA limit, iterate over the excess channels/samples
for (int sample = blockIdx.z; sample < nsamples; sample += gridDim.z) {
for (int channel = blockIdx.y; channel < nchannels; channel += gridDim.y) {
const float * xc = x + sample*s03 + channel*s02 + row*s01;
for (int i0 = 2*tid; i0 < ncols; i0 += 2*block_size) {
int ix0;
int ix1;
if (is_neox && i0 < n_dims) {
ix0 = i0/2;
ix1 = i0/2 + n_dims/2;
} else {
ix0 = i0 + 0;
ix1 = i0 + 1;
const uint32_t mul_row = fastmodulo(row, mul_nrows_packed);
const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed);
const uint32_t mul_sample = fastmodulo(sample, mul_nsamples_packed);
const float * mulc = mul + mul_sample*mul_s03 + mul_channel*mul_s02 + mul_row*mul_s01;
float tmp = 0.0f;
for (int col = tid; col < ncols; col += block_size) {
const float xi = xc[col];
tmp += xi * xi;
}
tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
const float scale = rsqrtf(tmp/ncols + eps);
int64_t idst = sample*s3 + channel*s2 + row*s1;
if (set_rows_stride != 0) {
idst = row*s1 + row_indices[channel]*set_rows_stride;
}
D * dstc = dst + idst;
for (int i0 = 2*tid; i0 < ncols; i0 += 2*block_size) {
int ix0;
int ix1;
if (is_neox && i0 < n_dims) {
ix0 = i0/2;
ix1 = i0/2 + n_dims/2;
} else {
ix0 = i0 + 0;
ix1 = i0 + 1;
}
const float x0 = scale * xc[ix0] * mulc[fastmodulo(ix0, mul_ncols_packed)];
const float x1 = scale * xc[ix1] * mulc[fastmodulo(ix1, mul_ncols_packed)];
if (i0 >= n_dims) {
dstc[ix0] = ggml_cuda_cast<D>(x0);
dstc[ix1] = ggml_cuda_cast<D>(x1);
continue;
}
const float theta_base = pos[channel]*powf(theta_scale, i0/2.0f);
const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f;
float cos_theta;
float sin_theta;
rope_yarn<true>(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta);
dstc[ix0] = ggml_cuda_cast<D>(x0*cos_theta - x1*sin_theta);
dstc[ix1] = ggml_cuda_cast<D>(x0*sin_theta + x1*cos_theta);
}
if constexpr (block_size > WARP_SIZE) {
// sync is needed as we reuse s_sum across block_reduce invocations, see #26385
__syncthreads();
}
}
const float x0 = scale * x[ix0] * mul[fastmodulo(ix0, mul_ncols_packed)];
const float x1 = scale * x[ix1] * mul[fastmodulo(ix1, mul_ncols_packed)];
if (i0 >= n_dims) {
dst[ix0] = ggml_cuda_cast<D>(x0);
dst[ix1] = ggml_cuda_cast<D>(x1);
continue;
}
const float theta_base = pos[channel]*powf(theta_scale, i0/2.0f);
const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f;
float cos_theta;
float sin_theta;
rope_yarn<true>(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta);
dst[ix0] = ggml_cuda_cast<D>(x0*cos_theta - x1*sin_theta);
dst[ix1] = ggml_cuda_cast<D>(x0*sin_theta + x1*cos_theta);
}
}
@ -806,7 +816,7 @@ static void rms_norm_mul_rope_cuda(
const bool is_neox, cudaStream_t stream) {
GGML_ASSERT(ncols % 2 == 0);
const dim3 blocks_num(nrows, nchannels, nsamples);
const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX));
const float theta_scale = powf(freq_base, -2.0f/n_dims);
@ -820,13 +830,13 @@ static void rms_norm_mul_rope_cuda(
const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32*sizeof(float), stream};
if (freq_factors == nullptr) {
ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<256, false, D>, launch_params,
x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
freq_factors, row_indices, set_rows_stride, is_neox);
} else {
ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<256, true, D>, launch_params,
x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
freq_factors, row_indices, set_rows_stride, is_neox);
@ -836,13 +846,13 @@ static void rms_norm_mul_rope_cuda(
const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32*sizeof(float), stream};
if (freq_factors == nullptr) {
ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<1024, false, D>, launch_params,
x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
freq_factors, row_indices, set_rows_stride, is_neox);
} else {
ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<1024, true, D>, launch_params,
x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
x, dst, ncols, nchannels, nsamples, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
freq_factors, row_indices, set_rows_stride, is_neox);

View file

@ -134,24 +134,67 @@ static void unary_cuda(const T * x, T * dst, const int k, cudaStream_t stream) {
ggml_cuda_kernel_launch(unary_op_kernel<op, T>, launch_params, x, dst, k);
}
template <float (*op)(float), typename T>
static __global__ void unary_op_kernel_strided(const T * x, T * dst, const int k,const int64_t ne00,const int64_t ne01,const int64_t ne02,const size_t nb00,const size_t nb01,const size_t nb02,const size_t nb03) {
ggml_cuda_pdl_lc();
const int i = blockDim.x*blockIdx.x + threadIdx.x;
if (i >= k) {
return;
}
int64_t rem = i;
const int64_t i0 = rem % ne00; rem /= ne00;
const int64_t i1 = rem % ne01; rem /= ne01;
const int64_t i2 = rem % ne02;
const int64_t i3 = rem / ne02;
const size_t src_byte_offset = i0 * nb00 + i1 * nb01 + i2 * nb02 + i3 * nb03;
const T * src_ptr = (const T *)((const char *)x + src_byte_offset);
ggml_cuda_pdl_sync();
dst[i] = ggml_cuda_cast<T>(op(ggml_cuda_cast<float>(*src_ptr)));
}
template <float (*op)(float), typename T>
static void unary_cuda_strided(const T * x, T * dst, const int k,const int64_t ne00,const int64_t ne01,const int64_t ne02,const size_t nb00,const size_t nb01,const size_t nb02,const size_t nb03, cudaStream_t stream) {
const int num_blocks = (k + CUDA_NEG_BLOCK_SIZE - 1) / CUDA_NEG_BLOCK_SIZE;
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params((dim3)num_blocks, CUDA_NEG_BLOCK_SIZE, 0, stream);
ggml_cuda_kernel_launch(unary_op_kernel_strided<op, T>, launch_params, x, dst, k, ne00,ne01,ne02,nb00,nb01,nb02,nb03);
}
template <float (*op)(float)>
void ggml_cuda_op_unary(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0];
const void * src0_d = src0->data;
void * dst_d = dst->data;
cudaStream_t stream = ctx.stream();
GGML_ASSERT(ggml_is_contiguous(src0));
cudaStream_t stream = ctx.stream();
GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16);
GGML_ASSERT(src0->type == dst->type);
if (src0->type == GGML_TYPE_F16) {
unary_cuda<op>((const half *)src0_d, (half *)dst_d, ggml_nelements(src0), stream);
} else if (src0->type == GGML_TYPE_BF16) {
unary_cuda<op>((const nv_bfloat16 *)src0_d, (nv_bfloat16 *)dst_d, ggml_nelements(src0), stream);
if (ggml_is_contiguous(src0)) {
if (src0->type == GGML_TYPE_F16) {
unary_cuda<op>((const half *)src0_d, (half *)dst_d, ggml_nelements(src0), stream);
} else if (src0->type == GGML_TYPE_BF16) {
unary_cuda<op>((const nv_bfloat16 *)src0_d, (nv_bfloat16 *)dst_d, ggml_nelements(src0), stream);
} else {
unary_cuda<op>((const float *)src0_d, (float *)dst_d, ggml_nelements(src0), stream);
}
} else {
unary_cuda<op>((const float *)src0_d, (float *)dst_d, ggml_nelements(src0), stream);
if (src0->type == GGML_TYPE_F16) {
unary_cuda_strided<op>((const half *)src0_d, (half *)dst_d, ggml_nelements(src0),
src0->ne[0], src0->ne[1], src0->ne[2],
src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3], stream);
} else if (src0->type == GGML_TYPE_BF16) {
unary_cuda_strided<op>((const nv_bfloat16 *)src0_d, (nv_bfloat16 *)dst_d, ggml_nelements(src0),
src0->ne[0], src0->ne[1], src0->ne[2],
src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3], stream);
} else {
unary_cuda_strided<op>((const float *)src0_d, (float *)dst_d, ggml_nelements(src0),
src0->ne[0], src0->ne[1], src0->ne[2],
src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3], stream);
}
}
}

View file

@ -80,6 +80,22 @@ static bool ggml_metal_mul_mv_mma_type_supported(enum ggml_type type) {
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
case GGML_TYPE_Q6_K:
case GGML_TYPE_BF16:
case GGML_TYPE_Q1_0:
case GGML_TYPE_Q2_0:
case GGML_TYPE_MXFP4:
case GGML_TYPE_Q2_K:
case GGML_TYPE_Q3_K:
case GGML_TYPE_IQ2_XXS:
case GGML_TYPE_IQ2_XS:
case GGML_TYPE_IQ2_S:
case GGML_TYPE_IQ3_XXS:
case GGML_TYPE_IQ3_S:
case GGML_TYPE_IQ1_S:
case GGML_TYPE_IQ1_M:
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_IQ4_XS:
case GGML_TYPE_TQ2_0:
return true;
default:
return false;
@ -105,10 +121,18 @@ static int64_t ggml_metal_mul_mv_mma_rows_min(enum ggml_type type) {
switch (type) {
case GGML_TYPE_F32:
return 6;
case GGML_TYPE_TQ2_0:
return 5;
case GGML_TYPE_BF16:
return 4;
case GGML_TYPE_F16:
case GGML_TYPE_Q2_0:
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_Q2_K:
case GGML_TYPE_IQ4_NL:
case GGML_TYPE_MXFP4:
return 3;
default:
return 2;

View file

@ -2476,7 +2476,8 @@ static int ggml_metal_op_mul_mat_mma(ggml_metal_op_t ctx, int idx) {
if (fuse_add) {
dst = ctx->node(idx + n_fuse - 1);
res = dst->src[0]->op == GGML_OP_MUL_MAT ? dst->src[1] : dst->src[0];
// the residual is the other operand of the ADD, by identity: it can itself be a MUL_MAT output
res = dst->src[0] == op ? dst->src[1] : dst->src[0];
}
auto pipeline = ggml_metal_library_get_pipeline_mul_mv_mma_auto(lib, op, fuse_add);

View file

@ -605,4 +605,23 @@ MUL_MV_MMA_GEN("q5_1", block_q5_1, 2, dequantize_q5_1)
MUL_MV_MMA_GEN("q4_K", block_q4_K, QK_NL, dequantize_q4_K)
MUL_MV_MMA_GEN("q6_K", block_q6_K, QK_NL, dequantize_q6_K)
#if defined(GGML_METAL_HAS_BF16)
MUL_MV_MMA_GEN("bf16", bfloat4x4, 1, dequantize_bf16)
#endif
MUL_MV_MMA_GEN("q1_0", block_q1_0, 8, dequantize_q1_0)
MUL_MV_MMA_GEN("q2_0", block_q2_0, 4, dequantize_q2_0)
MUL_MV_MMA_GEN("mxfp4", block_mxfp4, 2, dequantize_mxfp4)
MUL_MV_MMA_GEN("q2_K", block_q2_K, QK_NL, dequantize_q2_K)
MUL_MV_MMA_GEN("q3_K", block_q3_K, QK_NL, dequantize_q3_K)
MUL_MV_MMA_GEN("iq2_xxs", block_iq2_xxs, QK_NL, dequantize_iq2_xxs)
MUL_MV_MMA_GEN("iq2_xs", block_iq2_xs, QK_NL, dequantize_iq2_xs)
MUL_MV_MMA_GEN("iq2_s", block_iq2_s, QK_NL, dequantize_iq2_s)
MUL_MV_MMA_GEN("iq3_xxs", block_iq3_xxs, QK_NL, dequantize_iq3_xxs)
MUL_MV_MMA_GEN("iq3_s", block_iq3_s, QK_NL, dequantize_iq3_s)
MUL_MV_MMA_GEN("iq1_s", block_iq1_s, QK_NL, dequantize_iq1_s)
MUL_MV_MMA_GEN("iq1_m", block_iq1_m, QK_NL, dequantize_iq1_m)
MUL_MV_MMA_GEN("iq4_nl", block_iq4_nl, 2, dequantize_iq4_nl)
MUL_MV_MMA_GEN("iq4_xs", block_iq4_xs, QK_NL, dequantize_iq4_xs)
MUL_MV_MMA_GEN("tq2_0", block_tq2_0, QK_NL, dequantize_tq2_0)
#undef MUL_MV_MMA_GEN

View file

@ -8186,10 +8186,14 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const
// cm2 dense is fast, so it needs a larger reduction to win.
// With quantized K/V, sparse only breaks even around 16x (measured on RDNA3/RDNA4).
const int64_t min_ratio = tuning_params.path == FA_COOPMAT2 ? 4 : (kv_f16 ? 2 : 16);
// coopmat2 vector decode requires 8B strides.
auto sparse_gather_aligned = [](const ggml_tensor * t) {
return (t->type != GGML_TYPE_F16 && t->type != GGML_TYPE_BF16) ||
(t->nb[1] | t->nb[2] | t->nb[3]) % (4 * sizeof(ggml_fp16_t)) == 0;
};
const bool use_sparse = !disable_sparse && n_kv_max > 0 && mask &&
max_bias == 0.0f && logit_softcap == 0.0f &&
// the cm2 sparse gather only reads f16
(kv_f16 || tuning_params.path != FA_COOPMAT2) &&
(tuning_params.path != FA_COOPMAT2 || (sparse_gather_aligned(k) && sparse_gather_aligned(v))) &&
nem0 == KV &&
(int64_t)KV >= std::max<int64_t>(4096, min_ratio * (int64_t)n_kv_max) &&
(gqa_ratio > 1 || (tuning_params.path == FA_SCALAR && N == 1));

View file

@ -18,7 +18,8 @@
#ifdef GL_NV_cooperative_matrix_decode_vector
#extension GL_NV_cooperative_matrix_decode_vector : enable
#endif
#extension GL_EXT_buffer_reference : enable
#extension GL_EXT_buffer_reference2 : enable
#extension GL_EXT_shader_explicit_arithmetic_types_int64 : enable
#extension GL_KHR_shader_subgroup_ballot : enable
#extension GL_KHR_shader_subgroup_vote : enable
#extension GL_EXT_null_initializer : enable
@ -35,6 +36,10 @@
#define FA_GATHER_BS 1u
#endif
layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_Byte {
uint8_t raw;
};
// buffer_reference stride = sizeof(struct) = FaBlockBytesK/V.
layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_K {
uint8_t raw[FaBlockBytesK];
@ -113,48 +118,71 @@ layout (binding = 1) readonly buffer K {uint8_t data_k[];};
layout (binding = 2) readonly buffer V {uint8_t data_v[];};
layout (binding = 3) readonly buffer M {uint8_t data_m[];};
// f16 aliases for the sparse gather callbacks.
layout (binding = 1) readonly buffer KF16 {float16_t data_kf16[];};
layout (binding = 2) readonly buffer VF16 {float16_t data_vf16[];};
// Native 16-bit aliases for the sparse gather callbacks.
layout (binding = 1) readonly buffer K16 {FLOAT_TYPE data_k16[];};
layout (binding = 2) readonly buffer V16 {FLOAT_TYPE data_v16[];};
layout (binding = 3) readonly buffer MF16 {float16_t data_mf16[];};
#ifdef GL_NV_cooperative_matrix_decode_vector
layout (binding = 1) readonly buffer KF16V4 {f16vec4 data_kf16v4[];};
layout (binding = 2) readonly buffer VF16V4 {f16vec4 data_vf16v4[];};
layout (binding = 1) readonly buffer K16V4 {FLOAT_TYPEV4 data_k16v4[];};
layout (binding = 2) readonly buffer V16V4 {FLOAT_TYPEV4 data_v16v4[];};
#endif
// K/V/mask f16-element offsets for the current head/batch, set in main().
// K/V/mask offsets in 16-bit elements for the current head/batch, set in main().
uint32_t g_k_off_elem, g_v_off_elem, g_m_off_elem;
#if !defined(BFLOAT16)
// blockCoords are in block units: KV slot = blockCoords[0],
// head dim = blockCoords[1]*FA_GATHER_BS + coordInBlock[1].
float16_t faGatherK(const decodeBufFA_K unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return float16_t(0); }
FLOAT_TYPE faGatherK(const decodeBufFA_K bl_in, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return FLOAT_TYPE(0.0); }
const int r = data_sparse[sparse_base + blockCoords[0]];
return r < 0 ? float16_t(0) : data_kf16[g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
if (r < 0) { return FLOAT_TYPE(0.0); }
#if !defined(BFLOAT16)
if (USE_DECODE_K) {
decodeBufFA_K block = decodeBufFA_K(decodeBufFA_Byte(bl_in) + uint64_t(uint(r) - blockCoords[0]) * k_stride * FaBlockBytesK);
return faDecodeK(block, blockCoords, coordInBlock);
}
#endif
return data_k16[g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
}
float16_t faGatherV(const decodeBufFA_V unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return float16_t(0); }
FLOAT_TYPE faGatherV(const decodeBufFA_V bl_in, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return FLOAT_TYPE(0.0); }
const int r = data_sparse[sparse_base + blockCoords[0]];
return r < 0 ? float16_t(0) : data_vf16[g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
if (r < 0) { return FLOAT_TYPE(0.0); }
#if !defined(BFLOAT16)
if (USE_DECODE_V) {
decodeBufFA_V block = decodeBufFA_V(decodeBufFA_Byte(bl_in) + uint64_t(uint(r) - blockCoords[0]) * v_stride * FaBlockBytesV);
return faDecodeV(block, blockCoords, coordInBlock);
}
#endif
return data_v16[g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]];
}
#ifdef GL_NV_cooperative_matrix_decode_vector
f16vec4 faGatherKVector(const decodeBufFA_K unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return f16vec4(0); }
FLOAT_TYPEV4 faGatherKVector(const decodeBufFA_K bl_in, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return FLOAT_TYPEV4(0.0); }
const int r = data_sparse[sparse_base + blockCoords[0]];
if (r < 0) { return f16vec4(0); }
if (r < 0) { return FLOAT_TYPEV4(0.0); }
#if !defined(BFLOAT16)
if (USE_DECODE_K) {
decodeBufFA_K block = decodeBufFA_K(decodeBufFA_Byte(bl_in) + uint64_t(uint(r) - blockCoords[0]) * k_stride * FaBlockBytesK);
return faDecodeKVector(block, blockCoords, coordInBlock);
}
#endif
const uint32_t o = g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1];
return data_kf16v4[o / 4];
return data_k16v4[o / 4];
}
f16vec4 faGatherVVector(const decodeBufFA_V unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return f16vec4(0); }
FLOAT_TYPEV4 faGatherVVector(const decodeBufFA_V bl_in, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) {
if (blockCoords[0] >= p.split_kv) { return FLOAT_TYPEV4(0.0); }
const int r = data_sparse[sparse_base + blockCoords[0]];
if (r < 0) { return f16vec4(0); }
if (r < 0) { return FLOAT_TYPEV4(0.0); }
#if !defined(BFLOAT16)
if (USE_DECODE_V) {
decodeBufFA_V block = decodeBufFA_V(decodeBufFA_Byte(bl_in) + uint64_t(uint(r) - blockCoords[0]) * v_stride * FaBlockBytesV);
return faDecodeVVector(block, blockCoords, coordInBlock);
}
#endif
const uint32_t o = g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1];
return data_vf16v4[o / 4];
return data_v16v4[o / 4];
}
#define FAGATHERK , faGatherK, faGatherKVector
@ -163,7 +191,6 @@ f16vec4 faGatherVVector(const decodeBufFA_V unused, const uint32_t blockCoords[2
#define FAGATHERK , faGatherK
#define FAGATHERV , faGatherV
#endif
#endif
// Add gathered mask to S (slope==1 since sparse requires max_bias==0). col = slot in block jblk.
ACC_TYPE faAddSparseMask(const uint32_t row, const uint32_t col, const ACC_TYPE elem, const uint32_t jblk) {
@ -252,8 +279,8 @@ void main() {
tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0);
const uint bs_k = USE_SPARSE ? FA_GATHER_BS : fa_block_elems(FaTypeK);
const uint bs_v = USE_SPARSE ? FA_GATHER_BS : fa_block_elems(FaTypeV);
const uint bs_k = USE_SPARSE ? max(FA_GATHER_BS, BLOCK_SIZE_K) : BLOCK_SIZE_K;
const uint bs_v = USE_SPARSE ? max(FA_GATHER_BS, BLOCK_SIZE_V) : BLOCK_SIZE_V;
tensorLayoutK = setTensorLayoutBlockSizeNV(tensorLayoutK, 1, bs_k);
tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, bs_v);
@ -384,18 +411,15 @@ void main() {
uint32_t k_offset = ik2*p.nb12 + ik3*p.nb13;
// F16: bs_k==1 (direct load). F32: bs_k==4 (vec4 / dequantFuncF32). Quantized types: bs_k==32.
#if defined(BFLOAT16)
coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose);
#else
const bool k_use_decode = (bs_k > 1u);
if (USE_SPARSE) {
coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FAGATHERK);
} else if (k_use_decode) {
#if !defined(BFLOAT16)
} else if (USE_DECODE_K) {
coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FADECODEK);
#endif
} else {
coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose);
}
#endif
S = coopMatMulAdd(Qf16, K_T, S);
if (LOGIT_SOFTCAP) {
@ -458,18 +482,15 @@ void main() {
coopmat<FLOAT_TYPE, gl_ScopeWorkgroup, Bc, HSV_pad, gl_MatrixUseB> V;
uint32_t v_offset = iv2*p.nb22 + iv3*p.nb23;
#if defined(BFLOAT16)
coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad));
#else
const bool v_use_decode = (bs_v > 1u);
if (USE_SPARSE) {
coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FAGATHERV);
} else if (v_use_decode) {
#if !defined(BFLOAT16)
} else if (USE_DECODE_V) {
coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FADECODEV);
#endif
} else {
coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad));
}
#endif
L = eM*L + rowsum;

View file

@ -60,7 +60,13 @@ void topk(const uint row) {
if (gl_GlobalInvocationID.x < p.ncols_input) {
if (p.first_pass != 0) {
const uint row_offset = row * p.ncols_input;
dst_row[tid] = ivec2(gl_GlobalInvocationID.x, floatBitsToInt(data_a[row_offset + gl_GlobalInvocationID.x]));
// NaN ranks lowest, like -inf, so that every value has a place in
// the ordering the search below counts
float a = float(data_a[row_offset + gl_GlobalInvocationID.x]);
if (isnan(a)) {
a = uintBitsToFloat(0xFF800000);
}
dst_row[tid] = ivec2(gl_GlobalInvocationID.x, floatBitsToInt(a));
} else {
const uint row_offset = row * p.ncols_input;
dst_row[tid] = data_s[row_offset + gl_GlobalInvocationID.x];
@ -76,8 +82,10 @@ void topk(const uint row) {
if (tid < s) {
ivec2 a = dst_row[tid];
ivec2 b = dst_row[tid + s];
// compare as floats: the bit patterns of negative values
// order the other way as integers
if (a.x >= p.orig_ncols ||
b.x < p.orig_ncols && b.y > a.y) {
b.x < p.orig_ncols && intBitsToFloat(b.y) > intBitsToFloat(a.y)) {
dst_row[tid] = b;
}
}
@ -95,9 +103,11 @@ void topk(const uint row) {
int shift = 32 - SUBGROUP_SIZE_LOG2;
uint mask = ((1 << SUBGROUP_SIZE_LOG2) - 1) << shift;
// The current range.
// The current range, [range_min, range_max). It starts as every value
// (+inf maps to 0xFF800000 and NaN was replaced by -inf), so the
// buckets always hold at least limit values.
uint range_min = 0;
uint range_max = 0xFF800000;
uint range_max = 0xFFFFFFFF;
// How many are above the current range, and how many we need to find.
uint total = 0;
uint limit = min(p.k, p.ncols_input - gl_WorkGroupID.x * BLOCK_SIZE);
@ -138,8 +148,12 @@ void topk(const uint row) {
total = sh_total;
// Update the range, and break if we've found the K-th largest.
range_max = range_min + ((min_idx + 1) << shift);
range_min = range_min + (min_idx << shift);
// The end of the top bucket wraps past 2^32, clamp it instead.
range_min = range_min + (uint(min_idx) << shift);
range_max = range_min + (1u << shift);
if (range_max < range_min) {
range_max = 0xFFFFFFFF;
}
if (total == p.k) {
break;

View file

@ -318,6 +318,7 @@ class Keys:
class Classifier:
OUTPUT_LABELS = "{arch}.classifier.output_labels"
POOLING_TYPE = "{arch}.classifier.pooling_type"
ACTIVATION = "{arch}.classifier.activation"
class ShortConv:
L_CACHE = "{arch}.shortconv.l_cache"
@ -408,6 +409,7 @@ class Keys:
BLOCK_COUNT = "clip.vision.block_count"
IMAGE_MEAN = "clip.vision.image_mean"
IMAGE_STD = "clip.vision.image_std"
IMAGE_RESIZE_ALGO = "clip.vision.image_resize_algo"
SPATIAL_MERGE_SIZE = "clip.vision.spatial_merge_size"
SWIGLU_CLAMP = "clip.vision.swiglu_clamp"
EXPERT_COUNT_PER_LAYER = "clip.vision.expert_count_per_layer" # dots3note pyramid MoE, 0 = dense layer
@ -5171,6 +5173,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.ATTN_OUT,
MODEL_TENSOR.OUTPUT,
MODEL_TENSOR.DENSE_2_OUT, # LFM2-ColBert-350M
MODEL_TENSOR.TOKEN_TYPES, # decision head
MODEL_TENSOR.CLS,
MODEL_TENSOR.CLS_NORM,
MODEL_TENSOR.CLS_OUT,
],
MODEL_ARCH.LFM2MOE: [
MODEL_TENSOR.TOKEN_EMBD,
@ -6070,6 +6076,8 @@ class DecisionType:
NIMBLE = "nimble" # same as openjev, the prompt lists all the questions of the request
CLEF = "clef" # joint head over all questions, one score per option
PPLX_DECIDER = "pplx-decider" # same as openjev, label codes of 1 or 2 letters
LFM2_D1 = "lfm2-d1" # same as openjev, the labels depend on the question type
LFM2_D1_OMNI = "lfm2-d1-omni" # same head as laya on a bidirectional LFM2 trunk, other prompt layout
class VisionProjectorType:
@ -6130,6 +6138,9 @@ class VisionProjectorType:
MIMO_AUDIO = "mimo_audio"
GRANITE4_VISION = "granite4_vision"
MUSE_GLIMMER = "muse-glimmer"
COHERE2V = "cohere2v"
D1OMNI_V = "d1omni_v" # lfm2 vision, without separator tokens
D1OMNI_A = "d1omni_a" # lfm2a audio, with a residual block after the projector
# Items here are (block size, type size)

View file

@ -1340,6 +1340,9 @@ class GGUFWriter:
def add_classifier_pooling_type(self, value: PoolingType) -> None:
self.add_uint32(Keys.Classifier.POOLING_TYPE.format(arch=self.arch), value.value)
def add_classifier_activation(self, value: str) -> None:
self.add_string(Keys.Classifier.ACTIVATION.format(arch=self.arch), value)
def add_decision_type(self, value: str) -> None:
self.add_string(Keys.Decision.TYPE.format(arch=self.arch), value)
@ -1432,6 +1435,9 @@ class GGUFWriter:
def add_vision_image_mean(self, values: Sequence[float]) -> None:
self.add_array(Keys.ClipVision.IMAGE_MEAN, values)
def add_vision_image_resize_algo(self, value: str) -> None:
self.add_string(Keys.ClipVision.IMAGE_RESIZE_ALGO, value)
def add_vision_image_std(self, values: Sequence[float]) -> None:
self.add_array(Keys.ClipVision.IMAGE_STD, values)

View file

@ -255,6 +255,7 @@ class TensorNameMap:
"model.layers.{bid}.self_attn.language_expert_query_key_value", # cogvlm
"model.layers.{bid}.linear_attn.in_proj_qkv", # qwen3.5
"head.layers.{bid}.self_attn.in_proj", # laya
"layers.{bid}.linear_attn.in_proj_qkv", # qwen3.5 text
),
# Attention query
@ -397,6 +398,7 @@ class TensorNameMap:
MODEL_TENSOR.ATTN_GATE: (
"model.layers.{bid}.self_attn.gate_proj", # afmoe muse-glimmer
"model.layers.{bid}.linear_attn.in_proj_z", # qwen3.5
"layers.{bid}.linear_attn.in_proj_z", # qwen3.5 text
"model.layers.{bid}.self_attn.g_proj", # step3.5 head-wise attention gate
"model.layers.{bid}.self_attn.output_gate", # minimax-01
"model.layers.{bid}.self_attn.linear_gate", # hy-v4
@ -842,6 +844,7 @@ class TensorNameMap:
"model.layers.{bid}.mamba.conv1d", # jamba falcon-h1 granite-hybrid
"model.layers.layers.{bid}.mixer.conv1d", # plamo2
"model.layers.{bid}.linear_attn.conv1d", # qwen3next
"layers.{bid}.linear_attn.conv1d", # qwen3.5 text
),
MODEL_TENSOR.SSM_X: (
@ -857,6 +860,7 @@ class TensorNameMap:
"model.layers.{bid}.mamba.dt_proj", # jamba falcon-h1 granite-hybrid
"model.layers.layers.{bid}.mixer.dt_proj", # plamo2
"model.layers.{bid}.linear_attn.dt_proj", # qwen3next
"layers.{bid}.linear_attn.dt_proj", # qwen3.5 text
"backbone.layers.{bid}.mixer.dt", # nemotron-h-moe
"model.layers.{bid}.self_attn.dt_proj", # kimi
"model.layers.{bid}.attention.dt_proj", # bailingmoe3
@ -873,6 +877,7 @@ class TensorNameMap:
"model.layers.{bid}.mamba.A_log", # jamba falcon-h1 granite-hybrid
"model.layers.layers.{bid}.mixer.A_log", # plamo2
"model.layers.{bid}.linear_attn.A_log", # qwen3next
"layers.{bid}.linear_attn.A_log", # qwen3.5 text
"model.layers.{bid}.self_attn.A_log", # kimi
"model.layers.{bid}.attention.A_log", # bailingmoe3
),
@ -899,6 +904,7 @@ class TensorNameMap:
MODEL_TENSOR.SSM_NORM: (
"model.layers.{bid}.mamba.norm", # falcon-h1 granite-hybrid
"model.layers.{bid}.linear_attn.norm", # qwen3next
"layers.{bid}.linear_attn.norm", # qwen3.5 text
"backbone.layers.{bid}.mixer.norm", # mamba2
"model.layers.{bid}.self_attn.o_norm", # kimi
"model.layers.{bid}.attention.o_norm", # bailingmoe3
@ -909,11 +915,13 @@ class TensorNameMap:
"backbone.layers.{bid}.mixer.out_proj", # mamba
"model.layers.{bid}.mamba.out_proj", # jamba falcon-h1 granite-hybrid
"model.layers.{bid}.linear_attn.out_proj", # qwen3next
"layers.{bid}.linear_attn.out_proj", # qwen3.5 text
"model.layers.layers.{bid}.mixer.out_proj", # plamo2
),
MODEL_TENSOR.SSM_ALPHA: (
"model.layers.{bid}.linear_attn.in_proj_a", # qwen3.5
"layers.{bid}.linear_attn.in_proj_a", # qwen3.5 text
),
MODEL_TENSOR.SSM_BETA_ALPHA: (
@ -941,6 +949,7 @@ class TensorNameMap:
),
MODEL_TENSOR.SSM_BETA: (
"model.layers.{bid}.linear_attn.in_proj_b", # qwen3.5
"layers.{bid}.linear_attn.in_proj_b", # qwen3.5 text
"model.layers.{bid}.self_attn.b_proj", # Kimi Linear
"model.layers.{bid}.attention.b_proj", # bailingmoe3
),
@ -1622,6 +1631,7 @@ class TensorNameMap:
MODEL_TENSOR.V_MMPROJ: (
"aligner.w{bid}", # deepseek4v (w1 -> mm.1, w2 -> mm.2)
"multi_modal_projector.linear_{bid}",
"model.multi_modal_projector.linear_{bid}", # cohere2v
"mm_projector.proj.linear_{bid}", # Kimi-K2.5
"visual.merger.mlp.{bid}", # qwen2vl
"mlp_AR.linear_{bid}", # PaddleOCR-VL

View file

@ -399,6 +399,8 @@ extern "C" {
enum ggml_type type_k; // data type for K cache [EXPERIMENTAL]
enum ggml_type type_v; // data type for V cache [EXPERIMENTAL]
size_t moe_cache_size; // device cache in bytes for the experts kept in host memory, 0 = disabled [EXPERIMENTAL]
// Abort callback
// if it returns true, execution of llama_decode() will be aborted
// currently works only with CPU execution

View file

@ -6087,7 +6087,7 @@ std::string gpttype_parse_chat_tool_calls(const std::string & generated_text,
common_chat_msg parsed;
try
{
parsed = common_chat_parse(generated_text, is_partial, parser_params);
parsed = common_chat_parse(common_chat_input(generated_text), is_partial, parser_params);
}
catch(...)
{

View file

@ -370,6 +370,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_CLASSIFIER_OUTPUT_LABELS, "%s.classifier.output_labels" },
{ LLM_KV_CLASSIFIER_POOLING_TYPE, "%s.classifier.pooling_type" },
{ LLM_KV_CLASSIFIER_ACTIVATION, "%s.classifier.activation" },
{ LLM_KV_DECISION_BLOCK_COUNT, "%s.decision.block_count" },
{ LLM_KV_DECISION_ROUTING_BLOCK_COUNT, "%s.decision.routing_block_count" },

View file

@ -416,6 +416,7 @@ enum llm_kv {
LLM_KV_CLASSIFIER_OUTPUT_LABELS,
LLM_KV_CLASSIFIER_POOLING_TYPE,
LLM_KV_CLASSIFIER_ACTIVATION,
LLM_KV_DECISION_BLOCK_COUNT,
LLM_KV_DECISION_ROUTING_BLOCK_COUNT,

View file

@ -9,6 +9,7 @@
#include "llama-memory.h"
#include "llama-mmap.h"
#include "llama-model.h"
#include "llama-moe-cache.h"
#include "llama-ext.h"
#include "llama-sampler.h"
#include "llama.h"
@ -275,8 +276,9 @@ llama_context::llama_context(
}
}
cparams.op_offload = params.op_offload;
cparams.kv_unified = params.kv_unified;
cparams.op_offload = params.op_offload;
cparams.kv_unified = params.kv_unified;
cparams.moe_cache_size = params.moe_cache_size;
// initialized later
cparams.pipeline_parallel = false;
@ -470,6 +472,22 @@ llama_context::llama_context(
LLAMA_LOG_INFO("%s: pipeline parallelism enabled\n", __func__);
}
if (cparams.moe_cache_size > 0) {
if (cparams.pipeline_parallel || model.n_devices() > 1) {
throw std::runtime_error("MoE cache does not support multiple devices");
}
for (size_t i = 0; i < backend_ptrs.size(); ++i) {
const auto type = ggml_backend_dev_type(ggml_backend_get_device(backend_ptrs[i]));
if (type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU) {
moe_cache = std::make_unique<llama_moe_cache>(model, backend_ptrs[i], backend_buft[i], cparams.moe_cache_size);
break;
}
}
if (!moe_cache) {
throw std::runtime_error("MoE cache requires a GPU backend");
}
}
sched_reserve();
if (!cparams.flash_attn) {
@ -2615,6 +2633,7 @@ llm_graph_params llama_context::graph_params(
/*.loras =*/ loras.get(),
/*.mctx =*/ mctx,
/*.cross =*/ &cross,
/*.moe_cache =*/ moe_cache.get(),
/*.prec_policy =*/ &model.prec_policy,
/*.samplers =*/ sampling.samplers,
/*.n_outputs =*/ n_outputs,
@ -2655,7 +2674,14 @@ ggml_status llama_context::graph_compute(
}
bool llama_context::sched_copy_experts(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph, void * user_data) {
auto & st = static_cast<llama_context *>(user_data)->copy_experts;
auto * lctx = static_cast<llama_context *>(user_data);
// the slot maps of the MoE cache
if (lctx->moe_cache && lctx->moe_cache->copy(backend, src, dst, graph)) {
return true;
}
auto & st = lctx->copy_experts;
// the ids must be computed before the split starts, so only the first node of the split is considered
if (ggml_graph_n_nodes(graph) == 0) {
@ -2702,10 +2728,28 @@ bool llama_context::sched_copy_experts(ggml_backend_t backend, const ggml_tensor
last++;
}
// the experts in the MoE cache are copied from device memory, the others are uploaded
int64_t next = first;
for (int64_t e = first; e <= last && lctx->moe_cache; ) {
const int64_t n = lctx->moe_cache->copy_experts(backend, src, dst, e, last);
if (n == 0) {
e++;
continue;
}
if (next < e) {
ggml_backend_tensor_set_async(backend, dst, (const uint8_t *) src->data + next*expert_size, next*expert_size, (e - next)*expert_size);
}
e += n;
next = e;
}
// copy a bit extra to ensure there are no NaNs in the padding of the last expert, this is necessary for MMQ in the CUDA backend
const size_t offset = first*expert_size;
const size_t offset = next*expert_size;
const size_t padding = last < n_expert - 1 ? std::min<size_t>(expert_size, 512) : 0;
ggml_backend_tensor_set_async(backend, dst, (const uint8_t *) src->data + offset, offset, (last - first + 1)*expert_size + padding);
const size_t size = (last + 1 - next)*expert_size + padding;
if (size > 0) {
ggml_backend_tensor_set_async(backend, dst, (const uint8_t *) src->data + offset, offset, size);
}
first = last + 1;
}
@ -3572,6 +3616,11 @@ llama_memory_breakdown llama_context::memory_breakdown() const {
ret[buft].context += size;
}
}
if (moe_cache) {
for (const auto & [buft, size] : moe_cache->memory_breakdown()) {
ret[buft].context += size;
}
}
if (model.hparams.no_alloc) {
for (size_t i = 0; i < backends.size(); ++i) {
ggml_backend_t backend = backends[i].get();
@ -3861,6 +3910,7 @@ llama_context_params llama_context_default_params() {
/*.cb_eval_user_data =*/ nullptr,
/*.type_k =*/ GGML_TYPE_F16,
/*.type_v =*/ GGML_TYPE_F16,
/*.moe_cache_size =*/ 0,
/*.abort_callback =*/ nullptr,
/*.abort_callback_data =*/ nullptr,
/*.embeddings =*/ false,

View file

@ -7,6 +7,7 @@
#include "llama-adapter.h"
#include "llama-impl.h"
#include "llama-memory.h"
#include "llama-moe-cache.h"
#include "ggml-cpp.h"
#include "ggml-opt.h"
@ -17,6 +18,7 @@
struct llama_model;
class llama_batch_allocr;
class llama_moe_cache;
class llama_io_read_i;
class llama_io_write_i;
@ -271,7 +273,7 @@ private:
llm_graph_cb graph_get_cb() const;
// ggml_backend_sched copy callback, copies only the experts used by MUL_MAT_ID
// ggml_backend_sched copy callback, copies only the experts used by MUL_MAT_ID and updates the MoE cache
static bool sched_copy_experts(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph, void * user_data);
// disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device
@ -299,6 +301,7 @@ private:
llama_cross cross; // TODO: tmp for handling cross-attention - need something better probably
llama_memory_ptr memory;
llama_moe_cache_ptr moe_cache;
// decode output (2-dimensional array: [n_outputs][n_vocab])
buffer_view<float> logits = {nullptr, 0};

View file

@ -55,6 +55,8 @@ struct llama_cparams {
bool pipeline_parallel;
bool training; // set by llama_opt_init()
size_t moe_cache_size;
std::vector<bool> embeddings_layer_inp; // [n_layer()] extract input embeddings for layer
enum llama_context_type ctx_type;

View file

@ -2,6 +2,7 @@
#include "llama-impl.h"
#include "llama-model.h"
#include "llama-moe-cache.h"
#include "llama-batch.h"
#include "llama-cparams.h"
#include "llama-sampler.h"
@ -1524,6 +1525,7 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) :
loras (params.loras),
mctx (params.mctx),
cross (params.cross),
moe_cache (params.moe_cache),
prec_policy (params.prec_policy),
samplers (params.samplers),
cb_func (params.cb),
@ -1590,8 +1592,12 @@ ggml_tensor * llm_graph_context::build_lora_mm_id(
ggml_tensor * w, // ggml_tensor * as
ggml_tensor * cur, // ggml_tensor * b
ggml_tensor * ids,
ggml_tensor * w_s) const {
ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids);
ggml_tensor * w_s,
ggml_tensor * slots) const {
// the experts in the MoE cache are selected by their slots
ggml_tensor * res = slots == nullptr ?
ggml_mul_mat_id(ctx0, w, cur, ids) :
ggml_mul_mat_id(ctx0, moe_cache->get_experts(w), cur, slots);
if (prec_policy) {
prec_policy->apply(res);
@ -2206,6 +2212,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
//call early so that topk-moe can be used
ggml_build_forward_expand(gf, weights);
// the experts of host-resident layers may be read from the MoE cache
ggml_tensor * slots = build_moe_cache_slots(selected_experts, up_exps, gate_exps, down_exps, gate_up_exps, il);
cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens);
if (weight_before_ffn) {
@ -2220,7 +2229,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
if (gate_up_exps) {
// merged gate_up path: one mul_mat_id, then split into gate and up views
ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s); // [n_ff*2, n_expert_used, n_tokens]
ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s, slots); // [n_ff*2, n_expert_used, n_tokens]
cb(gate_up, "ffn_moe_gate_up", il);
if (up_exps_s) {
@ -2239,7 +2248,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
cb(up, "ffn_moe_up", il);
} else {
// separate gate and up path
up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s); // [n_ff, n_expert_used, n_tokens]
up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s, slots); // [n_ff, n_expert_used, n_tokens]
cb(up, "ffn_moe_up", il);
if (up_exps_s) {
@ -2252,7 +2261,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
}
if (gate_exps) {
cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s); // [n_ff, n_expert_used, n_tokens]
cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s, slots); // [n_ff, n_expert_used, n_tokens]
cb(cur, "ffn_moe_gate", il);
} else {
cur = up;
@ -2353,7 +2362,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
GGML_ABORT("fatal error");
}
experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s); // [n_embd, n_expert_used, n_tokens]
experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, slots); // [n_embd, n_expert_used, n_tokens]
if (arch == LLM_ARCH_MISTRAL4) {
// src1 can exceed F16 range
ggml_prec_set_src(experts, GGML_PREC_F32, 1);
@ -2410,6 +2419,45 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
return moe_out;
}
ggml_tensor * llm_graph_context::build_moe_cache_slots(
ggml_tensor * selected_experts,
ggml_tensor * up_exps,
ggml_tensor * gate_exps,
ggml_tensor * down_exps,
ggml_tensor * gate_up_exps,
int il) const {
if (moe_cache == nullptr) {
return nullptr;
}
ggml_tensor * slot_map = moe_cache->get_slot_map(il, selected_experts->ne[1], selected_experts->ne[0]);
if (slot_map == nullptr) {
return nullptr;
}
for (ggml_tensor * w : { up_exps, gate_exps, down_exps, gate_up_exps }) {
if (w != nullptr && moe_cache->get_experts(w) == nullptr) {
return nullptr;
}
}
ggml_tensor * ids = selected_experts;
if (!ggml_is_contiguous(ids)) {
ids = ggml_cont(ctx0, ids);
}
ids = ggml_reshape_1d(ctx0, ids, ggml_nelements(ids));
// the slot map is a host weight, so the scheduler starts a new split here and copies it with the copy callback
// the callback reads the selected experts, uploads the missing ones and updates the slot map
ggml_tensor * slots = ggml_get_rows(ctx0, slot_map, ids); // [1, n_expert_used*n_tokens]
if (!ggml_backend_supports_op(moe_cache->backend(), slots)) {
return nullptr;
}
ggml_backend_sched_set_tensor_backend(sched, slots, moe_cache->backend());
cb(slots, "ffn_moe_slots", il);
return ggml_reshape_2d(ctx0, slots, selected_experts->ne[0], selected_experts->ne[1]); // [n_expert_used, n_tokens]
}
// input embeddings with optional lora
ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd, float tok_scale) const {
const int64_t n_embd_inp = hparams.n_embd_inp();
@ -3856,11 +3904,7 @@ void llm_graph_context::build_pooling(
if (cls_b) {
cur = ggml_add(ctx0, cur, cls_b);
}
if (arch == LLM_ARCH_MODERN_BERT) {
cur = ggml_gelu(ctx0, cur);
} else {
cur = ggml_tanh(ctx0, cur);
}
cur = ggml_unary(ctx0, cur, hparams.act_cls);
if (cls_norm) {
// head norm
cur = build_norm(cur, cls_norm, NULL, LLM_NORM, -1);

View file

@ -21,6 +21,8 @@ struct llama_cparams;
struct llama_layer;
struct llama_prec_policy;
class llama_moe_cache;
struct llama_memory_context_i;
class llama_kv_cache_context;
@ -793,6 +795,7 @@ struct llm_graph_params {
const llama_adapter_loras * loras;
const llama_memory_context_i * mctx;
const llama_cross * cross;
const llama_moe_cache * moe_cache;
const llama_prec_policy * prec_policy = nullptr;
@ -1036,6 +1039,7 @@ struct llm_graph_context {
const llama_adapter_loras * loras;
const llama_memory_context_i * mctx;
const llama_cross * cross;
const llama_moe_cache * moe_cache;
const llama_prec_policy * prec_policy;
@ -1078,11 +1082,13 @@ struct llm_graph_context {
ggml_tensor * w_s = nullptr) const;
// do mat_mul_id, while optionally apply lora and per-expert scale
// if slots is set, the experts are read from the MoE cache at these slots (see build_moe_cache_slots)
ggml_tensor * build_lora_mm_id(
ggml_tensor * w, // ggml_tensor * as
ggml_tensor * cur, // ggml_tensor * b
ggml_tensor * ids,
ggml_tensor * w_s = nullptr) const;
ggml_tensor * w_s = nullptr,
ggml_tensor * slots = nullptr) const;
ggml_tensor * build_norm(
ggml_tensor * cur,
@ -1179,6 +1185,15 @@ struct llm_graph_context {
ggml_tensor * down_exps_s = nullptr,
ggml_tensor * selected_experts_in = nullptr) const;
// the slots of the selected experts in the MoE cache, nullptr if the experts of the layer are not read from the cache
ggml_tensor * build_moe_cache_slots(
ggml_tensor * selected_experts,
ggml_tensor * up_exps,
ggml_tensor * gate_exps,
ggml_tensor * down_exps,
ggml_tensor * gate_up_exps,
int il) const;
//
// inputs
//

View file

@ -371,6 +371,7 @@ struct llama_hparams {
// llm_ffn_op_type_from_string() in llama-model.cpp, mirroring how
// rope_scaling_type_train is handled.
enum llm_ffn_op_type llm_ffn_op;
enum ggml_unary_op act_cls = GGML_UNARY_OP_TANH; // activation of the classifier head (RANK)
// Step35: optional per-layer clamps for (Swi)GLU
std::array<float, LLAMA_MAX_LAYERS> swiglu_clamp_exp; // clamping for expert FFN

View file

@ -199,6 +199,10 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos
// partial rollback via per-token snapshot index (bounded by n_rs_seq)
if (0 < p0 && p0 <= cell.pos && p1 > cell.pos) {
// a cell shared by several sequences cannot move back for only one of them
if (cell.seq_id.size() > 1) {
return false;
}
// the filter kept no layer (e.g. an MTP draft context), so only the position moves back
if (is_empty()) {
cell.pos = p0 - 1;

View file

@ -1238,6 +1238,13 @@ static const std::map<std::string, llm_ffn_op_type> LLM_FFN_OP_TYPES_FROM_STRING
{ "reglu", LLM_FFN_REGLU },
};
// transformers names, "gelu" is the exact (erf) variant
static const std::map<std::string, ggml_unary_op> LLM_CLS_ACT_TYPES_FROM_STRING = {
{ "gelu", GGML_UNARY_OP_GELU_ERF },
{ "silu", GGML_UNARY_OP_SILU },
{ "tanh", GGML_UNARY_OP_TANH },
};
llm_ffn_op_type llm_ffn_op_type_from_string(const std::string & name, llm_ffn_op_type fallback) {
const auto it = LLM_FFN_OP_TYPES_FROM_STRING.find(name);
if (it != LLM_FFN_OP_TYPES_FROM_STRING.end()) {
@ -1496,6 +1503,12 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_CAUSAL, hparams.causal_attn, false);
ml.get_key(LLM_KV_POOLING_TYPE, hparams.pooling_type, false);
ml.get_key(LLM_KV_CLASSIFIER_POOLING_TYPE, hparams.pooling_type_cls, false);
std::string act_cls;
if (ml.get_key(LLM_KV_CLASSIFIER_ACTIVATION, act_cls, false)) {
const auto it = LLM_CLS_ACT_TYPES_FROM_STRING.find(act_cls);
GGML_ASSERT(it != LLM_CLS_ACT_TYPES_FROM_STRING.end() && "unsupported classifier activation");
hparams.act_cls = it->second;
}
ml.get_key(LLM_KV_BLOCK_COUNT, hparams.n_layer_all);
GGML_ASSERT(hparams.n_layer_all > 0 && hparams.n_layer_all <= LLAMA_MAX_LAYERS);
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
@ -2528,6 +2541,11 @@ ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int i
llama_memory_i * llama_model::create_memory(const llama_memory_params & params, const llama_cparams & cparams) const {
llama_memory_i * res;
// the non-causal LFM2 decision graph reads the whole prompt in one batch, nothing is kept
if (arch == LLM_ARCH_LFM2 && !hparams.causal_attn && hparams.n_layer_decision > 0) {
return nullptr;
}
switch (arch) {
// Models that need specific instantiation should be handled in the
// switch statement
@ -3621,6 +3639,27 @@ void llama_model_base::create_tensor_qkv(llama_layer & layer, int bid,
}
}
llama_model_base::nextn_flags_t llama_model_base::nextn_flags(llama_model_loader & ml, llm_tensor trunk_probe) const {
int trunk = 0;
int mtp = 0;
// a file without the first trunk layer is MTP-only, a file without the first NextN layer is trunk-only
if (hparams.n_layer_nextn > 0) {
if (ml.get_weight(tn(trunk_probe, "weight", 0).str().c_str()) == nullptr) {
trunk = TENSOR_NOT_REQUIRED;
}
if (ml.get_weight(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", hparams.n_layer()).str().c_str()) == nullptr) {
mtp = TENSOR_NOT_REQUIRED;
}
}
if (!ml.load_mtp) {
mtp |= TENSOR_SKIP;
}
return { trunk, mtp };
}
void llama_model_base::load_swa_pattern(llama_model_loader & ml, uint32_t n_pattern, bool dense_first) {
if (ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, false)) {
return;

View file

@ -849,6 +849,13 @@ struct llama_model_base : public llama_model {
int64_t n_embd_, int64_t n_embd_q_, int64_t n_embd_k_, int64_t n_embd_v_,
int flags);
// helper: tensor flags for a file that holds the trunk, the NextN layers, or both
struct nextn_flags_t {
int trunk; // TENSOR_NOT_REQUIRED when the file holds only the NextN layers
int mtp; // TENSOR_NOT_REQUIRED when the file holds only the trunk, TENSOR_SKIP when MTP is not loaded
};
nextn_flags_t nextn_flags(llama_model_loader & ml, llm_tensor trunk_probe = LLM_TENSOR_ATTN_NORM) const;
// helper: read the SWA pattern as one flag per layer, or as a period expanded by set_swa_pattern
void load_swa_pattern(llama_model_loader & ml, uint32_t n_pattern, bool dense_first = false);

550
src/llama-moe-cache.cpp Normal file
View file

@ -0,0 +1,550 @@
#include "llama-moe-cache.h"
#include "llama-impl.h"
#include "llama-model.h"
#include "ggml-cpp.h"
#include <algorithm>
#include <stdexcept>
#include <unordered_map>
#include <vector>
namespace {
// LRU of the experts of a group of layers, the slot of each expert is kept in the slot map of its layer
struct moe_cache_lru {
int32_t n_expert = 0;
int32_t n_slots = 0;
std::vector<int32_t *> slot_map; // [n_layer] data of the slot maps, -1 if the expert is not cached
std::vector<int32_t> key_of; // [n_slots] il*n_expert + expert, -1 if empty
// doubly linked list of the slots, head is the least recently used
std::vector<int32_t> prev;
std::vector<int32_t> next;
int32_t head = -1;
int32_t tail = -1;
std::vector<uint32_t> seen; // [n_expert]
uint32_t seen_gen = 0;
std::vector<int32_t> uniq;
void init(int32_t n_layer, int32_t n_expert, int32_t n_slots) {
this->n_expert = n_expert;
this->n_slots = n_slots;
slot_map.assign(n_layer, nullptr);
key_of.assign(n_slots, -1);
prev.resize(n_slots);
next.resize(n_slots);
for (int32_t s = 0; s < n_slots; ++s) {
prev[s] = s - 1;
next[s] = s + 1 < n_slots ? s + 1 : -1;
}
head = 0;
tail = n_slots - 1;
seen.assign(n_expert, 0);
}
// move slot s to the tail (most recently used)
void touch(int32_t s) {
if (s == tail) {
return;
}
if (prev[s] >= 0) {
next[prev[s]] = next[s];
} else {
head = next[s];
}
prev[next[s]] = prev[s];
prev[s] = tail;
next[s] = -1;
next[tail] = s;
tail = s;
}
struct fill {
int32_t expert;
int32_t slot;
};
// give a slot to each expert selected by ids in layer il, the misses evict the least recently used experts
// returns false if the ids select more distinct experts than there are slots
bool plan(int32_t il, const int32_t * ids, size_t n_ids, std::vector<fill> & fills, size_t & n_hit) {
fills.clear();
n_hit = 0;
if (++seen_gen == 0) {
std::fill(seen.begin(), seen.end(), 0);
seen_gen = 1;
}
uniq.clear();
for (size_t i = 0; i < n_ids; ++i) {
GGML_ASSERT(ids[i] >= 0 && ids[i] < n_expert);
if (seen[ids[i]] != seen_gen) {
seen[ids[i]] = seen_gen;
uniq.push_back(ids[i]);
}
}
if (uniq.size() > (size_t) n_slots) {
return false;
}
int32_t * slots = slot_map[il];
// hits go to the tail first, so the head can be evicted below
for (int32_t e : uniq) {
if (slots[e] >= 0) {
touch(slots[e]);
n_hit++;
}
}
// sorted misses usually get consecutive slots, so the uploads can be merged
std::sort(uniq.begin(), uniq.end());
for (int32_t e : uniq) {
if (slots[e] >= 0) {
continue;
}
const int32_t s = head;
if (key_of[s] >= 0) {
slot_map[key_of[s] / n_expert][key_of[s] % n_expert] = -1;
}
key_of[s] = il*n_expert + e;
slots[e] = s;
touch(s);
fills.push_back({ e, s });
}
return true;
}
};
// gate, up, down or gate_up, down
static std::vector<ggml_tensor *> llama_moe_cache_layer_experts(const llama_layer & layer) {
std::vector<ggml_tensor *> res;
for (ggml_tensor * t : { layer.ffn_gate_up_exps, layer.ffn_gate_exps, layer.ffn_up_exps, layer.ffn_down_exps }) {
if (t != nullptr) {
res.push_back(t);
}
}
return res;
}
static bool llama_moe_cache_same_layout(const std::vector<ggml_tensor *> & a, const std::vector<ggml_tensor *> & b) {
if (a.size() != b.size()) {
return false;
}
for (size_t i = 0; i < a.size(); ++i) {
if (a[i]->type != b[i]->type || !ggml_are_same_shape(a[i], b[i]) || a[i]->nb[2] != b[i]->nb[2]) {
return false;
}
}
return true;
}
static bool llama_moe_cache_is_host_weight(const ggml_tensor * t) {
return t->buffer != nullptr &&
ggml_backend_buffer_get_usage(t->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
ggml_backend_buffer_is_host(t->buffer);
}
}
struct llama_moe_cache::impl {
// layers with the same expert tensor layout share the banks and the LRU of a group
struct group {
std::vector<ggml_tensor *> ref; // expert tensors of the first layer
std::vector<int32_t> layers;
std::vector<ggml_tensor *> banks; // device storage of all slots, one per expert tensor
size_t host_bytes = 0;
int32_t n_slots = 0;
moe_cache_lru lru;
};
struct layer {
int32_t ig = -1; // -1 if the layer is not cached
ggml_tensor * slot_map = nullptr; // I32 [1, n_expert] in host memory
std::vector<ggml_tensor *> experts; // host expert tensors, in the order of the banks
};
struct binding {
int32_t il;
int32_t ip; // index of the bank
ggml_tensor * cached; // view of the bank used in place of the host experts
};
struct stats {
size_t hits = 0;
size_t misses = 0;
size_t bytes = 0;
};
static constexpr int64_t max_batch = 32;
ggml_backend_t backend;
int32_t n_expert_used;
stats stats_small; // up to 8 tokens per ubatch
stats stats_large;
stats stats_copy; // experts copied from the cache for large batches
std::vector<group> groups;
std::vector<layer> layers;
std::unordered_map<const ggml_tensor *, binding> bindings; // host experts -> cached experts
std::unordered_map<const ggml_tensor *, int32_t> layer_of; // slot map -> layer
std::vector<int32_t> ids;
std::vector<moe_cache_lru::fill> fills;
// banks and their views on the device
ggml_context_ptr ctx;
ggml_backend_buffer_ptr buf;
size_t buf_size = 0;
// slot maps in host memory
ggml_context_ptr ctx_host;
ggml_backend_buffer_ptr buf_host;
size_t buf_host_size = 0;
// views used by copy_experts
ggml_context_ptr ctx_views;
impl(const llama_model & model, ggml_backend_t backend, ggml_backend_buffer_type_t buft, size_t size) :
backend(backend), n_expert_used(model.hparams.n_expert_used_max()), layers(model.layers.size()) {
ggml_backend_dev_t dev = ggml_backend_get_device(backend);
const auto dev_type = ggml_backend_dev_type(dev);
if (dev_type != GGML_BACKEND_DEVICE_TYPE_GPU && dev_type != GGML_BACKEND_DEVICE_TYPE_IGPU) {
throw std::runtime_error("MoE cache requires a GPU backend");
}
if (model.split_mode() == LLAMA_SPLIT_MODE_TENSOR) {
throw std::runtime_error("MoE cache does not support tensor parallelism");
}
if (model.hparams.n_expert == 0 || n_expert_used == 0) {
throw std::runtime_error("MoE cache requires a MoE model");
}
// only cache layers that keep all of their experts in host memory
size_t host_bytes = 0;
for (size_t il = 0; il < model.layers.size(); ++il) {
auto experts = llama_moe_cache_layer_experts(model.layers[il]);
if (experts.empty() || model.dev_layer(il) != dev ||
!std::all_of(experts.begin(), experts.end(), llama_moe_cache_is_host_weight)) {
continue;
}
auto it = std::find_if(groups.begin(), groups.end(), [&](const group & g) { return llama_moe_cache_same_layout(g.ref, experts); });
if (it == groups.end()) {
groups.emplace_back();
it = groups.end() - 1;
it->ref = experts;
}
it->layers.push_back(il);
for (const ggml_tensor * t : experts) {
it->host_bytes += ggml_nbytes(t);
host_bytes += ggml_nbytes(t);
}
}
if (groups.empty()) {
LLAMA_LOG_WARN("%s: no layer has all of its experts in host memory, MoE cache is disabled\n", __func__);
return;
}
// one extra slot at the end, CUDA MMQ can read past the last expert
const size_t alignment = ggml_backend_buft_get_alignment(buft);
auto alloc_size = [&](const group & g, int32_t n_slots) {
size_t res = 0;
for (const ggml_tensor * t : g.ref) {
res += GGML_PAD(t->nb[2]*(n_slots + 1), alignment);
}
return res;
};
// split the budget by the size of the experts, so each group caches the same fraction of its experts
size_t n_tensors = 0;
size_t n_tensors_host = 0;
for (group & g : groups) {
const int32_t n_expert = g.ref[0]->ne[2];
const size_t budget = (size_t) ((double) size*g.host_bytes/host_bytes);
const int32_t max_slots = g.layers.size()*n_expert;
while (g.n_slots < max_slots && alloc_size(g, g.n_slots + 1) <= budget) {
g.n_slots++;
}
if (g.n_slots < n_expert_used) {
LLAMA_LOG_WARN("%s: MoE cache budget is too small for %zu layers, they are not cached\n", __func__, g.layers.size());
g.n_slots = 0;
continue;
}
g.lru.init(model.layers.size(), n_expert, g.n_slots);
n_tensors += g.ref.size()*(1 + g.layers.size());
n_tensors_host += g.layers.size();
}
if (n_tensors == 0) {
throw std::runtime_error("MoE cache is too small to hold the experts of one token");
}
auto init_ctx = [](size_t n_tensors) {
ggml_init_params params = {
/*.mem_size =*/ n_tensors*ggml_tensor_overhead(),
/*.mem_buffer =*/ nullptr,
/*.no_alloc =*/ true,
};
ggml_context_ptr res(ggml_init(params));
if (!res) {
throw std::runtime_error("failed to create the MoE cache context");
}
return res;
};
ctx = init_ctx(n_tensors);
ctx_host = init_ctx(n_tensors_host);
ctx_views = init_ctx(2);
ggml_backend_buffer_type_t buft_host = ggml_backend_cpu_buffer_type();
const size_t alignment_host = ggml_backend_buft_get_alignment(buft_host);
for (size_t ig = 0; ig < groups.size(); ++ig) {
group & g = groups[ig];
if (g.n_slots == 0) {
continue;
}
for (const ggml_tensor * t : g.ref) {
ggml_tensor * bank = ggml_new_tensor_3d(ctx.get(), t->type, t->ne[0], t->ne[1], g.n_slots + 1);
GGML_ASSERT(bank->nb[2] == t->nb[2]);
ggml_format_name(bank, "moe_cache.%zu.%s", ig, t->name);
g.banks.push_back(bank);
}
for (int32_t il : g.layers) {
layer & l = layers[il];
l.ig = (int32_t) ig;
l.experts = llama_moe_cache_layer_experts(model.layers[il]);
for (size_t ip = 0; ip < l.experts.size(); ++ip) {
ggml_tensor * bank = g.banks[ip];
ggml_tensor * cached = ggml_view_3d(ctx.get(), bank, bank->ne[0], bank->ne[1], g.n_slots, bank->nb[1], bank->nb[2], 0);
ggml_format_name(cached, "moe_cache.%s", l.experts[ip]->name);
bindings[l.experts[ip]] = { il, (int32_t) ip, cached };
}
l.slot_map = ggml_new_tensor_2d(ctx_host.get(), GGML_TYPE_I32, 1, g.ref[0]->ne[2]);
ggml_format_name(l.slot_map, "moe_cache.slot_map-%d", il);
layer_of[l.slot_map] = il;
buf_host_size += GGML_PAD(ggml_nbytes(l.slot_map), alignment_host);
}
buf_size += alloc_size(g, g.n_slots);
}
if (model.hparams.no_alloc) {
// only used to measure the memory use, see llama_context::memory_breakdown
buf.reset(ggml_backend_buft_alloc_buffer(buft, 0));
buf_host.reset(ggml_backend_buft_alloc_buffer(buft_host, 0));
for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {
t->buffer = buf.get();
}
for (ggml_tensor * t = ggml_get_first_tensor(ctx_host.get()); t != nullptr; t = ggml_get_next_tensor(ctx_host.get(), t)) {
t->buffer = buf_host.get();
}
} else {
buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft));
buf_host.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx_host.get(), buft_host));
if (!buf || !buf_host) {
throw std::runtime_error("failed to allocate the MoE cache buffers");
}
ggml_backend_buffer_clear(buf.get(), 0);
ggml_backend_buffer_clear(buf_host.get(), 0xff); // all slots are -1
buf_size = ggml_backend_buffer_get_size(buf.get());
buf_host_size = ggml_backend_buffer_get_size(buf_host.get());
for (group & g : groups) {
for (int32_t il : g.layers) {
if (layers[il].slot_map != nullptr) {
g.lru.slot_map[il] = (int32_t *) layers[il].slot_map->data;
}
}
}
}
// as weights, the ops that read the banks run on the device and the slot maps are copied with the copy callback
ggml_backend_buffer_set_usage(buf.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
ggml_backend_buffer_set_usage(buf_host.get(), GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
LLAMA_LOG_INFO("%s: %10s MoE cache size = %8.2f MiB for %.2f MiB of host experts\n", __func__,
ggml_backend_buft_name(buft), buf_size/1024.0/1024.0, host_bytes/1024.0/1024.0);
for (const group & g : groups) {
LLAMA_LOG_INFO("%s: %2zu layers, %s: %5d slots (%.1f%%)\n", __func__,
g.layers.size(), ggml_type_name(g.ref.back()->type), g.n_slots, 100.0*g.n_slots/(g.layers.size()*g.ref[0]->ne[2]));
}
}
~impl() {
log_stats();
}
ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const {
if (il < 0 || il >= (int32_t) layers.size() || layers[il].ig < 0) {
return nullptr;
}
const layer & l = layers[il];
// large batches use most experts of a layer, so they gain little from the cache and would evict the experts used in generation
if (n_tokens == 0 || n_tokens > max_batch || std::min(n_tokens*n_expert_used, l.slot_map->ne[1]) > groups[l.ig].n_slots) {
return nullptr;
}
return l.slot_map;
}
ggml_tensor * get_experts(const ggml_tensor * w) const {
const auto it = bindings.find(w);
return it != bindings.end() ? it->second.cached : nullptr;
}
int64_t copy_experts(ggml_backend_t backend, const ggml_tensor * w, ggml_tensor * dst, int64_t e, int64_t last) {
const auto it = bindings.find(w);
if (it == bindings.end() || backend != this->backend) {
return 0;
}
const binding & b = it->second;
const group & g = groups[layers[b.il].ig];
// large batches only read the cache, so the experts used in generation stay in it
const int32_t * slots = g.lru.slot_map[b.il];
if (slots == nullptr || slots[e] < 0) {
return 0;
}
int64_t n = 1;
while (e + n <= last && slots[e + n] == slots[e] + n) {
n++;
}
ggml_tensor * bank = g.banks[b.ip];
ggml_reset(ctx_views.get());
ggml_tensor * src_view = ggml_view_3d(ctx_views.get(), bank, bank->ne[0], bank->ne[1], n, bank->nb[1], bank->nb[2], slots[e]*bank->nb[2]);
ggml_tensor * dst_view = ggml_view_3d(ctx_views.get(), dst, dst->ne[0], dst->ne[1], n, dst->nb[1], dst->nb[2], e*dst->nb[2]);
ggml_backend_view_init(src_view);
ggml_backend_view_init(dst_view);
ggml_backend_tensor_copy_async(backend, backend, src_view, dst_view);
stats_copy.hits += n;
stats_copy.bytes += ggml_nbytes(src_view);
return n;
}
bool copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph) {
const auto it = layer_of.find(src);
if (it == layer_of.end()) {
return false;
}
const int32_t il = it->second;
const layer & l = layers[il];
group & g = groups[l.ig];
GGML_ASSERT(backend == this->backend);
// the get_rows that looks up the slots of the selected experts
const int n_nodes = ggml_graph_n_nodes(graph);
const ggml_tensor * lookup = nullptr;
for (int i = 0; i < n_nodes && lookup == nullptr; ++i) {
const ggml_tensor * node = ggml_graph_node(graph, i);
if (node->op == GGML_OP_GET_ROWS && node->src[0] == dst) {
lookup = node;
}
}
GGML_ASSERT(lookup != nullptr);
// the selected experts must be computed in an earlier split
// the scheduler starts a new split at the lookup because it reads a host weight, but only if the split already has inputs
const ggml_tensor * sel = lookup->src[1];
for (int i = 0; i < n_nodes; ++i) {
const ggml_tensor * node = ggml_graph_node(graph, i);
if (node == sel || node == sel->view_src) {
GGML_ABORT("the experts of layer %d are selected in the same split as their MoE cache lookup", il);
}
}
GGML_ASSERT(ggml_is_contiguous(sel));
ids.resize(ggml_nelements(sel));
ggml_backend_tensor_get_async(backend, sel, ids.data(), 0, ggml_nbytes(sel));
ggml_backend_synchronize(backend);
size_t n_hit = 0;
if (!g.lru.plan(il, ids.data(), ids.size(), fills, n_hit)) {
GGML_ABORT("the MoE cache is too small for the experts selected in layer %d", il);
}
// upload the missing experts, consecutive experts going to consecutive slots are uploaded together
size_t bytes = 0;
for (size_t ip = 0; ip < l.experts.size(); ++ip) {
const ggml_tensor * w = l.experts[ip];
ggml_tensor * bank = g.banks[ip];
const size_t expert_size = w->nb[2];
for (size_t i = 0; i < fills.size();) {
size_t n = 1;
while (i + n < fills.size() && fills[i + n].expert == fills[i].expert + (int32_t) n && fills[i + n].slot == fills[i].slot + (int32_t) n) {
n++;
}
ggml_backend_tensor_set_async(backend, bank, (const uint8_t *) w->data + fills[i].expert*expert_size, fills[i].slot*expert_size, n*expert_size);
bytes += n*expert_size;
i += n;
}
}
stats & st = ids.size() <= (size_t) 8*n_expert_used ? stats_small : stats_large;
st.hits += n_hit;
st.misses += fills.size();
st.bytes += bytes;
// the next copy synchronizes the backend before it changes the slot map again
ggml_backend_tensor_set_async(backend, dst, src->data, 0, ggml_nbytes(src));
return true;
}
void log_stats() const {
auto log = [](const char * name, const stats & st) {
const size_t n = st.hits + st.misses;
if (n == 0) {
return;
}
LLAMA_LOG_INFO("llama_moe_cache: %s: hits = %zu, misses = %zu, hit rate = %.2f%%, uploaded = %.2f MiB\n",
name, st.hits, st.misses, 100.0*st.hits/n, st.bytes/1024.0/1024.0);
};
log("ubatch <= 8", stats_small);
log("ubatch > 8", stats_large);
if (stats_copy.hits > 0) {
LLAMA_LOG_INFO("llama_moe_cache: large batches: %zu experts copied from the cache, %.2f MiB\n", stats_copy.hits, stats_copy.bytes/1024.0/1024.0);
}
}
};
llama_moe_cache::llama_moe_cache(const llama_model & model, ggml_backend_t backend, ggml_backend_buffer_type_t buft, size_t size) :
pimpl(new impl(model, backend, buft, size)) {
}
llama_moe_cache::~llama_moe_cache() = default;
ggml_backend_t llama_moe_cache::backend() const {
return pimpl->backend;
}
ggml_tensor * llama_moe_cache::get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const {
return pimpl->get_slot_map(il, n_tokens, n_expert_used);
}
ggml_tensor * llama_moe_cache::get_experts(const ggml_tensor * w) const {
return pimpl->get_experts(w);
}
bool llama_moe_cache::copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph) {
return pimpl->copy(backend, src, dst, graph);
}
int64_t llama_moe_cache::copy_experts(ggml_backend_t backend, const ggml_tensor * w, ggml_tensor * dst, int64_t e, int64_t last) {
return pimpl->copy_experts(backend, w, dst, e, last);
}
std::map<ggml_backend_buffer_type_t, size_t> llama_moe_cache::memory_breakdown() const {
std::map<ggml_backend_buffer_type_t, size_t> res;
if (pimpl->buf) {
res[ggml_backend_buffer_get_type(pimpl->buf.get())] += pimpl->buf_size;
}
if (pimpl->buf_host) {
res[ggml_backend_buffer_get_type(pimpl->buf_host.get())] += pimpl->buf_host_size;
}
return res;
}

39
src/llama-moe-cache.h Normal file
View file

@ -0,0 +1,39 @@
#pragma once
#include "ggml-backend.h"
#include <map>
#include <memory>
struct llama_model;
// keeps the most recently used experts of host-resident MoE layers in a device buffer
// each layer has a slot map in host memory: when the scheduler copies it to the device, the copy callback uploads the missing experts
class llama_moe_cache {
public:
llama_moe_cache(const llama_model & model, ggml_backend_t backend, ggml_backend_buffer_type_t buft, size_t size);
~llama_moe_cache();
ggml_backend_t backend() const;
// the slot map of layer il, if its experts can be read from the cache for n_tokens tokens, nullptr otherwise
ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const;
// the experts of w in the cache, nullptr if w is not cached
ggml_tensor * get_experts(const ggml_tensor * w) const;
// ggml_backend_sched copy callback, returns false if src is not a slot map
bool copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph);
// for large batches: copy the experts of w that are in the cache, starting at expert e and up to expert last, to the copy dst of w
// returns the number of experts copied, 0 if expert e is not in the cache
int64_t copy_experts(ggml_backend_t backend, const ggml_tensor * w, ggml_tensor * dst, int64_t e, int64_t last);
std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const;
private:
struct impl;
std::unique_ptr<impl> pimpl;
};
using llama_moe_cache_ptr = std::unique_ptr<llama_moe_cache>;

View file

@ -1086,6 +1086,12 @@ static void llama_sampler_greedy_backend_apply(
struct ggml_tensor * curl = ggml_argmax(ctx, logits);
ggml_set_name(curl, "greedy_argmax");
if (data->candidates != nullptr) {
struct ggml_tensor * candidates = ggml_reshape_2d(ctx, data->candidates, 1, ggml_nelements(data->candidates));
curl = ggml_get_rows(ctx, candidates, curl);
ggml_set_name(curl, "greedy_sampled_token");
}
data->sampled = curl;
}

View file

@ -3099,6 +3099,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|| t.first == "▁<PRE>" // CodeLlama
|| t.first == "<|code_prefix|>" // GLM-4.5
|| t.first == "<|prefix|>" // Falcon-H1-Tiny-Coder
|| t.first == "<|plamo:fim_prefix|>" // PLaMo-3
) {
special_fim_pre_id = t.second;
if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
@ -3120,6 +3121,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|| t.first == "▁<SUF>" // CodeLlama
|| t.first == "<|code_suffix|>" // GLM-4.5
|| t.first == "<|suffix|>" // Falcon-H1-Tiny-Coder
|| t.first == "<|plamo:fim_suffix|>" // PLaMo-3
) {
special_fim_suf_id = t.second;
if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
@ -3141,6 +3143,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
|| t.first == "▁<MID>" // CodeLlama
|| t.first == "<|code_middle|>" // GLM-4.5
|| t.first == "<|middle|>" // Falcon-H1-Tiny-Coder
|| t.first == "<|plamo:fim_middle|>" // PLaMo-3
) {
special_fim_mid_id = t.second;
if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {
@ -3191,6 +3194,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
if (special_fim_sep_id == LLAMA_TOKEN_NULL) {
if (false
|| t.first == "<|file_sep|>" // Qwen
|| t.first == "<|plamo:file_separator|>" // PLaMo-2/3
) {
special_fim_sep_id = t.second;
if ((attr & LLAMA_TOKEN_ATTR_CONTROL) == 0) {

View file

@ -22,6 +22,7 @@
#include "llama-memory-recurrent.cpp"
#include "llama-model-loader.cpp"
#include "llama-model-saver.cpp"
#include "llama-moe-cache.cpp"
#include "llama-quant.cpp"
#include "llama-hparams.cpp"
#include "llama-graph.cpp"

View file

@ -64,15 +64,9 @@ void llama_model_bailingmoe3::load_arch_tensors(llama_model_loader & ml) {
const int64_t qk_head_dim = hparams.n_embd_head_k_mla();
const int64_t v_head_dim = hparams.n_embd_head_v_mla();
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
for (int il = 0; il < n_layer; ++il) {
auto & layer = layers[il];

View file

@ -40,18 +40,9 @@ void llama_model_cohere2moe::load_arch_hparams(llama_model_loader & ml) {
void llama_model_cohere2moe::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
// Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
// tensors live in a separate file. Mark MTP tensors NOT_REQUIRED so the
// trunk loads cleanly.
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);

View file

@ -56,15 +56,9 @@ void llama_model_deepseek2::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const int64_t n_expert_shared = hparams.n_expert_shared;
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
const bool is_mla = hparams.is_mla();

View file

@ -46,15 +46,9 @@ void llama_model_deepseek32::load_arch_hparams(llama_model_loader & ml) {
void llama_model_deepseek32::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
const bool is_mla = hparams.is_mla();
if (!is_mla) {

View file

@ -93,9 +93,9 @@ void llama_model_deepseek4::load_arch_tensors(llama_model_loader & ml) {
const int64_t hc_dim = hc_mult * n_embd;
const int64_t hc_mix_dim = (2 + hc_mult) * hc_mult;
const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
const int mtp_flags = ml.load_mtp ? 0 : TENSOR_SKIP;
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);

View file

@ -75,19 +75,9 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const int64_t n_expert_shared = hparams.n_expert_shared;
// MTP-only: the GGUF carries only the NextN/MTP block(s) (user split target/draft).
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
// Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
// tensors live in a separate file (or were stripped at conversion). Mark
// MTP tensors NOT_REQUIRED so the trunk loads cleanly.
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
const bool is_mla = hparams.is_mla();
if (!is_mla) {

View file

@ -29,15 +29,9 @@ void llama_model_glm4_moe::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const int64_t n_expert_shared = hparams.n_expert_shared;
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
GGML_ASSERT(hparams.n_expert > 0 && "n_expert must be > 0 for GLM4_MOE MoE layers");
GGML_ASSERT(hparams.n_expert_used() > 0 && "n_expert_used must be > 0 for GLM4_MOE MoE layers");

View file

@ -64,12 +64,9 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
const int64_t hc = hparams.dsv4_hc_mult;
const int64_t hc_mix_dim = (2 + hc)*hc;
// the NextN block is loaded but only used by the MTP graph.
// Separated trunk_only/mtp_only handling TODO with DECODER_MTP graph in the MTP follow up
int mtp_flags = 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
@ -79,18 +76,19 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
for (int i = 0; i < n_layer_all; ++i) {
auto & layer = layers[i];
const int flags = (i >= n_layer) ? mtp_flags : 0;
const bool is_nextn = i >= n_layer;
const int flags = is_nextn ? mtp_flags : trunk_flags;
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
if (i < n_layer) {
layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc*n_embd, hc_mix_dim}, 0);
layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0);
layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);
layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc*n_embd, hc_mix_dim}, 0);
layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0);
layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0);
layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc*n_embd, hc_mix_dim}, flags);
layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, flags);
layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, flags);
layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc*n_embd, hc_mix_dim}, flags);
layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, flags);
layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, flags);
}
const int64_t head_dim = hparams.n_embd_head_kda;
@ -100,25 +98,25 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
if (hparams.is_recr(i)) {
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);
return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, flags);
};
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);
create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, flags);
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);
layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, head_dim}, flags);
layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, flags);
layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, flags);
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), {n_head}, 0);
layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0);
layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), {n_head}, flags);
layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, flags);
layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", i), {n_embd, head_dim}, 0);
layer.ssm_g_b = create_tensor(tn(LLM_TENSOR_SSM_G_B, "weight", i), {head_dim, 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);
layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", i), {n_embd, head_dim}, flags);
layer.ssm_g_b = create_tensor(tn(LLM_TENSOR_SSM_G_B, "weight", i), {head_dim, d_inner}, flags);
layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, flags);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, flags);
} else {
const int64_t q_lora_rank = hparams.n_lora_q;
const int64_t kv_lora_rank = hparams.n_lora_kv;
@ -155,9 +153,9 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
}
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);
layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, flags);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, flags);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags);
} else {
const int64_t n_ff_exp = hparams.n_ff_exp(i);
const int64_t n_expert_shared = hparams.n_expert_shared;
@ -187,7 +185,7 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
std::unique_ptr<llm_graph_context> llama_model_glm5_next::build_arch_graph(const llm_graph_params & params) const {
if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
throw std::runtime_error("GLM5-Next NextN graph not implemented yet");
return std::make_unique<graph_mtp>(*this, params);
}
return std::make_unique<graph>(*this, params);
}
@ -539,6 +537,136 @@ ggml_tensor * llama_model_glm5_next::graph::build_hc_post(
return out;
}
// construct the graph helpers without building the trunk, so graph_mtp can reuse them
llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_params & params, no_build) :
llm_build_delta_net_base(params), model(model) {
}
llama_model_glm5_next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) :
graph(model, params, no_build{}) {
GGML_ASSERT(hparams.n_layer_nextn > 0 && "GLM5-Next MTP requires n_layer_nextn > 0");
GGML_ASSERT(hparams.n_layer_nextn == 1 && "GLM5-Next MTP currently supports a single NextN block");
const int il = hparams.n_layer() + cparams.nextn_layer_offset;
GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
"nextn_layer_offset out of range [0, n_layer_nextn)");
const auto & layer = model.layers[il];
GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm &&
"GLM5-Next MTP block is missing - load the model with MTP enabled");
GGML_ASSERT(hparams.n_embd_out() == (uint32_t) n_embd && "GLM5-Next MTP hidden width mismatch");
const auto * mctx_hyb = static_cast<const llama_memory_hybrid_idx_context *>(mctx);
auto * inp_hyb = build_inp_mem_hybrid_k();
auto * inp_attn = inp_hyb->get_attn();
auto * inp_kpool = build_inp_kpool(mctx_hyb);
ggml_build_forward_expand(gf, inp_hyb->get_recr()->s_copy);
ggml_tensor * inp_out_ids = build_inp_out_ids();
auto inp = std::make_unique<llm_graph_input_embd_h>(n_embd);
inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
ggml_set_input(inp->tokens);
inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
ggml_set_input(inp->embd);
inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens);
ggml_set_input(inp->h);
ggml_set_name(inp->h, "mtp_h_input");
ggml_tensor * tok_embd = ggml_get_rows(ctx0,
layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd, inp->tokens);
cb(tok_embd, "mtp_tok_embd", il);
ggml_tensor * h = inp->h;
res->add_input(std::move(inp));
ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
cb(e_norm, "mtp_enorm", il);
ggml_tensor * h_norm = build_norm(h, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
cb(h_norm, "mtp_hnorm", il);
ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, ggml_concat(ctx0, e_norm, h_norm, 0), layer.nextn.eh_proj_s);
cb(cur, "mtp_eh_proj", il);
ggml_tensor * inpSA = cur;
cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_attn_norm", il);
ggml_tensor * prev_sel = nullptr;
cur = build_dsa_layer(cur, layer, mctx_hyb, inp_attn, inp_kpool, &prev_sel, il);
cb(cur, "mtp_attn_out", il);
// narrow to the output tokens before the position-wise FFN; unmasked nextn embeddings need all rows
if (crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
cb(ffn_inp, "mtp_ffn_inp", il);
cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "mtp_ffn_norm", il);
ggml_tensor * moe_out = build_moe_ffn(cur,
layer.ffn_gate_inp,
layer.ffn_up_exps,
layer.ffn_gate_exps,
layer.ffn_down_exps,
layer.ffn_exp_probs_b,
n_expert, n_expert_used,
LLM_FFN_SILU, hparams.expert_weights_norm,
hparams.expert_weights_scale,
(llama_expert_gating_func_type) hparams.expert_gating_func,
il);
cb(moe_out, "mtp_ffn_moe_out", il);
ggml_tensor * ffn_shexp = build_ffn(cur,
layer.ffn_up_shexp, nullptr, nullptr,
layer.ffn_gate_shexp, nullptr, nullptr,
layer.ffn_down_shexp, nullptr, nullptr,
nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
cb(ffn_shexp, "mtp_ffn_shexp", il);
cur = ggml_add(ctx0, moe_out, ffn_shexp);
cb(cur, "mtp_ffn_out", il);
cur = ggml_add(ctx0, cur, ffn_inp);
cb(cur, "mtp_post_ffn", il);
ggml_tensor * head_norm = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm;
GGML_ASSERT(head_norm && "GLM5-Next MTP: missing both nextn.shared_head_norm and output_norm");
cur = build_norm(cur, head_norm, nullptr, LLM_NORM_RMS, -1);
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
if (crop_after_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cb(cur, "mtp_shared_head_norm", -1);
ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
GGML_ASSERT(head_w && "GLM5-Next MTP: missing both nextn.shared_head_head and output");
cur = build_lora_mm(head_w, cur, head_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}
llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_params & params) :
llm_build_delta_net_base(params), model(model) {

View file

@ -22,18 +22,9 @@ void llama_model_hy_v3::load_arch_hparams(llama_model_loader & ml) {
void llama_model_hy_v3::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
// Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
// tensors live in a separate file (e.g. user split target/draft). Mark
// MTP tensors NOT_REQUIRED so the trunk loads cleanly.
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);

View file

@ -4,6 +4,9 @@
#include <algorithm>
// question types of a decision model: choice, score, noul
static const uint32_t N_DECISION_TYPES = 3;
void llama_model_lfm2::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_SHORTCONV_L_CACHE, hparams.n_shortconv_l_cache);
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
@ -23,6 +26,15 @@ void llama_model_lfm2::load_arch_hparams(llama_model_loader & ml) {
default: type = LLM_TYPE_UNKNOWN;
}
ml.get_key(LLM_KV_DECISION_BLOCK_COUNT, hparams.n_layer_decision, false);
if (hparams.n_layer_decision > 0) {
if (hparams.n_layer_decision >= hparams.n_layer() || hparams.causal_attn) {
throw std::runtime_error("invalid decision head");
}
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
hparams.n_embd_out_impl = N_DECISION_TYPES;
}
if (const auto is_swa = ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false); is_swa && hparams.n_swa > 0) {
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
@ -37,13 +49,49 @@ void llama_model_lfm2::load_arch_tensors(llama_model_loader &) {
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM_LFM2, "weight"), {n_embd}, 0);
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
if (output == NULL) {
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
if (hparams.n_layer_decision > 0) {
// decision head: plain pre-norm blocks with biases
for (int i = n_layer - (int) hparams.n_layer_decision; i < n_layer; ++i) {
auto & layer = layers[i];
const int64_t n_ff_head = hparams.n_ff(i);
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0);
layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, 3 * n_embd}, 0);
layer.wqkv_b = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias", i), {3 * n_embd}, 0);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0);
layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", i), {n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff_head}, 0);
layer.ffn_up_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", i), {n_ff_head}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff_head, n_embd}, 0);
layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i), {n_embd}, 0);
}
if (n_token_types != N_DECISION_TYPES) {
throw std::runtime_error("decision model must have one token type per question type");
}
type_embd = create_tensor(tn(LLM_TENSOR_TOKEN_TYPES, "weight"), {n_embd, n_token_types}, 0);
cls_norm = create_tensor(tn(LLM_TENSOR_CLS_NORM, "weight"), {n_embd}, 0);
cls_norm_b = create_tensor(tn(LLM_TENSOR_CLS_NORM, "bias"), {n_embd}, 0);
cls = create_tensor(tn(LLM_TENSOR_CLS, "weight"), {n_embd, n_embd}, 0);
cls_b = create_tensor(tn(LLM_TENSOR_CLS, "bias"), {n_embd}, 0);
cls_out = create_tensor(tn(LLM_TENSOR_CLS_OUT, "weight"), {n_embd, 1}, 0);
cls_out_b = create_tensor(tn(LLM_TENSOR_CLS_OUT, "bias"), {1}, 0);
} else {
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
if (output == NULL) {
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
}
}
for (int i = 0; i < n_layer; ++i) {
for (int i = 0; i < n_layer - (int) hparams.n_layer_decision; ++i) {
auto & layer = layers[i];
const bool is_moe_layer = i >= static_cast<int>(hparams.n_layer_dense_lead);
@ -87,6 +135,9 @@ void llama_model_lfm2::load_arch_tensors(llama_model_loader &) {
}
std::unique_ptr<llm_graph_context> llama_model_lfm2::build_arch_graph(const llm_graph_params & params) const {
if (hparams.n_layer_decision > 0) {
return std::make_unique<graph_decision>(*this, params);
}
if (hparams.swa_type == LLAMA_SWA_TYPE_STANDARD) {
return std::make_unique<graph<true>>(*this, params);
} else {
@ -294,6 +345,242 @@ llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_
ggml_build_forward_expand(gf, cur);
}
// media entries (an image or audio prefix) are embeddings, text entries are tokens
static bool lfm2_is_media(const llama_ubatch & ubatch, int64_t i) {
return ubatch.is_mixed() ? ubatch.type[i] != 0 : ubatch.token == nullptr;
}
// non-causal within a sequence, the media never reads the text, so it is a function of the media alone
// in the head, the text and the media only read their own kind
class llm_graph_input_attn_media : public llm_graph_input_attn_no_cache {
public:
llm_graph_input_attn_media(const llama_hparams & hparams, const llama_cparams & cparams, bool is_head) :
llm_graph_input_attn_no_cache(hparams, cparams), is_head(is_head) {}
void set_input(const llama_ubatch * ubatch) override {
const int64_t n_tokens = ubatch->n_tokens;
std::vector<bool> is_media(n_tokens);
for (int64_t i = 0; i < n_tokens; ++i) {
is_media[i] = lfm2_is_media(*ubatch, i);
}
const auto fill_mask = [&](auto * data, auto zero, auto ninf) {
for (int64_t i1 = 0; i1 < n_tokens; ++i1) {
for (int64_t i0 = 0; i0 < n_tokens; ++i0) {
bool visible = ubatch->seq_id[i0][0] == ubatch->seq_id[i1][0];
if (is_head) {
visible = visible && is_media[i0] == is_media[i1];
} else {
visible = visible && !(is_media[i1] && !is_media[i0]);
}
data[i1 * n_tokens + i0] = visible ? zero : ninf;
}
}
};
GGML_ASSERT(ggml_backend_buffer_is_host(self_kq_mask->buffer));
if (self_kq_mask->type == GGML_TYPE_F16) {
fill_mask((ggml_fp16_t *) self_kq_mask->data, ggml_fp32_to_fp16(0.0f), ggml_fp32_to_fp16(-INFINITY));
} else {
fill_mask((float *) self_kq_mask->data, 0.0f, -INFINITY);
}
}
const bool is_head;
};
// 1 where the previous (next) token is the left (right) neighbor in the same sequence
// the last media entry does not read the text on its right
class llm_graph_input_conv_mask : public llm_graph_input_i {
public:
void set_input(const llama_ubatch * ubatch) override {
const int64_t n_tokens = ubatch->n_tokens;
std::vector<float> data_left(n_tokens, 0.0f);
std::vector<float> data_right(n_tokens, 0.0f);
for (int64_t i = 0; i + 1 < n_tokens; ++i) {
const bool is_next = ubatch->seq_id[i][0] == ubatch->seq_id[i + 1][0] && ubatch->pos[i] + 1 == ubatch->pos[i + 1];
data_right[i] = is_next && !(lfm2_is_media(*ubatch, i) && !lfm2_is_media(*ubatch, i + 1));
data_left[i + 1] = is_next;
}
ggml_backend_tensor_set(left, data_left.data(), 0, ggml_nbytes(left));
ggml_backend_tensor_set(right, data_right.data(), 0, ggml_nbytes(right));
}
ggml_tensor * left = nullptr; // F32 [1, n_tokens]
ggml_tensor * right = nullptr; // F32 [1, n_tokens]
};
llama_model_lfm2::graph_decision::graph_decision(const llama_model & model, const llm_graph_params & params) :
llm_graph_context(params) {
const int64_t n_embd_head = hparams.n_embd_head_v();
const int n_layer_enc = n_layer - hparams.n_layer_decision;
ggml_tensor * cur = build_inp_embd(model.tok_embd);
cb(cur, "model.embed_tokens", -1);
ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = build_inp_out_ids();
const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32;
llm_graph_input_attn_no_cache * inp_attn[2];
for (bool is_head : {false, true}) {
auto inp = std::make_unique<llm_graph_input_attn_media>(hparams, cparams, is_head);
inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1);
ggml_set_input(inp->self_kq_mask);
inp->self_kq_mask_cnv = inp->self_kq_mask;
inp_attn[is_head] = (llm_graph_input_attn_no_cache *) res->add_input(std::move(inp));
}
auto inp_conv = std::make_unique<llm_graph_input_conv_mask>();
inp_conv->left = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, n_tokens);
inp_conv->right = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, n_tokens);
ggml_set_input(inp_conv->left);
ggml_set_input(inp_conv->right);
ggml_tensor * conv_left = inp_conv->left;
ggml_tensor * conv_right = inp_conv->right;
res->add_input(std::move(inp_conv));
for (int il = 0; il < n_layer_enc; ++il) {
const auto & layer = model.layers[il];
ggml_tensor * inpL = cur;
cur = build_norm(cur, layer.attn_norm, NULL, LLM_NORM_RMS, il);
cb(cur, "model.layers.{}.operator_norm", il);
if (hparams.is_recr(il)) {
ggml_tensor * bcx = build_lora_mm(layer.shortconv.in_proj, cur);
cb(bcx, "model.layers.{}.conv.in_proj", il);
ggml_tensor * b = ggml_view_2d(ctx0, bcx, n_embd, n_tokens, bcx->nb[1], 0 * n_embd * ggml_element_size(bcx));
ggml_tensor * c = ggml_view_2d(ctx0, bcx, n_embd, n_tokens, bcx->nb[1], 1 * n_embd * ggml_element_size(bcx));
ggml_tensor * x = ggml_view_2d(ctx0, bcx, n_embd, n_tokens, bcx->nb[1], 2 * n_embd * ggml_element_size(bcx));
// centred 3-tap conv, a tap outside the sequence reads 0
ggml_tensor * bx = ggml_mul(ctx0, b, x);
ggml_tensor * bxp = ggml_pad_ext(ctx0, bx, 0, 0, 1, 1, 0, 0, 0, 0);
ggml_tensor * prv = ggml_view_2d(ctx0, bxp, n_embd, n_tokens, bxp->nb[1], 0);
ggml_tensor * nxt = ggml_view_2d(ctx0, bxp, n_embd, n_tokens, bxp->nb[1], 2 * bxp->nb[1]);
GGML_ASSERT(hparams.n_shortconv_l_cache == 3);
ggml_tensor * taps = ggml_cont(ctx0, ggml_transpose(ctx0, layer.shortconv.conv));
ggml_tensor * tap0 = ggml_view_1d(ctx0, taps, n_embd, 0 * taps->nb[1]);
ggml_tensor * tap1 = ggml_view_1d(ctx0, taps, n_embd, 1 * taps->nb[1]);
ggml_tensor * tap2 = ggml_view_1d(ctx0, taps, n_embd, 2 * taps->nb[1]);
ggml_tensor * y = ggml_mul(ctx0, bx, tap1);
y = ggml_add(ctx0, y, ggml_mul(ctx0, ggml_mul(ctx0, prv, tap0), conv_left));
y = ggml_add(ctx0, y, ggml_mul(ctx0, ggml_mul(ctx0, nxt, tap2), conv_right));
cb(y, "model.layers.{}.conv.conv", il);
cur = build_lora_mm(layer.shortconv.out_proj, ggml_mul(ctx0, c, y));
cb(cur, "model.layers.{}.conv.out_proj", il);
} else {
auto [q, k, v] = build_qkv(layer, cur, n_embd_head, n_head, hparams.n_head_kv(il), il);
q = build_norm(q, layer.attn_q_norm, NULL, LLM_NORM_RMS, il);
k = build_norm(k, layer.attn_k_norm, NULL, LLM_NORM_RMS, il);
q = ggml_rope_ext(ctx0, q, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor,
attn_factor, beta_fast, beta_slow);
k = ggml_rope_ext(ctx0, k, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor,
attn_factor, beta_fast, beta_slow);
cur = build_attn(inp_attn[0],
layer.wo, NULL, layer.wo_s,
q, k, v, nullptr, nullptr, nullptr, 1.0f / sqrtf(float(n_embd_head)), il);
cb(cur, "model.layers.{}.self_attn.out_proj", il);
}
cur = ggml_add(ctx0, cur, inpL);
ggml_tensor * ffn_out = build_norm(cur, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
ffn_out = build_ffn(ffn_out,
layer.ffn_up, NULL, NULL,
layer.ffn_gate, NULL, NULL,
layer.ffn_down, NULL, NULL,
NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
cur = ggml_add(ctx0, cur, ffn_out);
cb(cur, "l_out", il);
}
cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
cb(cur, "result_norm", -1);
cur = build_decision_head(model, cur, inp_attn[1], inp_out_ids);
res->t_embd = cur;
ggml_build_forward_expand(gf, cur);
}
// same as llama_model_modern_bert::graph::build_decision_head(), with the head counts of the head layers
ggml_tensor * llama_model_lfm2::graph_decision::build_decision_head(
const llama_model & model,
ggml_tensor * inp,
llm_graph_input_attn_no_cache * inp_attn,
ggml_tensor * inp_out_ids) {
const int64_t n_embd_head = hparams.n_embd_head_v();
const int n_layer_enc = n_layer - hparams.n_layer_decision;
ggml_tensor * scores = nullptr;
// the question type is not a graph input, so the head is evaluated for each of them
for (uint32_t it = 0; it < N_DECISION_TYPES; ++it) {
ggml_tensor * type_row = ggml_view_1d(ctx0, model.type_embd, n_embd, it * model.type_embd->nb[1]);
ggml_tensor * inpL = ggml_add(ctx0, inp, type_row);
for (int il = n_layer_enc; il < n_layer; ++il) {
const auto & layer = model.layers[il];
ggml_tensor * cur = build_norm(inpL, layer.attn_norm, layer.attn_norm_b, LLM_NORM, il);
cb(cur, "attn_norm", il);
// no positional encoding in the head
auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);
cur = build_attn(inp_attn,
layer.wo, layer.wo_b, layer.wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
cb(cur, "kqv_out", il);
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
cb(ffn_inp, "ffn_inp", il);
cur = build_norm(ffn_inp, layer.ffn_norm, layer.ffn_norm_b, LLM_NORM, il);
cb(cur, "ffn_norm", il);
cur = build_ffn(cur,
layer.ffn_up, layer.ffn_up_b, NULL,
NULL, NULL, NULL,
layer.ffn_down, layer.ffn_down_b, NULL,
NULL,
LLM_FFN_RELU,
LLM_FFN_SEQ, il);
inpL = ggml_add(ctx0, cur, ffn_inp);
}
// scorer
ggml_tensor * cur = build_norm(inpL, model.cls_norm, model.cls_norm_b, LLM_NORM, -1);
cur = ggml_add(ctx0, build_lora_mm(model.cls, cur), model.cls_b);
cur = ggml_gelu_erf(ctx0, cur);
cur = ggml_add(ctx0, build_lora_mm(model.cls_out, cur), model.cls_out_b);
scores = scores ? ggml_concat(ctx0, scores, cur, 0) : cur;
}
cb(scores, "decision_scores", -1);
return scores;
}
// Explicit template instantiations
template struct llama_model_lfm2::graph<true>;
template struct llama_model_lfm2::graph<false>;

View file

@ -25,15 +25,9 @@ void llama_model_mimo2::load_arch_hparams(llama_model_loader & ml) {
void llama_model_mimo2::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);

View file

@ -2186,6 +2186,17 @@ struct llama_model_lfm2 : public llama_model_base {
graph(const llama_model & model, const llm_graph_params & params);
};
// non-causal trunk without memory, then the decision head
struct graph_decision : public llm_graph_context {
graph_decision(const llama_model & model, const llm_graph_params & params);
ggml_tensor * build_decision_head(
const llama_model & model,
ggml_tensor * inp,
llm_graph_input_attn_no_cache * inp_attn,
ggml_tensor * inp_out_ids);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
@ -2741,6 +2752,10 @@ struct llama_model_glm5_next : public llama_model_base {
struct graph : public llm_build_delta_net_base {
graph(const llama_model & model, const llm_graph_params & params);
// build the helpers without the trunk, so graph_mtp can reuse them
struct no_build {};
graph(const llama_model & model, const llm_graph_params & params, no_build);
// collapse the hc streams with per-stream weights
ggml_tensor * build_hc_pre(
ggml_tensor * x,
@ -2789,6 +2804,11 @@ struct llama_model_glm5_next : public llama_model_base {
};
// the NextN / MTP draft head: one DSA block appended after the trunk
struct graph_mtp : public graph {
graph_mtp(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};

View file

@ -1,7 +1,7 @@
#include "models.h"
// question types of a decision model: choice, score, noul
static const uint32_t N_DECISION_TYPES = 3;
// static const uint32_t N_DECISION_TYPES = 3; //kcpp: already defined in lfm2.cpp
void llama_model_modern_bert::load_arch_hparams(llama_model_loader & ml) {
const bool found_swa = ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
@ -28,6 +28,12 @@ void llama_model_modern_bert::load_arch_hparams(llama_model_loader & ml) {
hparams.pooling_type_cls = LLAMA_POOLING_TYPE_MEAN;
}
// GGUFs without a classifier activation use gelu, the transformers default
std::string act_cls;
if (!ml.get_key(LLM_KV_CLASSIFIER_ACTIVATION, act_cls, false)) {
hparams.act_cls = GGML_UNARY_OP_GELU_ERF;
}
ml.get_key(LLM_KV_DECISION_BLOCK_COUNT, hparams.n_layer_decision, false);
if (hparams.n_layer_decision > 0) {
if (hparams.n_layer_decision >= hparams.n_layer()) {

View file

@ -49,9 +49,9 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) {
void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.attn_norm.weight") == nullptr;
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
const int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
// mamba2 Mixer SSM params
// NOTE: int64_t for tensor dimensions

View file

@ -33,9 +33,9 @@ void llama_model_qwen35::load_arch_hparams(llama_model_loader & ml) {
void llama_model_qwen35::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);

View file

@ -36,9 +36,9 @@ void llama_model_qwen35moe::load_arch_hparams(llama_model_loader & ml) {
void llama_model_qwen35moe::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);

View file

@ -35,9 +35,9 @@ void llama_model_qwen3next::load_arch_tensors(llama_model_loader & ml) {
throw std::runtime_error(arch_name() + " model cannot have zero experts");
}
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0;
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);

View file

@ -175,10 +175,9 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) {
const int64_t hc_dim = hc * n_embd;
const int64_t hc_lr = hparams.hc_low_rank;
// an MTP-only file carries the MTP block, the embeddings and the LM head, but no trunk
const bool mtp_only = n_layer_nextn > 0 && ml.get_weight(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", 0).str().c_str()) == nullptr;
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
const int mtp_flags = ml.load_mtp ? 0 : TENSOR_SKIP;
const auto nf = nextn_flags(ml, LLM_TENSOR_HC_ATTN_NORM);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);

View file

@ -37,18 +37,9 @@ void llama_model_step35::load_arch_hparams(llama_model_loader & ml) {
void llama_model_step35::load_arch_tensors(llama_model_loader & ml) {
LLAMA_LOAD_LOCALS;
const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
// Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP
// tensors live in a separate file (e.g. user split target/draft). Mark
// MTP tensors NOT_REQUIRED so the trunk loads cleanly.
const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0;
if (!ml.load_mtp) {
mtp_flags |= TENSOR_SKIP;
}
const auto nf = nextn_flags(ml);
const int trunk_flags = nf.trunk;
const int mtp_flags = nf.mtp;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);

View file

@ -58,6 +58,7 @@
#define KEY_PATCH_SIZE "clip.vision.patch_size"
#define KEY_IMAGE_MEAN "clip.vision.image_mean"
#define KEY_IMAGE_STD "clip.vision.image_std"
#define KEY_IMAGE_RESIZE_ALGO "clip.vision.image_resize_algo"
#define KEY_PROJ_SCALE_FACTOR "clip.vision.projector.scale_factor"
#define KEY_PROJ_SAMPLE_QUERY_SIDE "clip.vision.projector.query_side"
#define KEY_PROJ_SAMPLE_WINDOW_SIDE "clip.vision.projector.window_side"
@ -475,6 +476,7 @@ enum projector_type {
PROJECTOR_TYPE_MERALION,
PROJECTOR_TYPE_MUSIC_FLAMINGO,
PROJECTOR_TYPE_LFM2,
PROJECTOR_TYPE_D1OMNI_V,
PROJECTOR_TYPE_KIMIVL,
PROJECTOR_TYPE_PADDLEOCR,
PROJECTOR_TYPE_LIGHTONOCR,
@ -487,6 +489,7 @@ enum projector_type {
PROJECTOR_TYPE_DEEPSEEKOCR2,
PROJECTOR_TYPE_DEEPSEEK4V,
PROJECTOR_TYPE_LFM2A,
PROJECTOR_TYPE_D1OMNI_A,
PROJECTOR_TYPE_GLM4V,
PROJECTOR_TYPE_GLM5V,
PROJECTOR_TYPE_YOUTUVL,
@ -507,6 +510,7 @@ enum projector_type {
PROJECTOR_TYPE_POCKETTTS_SPKENC,
PROJECTOR_TYPE_POCKETTTS_GEN,
PROJECTOR_TYPE_MUSE_GLIMMER,
PROJECTOR_TYPE_COHERE2V,
PROJECTOR_TYPE_UNKNOWN,
};
@ -542,6 +546,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
{ PROJECTOR_TYPE_MERALION, "meralion"},
{ PROJECTOR_TYPE_MUSIC_FLAMINGO, "musicflamingo"},
{ PROJECTOR_TYPE_LFM2, "lfm2"},
{ PROJECTOR_TYPE_D1OMNI_V, "d1omni_v"},
{ PROJECTOR_TYPE_KIMIVL, "kimivl"},
{ PROJECTOR_TYPE_PADDLEOCR, "paddleocr"},
{ PROJECTOR_TYPE_LIGHTONOCR, "lightonocr"},
@ -554,6 +559,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
{ PROJECTOR_TYPE_DEEPSEEKOCR2, "deepseekocr2"},
{ PROJECTOR_TYPE_DEEPSEEK4V, "deepseek4v"},
{ PROJECTOR_TYPE_LFM2A, "lfm2a"},
{ PROJECTOR_TYPE_D1OMNI_A, "d1omni_a"},
{ PROJECTOR_TYPE_GLM4V, "glm4v"},
{ PROJECTOR_TYPE_GLM5V, "glm5v"},
{ PROJECTOR_TYPE_YOUTUVL, "youtuvl"},
@ -574,6 +580,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
{ PROJECTOR_TYPE_POCKETTTS_SPKENC, "pockettts_spkenc"},
{ PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"},
{ PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"},
{ PROJECTOR_TYPE_COHERE2V, "cohere2v"},
};
static projector_type clip_projector_type_from_string(const std::string & str) {

View file

@ -623,6 +623,10 @@ struct clip_model {
ggml_tensor * mm_3_b = nullptr;
ggml_tensor * mm_4_w = nullptr;
ggml_tensor * mm_4_b = nullptr;
ggml_tensor * mm_5_w = nullptr;
ggml_tensor * mm_5_b = nullptr;
ggml_tensor * mm_6_w = nullptr;
ggml_tensor * mm_6_b = nullptr;
// GLMV-Edge projection
ggml_tensor * mm_model_adapter_conv_w = nullptr;

View file

@ -997,7 +997,9 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
switch (ctx->proj_type()) {
case PROJECTOR_TYPE_GEMMA3:
case PROJECTOR_TYPE_IDEFICS3:
case PROJECTOR_TYPE_COHERE2V:
case PROJECTOR_TYPE_LFM2:
case PROJECTOR_TYPE_D1OMNI_V:
case PROJECTOR_TYPE_JANUS_PRO:
case PROJECTOR_TYPE_PHI4:
{
@ -1132,6 +1134,7 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
builder = std::make_unique<clip_graph_deepseekocr2>(ctx, img);
} break;
case PROJECTOR_TYPE_LFM2A:
case PROJECTOR_TYPE_D1OMNI_A:
{
builder = std::make_unique<clip_graph_conformer>(ctx, img);
} break;
@ -1597,11 +1600,36 @@ struct clip_model_loader {
get_u32(KEY_PREPROC_IMAGE_SIZE, hparams.image_longest_edge, false);
hparams.set_limit_image_tokens();
} break;
case PROJECTOR_TYPE_LFM2:
case PROJECTOR_TYPE_COHERE2V:
{
hparams.image_resize_algo = RESIZE_ALGO_BILINEAR;
hparams.image_resize_algo_rf = RESIZE_ALGO_BILINEAR;
hparams.image_resize_algo_ov = RESIZE_ALGO_BILINEAR;
hparams.image_pad_rf = PAD_NONE;
get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge);
get_u32(KEY_PREPROC_MAX_TILES, hparams.preproc_max_tiles);
if (hparams.preproc_max_tiles <= 0 || hparams.preproc_max_tiles > 256) {
throw std::runtime_error(string_format("%s: preproc_max_tiles (%d) must be in range [1, 256]\n", __func__, hparams.preproc_max_tiles));
}
} break;
case PROJECTOR_TYPE_LFM2:
case PROJECTOR_TYPE_D1OMNI_V:
{
// default for older GGUFs
std::string resize_algo = "bilinear";
get_string(KEY_IMAGE_RESIZE_ALGO, resize_algo, false);
if (resize_algo == "bilinear") {
hparams.image_resize_algo = RESIZE_ALGO_BILINEAR;
} else if (resize_algo == "bicubic") {
hparams.image_resize_algo = RESIZE_ALGO_BICUBIC;
} else if (resize_algo == "lanczos") {
hparams.image_resize_algo = RESIZE_ALGO_LANCZOS;
} else {
throw std::runtime_error(string_format("%s: unsupported image resize algo: %s\n", __func__, resize_algo.c_str()));
}
hparams.image_resize_algo_rf = hparams.image_resize_algo;
hparams.image_resize_algo_ov = hparams.image_resize_algo;
// the tiles stretch the image to the grid (d1-omni vision.py)
if (model.proj_type == PROJECTOR_TYPE_D1OMNI_V) {
hparams.image_pad_rf = PAD_NONE;
}
get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false);
// ref: https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B/blob/main/processor_config.json
hparams.set_limit_image_tokens(64, 256);
@ -2047,6 +2075,7 @@ struct clip_model_loader {
hparams.set_warmup_n_tokens(32*32);
} break;
case PROJECTOR_TYPE_LFM2A:
case PROJECTOR_TYPE_D1OMNI_A:
{
// audio preprocessing params
hparams.audio_chunk_len = 1; // in seconds
@ -2863,7 +2892,15 @@ struct clip_model_loader {
{
model.mm_fc_w = get_tensor(string_format(TN_MM_PROJECTOR, "weight"));
} break;
case PROJECTOR_TYPE_COHERE2V:
{
model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight"));
model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 1, "bias"));
model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight"));
model.mm_2_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"));
} break;
case PROJECTOR_TYPE_LFM2:
case PROJECTOR_TYPE_D1OMNI_V:
{
model.mm_input_norm_w = get_tensor(TN_MM_INP_NORM, false);
model.mm_input_norm_b = get_tensor(TN_MM_INP_NORM_B, false);
@ -3476,6 +3513,7 @@ struct clip_model_loader {
model.mm_input_proj_w = get_tensor(string_format(TN_A_MM_INP_PROJ, "weight"));
} break;
case PROJECTOR_TYPE_LFM2A:
case PROJECTOR_TYPE_D1OMNI_A:
{
for (int i : {0, 2, 3, 5, 6}) {
model.pre_encode_conv_X_w[i] = get_tensor(string_format(TN_CONV1D, i, "weight"));
@ -3491,6 +3529,16 @@ struct clip_model_loader {
model.mm_3_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 3, "weight"));
model.mm_3_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 3, "bias"));
// residual block after the projector: norm, down, up
if (model.proj_type == PROJECTOR_TYPE_D1OMNI_A) {
model.mm_4_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 4, "weight"));
model.mm_4_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 4, "bias"));
model.mm_5_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 5, "weight"));
model.mm_5_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 5, "bias"));
model.mm_6_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 6, "weight"));
model.mm_6_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 6, "bias"));
}
for (int il = 0; il < hparams.n_layer; ++il) {
auto & layer = model.layers[il];
@ -4326,6 +4374,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
case PROJECTOR_TYPE_GEMMA4V:
case PROJECTOR_TYPE_GEMMA4UV:
case PROJECTOR_TYPE_IDEFICS3:
case PROJECTOR_TYPE_COHERE2V:
case PROJECTOR_TYPE_INTERNVL:
case PROJECTOR_TYPE_NEMOTRON_V2_VL:
case PROJECTOR_TYPE_LLAMA4:
@ -4341,6 +4390,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
n_patches = ctx->model.hparams.image_size / ctx->model.hparams.patch_size;
} break;
case PROJECTOR_TYPE_LFM2:
case PROJECTOR_TYPE_D1OMNI_V:
case PROJECTOR_TYPE_KIMIVL:
case PROJECTOR_TYPE_KIMIK25:
{
@ -4464,6 +4514,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
}
} break;
case PROJECTOR_TYPE_LFM2A:
case PROJECTOR_TYPE_D1OMNI_A:
{
n_patches = ((((img->nx() + 1) / 2) + 1) / 2 + 1) / 2;
} break;
@ -5414,6 +5465,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
case PROJECTOR_TYPE_GEMMA3:
case PROJECTOR_TYPE_GEMMA3NV:
case PROJECTOR_TYPE_IDEFICS3:
case PROJECTOR_TYPE_COHERE2V:
case PROJECTOR_TYPE_INTERNVL:
case PROJECTOR_TYPE_NEMOTRON_V2_VL:
case PROJECTOR_TYPE_QWEN2A:
@ -5421,6 +5473,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
case PROJECTOR_TYPE_GLMA:
case PROJECTOR_TYPE_ULTRAVOX:
case PROJECTOR_TYPE_LFM2:
case PROJECTOR_TYPE_D1OMNI_V:
case PROJECTOR_TYPE_VOXTRAL:
case PROJECTOR_TYPE_MERALION:
case PROJECTOR_TYPE_MUSIC_FLAMINGO:
@ -5680,6 +5733,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
}
} break;
case PROJECTOR_TYPE_LFM2A:
case PROJECTOR_TYPE_D1OMNI_A:
{
GGML_ASSERT(imgs.entries.size() == 1);
const auto n_frames = clip_n_output_tokens(ctx, &imgs.entries.front());
@ -6143,11 +6197,13 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
return ctx->model.mm_2_w->ne[1];
case PROJECTOR_TYPE_GLMA:
case PROJECTOR_TYPE_LFM2:
case PROJECTOR_TYPE_D1OMNI_V:
case PROJECTOR_TYPE_KIMIVL:
case PROJECTOR_TYPE_PADDLEOCR:
case PROJECTOR_TYPE_KIMIK25:
case PROJECTOR_TYPE_YASA2:
case PROJECTOR_TYPE_DEEPSEEK4V:
case PROJECTOR_TYPE_COHERE2V:
return ctx->model.mm_2_w->ne[1];
case PROJECTOR_TYPE_HUNYUANVL:
return ctx->model.mm_model_proj->ne[1];
@ -6158,6 +6214,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
return ctx->model.mm_fc_w->ne[1];
case PROJECTOR_TYPE_LFM2A:
return ctx->model.position_embeddings->ne[0];
case PROJECTOR_TYPE_D1OMNI_A:
return ctx->model.mm_3_w->ne[1];
case PROJECTOR_TYPE_GRANITE_SPEECH:
return ctx->model.qf_proj_blocks[0].qf_proj_linear_w->ne[1];
case PROJECTOR_TYPE_GRANITE4_VISION:

View file

@ -4,7 +4,7 @@ ggml_cgraph * clip_graph_conformer::build() {
const int n_frames = img.nx();
const int n_pos = n_frames / 2;
const int n_pos_embd = (((((n_frames + 1) / 2) + 1) / 2 + 1) / 2) * 2 - 1;
GGML_ASSERT(model.position_embeddings->ne[1] >= n_pos);
GGML_ASSERT(!model.position_embeddings || model.position_embeddings->ne[1] >= n_pos);
ggml_tensor * pos_emb = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 512, n_pos_embd);
ggml_set_name(pos_emb, "pos_emb");
@ -164,7 +164,7 @@ ggml_cgraph * clip_graph_conformer::build() {
// TODO @ngxson : support this ops in ggml
{
int64_t d = x->ne[0] / 2;
ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_view_2d(ctx0, x, d, x->ne[1], x->nb[1], d * x->nb[0]));
ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, x, d, x->ne[1], x->nb[1], d * x->nb[0])));
x = ggml_mul(ctx0, ggml_view_2d(ctx0, x, d, x->ne[1], x->nb[1], 0), gate);
x = ggml_cont(ctx0, ggml_transpose(ctx0, x));
}
@ -207,6 +207,13 @@ ggml_cgraph * clip_graph_conformer::build() {
cb(cur, "audio_adapter.model.{}", 0);
cur = build_ffn(cur, model.mm_1_w, model.mm_1_b, nullptr, nullptr, model.mm_3_w, model.mm_3_b, FFN_GELU_ERF, -1);
// d1omni_a: residual block after the projector
if (model.mm_4_w) {
ggml_tensor * x = build_norm(cur, model.mm_4_w, model.mm_4_b, NORM_TYPE_NORMAL, 1e-5, -1);
x = build_ffn(x, model.mm_5_w, model.mm_5_b, nullptr, nullptr, model.mm_6_w, model.mm_6_b, FFN_GELU_ERF, -1);
cur = ggml_add(ctx0, cur, x);
}
cb(cur, "projected", -1);
ggml_build_forward_expand(gf, cur);

View file

@ -4,7 +4,7 @@ ggml_cgraph * clip_graph_siglip::build() {
ggml_tensor * inp = build_inp();
ggml_tensor * learned_pos_embd = model.position_embeddings;
if (proj_type == PROJECTOR_TYPE_LFM2 || proj_type == PROJECTOR_TYPE_PHI4) {
if (proj_type == PROJECTOR_TYPE_LFM2 || proj_type == PROJECTOR_TYPE_D1OMNI_V || proj_type == PROJECTOR_TYPE_PHI4) {
learned_pos_embd = resize_position_embeddings();
}
@ -45,7 +45,17 @@ ggml_cgraph * clip_graph_siglip::build() {
cur = build_patch_merge_permute(cur, scale_factor);
cur = build_mm(model.mm_fc_w, cur);
} else if (proj_type == PROJECTOR_TYPE_LFM2) {
} else if (proj_type == PROJECTOR_TYPE_COHERE2V) {
// tiles are square, so the pixel shuffle is the same as Idefics3
cur = build_patch_merge_permute(cur, model.hparams.n_merge);
cur = build_mm(model.mm_1_w, cur);
cur = ggml_add(ctx0, cur, model.mm_1_b);
// linear_1 output is [x, gate], HF computes silu(gate) * x
cur = ggml_swiglu_swapped(ctx0, cur);
cur = build_mm(model.mm_2_w, cur);
cur = ggml_add(ctx0, cur, model.mm_2_b);
} else if (proj_type == PROJECTOR_TYPE_LFM2 || proj_type == PROJECTOR_TYPE_D1OMNI_V) {
// pixel unshuffle block
const int scale_factor = model.hparams.n_merge;
cur = build_patch_merge_permute(cur, scale_factor);
@ -60,11 +70,12 @@ ggml_cgraph * clip_graph_siglip::build() {
cur = ggml_add(ctx0, cur, model.mm_input_norm_b);
}
// d1-omni uses the exact gelu in the projector
cur = build_ffn(cur,
model.mm_1_w, model.mm_1_b,
nullptr, nullptr,
model.mm_2_w, model.mm_2_b,
FFN_GELU,
proj_type == PROJECTOR_TYPE_D1OMNI_V ? FFN_GELU_ERF : FFN_GELU,
-1);
} else if (proj_type == PROJECTOR_TYPE_JANUS_PRO) {

View file

@ -994,6 +994,42 @@ bool mtmd_audio_preprocessor_conformer::preprocess(const float *
return true;
}
//
// mtmd_audio_preprocessor_d1omni
//
bool mtmd_audio_preprocessor_d1omni::preprocess(const float * samples,
size_t n_samples,
std::vector<mtmd_audio_mel> & output) const {
if (n_samples == 0) {
return false;
}
const size_t n_max = 30 * hparams.audio_sample_rate;
const size_t n_min = hparams.audio_sample_rate / 2;
std::vector<float> buf(samples, samples + std::min(n_samples, n_max));
buf.resize(std::max(buf.size(), n_min), 0.0f);
if (!mtmd_audio_preprocessor_conformer::preprocess(buf.data(), buf.size(), output)) {
return false;
}
// the encoder reads one frame per hop, not the extra frame of the centre padding (NeMo: seq_len)
const int64_t n_frames = buf.size() / hparams.audio_hop_len;
for (auto & mel : output) {
if (mel.n_len <= n_frames) {
continue;
}
std::vector<float> data((size_t) mel.n_mel * n_frames);
for (int64_t j = 0; j < mel.n_mel; ++j) {
std::copy_n(mel.data.begin() + (size_t) j * mel.n_len, n_frames, data.begin() + (size_t) j * n_frames);
}
mel.n_len = n_frames;
mel.data = std::move(data);
}
return true;
}
//
// mtmd_audio_preprocessor_granite_speech
//

View file

@ -80,6 +80,12 @@ struct mtmd_audio_preprocessor_conformer : mtmd_audio_preprocessor {
mtmd_audio_cache cache;
};
// same as conformer, the audio is cut to 30 s and padded to 0.5 s (d1-omni audio.py)
struct mtmd_audio_preprocessor_d1omni : mtmd_audio_preprocessor_conformer {
using mtmd_audio_preprocessor_conformer::mtmd_audio_preprocessor_conformer;
bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) const override;
};
struct mtmd_audio_preprocessor_granite_speech : mtmd_audio_preprocessor {
mtmd_audio_preprocessor_granite_speech(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {}
void initialize() override;

Some files were not shown because too many files have changed in this diff Show more