diff --git a/common/arg.cpp b/common/arg.cpp index de60522bb..e1e9ef36e 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3647,6 +3647,18 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_REASONING")); + add_opt(common_arg( + {"--reasoning-effort"}, "LEVEL", + "reasoning effort level given to the chat template: 'default' to keep the template default,\n" + "or a level such as 'minimal', 'low', 'medium', 'high', 'xhigh' or 'max' (default: default)", + [](common_params & params, const std::string & value) { + if (value == "default") { + params.default_template_kwargs.erase("reasoning_effort"); + } else { + params.default_template_kwargs["reasoning_effort"] = json(value).dump(); + } + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_REASONING_EFFORT")); add_opt(common_arg( {"--reasoning-budget"}, "N", "token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)", @@ -4066,6 +4078,9 @@ common_params_context common_params_parser_init(common_params & params, llama_ex {"--spec-draft-n-max"}, "N", string_format("number of tokens to draft for speculative decoding (default: %d)", params.speculative.draft.n_max), [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } params.speculative.draft.n_max = value; } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MAX")); diff --git a/common/chat.cpp b/common/chat.cpp index ec374e91f..803b4021b 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -934,6 +934,10 @@ static std::string common_chat_template_direct_apply_impl( bool enabled = inp["preserve_reasoning"].get(); jinja::caps_apply_preserve_reasoning(ctx, enabled); } + if (inp.contains("reasoning_effort") && inp["reasoning_effort"].is_string() && !inp["reasoning_effort"].empty()) { + std::string reasoning_effort = inp["reasoning_effort"].get(); + jinja::caps_apply_reasoning_effort(ctx, reasoning_effort); + } jinja::global_from_json(ctx, inp, inputs.mark_input); diff --git a/common/imatrix-loader.cpp b/common/imatrix-loader.cpp index efe9aecee..71d3b500f 100644 --- a/common/imatrix-loader.cpp +++ b/common/imatrix-loader.cpp @@ -102,7 +102,8 @@ bool common_imatrix_load(const std::string & fname, common_imatrix & imatrix) { const int64_t chunk_count_key = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_COUNT); const int64_t chunk_size_key = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_SIZE); - if (datasets_key != -1 && gguf_get_arr_type(ctx_gguf, datasets_key) == GGUF_TYPE_STRING) { + if (datasets_key != -1 && gguf_get_kv_type(ctx_gguf, datasets_key) == GGUF_TYPE_ARRAY && + gguf_get_arr_type(ctx_gguf, datasets_key) == GGUF_TYPE_STRING) { const int64_t n = gguf_get_arr_n(ctx_gguf, datasets_key); imatrix.datasets.reserve(imatrix.datasets.size() + n); for (int64_t i = 0; i < n; ++i) { @@ -143,6 +144,13 @@ bool common_imatrix_load(const std::string & fname, common_imatrix & imatrix) { return false; } + if (in_sum2->type != GGML_TYPE_F32 || counts->type != GGML_TYPE_F32) { + LOG_ERR("%s: sums and counts for %s must be F32\n", __func__, name.c_str()); + gguf_free(ctx_gguf); + ggml_free(ctx); + return false; + } + auto & e = imatrix.entries[name]; const int64_t nval = ggml_nelements(in_sum2); diff --git a/common/jinja/caps.cpp b/common/jinja/caps.cpp index 6f4364072..00b02119a 100644 --- a/common/jinja/caps.cpp +++ b/common/jinja/caps.cpp @@ -20,7 +20,7 @@ namespace jinja { using caps_json_fn = std::function; using caps_ctx_fn = std::function; -using caps_analyze_fn = std::function; +using caps_analyze_fn = std::function; void caps_apply_preserve_reasoning(jinja::context & ctx, bool enabled) { ctx.set_val("preserve_thinking", mk_val(enabled)); @@ -29,6 +29,12 @@ void caps_apply_preserve_reasoning(jinja::context & ctx, bool enabled) { ctx.set_val("drop_thinking", mk_val(!enabled)); } +void caps_apply_reasoning_effort(jinja::context & ctx, const std::string & effort) { + value var = mk_val(effort); // bind to the same value for stats + ctx.set_val("reasoning_effort", var); + ctx.set_val("reasoning_strength", var); +} + static void caps_try_execute(jinja::program & prog, const caps_json_fn & messages_fn, const caps_ctx_fn & ctx_fn, @@ -65,7 +71,7 @@ static void caps_try_execute(jinja::program & prog, // ignore exceptions during capability analysis } - analyze_fn(success, messages, tools, result); + analyze_fn(ctx, success, messages, tools, result); } // for debugging only @@ -90,6 +96,7 @@ std::map caps::to_map() const { {"supports_parallel_tool_calls", supports_parallel_tool_calls}, {"supports_system_role", supports_system_role}, {"supports_preserve_reasoning", supports_preserve_reasoning}, + {"supports_reasoning_effort", supports_reasoning_effort}, {"supports_object_arguments", supports_object_arguments}, }; } @@ -127,7 +134,7 @@ caps caps_get(jinja::program & prog) { }, nullptr, // ctx_fn nullptr, // tools_fn - [&](bool success, value & messages, value &, const std::string &) { + [&](context &, bool success, value & messages, value &, const std::string &) { auto & content = messages->at(0)->at("content"); caps_print_stats(content, "messages[0].content"); if (has_op(content, "selectattr") || has_op(content, "array_access")) { @@ -161,7 +168,7 @@ caps caps_get(jinja::program & prog) { }, nullptr, // ctx_fn nullptr, // tools_fn - [&](bool, value & messages, value &, const std::string &) { + [&](context &, bool, value & messages, value &, const std::string &) { auto & content = messages->at(0)->at("content"); caps_print_stats(content, "messages[0].content"); if (!content->stats.used) { @@ -237,7 +244,7 @@ caps caps_get(jinja::program & prog) { }, }); }, - [&](bool success, value & messages, value & tools, const std::string &) { + [&](context &, bool success, value & messages, value & tools, const std::string &) { if (!success) { return; // Nothing can be inferred } @@ -330,7 +337,7 @@ caps caps_get(jinja::program & prog) { }, }); }, - [&](bool success, value & messages, value & tools, const std::string &) { + [&](context &, bool success, value & messages, value & tools, const std::string &) { if (!success) { result.supports_tool_calls = false; result.supports_tools = false; @@ -432,7 +439,7 @@ caps caps_get(jinja::program & prog) { }, }); }, - [&](bool success, value & messages, value &, const std::string &) { + [&](context &, bool success, value & messages, value &, const std::string &) { if (!success) { result.supports_parallel_tool_calls = false; return; @@ -489,7 +496,7 @@ caps caps_get(jinja::program & prog) { caps_apply_preserve_reasoning(ctx, true); }, nullptr, // tools_fn - [&](bool, value &, value &, const std::string & output) { + [&](context &, bool, value &, value &, const std::string & output) { // note: we cannot use stats here because the reasoning_content may be used for "if" condition test, but not actually outputted in the final result if (output.find(reasoning_placeholder) != std::string::npos) { result.supports_preserve_reasoning = true; @@ -497,6 +504,32 @@ caps caps_get(jinja::program & prog) { } ); + JJ_DEBUG("%s\n", ">>> Running capability check: reasoning effort"); + + // case: reasoning effort level + caps_try_execute( + prog, + [&]() { + // messages + return json::array({ + { + {"role", "user"}, + {"content", "User message"} + }, + }); + }, + [&](context & ctx) { + ctx.set_val("enable_thinking", mk_val(true)); + caps_apply_reasoning_effort(ctx, "low"); + }, + nullptr, // tools_fn + [&](context & ctx, bool, value &, value &, const std::string &) { + value effort = ctx.get_val("reasoning_effort"); + caps_print_stats(effort, "reasoning_effort"); + result.supports_reasoning_effort = effort->stats.used; + } + ); + JJ_DEBUG("%s\n", result.to_string().c_str()); return result; diff --git a/common/jinja/caps.h b/common/jinja/caps.h index a290cd7da..b81dd95f2 100644 --- a/common/jinja/caps.h +++ b/common/jinja/caps.h @@ -16,6 +16,9 @@ struct caps { // supports preserve reasoning trace in the full history, not just the last assistant message bool supports_preserve_reasoning = false; + // supports reasoning effort levels + bool supports_reasoning_effort = false; + // one of the 2 content capabilities must be true bool supports_string_content = true; bool supports_typed_content = false; @@ -32,5 +35,6 @@ struct caps { caps caps_get(jinja::program & prog); void caps_apply_preserve_reasoning(jinja::context & ctx, bool enabled); +void caps_apply_reasoning_effort(jinja::context & ctx, const std::string & effort); } // namespace jinja diff --git a/common/jinja/runtime.cpp b/common/jinja/runtime.cpp index afba0025f..9cc5cfa46 100644 --- a/common/jinja/runtime.cpp +++ b/common/jinja/runtime.cpp @@ -266,7 +266,7 @@ value binary_expression::execute_impl(context & ctx) { return res; } for (int64_t i = 0; i < repeat; ++i) { - res->val_str = res->val_str.append(str); + res->val_str.append(str); } return res; } diff --git a/common/jinja/runtime.h b/common/jinja/runtime.h index 0884a1592..69bd683c6 100644 --- a/common/jinja/runtime.h +++ b/common/jinja/runtime.h @@ -763,14 +763,22 @@ struct runtime { gather_string_parts_recursive(val, parts); // join consecutive parts with the same type auto & p = parts->val_str.parts; - for (size_t i = 1; i < p.size(); ) { - if (p[i].is_input == p[i - 1].is_input) { - p[i - 1].val += p[i].val; - p.erase(p.begin() + i); + if (p.empty()) { + return parts; + } + size_t w = 0; + for (size_t r = 1; r < p.size(); r++) { + if (p[w].is_input == p[r].is_input) { + p[w].val += p[r].val; } else { - i++; + w++; + if (w != r) { + // the guard is needed, self-move leaves the string in an unspecified state + p[w] = std::move(p[r]); + } } } + p.resize(w + 1); return parts; } diff --git a/common/jinja/string.cpp b/common/jinja/string.cpp index 8087e15b3..bde679e4e 100644 --- a/common/jinja/string.cpp +++ b/common/jinja/string.cpp @@ -103,7 +103,7 @@ void string::mark_input_based_on(const string & other) { } } -string string::append(const string & other) { +string & string::append(const string & other) { for (const auto & part : other.parts) { parts.push_back(part); } diff --git a/common/jinja/string.h b/common/jinja/string.h index c4963000a..669afb8f1 100644 --- a/common/jinja/string.h +++ b/common/jinja/string.h @@ -47,7 +47,7 @@ struct string { // mark this string as input if other has ALL parts as input void mark_input_based_on(const string & other); - string append(const string & other); + string & append(const string & other); // in-place transformations diff --git a/common/preset.cpp b/common/preset.cpp index 0b29af883..4c61e93ee 100644 --- a/common/preset.cpp +++ b/common/preset.cpp @@ -365,8 +365,25 @@ struct local_model { std::string name; std::string path; std::string path_mmproj; + std::string path_draft; }; +// TODO @ngxson: handle "eagle3-" when it's supported by common_speculative_types_from_gguf() +static const char * draft_prefixes[] = { "mtp-", "dspark-", "dflash-" }; + +static bool is_mmproj_file(const std::string & fname) { + return fname.find("mmproj") != std::string::npos; +} + +static bool is_draft_file(const std::string & fname) { + for (const auto & prefix : draft_prefixes) { + if (fname.rfind(prefix, 0) == 0) { + return true; + } + } + return false; +} + common_presets common_preset_context::load_from_models_dir(const std::string & models_dir) const { if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) { throw std::runtime_error(string_format("error: '%s' does not exist or is not a directory\n", models_dir.c_str())); @@ -378,10 +395,15 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m common_file_info model_file; common_file_info first_shard_file; common_file_info mmproj_file; + common_file_info draft_file; for (const auto & file : files) { if (string_ends_with(file.name, ".gguf")) { - if (file.name.find("mmproj") != std::string::npos) { + if (is_mmproj_file(file.name)) { mmproj_file = file; + } else if (is_draft_file(file.name)) { + if (draft_file.path.empty()) { + draft_file = file; // first sidecar found wins + } } else if (file.name.find("-00001-of-") != std::string::npos) { first_shard_file = file; } else { @@ -393,7 +415,8 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m local_model model{ /* name */ name, /* path */ first_shard_file.path.empty() ? model_file.path : first_shard_file.path, - /* path_mmproj */ mmproj_file.path // can be empty + /* path_mmproj */ mmproj_file.path, // can be empty + /* path_draft */ draft_file.path // can be empty }; if (!model.path.empty()) { models.push_back(model); @@ -405,13 +428,17 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m if (file.is_dir) { scan_subdir(file.path, file.name); } else if (string_ends_with(file.name, ".gguf")) { + if (is_mmproj_file(file.name) || is_draft_file(file.name)) { + continue; // companion file, cannot be loaded as a model on its own + } // single file model std::string name = file.name; string_replace_all(name, ".gguf", ""); local_model model{ /* name */ name, /* path */ file.path, - /* path_mmproj */ "" + /* path_mmproj */ "", + /* path_draft */ "" }; models.push_back(model); } @@ -426,6 +453,9 @@ common_presets common_preset_context::load_from_models_dir(const std::string & m if (!model.path_mmproj.empty()) { preset.set_option(*this, "LLAMA_ARG_MMPROJ", model.path_mmproj); } + if (!model.path_draft.empty()) { + preset.set_option(*this, "LLAMA_ARG_SPEC_DRAFT_MODEL", model.path_draft); + } out[preset.name] = preset; } diff --git a/conversion/__init__.py b/conversion/__init__.py index 695289b73..f4d475de7 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -161,6 +161,8 @@ TEXT_MODEL_MAP: dict[str, str] = { "MiniCPM3ForCausalLM": "minicpm", "MiniCPMForCausalLM": "minicpm", "MiniCPMV4_6ForConditionalGeneration": "minicpm", + "MiniMaxText01ForCausalLM": "minimax", + "MiniMaxM1ForCausalLM": "minimax", "MiniMaxM2ForCausalLM": "minimax", "MiniMaxM3SparseForCausalLM": "minimax", "MiniMaxM3SparseForConditionalGeneration": "minimax", diff --git a/conversion/minimax.py b/conversion/minimax.py index c2175cc93..d7a00bac9 100644 --- a/conversion/minimax.py +++ b/conversion/minimax.py @@ -1,13 +1,121 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import Iterable, Sequence, TYPE_CHECKING import torch if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, TextModel, MmprojModel, gguf +from .base import ModelBase, TextModel, MmprojModel, gguf, logger + + +@ModelBase.register("MiniMaxText01ForCausalLM") +@ModelBase.register("MiniMaxM1ForCausalLM") +class MiniMaxText01Model(TextModel): + model_arch = gguf.MODEL_ARCH.MINIMAX01 + + def _get_suppress_tokens(self) -> Sequence[int] | None: + import json + from transformers import AutoTokenizer + from .base import LazyTorchTensor + + # check added tokens embeddings in embeddings tensor for zero-valued embeddings + # they get in the way of the token sampling process and must be suppressed + + tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) + tokenizer_vocab_size = tokenizer.vocab_size + + with open(self.dir_model / "model.safetensors.index.json", "r", encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] + + embeddings_tensor_name = "model.embed_tokens.weight" + embeddings_shard_name = weight_map[embeddings_tensor_name] + with gguf.utility.SafetensorsLocal(self.dir_model / embeddings_shard_name) as model_shard: + embeddings_data = model_shard[embeddings_tensor_name] + + embeddings_weights_dtype = LazyTorchTensor._dtype_str_map[embeddings_data.dtype] + embeddings_weights = torch.from_numpy(embeddings_data.mmap_bytes()).view(embeddings_weights_dtype).reshape(embeddings_data.shape) + embeddings_vocab_size = embeddings_weights.shape[0] + + embeddings_added_tokens = embeddings_weights[tokenizer_vocab_size:embeddings_vocab_size] + embeddings_zero_rows = torch.all(embeddings_added_tokens == 0, dim=1) + tokens_zero_embeddings_ids = (torch.nonzero(embeddings_zero_rows, as_tuple=False).flatten() + tokenizer_vocab_size).tolist() + + return tokens_zero_embeddings_ids + + def set_vocab(self) -> None: + from pathlib import Path + + self._set_vocab_gpt2() + + for tmpl_file in [ + self.dir_model / "chat_template.jinja", + Path(__file__).parent.parent / "models" / "templates" / "MiniMax-M1.jinja" + ]: + if tmpl_file.is_file(): + self.gguf_writer.add_chat_template(tmpl_file.read_text(encoding="utf-8")) + logger.info(f"Chat template overridden with {tmpl_file}.") + break + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + suppress_tokens = self._get_suppress_tokens() + if suppress_tokens: + logger.info(f"Suppressing tokens with zero embeddings {suppress_tokens}") + self.gguf_writer.add_suppress_tokens(suppress_tokens) + + layernorm_full_attention_alpha = self.hparams["layernorm_full_attention_alpha"] + layernorm_full_attention_beta = self.hparams["layernorm_full_attention_beta"] + layernorm_linear_attention_alpha = self.hparams["layernorm_linear_attention_alpha"] + layernorm_linear_attention_beta = self.hparams["layernorm_linear_attention_beta"] + layernorm_mlp_alpha = self.hparams["layernorm_mlp_alpha"] + layernorm_mlp_beta = self.hparams["layernorm_mlp_beta"] + assert layernorm_full_attention_alpha == layernorm_linear_attention_alpha == layernorm_mlp_alpha + assert layernorm_full_attention_beta == layernorm_linear_attention_beta == layernorm_mlp_beta == 1.0 + # we do not store the layernorm betas as they are all 1.0 + # layernorm alphas are stored as single residual_scale hparam + self.gguf_writer.add_residual_scale(layernorm_full_attention_alpha) + + self.gguf_writer.add_rope_dimension_count(self.hparams["rotary_dim"]) + + _experts: list[dict[str, Tensor]] | None = None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # process the experts separately + if name.find("block_sparse_moe.experts") != -1: + n_experts = self.hparams["num_local_experts"] + + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + # merge the experts into a single 3d tensor + for wid in ["w1", "w2", "w3"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + + merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight" + + new_name = self.map_tensor_name(merged_name) + + yield from super().modify_tensors(data_torch, new_name, bid) + return + else: + return + + yield from super().modify_tensors(data_torch, name, bid) @ModelBase.register("MiniMaxM2ForCausalLM") diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 405f1e864..604cc6902 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2477,7 +2477,8 @@ extern "C" { struct ggml_tensor * A, struct ggml_tensor * B, struct ggml_tensor * C, - struct ggml_tensor * ids); + struct ggml_tensor * ids, + int64_t K); // partition into non-overlapping windows with padding if needed // example: diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 50b2124c2..ac91080e6 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -472,6 +472,8 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; case GGML_OP_CONV_2D: return ggml_is_contiguous(op->src[0]); + case GGML_OP_SSM_SCAN: + return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1; default: return true; } diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 25bb74383..001e1ae85 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -9644,11 +9644,13 @@ static void ggml_compute_forward_ssm_scan_f32( const int64_t ng = src4->ne[1]; const int64_t nt = src1->ne[2]; // number of tokens per sequence const int64_t ns = src1->ne[3]; // number of sequences in the batch + const int64_t K = ggml_get_op_params_i32(dst, 0); // can't use ggml_nbytes because src1 is not necessarily contiguous const int64_t s_off = ggml_nelements(src1) * ggml_element_size(src1); - GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*ns == ggml_nelements(dst)); + GGML_ASSERT(K >= 1); + GGML_ASSERT(ggml_nelements(src1) + K*nc*nr*nh*ns == ggml_nelements(dst)); GGML_ASSERT(src0->nb[0] == sizeof(float)); GGML_ASSERT(src1->nb[0] == sizeof(float)); GGML_ASSERT(src2->nb[0] == sizeof(float)); @@ -9657,6 +9659,7 @@ static void ggml_compute_forward_ssm_scan_f32( GGML_ASSERT(src5->nb[0] == sizeof(float)); GGML_ASSERT(src6->nb[0] == sizeof(int32_t)); GGML_ASSERT(nh % ng == 0); + GGML_ASSERT(src3->ne[0] == 1 || K == 1); // heads per thread const int dh = (nh + nth - 1)/nth; @@ -9831,6 +9834,13 @@ static void ggml_compute_forward_ssm_scan_f32( } } } + const int64_t slot = nt - 1 - i2; + if (K > 1 && slot > 0 && slot < K) { + float * s_snapshot = (float *) ((char *) dst->data + s_off + (slot*ns + i3)*(src0->nb[3])); + for (int h = ih0; h < ih1; ++h) { + memcpy((char *) s_snapshot + h*src0->nb[2], (char *) s + h*src0->nb[2], src0->nb[2]); + } + } // use the output as the source when it's not the first token-wise iteration s0 = s; } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 991d4a187..51e11ee68 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5202,11 +5202,17 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16); case GGML_OP_SSM_SCAN: { + const int32_t K = ggml_get_op_params_i32(op, 0); + if (op->src[3]->ne[0] == 1) { // Mamba2 // (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0) return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0; } else { + if (K > 1) { + return false; + } + // Mamba // (kernel only supports d_state == 16, d_head == 1, n_head % 128 == 0, n_group == 1) return op->src[0]->ne[0] == 16 && op->src[0]->ne[1] == 1 && op->src[0]->ne[2] % 128 == 0 && op->src[4]->ne[1] == 1; diff --git a/ggml/src/ggml-cuda/ssm-scan.cu b/ggml/src/ggml-cuda/ssm-scan.cu index f3418c2af..ef342f01f 100644 --- a/ggml/src/ggml-cuda/ssm-scan.cu +++ b/ggml/src/ggml-cuda/ssm-scan.cu @@ -149,7 +149,7 @@ __global__ void __launch_bounds__(d_state, 1) const int src0_nb2, const int src0_nb3, const int src1_nb2, const int src1_nb3, const int src2_nb1, const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3, - const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok) { + const int64_t s_off, const int64_t n_head, const int64_t d_head, const int64_t n_group, const int64_t n_tok, const int64_t K) { const float * GGML_CUDA_RESTRICT src0 = src0_ptr; const float * GGML_CUDA_RESTRICT src1 = src1_ptr; const float * GGML_CUDA_RESTRICT src2 = src2_ptr; @@ -217,6 +217,16 @@ __global__ void __launch_bounds__(d_state, 1) if (lane == 0) { y_warp[i * stride_y] = state_sum; } + + // Slot 0 is the final state written below; slots 1..K-1 are rollback snapshots. + const int64_t slot = n_tok - 1 - i; + if (K > 1 && slot > 0 && slot < K) { + float * s_snapshot_warp = (float *) ((char *) dst + s_off + (slot * gridDim.y + seq_idx) * src0_nb3 + head_idx * src0_nb2 + head_off * d_state); +#pragma unroll + for (int j = 0; j < c_factor; j++) { + s_snapshot_warp[WARP_SIZE * j + lane] = state[j]; + } + } } // write back the state @@ -232,7 +242,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa const int src2_nb2, const int src3_nb1, const int src4_nb2, const int src4_nb3, const int src5_nb2, const int src5_nb3, const int64_t s_off, const int64_t d_state, const int64_t head_dim, const int64_t n_head, const int64_t n_group, const int64_t n_tok, const int64_t n_seq, - cudaStream_t stream) { + const int64_t K, cudaStream_t stream) { // NOTE: if you change conditions here, be sure to update the corresponding supports_op condition! if (src3_nb1 == sizeof(float)) { // Mamba-2 @@ -245,7 +255,7 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa ggml_cuda_kernel_launch(ssm_scan_f32_group<128/WARP_SIZE, 128>, launch_params, src0, src1, src2, src3, src4, src5, src6, dst, src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1, - src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok); + src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K); } else if (d_state == 256) { // Falcon-H1 constexpr int threads = 256; constexpr int num_warps = threads/WARP_SIZE; @@ -255,12 +265,13 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa ggml_cuda_kernel_launch(ssm_scan_f32_group<256/WARP_SIZE, 256>, launch_params, src0, src1, src2, src3, src4, src5, src6, dst, src0_nb2, src0_nb3, src1_nb2, src1_nb3, src2_nb1, src2_nb2, src3_nb1, - src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok); + src4_nb2, src4_nb3, src5_nb2, src5_nb3, s_off, n_head, head_dim, n_group, n_tok, K); } else { GGML_ABORT("doesn't support d_state!=(128 or 256)."); } } else { // Mamba-1 + GGML_ASSERT(K == 1); constexpr int threads = 128; GGML_ASSERT(n_head % threads == 0); GGML_ASSERT(head_dim == 1); @@ -769,10 +780,12 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int64_t ng = src4->ne[1]; // n_group const int64_t n_t = src1->ne[2]; // number of tokens per sequence const int64_t n_s = src1->ne[3]; // number of sequences in the batch + const int32_t K_param = ggml_get_op_params_i32(dst, 0); + const int64_t K = K_param > 0 ? K_param : 1; const int64_t s_off = ggml_nelements(src1) * sizeof(float); - GGML_ASSERT(ggml_nelements(src1) + nc*nr*nh*n_s == ggml_nelements(dst)); + GGML_ASSERT(ggml_nelements(src1) + K*nc*nr*nh*n_s == ggml_nelements(dst)); GGML_ASSERT(src0->nb[0] == sizeof(float)); GGML_ASSERT(src1->nb[0] == sizeof(float)); GGML_ASSERT(src2->nb[0] == sizeof(float)); @@ -780,6 +793,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { GGML_ASSERT(src4->nb[0] == sizeof(float)); GGML_ASSERT(src5->nb[0] == sizeof(float)); GGML_ASSERT(src6->nb[0] == sizeof(int32_t)); + GGML_ASSERT(src3->ne[0] == 1 || K == 1); const float * src0_d = (const float *) src0->data; const float * src1_d = (const float *) src1->data; @@ -814,6 +828,7 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const bool is_mamba2 = (src3->nb[1] == sizeof(float)); const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const bool use_ssd = is_mamba2 && n_t > SSM_SSD_MIN_TOKENS + && K == 1 && n_t <= SSM_SSD_MAX_TOKENS && GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_TURING @@ -841,5 +856,5 @@ void ggml_cuda_op_ssm_scan(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ssm_scan_f32_cuda(src0_d, src1_d, src2_d, src3_d, src4_d, src5_d, src6_d, dst_d, src0->nb[2], src0->nb[3], src1->nb[2], src1->nb[3], src2->nb[1], src2->nb[2], src3->nb[1], src4->nb[2], src4->nb[3], src5->nb[2], src5->nb[3], - s_off, nc, nr, nh, ng, n_t, n_s, stream); + s_off, nc, nr, nh, ng, n_t, n_s, K, stream); } diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index c91aa16de..bccc86c9e 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1382,9 +1382,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te ggml_is_contiguous_rows(op->src[1]) && ggml_is_contiguous_rows(op->src[2]) && ggml_is_contiguous_rows(op->src[3]); - case GGML_OP_SSM_CONV: case GGML_OP_SSM_SCAN: return has_simdgroup_reduction; + case GGML_OP_SSM_CONV: + return has_simdgroup_reduction; case GGML_OP_RWKV_WKV6: case GGML_OP_RWKV_WKV7: return true; diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index cf32c5c5b..1f6e8c48b 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -880,6 +880,7 @@ typedef struct { int64_t n_group; int64_t n_seq_tokens; int64_t n_seqs; + int64_t K; uint64_t s_off; uint64_t nb00; uint64_t nb01; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 6d324056d..b7f9b2d0d 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -1710,6 +1710,10 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) { const int64_t n_group = ne41; const int64_t n_seq_tokens = ne12; const int64_t n_seqs = ne13; + const int64_t K = ggml_get_op_params_i32(op, 0); + + GGML_ASSERT(K >= 1); + GGML_ASSERT(ggml_nelements(op->src[1]) + K*d_state*d_inner*n_head*n_seqs == ggml_nelements(op)); ggml_metal_kargs_ssm_scan args = { /*.d_state =*/ d_state, @@ -1718,6 +1722,7 @@ int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) { /*.n_group =*/ n_group, /*.n_seq_tokens =*/ n_seq_tokens, /*.n_seqs =*/ n_seqs, + /*.K =*/ K, /*.s_off =*/ ggml_nelements(op->src[1]) * sizeof(float), /*.nb00 =*/ nb00, /*.nb01 =*/ nb01, diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index b38b23edc..243c997fc 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -2429,6 +2429,8 @@ kernel void kernel_ssm_scan_f32( const int32_t nh = args.n_head; const int32_t ng = args.n_group; const int32_t n_t = args.n_seq_tokens; + const int32_t n_s = args.n_seqs; + const int32_t K = args.K; const int32_t s_off = args.s_off; @@ -2487,6 +2489,12 @@ kernel void kernel_ssm_scan_f32( // recurse s0 = s; + const int32_t slot = n_t - 1 - (i2 + t); + if (slot > 0 && slot < K) { + device float * s_snapshot = (device float *) ((device char *) s_buff + (int64_t) slot*n_s*args.nb03); + s_snapshot[i] = s; + } + B += args.ns42; C += args.ns52; } diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index ed3e9ee68..ff86fe2a1 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1867,6 +1867,7 @@ struct vk_op_ssm_scan_push_constants { uint32_t nb42, nb43, nb52, nb53; uint32_t s_off; uint32_t n_head, d_head, n_group, n_tok; + uint32_t n_seq, K; }; struct vk_op_ssm_conv_push_constants { uint32_t nb01, nb02; @@ -2071,7 +2072,7 @@ struct ggml_vk_garbage_collector { static void ggml_vk_preallocate_buffers(ggml_backend_vk_context * ctx, vk_context subctx); static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested = nullptr); static void ggml_pipeline_allocate_descriptor_sets(ggml_backend_vk_context * ctx); -static bool ggml_vk_intel_windows_driver_equals_or_newer_than(uint32_t driver_version, uint32_t threshold_major, uint32_t threshold_minor); +static bool ggml_vk_intel_windows_driver_in_range(uint32_t driver_version, uint32_t lower_major, uint32_t lower_minor, uint32_t upper_major, uint32_t upper_minor); static bool vk_memory_logger_enabled = false; @@ -3967,7 +3968,10 @@ static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vec } // Needs to be kept up to date on shader changes - const uint32_t bank_conflict_offset = device->coopmat_support ? 8 : 1; + // Needs to stay aligned with ggml_vk_mul_mm_spec. + const bool intel_shmem_stride_pad_zero = device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && + device->driver_id == vk::DriverId::eIntelProprietaryWindows; + const uint32_t bank_conflict_offset = intel_shmem_stride_pad_zero ? 0 : (device->coopmat_support ? 8 : 1); const uint32_t type_size = device->fp16 ? sizeof(ggml_fp16_t) : sizeof(float); const uint32_t warps = warptile[0] / warptile[10]; @@ -4584,8 +4588,13 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } #endif - auto const &ggml_vk_mul_mm_spec = [](std::vector spec, bool aligned) { - spec.push_back(aligned ? 1u : 0u); + auto const &ggml_vk_mul_mm_spec = [&device](std::vector spec, bool aligned) { + spec.push_back(aligned ? 1u : 0u); // constantID=11: ALIGNED + if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && + device->driver_id == vk::DriverId::eIntelProprietaryWindows) { + spec.push_back(0u); // constantID=12: SHMEM_STRIDE_PAD = 0 + spec.push_back(1u); // constantID=13: APPLY_SLM_A_RESHAPE = true + } return spec; }; @@ -5747,10 +5756,9 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_argmax_f32, "argmax_f32", argmax_f32_len, argmax_f32_data, "main", 2, sizeof(vk_op_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); ggml_vk_create_pipeline(device, device->pipeline_sum_rows_f32, "sum_rows_f32", sum_rows_f32_len, sum_rows_f32_data, "main", 2, sizeof(vk_op_sum_rows_push_constants), {1, 1, 1}, { device->subgroup_size }, 1); - // Intel Windows driver older than 32.0.101.8860 will crash when using fwht kernels on Xe2+ GPUS so we gate that here + // Intel Windows driver in range [32.0.101.8509, 32.0.101.8860) will crash when using fwht kernels so we gate that here const bool can_use_fwht = device->driver_id != vk::DriverId::eIntelProprietaryWindows || - device->architecture != vk_device_architecture::INTEL_XE2 || - (device->architecture == vk_device_architecture::INTEL_XE2 && ggml_vk_intel_windows_driver_equals_or_newer_than(device->properties.driverVersion, 101, 8860)); + !ggml_vk_intel_windows_driver_in_range(device->properties.driverVersion, 101, 8509, 101, 8860); if (can_use_fwht && device->subgroup_basic && device->subgroup_shuffle) { int idx = 0; for (uint32_t n : {64, 128, 256, 512}) { @@ -12764,7 +12772,8 @@ static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx, (uint32_t)src4->nb[2], (uint32_t)src4->nb[3], (uint32_t)src5->nb[2], (uint32_t)src5->nb[3], (uint32_t)s_off, - n_head, head_dim, n_group, n_tok + n_head, head_dim, n_group, n_tok, + n_seq, (uint32_t) ggml_get_op_params_i32(dst, 0) }; vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); @@ -18902,17 +18911,23 @@ static uint32_t ggml_vk_intel_shader_core_count(const vk::PhysicalDevice& vkdev) } } -static bool ggml_vk_intel_windows_driver_equals_or_newer_than(uint32_t driver_version, uint32_t threshold_major, uint32_t threshold_minor) { +// checks whether lower <= driver_version < upper, with each bound given as xxx.yyyy +static bool ggml_vk_intel_windows_driver_in_range(uint32_t driver_version, uint32_t lower_major, uint32_t lower_minor, uint32_t upper_major, uint32_t upper_minor) { #if defined(_WIN32) // Intel Windows encodes xxx.yyyy as [31:14].[13:0]. const uint32_t major = driver_version >> 14; const uint32_t minor = driver_version & 0x3fff; - return major > threshold_major || (major == threshold_major && minor >= threshold_minor); + const bool ge_lower = major > lower_major || (major == lower_major && minor >= lower_minor); + const bool lt_upper = major < upper_major || (major == upper_major && minor < upper_minor); + + return ge_lower && lt_upper; #else GGML_UNUSED(driver_version); - GGML_UNUSED(threshold_major); - GGML_UNUSED(threshold_minor); + GGML_UNUSED(lower_major); + GGML_UNUSED(lower_minor); + GGML_UNUSED(upper_major); + GGML_UNUSED(upper_minor); return true; #endif } @@ -19450,8 +19465,9 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * } else if (tensor->op == GGML_OP_ADD_ID) { tensor_clone = ggml_add_id(ggml_ctx, src_clone[0], src_clone[1], src_clone[2]); } else if (tensor->op == GGML_OP_SSM_SCAN) { + const int32_t K = ggml_get_op_params_i32(tensor, 0); tensor_clone = ggml_ssm_scan(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], - src_clone[3], src_clone[4], src_clone[5], src_clone[6]); + src_clone[3], src_clone[4], src_clone[5], src_clone[6], K); } else if (tensor->op == GGML_OP_SSM_CONV) { tensor_clone = ggml_ssm_conv(ggml_ctx, src_clone[0], src_clone[1]); } else if (tensor->op == GGML_OP_ROLL) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index 57c0410e4..3df88044a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -119,10 +119,13 @@ layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working wit #endif #ifdef COOPMAT -#define SHMEM_STRIDE (BK / 2 + 4) +layout(constant_id = 12) const uint SHMEM_STRIDE_PAD = 4; +layout(constant_id = 13) const bool APPLY_SLM_A_RESHAPE = false; #else -#define SHMEM_STRIDE (BK / 2 + 1) +const uint SHMEM_STRIDE_PAD = 1; +const bool APPLY_SLM_A_RESHAPE = false; #endif +#define SHMEM_STRIDE (BK / 2 + SHMEM_STRIDE_PAD) shared FLOAT_TYPEV2 buf_a[BM * SHMEM_STRIDE]; shared FLOAT_TYPEV2 buf_b[BN * SHMEM_STRIDE]; @@ -302,7 +305,7 @@ void main() { [[unroll]] for (uint i = 0; i < BK; i += TK) { [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { // Load from shared into cache - coopMatLoad(cache_a, buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(cache_a, buf_a, a_shmem_index(warp_r * WM + cm_row * TM, i / 2), a_shmem_stride(), gl_CooperativeMatrixLayoutRowMajor); [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { coopMatLoad(cache_b, buf_b, (warp_c * WN + cm_col * TN) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl index 63af2ce68..7d852dced 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl @@ -1,60 +1,76 @@ +// k_pair is the K coordinate measured in FLOAT_TYPEV2 elements. +uint a_shmem_index(uint m, uint k_pair) { + if (APPLY_SLM_A_RESHAPE) { + const uint tile_width = TK / 2; + return (k_pair / tile_width) * BM * tile_width + + m * tile_width + + k_pair % tile_width; + } + return m * SHMEM_STRIDE + k_pair; +} + +uint a_shmem_stride() { + return APPLY_SLM_A_RESHAPE ? TK / 2 : SHMEM_STRIDE; +} + +void store_a(uint m, uint k_pair, FLOAT_TYPEV2 value) { + buf_a[a_shmem_index(m, k_pair)] = value; +} + void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uint idx_m, const uint block, const uint end_k) { #if defined(DATA_A_F32) || defined(DATA_A_F16) #if LOAD_VEC_A == 8 if (ALIGNED != 0) { const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + const uint k_pair = row * LOAD_VEC_A / 2; FLOAT_TYPEV8 aa = FLOAT_TYPEV8(data_a[idx]); - buf_a[buf_idx ] = aa[0].xy; - buf_a[buf_idx + 1] = aa[0].zw; - buf_a[buf_idx + 2] = aa[1].xy; - buf_a[buf_idx + 3] = aa[1].zw; + store_a(col, k_pair, aa[0].xy); + store_a(col, k_pair + 1, aa[0].zw); + store_a(col, k_pair + 2, aa[1].xy); + store_a(col, k_pair + 3, aa[1].zw); return; } #elif LOAD_VEC_A == 4 if (ALIGNED != 0) { const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + const uint k_pair = row * LOAD_VEC_A / 2; FLOAT_TYPEV4 aa = FLOAT_TYPEV4(data_a[idx]); - buf_a[buf_idx ] = aa.xy; - buf_a[buf_idx + 1] = aa.zw; + store_a(col, k_pair, aa.xy); + store_a(col, k_pair + 1, aa.zw); return; } #endif const uint idx = pos_a + col * p.stride_a + row * 2; - const uint buf_idx = col * SHMEM_STRIDE + row; if (idx_m < p.M && block + row * 2 + 1 < end_k) { - buf_a[buf_idx] = FLOAT_TYPEV2(data_a_scalar[idx], - data_a_scalar[idx + 1]); + store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx], + data_a_scalar[idx + 1])); } else if (idx_m < p.M && block + row * 2 < end_k) { - buf_a[buf_idx] = FLOAT_TYPEV2(data_a_scalar[idx], 0.0f); + store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx], 0.0f)); } else { - buf_a[buf_idx] = FLOAT_TYPEV2(0.0f); + store_a(col, row, FLOAT_TYPEV2(0.0f)); } #elif defined(DATA_A_BF16) #if LOAD_VEC_A == 4 if (ALIGNED != 0) { const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + const uint k_pair = row * LOAD_VEC_A / 2; FLOAT_TYPEV4 aa = FLOAT_TYPEV4(TO_FLOAT_TYPE(data_a[idx])); - buf_a[buf_idx ] = aa.xy; - buf_a[buf_idx + 1] = aa.zw; + store_a(col, k_pair, aa.xy); + store_a(col, k_pair + 1, aa.zw); return; } #endif const uint idx = pos_a + col * p.stride_a + row * 2; - const uint buf_idx = col * SHMEM_STRIDE + row; if (idx_m < p.M && block + row * 2 + 1 < end_k) { - buf_a[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), - TO_FLOAT_TYPE(data_a_scalar[idx + 1])); + store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), + TO_FLOAT_TYPE(data_a_scalar[idx + 1]))); } else if (idx_m < p.M && block + row * 2 < end_k) { - buf_a[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), 0.0f); + store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), 0.0f)); } else { - buf_a[buf_idx] = FLOAT_TYPEV2(0.0f); + store_a(col, row, FLOAT_TYPEV2(0.0f)); } #elif defined(DATA_A_Q4_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 4; const uint iqs = idx & 0x03; @@ -64,13 +80,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v0 = (vec4(unpack8(vui & 0x0F0F0F0F)) - 8.0f) * d; const vec4 v1 = (vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) - 8.0f) * d; - buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xy); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v0.zw); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(v1.xy); - buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.zw); + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, FLOAT_TYPEV2(v0.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw)); + store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy)); + store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw)); #elif defined(DATA_A_Q4_1) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 4; const uint iqs = idx & 0x03; @@ -80,13 +96,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v0 = vec4(unpack8(vui & 0x0F0F0F0F)) * dm.x + dm.y; const vec4 v1 = vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) * dm.x + dm.y; - buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xy); - buf_a[buf_idx + 1 ] = FLOAT_TYPEV2(v0.zw); - buf_a[buf_idx + 8 ] = FLOAT_TYPEV2(v1.xy); - buf_a[buf_idx + 9 ] = FLOAT_TYPEV2(v1.zw); + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, FLOAT_TYPEV2(v0.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw)); + store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy)); + store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw)); #elif defined(DATA_A_Q5_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -98,12 +114,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint vui = uint(data_a_packed16[ib].qs[iqs]); const vec4 v = (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f) * d; - - buf_a[buf_idx ] = FLOAT_TYPEV2(v.xz); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(v.yw); + store_a(col, row, FLOAT_TYPEV2(v.xz)); + store_a(col, row + 8, FLOAT_TYPEV2(v.yw)); #elif defined(DATA_A_Q5_1) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 4; const uint iqs = idx & 0x03; @@ -119,13 +133,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v0 = vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, ((vui >> 12) & 0xF) | qh1.y) * dm.x + dm.y; const vec4 v1 = vec4(((vui >> 16) & 0xF) | qh2.x, ((vui >> 20) & 0xF) | qh2.y, ((vui >> 24) & 0xF) | qh3.x, ((vui >> 28) & 0xF) | qh3.y) * dm.x + dm.y; - buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xz); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v1.xz); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(v0.yw); - buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.yw); + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, FLOAT_TYPEV2(v0.xz)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v1.xz)); + store_a(col, k_pair + 8, FLOAT_TYPEV2(v0.yw)); + store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.yw)); #elif defined(DATA_A_Q8_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -135,11 +149,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const i8vec2 v1 = unpack8(int32_t(data_a_packed16[ib].qs[2*iqs + 1])).xy; const vec4 v = vec4(v0.x, v0.y, v1.x, v1.y) * d; - buf_a[buf_idx ] = FLOAT_TYPEV2(v.xy); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); #elif defined(DATA_A_Q1_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 16; const uint iqs = idx & 0xfu; @@ -147,13 +161,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const float d = float(data_a[ib].d); const uint bits = uint(data_a[ib].qs[iqs]); - buf_a[buf_idx ] = FLOAT_TYPEV2((bits & 0x01u) != 0u ? d : -d, (bits & 0x02u) != 0u ? d : -d); - buf_a[buf_idx + 1] = FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d); - buf_a[buf_idx + 2] = FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d); - buf_a[buf_idx + 3] = FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2((bits & 0x01u) != 0u ? d : -d, (bits & 0x02u) != 0u ? d : -d)); + store_a(col, k_pair + 1, FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d)); + store_a(col, k_pair + 2, FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d)); + store_a(col, k_pair + 3, FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d)); #elif defined(DATA_A_Q2_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 16; const uint iqs = idx & 0xfu; @@ -161,11 +175,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib].d); const uint bits = uint(data_a[ib].qs[iqs]); - buf_a[buf_idx ] = d * (FLOAT_TYPEV2(bits & 3u, (bits >> 2u) & 3u) - FLOAT_TYPEV2(1.0f)); - buf_a[buf_idx + 1] = d * (FLOAT_TYPEV2((bits >> 4u) & 3u, bits >> 6u) - FLOAT_TYPEV2(1.0f)); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, d * (FLOAT_TYPEV2(bits & 3u, (bits >> 2u) & 3u) - FLOAT_TYPEV2(1.0f))); + store_a(col, k_pair + 1, d * (FLOAT_TYPEV2((bits >> 4u) & 3u, bits >> 6u) - FLOAT_TYPEV2(1.0f))); #elif defined(DATA_A_Q2_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -180,11 +194,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v = dm.x * float(scales & 0xF) * qs - dm.y * float(scales >> 4); - buf_a[buf_idx ] = FLOAT_TYPEV2(v.xy); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); #elif defined(DATA_A_TQ2_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 128; // 2 values per idx const uint iqs = (idx % 128) * 2; // elem 0,2,4..254 @@ -197,10 +211,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec2 v = d * (vec2((qs >> shift) & 3) - 1.0); - buf_a[buf_idx] = FLOAT_TYPEV2(v.xy); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); #elif defined(DATA_A_Q3_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 128; // 2 values per idx const uint iqs = idx % 128; // 0..127 @@ -220,11 +234,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec2 qs = vec2(unpack8((uint(data_a_packed16[ib].qs[qsi / 2]) >> qsshift) & 0x0303).xy); const vec2 hm = vec2(unpack8(((uint(data_a_packed16[ib].hmask[hmi / 2]) >> (4 * n + halfsplit)) & 0x0101 ^ 0x0101) << 2).xy); - buf_a[buf_idx] = FLOAT_TYPEV2(dl * (qs.x - hm.x), - dl * (qs.y - hm.y)); + store_a(col, row * LOAD_VEC_A / 2, FLOAT_TYPEV2(dl * (qs.x - hm.x), + dl * (qs.y - hm.y))); #elif defined(DATA_A_Q4_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -256,11 +269,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 q = vec4(unpack8((data_a_packed32[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F)); - buf_a[buf_idx ] = FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m)); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m))); + store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m))); #elif defined(DATA_A_Q5_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -295,11 +308,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint qh = ((data_a_packed32[ib].qh[qhi / 4] >> (iqs / 16)) & 0x01010101) << 4; const vec4 q = vec4(unpack8(qs | qh)); - buf_a[buf_idx ] = FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m)); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m))); + store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m))); #elif defined(DATA_A_Q6_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 128; // 2 values per idx const uint iqs = idx % 128; // 0..127 @@ -318,10 +331,9 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint qh = (uint(data_a_packed16[ib].qh[qhi]) >> qhshift) & 0x0303; const vec2 q = (vec2(unpack8(ql | (qh << 4)).xy) - 32) * dscale; - buf_a[buf_idx] = FLOAT_TYPEV2(q.x, q.y); + store_a(col, row * LOAD_VEC_A / 2, FLOAT_TYPEV2(q.x, q.y)); #elif defined(DATA_A_IQ1_S) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -334,13 +346,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA; const int16_t grid = int16_t(iq1s_grid[qs | (bitfieldExtract(qh, 3 * int(ib8 & 3), 3) << 8)]); + const uint k_pair = row * LOAD_VEC_A / 2; [[unroll]] for (int k = 0; k < 4; ++k) { - buf_a[buf_idx + k] = FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), - dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta)); + store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), + dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta))); } #elif defined(DATA_A_IQ1_M) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 32; // 8 values per idx const uint ib8 = idx % 32; @@ -356,13 +368,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA; const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]); + const uint k_pair = row * LOAD_VEC_A / 2; [[unroll]] for (int k = 0; k < 4; ++k) { - buf_a[buf_idx + k] = FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), - dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta)); + store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), + dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta))); } #elif defined(DATA_A_IQ2_XXS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -383,17 +395,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 grid0 = vec4(unpack8(grid.x)); const vec4 grid1 = vec4(unpack8(grid.y)); - buf_a[buf_idx ] = db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, - (sign & 2) != 0 ? -grid0.y : grid0.y); - buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, - (sign & 8) != 0 ? -grid0.w : grid0.w); - buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, - (sign & 32) != 0 ? -grid1.y : grid1.y); - buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, - (sign & 128) != 0 ? -grid1.w : grid1.w); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, + (sign & 2) != 0 ? -grid0.y : grid0.y)); + store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, + (sign & 8) != 0 ? -grid0.w : grid0.w)); + store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, + (sign & 32) != 0 ? -grid1.y : grid1.y)); + store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, + (sign & 128) != 0 ? -grid1.w : grid1.w)); #elif defined(DATA_A_IQ2_XS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -409,17 +421,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 grid0 = vec4(unpack8(grid.x)); const vec4 grid1 = vec4(unpack8(grid.y)); - buf_a[buf_idx ] = db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, - (sign & 2) != 0 ? -grid0.y : grid0.y); - buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, - (sign & 8) != 0 ? -grid0.w : grid0.w); - buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, - (sign & 32) != 0 ? -grid1.y : grid1.y); - buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, - (sign & 128) != 0 ? -grid1.w : grid1.w); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, + (sign & 2) != 0 ? -grid0.y : grid0.y)); + store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, + (sign & 8) != 0 ? -grid0.w : grid0.w)); + store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, + (sign & 32) != 0 ? -grid1.y : grid1.y)); + store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, + (sign & 128) != 0 ? -grid1.w : grid1.w)); #elif defined(DATA_A_IQ2_S) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 32; // 8 values per idx const uint ib8 = idx % 32; // 0..31 @@ -437,17 +449,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 grid0 = vec4(unpack8(grid.x)); const vec4 grid1 = vec4(unpack8(grid.y)); - buf_a[buf_idx ] = db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, - (sign & 2) != 0 ? -grid0.y : grid0.y); - buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, - (sign & 8) != 0 ? -grid0.w : grid0.w); - buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, - (sign & 32) != 0 ? -grid1.y : grid1.y); - buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, - (sign & 128) != 0 ? -grid1.w : grid1.w); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, + (sign & 2) != 0 ? -grid0.y : grid0.y)); + store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, + (sign & 8) != 0 ? -grid0.w : grid0.w)); + store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, + (sign & 32) != 0 ? -grid1.y : grid1.y)); + store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, + (sign & 128) != 0 ? -grid1.w : grid1.w)); #elif defined(DATA_A_IQ3_XXS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 64; // 4 values per idx const uint iqs = idx % 64; // 0..63 @@ -465,13 +477,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint grid = iq3xxs_grid[qs]; const vec4 v = db * vec4(unpack8(grid)); - buf_a[buf_idx ] = FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, - (sign & 2) != 0 ? -v.y : v.y); - buf_a[buf_idx + 1] = FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, - (sign & 8) != 0 ? -v.w : v.w); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, + (sign & 2) != 0 ? -v.y : v.y)); + store_a(col, k_pair + 1, FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, + (sign & 8) != 0 ? -v.w : v.w)); #elif defined(DATA_A_IQ3_S) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 64; // 4 values per idx const uint iqs = idx % 64; // 0..63 @@ -487,13 +499,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint32_t grid = iq3s_grid[qs | ((qh << (8 - (iqs % 8))) & 256)]; const vec4 v = db * vec4(unpack8(grid)); - buf_a[buf_idx ] = FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, - (sign & 2) != 0 ? -v.y : v.y); - buf_a[buf_idx + 1] = FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, - (sign & 8) != 0 ? -v.w : v.w); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, + (sign & 2) != 0 ? -v.y : v.y)); + store_a(col, k_pair + 1, FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, + (sign & 8) != 0 ? -v.w : v.w)); #elif defined(DATA_A_IQ4_XS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 64; // 4 values per idx const uint ib32 = (idx % 64) / 8; // 0..7 @@ -507,11 +519,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const float d = float(data_a[ib].d); const vec4 v = d * float(int(sl | (sh << 4)) - 32) * vec4(kvalues_iq4nl[qs.x], kvalues_iq4nl[qs.y], kvalues_iq4nl[qs.z], kvalues_iq4nl[qs.w]); - buf_a[buf_idx ] = FLOAT_TYPEV2(v.xy); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); #elif defined(DATA_A_IQ4_NL) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -519,13 +531,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const FLOAT_TYPE d = FLOAT_TYPE(data_a_packed16[ib].d); const uint vui = uint(data_a_packed16[ib].qs[iqs]); - buf_a[buf_idx ] = d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF], - kvalues_iq4nl[bitfieldExtract(vui, 8, 4)]); - buf_a[buf_idx + 8] = d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)], - kvalues_iq4nl[vui >> 12]); + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF], + kvalues_iq4nl[bitfieldExtract(vui, 8, 4)])); + store_a(col, k_pair + 8, d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)], + kvalues_iq4nl[vui >> 12])); #elif defined(DATA_A_MXFP4) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 8; const uint iqs = (idx & 0x07) * 2; @@ -536,38 +548,37 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin #ifdef USE_OCP_FP4 const float d = e8m0_to_fp32(data_a[ib].e); const u8vec2 packed = u8vec2(vui, vui2); - buf_a[buf_idx ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * FLOAT_TYPE(d); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * FLOAT_TYPE(d); + store_a(col, row, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * FLOAT_TYPE(d)); + store_a(col, row + 8, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * FLOAT_TYPE(d)); #else const float d = e8m0_to_fp32(data_a[ib].e) * 0.5; - buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, - kvalues_mxfp4[vui2 & 0xF] * d); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, - kvalues_mxfp4[vui2 >> 4] * d); + store_a(col, row, FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, + kvalues_mxfp4[vui2 & 0xF] * d)); + store_a(col, row + 8, FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, + kvalues_mxfp4[vui2 >> 4] * d)); #endif #elif defined(DATA_A_NVFP4) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - // lo and hi nibbles are 8 elements apart, which doesn't quite line up with - // how the thread mapping and buf_idx calculation works for other types. - const uint buf_idx = col * SHMEM_STRIDE + (row & 3) + (row & ~3) * 2; - const uint ib = idx / 16u; const uint sub = (idx & 0xC) >> 2; const uint iqs = (idx & 0xF) * 2; const uint vui = uint(data_a[ib].qs[iqs]); const uint vui2 = uint(data_a[ib].qs[iqs+1]); + // lo and hi nibbles are 8 elements apart, which doesn't quite line up with + // how the thread mapping and buf_idx calculation works for other types. + const uint eff_row = (row & 3) + (row & ~3) * 2; #ifdef USE_OCP_FP4 const FLOAT_TYPE d = FLOAT_TYPE(ue4m3_from_bits(data_a[ib].d[sub])); const u8vec2 packed = u8vec2(vui, vui2); - buf_a[buf_idx ] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * d; - buf_a[buf_idx + 4] = FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * d; + store_a(col, eff_row, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * d); + store_a(col, eff_row + 4, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * d); #else const float d = ue4m3_to_fp32(data_a[ib].d[sub]) * 0.5; - buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, - kvalues_mxfp4[vui2 & 0xF] * d); - buf_a[buf_idx + 4] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, - kvalues_mxfp4[vui2 >> 4] * d); + store_a(col, eff_row, FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, + kvalues_mxfp4[vui2 & 0xF] * d)); + store_a(col, eff_row + 4, FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, + kvalues_mxfp4[vui2 >> 4] * d)); #endif #endif } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp b/ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp index c7416206d..4fecb3aa5 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/ssm_scan.comp @@ -33,6 +33,8 @@ layout(push_constant) uniform PushConstants { uint d_head; uint n_group; uint n_tok; + uint n_seq; + uint K; }; float softplus(float x) { @@ -114,6 +116,14 @@ void main() { if (lane == 0) { d[y_base_idx + i * stride_y] = state_sum; } + + const uint slot = n_tok - 1u - i; + if (slot > 0u && slot < K) { + const uint snapshot_base_idx = s_base_idx + slot * n_seq * (nb03 / 4u); + [[unroll]] for (uint j = 0; j < c_factor; j++) { + d[snapshot_base_idx + SUBGROUP_SIZE * j + lane] = state[j]; + } + } } // write back the state diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 8ce31073f..8cf8ab342 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5604,7 +5604,10 @@ struct ggml_tensor * ggml_ssm_scan( struct ggml_tensor * A, struct ggml_tensor * B, struct ggml_tensor * C, - struct ggml_tensor * ids) { + struct ggml_tensor * ids, + int64_t K) { + GGML_ASSERT(K >= 1); + GGML_ASSERT(K <= INT32_MAX); GGML_ASSERT(ggml_is_contiguous(s)); GGML_ASSERT(ggml_is_contiguous(dt)); GGML_ASSERT(ggml_is_contiguous(A)); @@ -5641,11 +5644,12 @@ struct ggml_tensor * ggml_ssm_scan( if (A->ne[0] != 1) { // Mamba-1 has more granular decay factors GGML_ASSERT(A->ne[0] == d_state); + GGML_ASSERT(K == 1); } } // concatenated y + ssm_states - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]); + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + K*s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]); result->op = GGML_OP_SSM_SCAN; result->src[0] = s; @@ -5656,6 +5660,8 @@ struct ggml_tensor * ggml_ssm_scan( result->src[5] = C; result->src[6] = ids; + ggml_set_op_params_i32(result, 0, (int32_t) K); + return result; } diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 98c4fa169..a312c8aa0 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -565,6 +565,7 @@ class MODEL_ARCH(IntEnum): GROVEMOE = auto() APERTUS = auto() COGVLM = auto() + MINIMAX01 = auto() MINIMAXM2 = auto() MINIMAXM3 = auto() RND1 = auto() @@ -1271,6 +1272,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.SEED_OSS: "seed_oss", MODEL_ARCH.GROVEMOE: "grovemoe", MODEL_ARCH.APERTUS: "apertus", + MODEL_ARCH.MINIMAX01: "minimax-01", MODEL_ARCH.MINIMAXM2: "minimax-m2", MODEL_ARCH.MINIMAXM3: "minimax-m3", MODEL_ARCH.COGVLM: "cogvlm", @@ -4592,6 +4594,24 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN_CHEXP, MODEL_TENSOR.FFN_UP_CHEXP, ], + MODEL_ARCH.MINIMAX01: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_NORM_2, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], MODEL_ARCH.MINIMAXM2: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 79d270ab8..ee3a8d3aa 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -225,6 +225,7 @@ class TensorNameMap: "rwkv.blocks.{bid}.ln2", # rwkv6 "model.layers.{bid}.ln2", # rwkv7 "model.layers.{bid}.post_attention_layernorm", # cogvlm + "model.layers.{bid}.self_attn.norm", # minimax-01 ), # Attention query-key-value @@ -321,7 +322,7 @@ class TensorNameMap: "h.{bid}.self_attention.dense", # bloom "model.layers.{bid}.self_attn.o_proj", # llama-hf nemotron olmoe olmo2 phimoe "layers.{bid}.self_attn.o_proj", # embeddinggemma - "model.layers.{bid}.self_attn.out_proj", # lfm2 + "model.layers.{bid}.self_attn.out_proj", # lfm2 minimax-01 "model.layers.{bid}.self_attn.linear_attn", # deci "layers.{bid}.attention.wo", # llama-pth "encoder.layer.{bid}.attention.output.dense", # bert @@ -385,6 +386,7 @@ class TensorNameMap: "model.layers.{bid}.self_attn.gate_proj", # afmoe muse-glimmer "model.layers.{bid}.linear_attn.in_proj_z", # qwen3.5 "model.layers.{bid}.self_attn.g_proj", # step3.5 head-wise attention gate + "model.layers.{bid}.self_attn.output_gate", # minimax-01 ), # Feed-forward norm diff --git a/skills/code-review/SKILL.md b/skills/code-review/SKILL.md index b9372ddda..a17c11d7c 100644 --- a/skills/code-review/SKILL.md +++ b/skills/code-review/SKILL.md @@ -46,6 +46,8 @@ Mandatory on every review; any finding here is **blocking**. Rule of thumb: GGUF - **Sizes/counts from tensor dims:** validate before allocating. Products like `ne[i]*nb[i]`/nbytes can overflow on crafted dims into an undersized alloc then heap overflow. Overflow checks must run BEFORE the arithmetic they guard - padding/alignment macros wrap to 0 near `SIZE_MAX`, so a guard after the pad passes. - **GGUF strings/arrays:** cap declared lengths and element counts before using them to size a loop or buffer; validate element type and length before casting an array to a pointer or reading fixed indices (`[i+1]`, `[0..2]`). +- **Element-type confusion:** casting `gguf_get_arr_data()` or `tensor->data` to `float *`/`int32_t *` needs an element-type check first (`gguf_get_kv_type() == GGUF_TYPE_ARRAY` then `gguf_get_arr_type()`; `type == GGML_TYPE_F32` for tensors). A `UINT8` array or `I8` tensor passes every length check, then gets read 4 bytes per element - a nearby length check is not a type check. +- **Loaders:** `GGML_ASSERT` on a file-derived value aborts the process; throw instead where the caller already catches (vocab, model loader, clip). - **File-supplied counts indexing fixed arrays:** bound any count (e.g. layer/block count into a `LLAMA_MAX_*` array) before indexing; watch checks that only fire when an optional key is present. - **Declared vs actual array length:** check the declared length of a GGUF array against the count actually read, not just against a buffer size. - **Bounds comparisons:** flag narrowing casts (`size_t`->`int32_t`) and signed/unsigned mixing that can bypass a length check and copy past a buffer. diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 8ed9391d7..1cd9b758f 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -128,6 +128,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_SEED_OSS, "seed_oss" }, { LLM_ARCH_GROVEMOE, "grovemoe" }, { LLM_ARCH_APERTUS, "apertus" }, + { LLM_ARCH_MINIMAX_01, "minimax-01" }, { LLM_ARCH_MINIMAX_M2, "minimax-m2" }, { LLM_ARCH_MINIMAX_M3, "minimax-m3" }, { LLM_ARCH_COGVLM, "cogvlm" }, @@ -978,6 +979,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_MINIMAX_01: return true; default: return false; @@ -1001,6 +1003,8 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_NEMOTRON_H: + case LLM_ARCH_NEMOTRON_H_MOE: return true; default: return false; @@ -1031,6 +1035,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_GRANITE_HYBRID: case LLM_ARCH_LFM2: case LLM_ARCH_LFM2MOE: + case LLM_ARCH_MINIMAX_01: case LLM_ARCH_MINIMAX_M2: case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_MISTRAL4: diff --git a/src/llama-arch.h b/src/llama-arch.h index 18d9de186..24252b9da 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -153,6 +153,7 @@ enum llm_arch { LLM_ARCH_NANBEIGE, LLM_ARCH_QWEN3TTS, LLM_ARCH_POCKETTTS, + LLM_ARCH_MINIMAX_01, LLM_ARCH_UNKNOWN, }; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 3f105cfb1..cebdd9058 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -106,7 +106,7 @@ llama_context::llama_context( cparams.n_rs_seq = params.n_rs_seq; if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) { - LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model arch does not support recurrent partial rollback; clamping to 0\n", + LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n", __func__, cparams.n_rs_seq); cparams.n_rs_seq = 0; } @@ -2310,6 +2310,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_DEEPSEEK4 || (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || + model.arch == LLM_ARCH_MINIMAX_01 || model.arch == LLM_ARCH_MINIMAX_M3) { res = std::max(n_tokens * 40, 32u * model.n_tensors()); } else { diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 781277f3f..e3f0cf0ed 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -217,6 +217,13 @@ uint32_t llama_hparams::n_embd_s() const { return n_embd_head_kda * n_embd_head_kda * n_head(); // 128 * 128 * 32 = 524288 } + if (n_embd_head_la != 0) { + // for MiniMax-Text-01 linear attention layers + // Full recurrent state: head_dim * head_dim * n_head + // tensor shape for linear attention: [head_dim, head_dim, n_head] + return n_embd_head_la * n_embd_head_la * n_head(); // 128 * 128 * 64 = 1048576 + } + // corresponds to Mamba's ssm_states size return ssm_d_state * ssm_d_inner; } diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 57de80824..d9ac112b7 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -164,6 +164,9 @@ struct llama_hparams { uint32_t ssm_dt_rank = 0; uint32_t ssm_n_group = 0; + // for MiniMax-Text-01 linear attention + uint32_t n_embd_head_la = 0; + // for Kimi Linear KDA uint32_t n_embd_head_kda = 0; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 7dd140f7c..6481b62e4 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -316,15 +316,19 @@ namespace GGUFMeta { struct GGUFMeta::ArrayInfo arr_info = GGUFMeta::GKV::get_kv(ctx, kid); + bool type_ok = false; switch (arr_info.gt) { case GGUF_TYPE_UINT32: - case GGUF_TYPE_INT32: GGML_ASSERT((std::is_same::value) || - (std::is_same::value)); break; - case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same::value)); break; - case GGUF_TYPE_STRING: GGML_ASSERT((std::is_same::value)); break; + case GGUF_TYPE_INT32: type_ok = (std::is_same::value) || + (std::is_same::value); break; + case GGUF_TYPE_FLOAT32: type_ok = (std::is_same::value); break; + case GGUF_TYPE_STRING: type_ok = (std::is_same::value); break; default: throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str())); } + if (!type_ok) { + throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt))); + } if constexpr (std::is_same::value) { const size_t n_items = gguf_get_arr_n(ctx, kid); @@ -357,16 +361,20 @@ namespace GGUFMeta { struct GGUFMeta::ArrayInfo arr_info = GGUFMeta::GKV::get_kv(ctx, kid); + bool type_ok = false; switch (arr_info.gt) { case GGUF_TYPE_BOOL: case GGUF_TYPE_UINT32: - case GGUF_TYPE_INT32: GGML_ASSERT((std::is_same::value) || - (std::is_same::value)); break; - case GGUF_TYPE_FLOAT32: GGML_ASSERT((std::is_same::value)); break; - case GGUF_TYPE_STRING: GGML_ASSERT((std::is_same::value)); break; + case GGUF_TYPE_INT32: type_ok = (std::is_same::value) || + (std::is_same::value); break; + case GGUF_TYPE_FLOAT32: type_ok = (std::is_same::value); break; + case GGUF_TYPE_STRING: type_ok = (std::is_same::value); break; default: throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str())); } + if (!type_ok) { + throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt))); + } if (arr_info.length > N_MAX) { throw std::runtime_error(format("array length %u for key %s exceeds max %u", (uint32_t) arr_info.length, key.c_str(), (uint32_t) N_MAX)); @@ -1003,7 +1011,7 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w ggml_tensor * B = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs); ggml_tensor * C = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs); ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs); - op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids); + op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids, /*K=*/1); } break; case GGML_OP_RWKV_WKV6: { @@ -1179,7 +1187,7 @@ struct ggml_tensor * llama_model_loader::create_tensor( if (use_mmap) { static std::once_flag once; std::call_once(once, [] { - LLAMA_LOG_WARN("llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --no-mmap for better performance\n"); + LLAMA_LOG_WARN("llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --load-mode none for better performance\n"); }); } } else { diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 066518e30..587300757 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -122,6 +122,7 @@ #include "models/mimo2.cpp" #include "models/minicpm.cpp" #include "models/minicpm3.cpp" +#include "models/minimax-01.cpp" #include "models/minimax-m2.cpp" #include "models/minimax-m3.cpp" #include "models/mistral3.cpp" @@ -442,6 +443,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_grovemoe(params); case LLM_ARCH_APERTUS: return new llama_model_apertus(params); + case LLM_ARCH_MINIMAX_01: + return new llama_model_minimax_01(params); case LLM_ARCH_MINIMAX_M2: return new llama_model_minimax_m2(params); case LLM_ARCH_MINIMAX_M3: @@ -944,6 +947,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_290B: return "290B"; case LLM_TYPE_314B: return "314B"; case LLM_TYPE_405B: return "405B"; + case LLM_TYPE_456B: return "456B"; case LLM_TYPE_671B: return "671B"; case LLM_TYPE_SMALL: return "0.1B"; case LLM_TYPE_MEDIUM: return "0.4B"; @@ -2429,7 +2433,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_recr = [&](uint32_t il) { return hparams.is_recr(il) && hparams.n_ff(il) == 0; }; - } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE) { + } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_MINIMAX_01) { filter_attn = [&](uint32_t il) { return il < hparams.n_layer() && !hparams.is_recr(il); }; @@ -2850,6 +2854,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_SEED_OSS: case LLM_ARCH_GROVEMOE: case LLM_ARCH_APERTUS: + case LLM_ARCH_MINIMAX_01: case LLM_ARCH_MINIMAX_M2: case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_COGVLM: diff --git a/src/llama-model.h b/src/llama-model.h index 341cb66fb..510ed8498 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -99,6 +99,7 @@ enum llm_type { LLM_TYPE_290B, LLM_TYPE_314B, LLM_TYPE_405B, + LLM_TYPE_456B, LLM_TYPE_671B, LLM_TYPE_SMALL, LLM_TYPE_MEDIUM, @@ -271,6 +272,7 @@ struct llama_layer { struct ggml_tensor * wv = nullptr; struct ggml_tensor * wo = nullptr; struct ggml_tensor * wqkv = nullptr; + struct ggml_tensor * wg = nullptr; struct ggml_tensor * wq_a = nullptr; struct ggml_tensor * wq_b = nullptr; struct ggml_tensor * wkv_a_mqa = nullptr; diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 39b1e672d..d25d82641 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2215,6 +2215,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { // Kimi-K2 doesn't need merges, skip LLAMA_LOG_INFO("%s: Kimi-K2 tokenizer detected, skipping BPE merges\n", __func__); } else { + if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str())); + } const int n_merges = gguf_get_arr_n(ctx, merges_keyidx); for (int i = 0; i < n_merges; i++) { const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i); @@ -2264,8 +2268,13 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { const int precompiled_charsmap_keyidx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str()); if (precompiled_charsmap_keyidx != -1) { + if (gguf_get_kv_type(ctx, precompiled_charsmap_keyidx) != GGUF_TYPE_ARRAY) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str())); + } const gguf_type pc_type = gguf_get_arr_type(ctx, precompiled_charsmap_keyidx); - GGML_ASSERT(pc_type == GGUF_TYPE_INT8 || pc_type == GGUF_TYPE_UINT8); + if (pc_type != GGUF_TYPE_INT8 && pc_type != GGUF_TYPE_UINT8) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_PRECOMPILED_CHARSMAP).c_str())); + } const size_t n_precompiled_charsmap = gguf_get_arr_n(ctx, precompiled_charsmap_keyidx); const char * pc = (const char *) gguf_get_arr_data(ctx, precompiled_charsmap_keyidx); @@ -2317,6 +2326,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { throw std::runtime_error("cannot find tokenizer merges in model file\n"); } { + if (gguf_get_kv_type(ctx, merges_keyidx) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, merges_keyidx) != GGUF_TYPE_STRING) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_MERGES).c_str())); + } const int n_merges = gguf_get_arr_n(ctx, merges_keyidx); for (int i = 0; i < n_merges; i++) { const std::string word = gguf_get_arr_str(ctx, merges_keyidx, i); @@ -2643,11 +2656,20 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { throw std::runtime_error("cannot find tokenizer vocab in model file\n"); } + if (gguf_get_kv_type(ctx, token_idx) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, token_idx) != GGUF_TYPE_STRING) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_LIST).c_str())); + } + const uint32_t n_tokens = gguf_get_arr_n(ctx, token_idx); const float * scores = nullptr; const int score_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SCORES).c_str()); if (score_idx != -1) { + if (gguf_get_kv_type(ctx, score_idx) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, score_idx) != GGUF_TYPE_FLOAT32) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SCORES).c_str())); + } const uint32_t n_scores = gguf_get_arr_n(ctx, score_idx); if (n_scores < n_tokens) { throw std::runtime_error("Index out of array bounds for scores (" + std::to_string(n_scores) + " < " + std::to_string(n_tokens) + ")\n"); @@ -2658,6 +2680,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { const int * toktypes = nullptr; const int toktype_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str()); if (toktype_idx != -1) { + if (gguf_get_kv_type(ctx, toktype_idx) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, toktype_idx) != GGUF_TYPE_INT32) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_TOKEN_TYPE).c_str())); + } const uint32_t n_toktypes = gguf_get_arr_n(ctx, toktype_idx); if (n_toktypes < n_tokens) { throw std::runtime_error("Index out of array bounds for toktypes (" + std::to_string(n_toktypes) + " < " + std::to_string(n_tokens) + ")\n"); @@ -2823,6 +2849,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { { const int suppress_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str()); if (suppress_idx != -1) { + if (gguf_get_kv_type(ctx, suppress_idx) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(ctx, suppress_idx) != GGUF_TYPE_INT32) { + throw std::runtime_error(format("invalid gguf type for %s", kv(LLM_KV_TOKENIZER_SUPPRESS_TOKENS).c_str())); + } const int n = gguf_get_arr_n(ctx, suppress_idx); const int32_t * data = (const int32_t *) gguf_get_arr_data(ctx, suppress_idx); // drop out-of-range ids diff --git a/src/llama.cpp b/src/llama.cpp index 81eb92e24..cb4b9a640 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -279,7 +279,11 @@ static bool llama_prepare_model_devices(const llama_model_params & params, llama } case GGML_BACKEND_DEVICE_TYPE_IGPU: - if (igpus.empty()) { + // igpus.empty() - workaround for integrated devices seen by multiple backends + // ref: https://github.com/ggml-org/llama.cpp/pull/23897 + // ggml_backend_dev_backend_reg - allow devices of the same backend regardless if integrated + // ref: https://github.com/ggml-org/llama.cpp/pull/23897#issuecomment-5264222997 + if (igpus.empty() || ggml_backend_dev_backend_reg(dev) == ggml_backend_dev_backend_reg(igpus.back().dev)) { igpus.push_back({false, dev}); } break; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index daaa20826..39e4f5dd8 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -43,6 +43,8 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); ml.get_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, false); + GGML_ASSERT(hparams.dsv4_o_group_count > 0); // avoid div by zero + if (hparams.expert_gating_func != LLAMA_EXPERT_GATING_FUNC_TYPE_SQRT_SOFTPLUS) { throw std::runtime_error("DSpark DSV4 draft expects sqrtsoftplus MoE scoring"); } diff --git a/src/models/mamba-base.cpp b/src/models/mamba-base.cpp index fd3fe3f03..1f994ae0a 100644 --- a/src/models/mamba-base.cpp +++ b/src/models/mamba-base.cpp @@ -2,6 +2,8 @@ #include "llama-memory-recurrent.h" +#include + llm_build_mamba_base::llm_build_mamba_base(const llm_graph_params & params) : llm_graph_context(params) {} ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp, @@ -118,7 +120,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba_layer(llm_graph_input_rs * inp, // Custom operator to optimize the parallel associative scan // as described in the Annex D of the Mamba paper. // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs} - return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids); + return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1); }; ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows); @@ -153,7 +155,8 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, int il) const { const auto * mctx_cur = inp->mctx; - const auto kv_head = mctx_cur->get_head(); + const auto kv_head = mctx_cur->get_head(); + const auto mem_size = mctx_cur->get_size(); const int64_t d_conv = hparams.ssm_d_conv; const int64_t d_inner = hparams.ssm_d_inner; @@ -164,6 +167,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, const int64_t n_seqs = ubatch.n_seqs; const int64_t n_seq_tokens = ubatch.n_seq_tokens; + const int64_t K = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1; GGML_ASSERT(n_seqs != 0); GGML_ASSERT(ubatch.equal_seqs()); @@ -173,6 +177,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + const int64_t state_slots = ssm_states_all->ne[1]; ggml_tensor * conv = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs); conv = ggml_reshape_3d(ctx0, conv, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs); @@ -198,15 +203,19 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, // => {d_conv - 1 + n_seq_tokens, d_inner + 2*n_group*d_state, n_seqs} ggml_tensor * conv_x = ggml_concat(ctx0, conv, ggml_transpose(ctx0, xBC), 0); - // copy last (d_conv - 1) columns back into the state cache - ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs, - conv_x->nb[1], conv_x->nb[2], n_seq_tokens * (conv_x->nb[0])); + const int64_t row_count = (d_conv - 1) * (d_inner + 2 * n_group * d_state); + const size_t row_size = ggml_row_size(conv_states_all->type, row_count); + const int64_t n_written = std::min(n_seq_tokens, K); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv, - ggml_view_1d(ctx0, conv_states_all, - (d_conv - 1) * (d_inner + 2 * n_group * d_state) * (n_seqs), - kv_head * (d_conv - 1) * (d_inner + 2 * n_group * d_state) * - ggml_element_size(conv_states_all)))); + for (int64_t slot = 0; slot < n_written; ++slot) { + ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs, + conv_x->nb[1], conv_x->nb[2], (n_seq_tokens - slot) * conv_x->nb[0]); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv, + ggml_view_2d(ctx0, conv_states_all, row_count, n_seqs, + conv_states_all->nb[1], + ((size_t) slot * mem_size + kv_head) * row_size))); + } // 1D convolution // The equivalent is to make a self-overlapping view of conv_x @@ -244,20 +253,27 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, // (this is necessary in order to properly use the states before they are overwritten, // while avoiding to make unnecessary copies of the states) auto get_ssm_rows = [&](ggml_context * ctx, ggml_tensor * states, ggml_tensor * ids) { - ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, mctx_cur->get_size()); + ggml_tensor * ssm = ggml_reshape_4d(ctx, states, d_state, head_dim, n_head, state_slots); // TODO: use semistructured matrices to implement state-space duality // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs} - return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids); + // K > 1 asks the backend to return rollback snapshots in addition to the final state. + return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, K); }; ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows); + const int64_t D = d_state * d_inner; + const int64_t n_written = std::min(n_seq_tokens, K); + const size_t row_size = ggml_row_size(ssm_states_all->type, D); + const size_t y_row_size = ggml_row_size(y_ssm->type, D); + const size_t state_offset = ggml_nelements(x) * ggml_element_size(x); - // store last states ggml_build_forward_expand( - gf, ggml_cpy(ctx0, ggml_view_1d(ctx0, y_ssm, d_state * d_inner * n_seqs, ggml_nelements(x) * x->nb[0]), - ggml_view_1d(ctx0, ssm_states_all, d_state * d_inner * n_seqs, - kv_head * d_state * d_inner * ggml_element_size(ssm_states_all)))); + gf, ggml_cpy(ctx0, + ggml_view_3d(ctx0, y_ssm, D, n_seqs, n_written, + y_row_size, y_row_size * n_seqs, state_offset), + ggml_view_3d(ctx0, ssm_states_all, D, n_seqs, n_written, + ssm_states_all->nb[1], (size_t) mem_size * row_size, kv_head * row_size))); ggml_tensor * y = ggml_view_4d(ctx0, y_ssm, head_dim, n_head, n_seq_tokens, n_seqs, x->nb[1], n_head * x->nb[1], n_seq_tokens * n_head * x->nb[1], 0); diff --git a/src/models/minimax-01.cpp b/src/models/minimax-01.cpp new file mode 100644 index 000000000..a6ccee191 --- /dev/null +++ b/src/models/minimax-01.cpp @@ -0,0 +1,520 @@ +#include "models.h" +#include "llama-memory-recurrent.h" + +void llama_model_minimax_01::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale); + + // we use n_embd_head_la to set recurrent memory n_embd_s + hparams.n_embd_head_la = hparams.n_embd_head_k_full; + + // Mark recurrent layers (lightning attention layers). + if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) { + uint32_t full_attn_interval = 8; + ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); + for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { + hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0); + } + } + + switch (hparams.n_layer()) { + case 80: type = LLM_TYPE_456B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_minimax_01::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + // output + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); + + // if output is NULL, init from the input tok embed + 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) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + if (!hparams.is_recr(i)) { + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0); + } else { + layer.attn_norm_2 = create_tensor(tn(LLM_TENSOR_ATTN_NORM_2, "weight", i), {n_embd_head_k * n_head}, 0); + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, 3 * n_embd_head_k * n_head}, 0); + layer.wg = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0); + } + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff, n_expert}, TENSOR_NOT_REQUIRED); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff, n_embd, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff, n_expert}, 0); + } +} + +std::unique_ptr llama_model_minimax_01::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +class llm_graph_input_la : public llm_graph_input_i { +public: + llm_graph_input_la(const llama_hparams & hparams) : hparams(hparams) {} + + void set_input(const llama_ubatch * ubatch) override { + // this operates on assumption that we have an equal ubatch split + + const int64_t n_head = hparams.n_head(); + const int32_t n_seqs = ubatch->n_seqs; + const int32_t n_seqs_unq = ubatch->n_seqs_unq; + const int32_t n_tokens = ubatch->n_tokens; + const int32_t n_seq_tokens = ubatch->n_seq_tokens; + + std::vector p0(n_seqs_unq); + std::fill(p0.begin(), p0.end(), std::numeric_limits::max()); + + // get lowest token position in a ubatch for each stream + for (int i = 0; i < n_tokens; ++i) { + llama_seq_id seq_id = ubatch->seq_id[i][0]; + int32_t seq_idx = ubatch->seq_idx[seq_id]; + llama_pos pos = ubatch->pos[i]; + if (p0[seq_idx] > pos) { + p0[seq_idx] = pos; + } + } + + if (inp_slopes) { + GGML_ASSERT(ggml_backend_buffer_is_host(inp_slopes->buffer)); + + float * data = (float *) inp_slopes->data; + + float start = powf(2, -powf(2, -(log2f(n_head) - 3))); + float ratio = start; + + for (int h = 0; h < n_head; ++h) { + data[h] = start * powf(ratio, h); + } + } + + if (inp_q_decay) { + GGML_ASSERT(ggml_backend_buffer_is_host(inp_q_decay->buffer)); + + float * slopes = (float *) inp_slopes->data; + float * data = (float *) inp_q_decay->data; + + for (int s = 0; s < n_seqs; ++s) { + for (int i = 0; i < n_seq_tokens; ++i) { + llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + i][0]; + int32_t seq_idx = ubatch->seq_idx[seq_id]; + llama_pos pos = ubatch->pos[s * n_seq_tokens + i]; + int pos_rel = pos - p0[seq_idx]; + + for (int h = 0; h < n_head; ++h) { + data[seq_idx * n_head * n_seq_tokens + i * n_head + h] = -slopes[h] * (pos_rel + 1); + } + } + } + } + + if (inp_k_decay) { + GGML_ASSERT(ggml_backend_buffer_is_host(inp_k_decay->buffer)); + + float * slopes = (float *) inp_slopes->data; + float * data = (float *) inp_k_decay->data; + + for (int s = 0; s < n_seqs; ++s) { + for (int i = 0; i < n_seq_tokens; ++i) { + llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + i][0]; + int32_t seq_idx = ubatch->seq_idx[seq_id]; + llama_pos pos = ubatch->pos[s * n_seq_tokens + i]; + int pos_rel = pos - p0[seq_idx]; + + for (int h = 0; h < n_head; ++h) { + data[seq_idx * n_head * n_seq_tokens + i * n_head + h] = -slopes[h] * (n_seq_tokens - pos_rel - 1); + } + } + } + } + + if (inp_diag_decay) { + GGML_ASSERT(ggml_backend_buffer_is_host(inp_diag_decay->buffer)); + + float * slopes = (float *) inp_slopes->data; + float * data = (float *) inp_diag_decay->data; + + for (int s = 0; s < n_seqs; ++s) { + for (int h = 0; h < n_head; ++h) { + for (int j = 0; j < n_seq_tokens; ++j) { + llama_seq_id seq_id = ubatch->seq_id[s * n_seq_tokens + j][0]; + int32_t seq_idx = ubatch->seq_idx[seq_id]; + llama_pos pos_j = ubatch->pos[s * n_seq_tokens + j]; + int pos_rel_j = pos_j - p0[seq_idx]; + + for (int i = 0; i < n_seq_tokens; ++i) { + llama_pos pos_i = ubatch->pos[s * n_seq_tokens + i]; + int pos_rel_i = pos_i - p0[seq_idx]; + + int index = pos_rel_j - pos_rel_i; + float s_index = index >= 0 ? -slopes[h] * index : -INFINITY; + data[seq_idx * n_head * n_seq_tokens * n_seq_tokens + h * n_seq_tokens * n_seq_tokens + j * n_seq_tokens + i] = s_index; + } + } + } + } + } + } + + bool can_reuse(const llm_graph_params & params) override { + bool res = true; + + if (params.ubatch.n_seq_tokens > 1) { + res &= ( inp_q_decay && inp_q_decay->ne[2] == params.ubatch.n_seq_tokens); + res &= ( inp_k_decay && inp_k_decay->ne[2] == params.ubatch.n_seq_tokens); + res &= (inp_diag_decay && inp_diag_decay->ne[1] == params.ubatch.n_seq_tokens); + } + + return res; + } + + const llama_hparams & hparams; + + ggml_tensor * inp_slopes = nullptr; // F32 [n_head] + ggml_tensor * inp_q_decay = nullptr; // F32 [1, n_head, n_batch] + ggml_tensor * inp_k_decay = nullptr; // F32 [1, n_head, n_batch] + ggml_tensor * inp_diag_decay = nullptr; // F32 [n_batch, n_batch, n_head] +}; + +llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + // GGML_ASSERT(n_embd_head == n_rot); this is wrong in case of minimax, head_dim = 128, n_rot = 64 + + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + auto * inp_hybrid = build_inp_mem_hybrid(); + auto * inp_rs = inp_hybrid->get_recr(); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + llm_graph_input_la * la = nullptr; + + auto inp = std::make_unique(hparams); + + inp->inp_slopes = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_head); + ggml_set_input(inp->inp_slopes); + cb(inp->inp_slopes, "slopes", -1); + + if (n_seq_tokens != 1) { + inp->inp_q_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); + ggml_set_input(inp->inp_q_decay); + cb(inp->inp_q_decay, "q_decay_exp", -1); + + inp->inp_k_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); + ggml_set_input(inp->inp_k_decay); + cb(inp->inp_k_decay, "k_decay_exp", -1); + + inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs); + ggml_set_input(inp->inp_diag_decay); + cb(inp->inp_diag_decay, "diag_decay_exp", -1); + } + + la = (llm_graph_input_la *) res->add_input(std::move(inp)); + + ggml_tensor * slopes = la->inp_slopes; + + for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = inpL; + + ggml_tensor * inpSA = inpL; + + cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + ggml_tensor * residual = cur; + + // self_attention + if (!hparams.is_recr(il)) { + // softmax attention layer + + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, n_head_kv, il); + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + cur = build_attn(inp_hybrid->get_attn(), + model.layers[il].wo, NULL, model.layers[il].wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il); + } else { + // lightning attention layer + + const auto * mctx_cur = inp_rs->mctx; + const auto kv_head = mctx_cur->get_head(); + + // TODO unneeded - any way to make conv states optional in recurrent memory? + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + ggml_build_forward_expand(gf, conv_state_all); + + float slope_scale = 1.0 - 1.0 * il / (n_layer - 1) + 1e-5; + ggml_tensor * slope_rate = ggml_scale(ctx0, slopes, slope_scale); + cb(slope_rate, "slope_rate", il); + + cur = ggml_reshape_4d(ctx0, cur, cur->ne[0], n_seq_tokens, 1, n_seqs); + + ggml_tensor * QKVcur = build_lora_mm(model.layers[il].wqkv, cur); + cb(QKVcur, "QKVcur", il); + + QKVcur = ggml_silu(ctx0, QKVcur); + cb(QKVcur, "QKVcur_silu", il); + + QKVcur = ggml_reshape_4d(ctx0, QKVcur, n_embd_head * 3, n_head, n_seq_tokens, n_seqs); + + ggml_tensor * Qcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 0*ggml_element_size(QKVcur)*n_embd_head); + ggml_tensor * Kcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 1*ggml_element_size(QKVcur)*n_embd_head); + ggml_tensor * Vcur = ggml_view_4d(ctx0, QKVcur, n_embd_head, n_head, n_seq_tokens, n_seqs, QKVcur->nb[1], QKVcur->nb[2], QKVcur->nb[3], 2*ggml_element_size(QKVcur)*n_embd_head); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + // get previous KV + ggml_tensor * la_states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, la_states_all, hparams.n_embd_s(), n_seqs); + + ggml_tensor * kv_old = ggml_reshape_4d(ctx0, state, n_embd_head, n_embd_head, n_head, n_seqs); + cb(kv_old, "kv_old", il); + + ggml_tensor * qkv = nullptr; + ggml_tensor * kv_new = nullptr; + + if (n_seq_tokens == 1) { + // lightning attention - optimized single token case for TG + + ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0); + cb(slopes_neg, "slopes_neg", il); + + ggml_tensor * ratio = ggml_exp(ctx0, slopes_neg); + cb(ratio, "ratio", il); + + ggml_tensor * ratio_3d = ggml_reshape_3d(ctx0, ratio, 1, 1, n_head); + cb(ratio_3d, "ratio3d", il); + + ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3)); + cb(v_trans, "v_trans", il); + + ggml_tensor * k_trans = ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 1, 2, 0, 3)); + cb(k_trans, "k_trans", il); + + ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_trans, v_trans); + cb(kv_cur, "kv_cur", il); + + ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, ratio_3d); + cb(kv_old_s, "kv_old_s", il); + + kv_new = ggml_add(ctx0, kv_old_s, kv_cur); + cb(kv_new, "kv_new", il); + + ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); + cb(q_trans, "q_trans", il); + + qkv = ggml_mul_mat(ctx0, kv_new, q_trans); + cb(qkv, "qkv", il); + } else if(n_seq_tokens > 1) { + // lightning attention - general multi token case for PP + + ggml_tensor * q_decay_exp = la->inp_q_decay; + ggml_tensor * k_decay_exp = la->inp_k_decay; + ggml_tensor * diag_decay_exp = la->inp_diag_decay; + + ggml_tensor * q_decay = ggml_exp(ctx0, ggml_scale(ctx0, q_decay_exp, slope_scale)); + cb(q_decay, "q_decay", il); + ggml_tensor * k_decay = ggml_exp(ctx0, ggml_scale(ctx0, k_decay_exp, slope_scale)); + cb(k_decay, "k_decay", il); + ggml_tensor * diag_decay = ggml_exp(ctx0, ggml_scale(ctx0, diag_decay_exp, slope_scale)); + cb(diag_decay, "diag_decay", il); + + ggml_tensor * q_s = ggml_mul(ctx0, Qcur, q_decay); + cb(q_s, "q_s", il); + + ggml_tensor * q_s_trans = ggml_permute(ctx0, q_s, 0, 2, 1, 3); + cb(q_s_trans, "q_s_trans", il); + + ggml_tensor * qkv_none_diag = ggml_mul_mat(ctx0, kv_old, q_s_trans); + cb(qkv_none_diag, "qkv_none_diag", il); + + ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); + cb(q_trans, "q_trans", il); + + ggml_tensor * k_trans = ggml_permute(ctx0, Kcur, 0, 2, 1, 3); + cb(k_trans, "k_trans", il); + + ggml_tensor * qk = ggml_mul_mat(ctx0, k_trans, q_trans); + cb(qk, "qk", il); + + qk = ggml_mul(ctx0, qk, diag_decay); + cb(qk, "qk_s", il); + + ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3)); + cb(v_trans, "v_trans", il); + + ggml_tensor * qkv_diag = ggml_mul_mat(ctx0, v_trans, qk); + cb(qkv_diag, "qkv_diag", il); + + qkv = ggml_add(ctx0, qkv_none_diag, qkv_diag); + cb(qkv, "qkv", il); + + ggml_build_forward_expand(gf, qkv); + + ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0*n_seq_tokens); + cb(slopes_neg, "slopes_neg", il); + + ggml_tensor * block_decay = ggml_exp(ctx0, slopes_neg); + cb(block_decay, "block_decay", il); + + ggml_tensor * block_decay_3d = ggml_reshape_3d(ctx0, block_decay, 1, 1, n_head); + cb(block_decay_3d, "block_decay_3d", il); + + ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, block_decay_3d); + cb(kv_old_s, "kv_old_s", il); + + ggml_tensor * k_after_decay = ggml_mul(ctx0, Kcur, k_decay); + cb(k_after_decay, "k_after_decay", il); + + ggml_tensor * k_after_decay_trans = ggml_cont(ctx0, ggml_permute(ctx0, k_after_decay, 1, 2, 0, 3)); + cb(k_after_decay_trans, "k_after_decay_trans", il); + + ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_after_decay_trans, v_trans); + cb(kv_cur, "kv_cur", il); + + kv_new = ggml_add(ctx0, kv_old_s, kv_cur); + cb(kv_new, "kv_new", il); + } + + // store new KV + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, kv_new, + ggml_view_1d(ctx0, la_states_all, hparams.n_embd_s() * n_seqs, + kv_head * hparams.n_embd_s() * ggml_element_size(la_states_all)))); + + qkv = ggml_cont(ctx0, ggml_permute(ctx0, qkv, 0, 2, 1, 3)); + cb(qkv, "qkv_permuted", il); + + qkv = ggml_reshape_4d(ctx0, qkv, qkv->ne[0]*qkv->ne[1], qkv->ne[2], 1, qkv->ne[3]); + + // norm + ggml_tensor * qkv_norm = build_norm(qkv, + model.layers[il].attn_norm_2, NULL, + LLM_NORM_RMS, il); + cb(qkv_norm, "qkv_norm", il); + + ggml_tensor * g = build_lora_mm(model.layers[il].wg, cur); + cb(g, "g", il); + + g = ggml_sigmoid(ctx0, g); + cb(g, "g_sigm", il); + + cur = ggml_mul(ctx0, g, qkv_norm); + + cur = build_lora_mm(model.layers[il].wo, cur); + cb(cur, "attn_out", il); + + cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], n_seq_tokens*n_seqs); + cb(cur, "attn_out", il); + } + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + residual = ggml_get_rows(ctx0, residual, inp_out_ids); + } + + residual = ggml_scale(ctx0, residual, hparams.f_residual_scale); + cb(residual, "residual_scaled_attn", il); + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, residual); + cb(ffn_inp, "ffn_inp", il); + + // MoE branch + cur = build_norm(ffn_inp, + model.layers[il].ffn_norm, NULL, + LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + residual = cur; + + cur = build_moe_ffn(cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + model.layers[il].ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, true, + hparams.expert_weights_scale, + LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, + il); + cb(cur, "ffn_moe_out", il); + + residual = ggml_scale(ctx0, residual, hparams.f_residual_scale); + cb(residual, "residual_scaled_ffn", il); + + cur = ggml_add(ctx0, cur, residual); + cb(cur, "ffn_out", il); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + // input for next layer + inpL = cur; + } + + cur = inpL; + + cur = build_norm(cur, + model.output_norm, NULL, + LLM_NORM_RMS, -1); + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + // lm_head + cur = build_lora_mm(model.output, cur, model.output_s); + + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp index 854d5aed0..1ba699d01 100644 --- a/src/models/minimax-m3.cpp +++ b/src/models/minimax-m3.cpp @@ -25,6 +25,8 @@ void llama_model_minimax_m3::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks); msa_p = { (int) hparams.indexer_block_size, (int) hparams.indexer_top_k, (int) hparams.indexer_local_blocks }; + GGML_ASSERT(hparams.indexer_block_size > 0); // avoid div by zero + switch (hparams.n_layer()) { case 60: type = LLM_TYPE_428B_A23B; break; default: type = LLM_TYPE_UNKNOWN; diff --git a/src/models/models.h b/src/models/models.h index ddb9ae2f1..cefaf56eb 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2043,6 +2043,19 @@ struct llama_model_apertus : public llama_model_base { }; +struct llama_model_minimax_01 : public llama_model_base { + llama_model_minimax_01(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_minimax_m2 : public llama_model_base { llama_model_minimax_m2(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/src/models/plamo2.cpp b/src/models/plamo2.cpp index 0b81513c3..d946b3cff 100644 --- a/src/models/plamo2.cpp +++ b/src/models/plamo2.cpp @@ -382,7 +382,7 @@ ggml_tensor * llama_model_plamo2::graph::build_plamo2_mamba_layer(llm_graph_inpu // Custom operator to optimize the parallel associative scan // as described in the Annex D of the Mamba paper. // => {d_inner, n_seq_tokens, n_seqs} and {d_state, d_inner, n_seqs} - return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids); + return ggml_ssm_scan(ctx, ssm, x, dt, A, B, C, ids, /*K=*/1); }; ggml_tensor * y_ssm = build_rs(inp, ssm_states_all, hparams.n_embd_s(), ubatch.n_seqs, get_ssm_rows); diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index ae164807b..86c2c1c94 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -603,7 +604,7 @@ struct clip_image_u8 { // return a dummy value, so that legacy code can still process image without errors return { 0, 0, 0 }; } - int idx = (y * nx + x) * 3; + size_t idx = ((size_t) y * (size_t) nx + (size_t) x) * 3; return { buf[idx], buf[idx + 1], buf[idx + 2] }; } @@ -611,8 +612,8 @@ struct clip_image_u8 { if (is_placeholder()) { return; // no-op } - int idx = (y * nx + x) * 3; - buf[idx] = rgb[0]; + size_t idx = ((size_t) y * (size_t) nx + (size_t) x) * 3; + buf[idx] = rgb[0]; buf[idx + 1] = rgb[1]; buf[idx + 2] = rgb[2]; } @@ -642,9 +643,25 @@ struct mtmd_serialization; // forward declaration struct clip_image_f32 { // marks the global view in e.g., DeepSeek-OCR Models bool add_viewsep = false; - // whether a learned newline (or EOI) token should be appended after the image (eg Granite4 Vision) + // appends a learned newline (or EOI) token after the image + // no model uses it now (Granite4 Vision moved to anyres), kept for future models bool add_newline = false; + // llava-next "anyres" tiling, used by Granite4 Vision + // the whole grid is encoded and assembled in a single graph + // NOTE: excluded from serialized: a deserialized image is always a placeholder, which is never encoded + struct anyres_info { + int grid_x = 0; // tiles per row, 0 means the image is not tiled + int grid_y = 0; // tiles per column + int orig_nx = 0; // size of the source image, used to drop the padding tokens + int orig_ny = 0; + + bool is_tiled() const { + return grid_x > 0 && grid_y > 0; + } + }; + anyres_info anyres; + clip_image_size get_size() const { return { nx_, ny_ }; } @@ -726,6 +743,25 @@ struct clip_image_f32 { } }; +// token area kept after removing the padding added by the anyres resize +// ref: https://github.com/huggingface/transformers/blob/v5.0.0/src/transformers/models/llava_next/modeling_llava_next.py#L109 +static inline void clip_anyres_unpad(int cur_w, int cur_h, int orig_w, int orig_h, + int & off_x, int & off_y, int & out_w, int & out_h) { + off_x = 0; + off_y = 0; + out_w = cur_w; + out_h = cur_h; + if ((float) orig_w / orig_h > (float) cur_w / cur_h) { + const int new_h = (int) std::floor((double) orig_h * cur_w / orig_w + 1e-7); + off_y = (cur_h - new_h) / 2; + out_h = cur_h - 2 * off_y; + } else { + const int new_w = (int) std::floor((double) orig_w * cur_h / orig_h + 1e-7); + off_x = (cur_w - new_w) / 2; + out_w = cur_w - 2 * off_x; + } +} + // // logging // diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index c06a2d65f..c876d3342 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1680,6 +1680,9 @@ struct clip_model_loader { hparams.image_resize_algo = RESIZE_ALGO_BICUBIC_PILLOW; hparams.image_resize_pad = PAD_NONE; get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false); + // n_merge is used as a divisor in clip_image_batch_encode + // (gh / n_merge); reject 0 to avoid int div-by-zero (DoS). + GGML_ASSERT(hparams.n_merge > 0); hparams.rope_theta = 10000.0f; // vision_config.rope_theta // MiniMax-M3: max_pixels 451584 (=672^2) -> 576 merged tokens (image_seq_length) hparams.set_limit_image_tokens(8, 576); @@ -1908,7 +1911,9 @@ struct clip_model_loader { // unlimited-ocr shares the v1 projector but tiles up to 32 get_u32(KEY_PREPROC_MIN_TILES, hparams.preproc_min_tiles, false); get_u32(KEY_PREPROC_MAX_TILES, hparams.preproc_max_tiles, false); - GGML_ASSERT(hparams.preproc_min_tiles <= hparams.preproc_max_tiles); + GGML_ASSERT(hparams.preproc_min_tiles >= 0 + && hparams.preproc_min_tiles <= hparams.preproc_max_tiles + && hparams.preproc_max_tiles <= 256); } break; case PROJECTOR_TYPE_HUNYUANVL: { @@ -1973,6 +1978,9 @@ struct clip_model_loader { hparams.audio_window_len = 400; hparams.audio_hop_len = 160; get_u32(KEY_A_CHUNK_SIZE, hparams.audio_chunk_size); + // context_size is squared for the attn_dists/mask buffers; cap to prevent int32 overflow + // (legitimate values are small, e.g. 12-200; 8192^2 = 67M still fits int32) + GGML_ASSERT(hparams.audio_chunk_size > 0 && hparams.audio_chunk_size <= 8192); get_u32(KEY_A_CONV_KERNEL_SIZE, hparams.audio_conv_kernel_size); get_u32(KEY_A_MAX_POS_EMB, hparams.audio_max_pos_emb); get_u32(KEY_A_PROJ_WINDOW_SIZE, hparams.audio_proj_window_size); @@ -2012,8 +2020,9 @@ struct clip_model_loader { // note: some models having hparams.image_size == 0, which means the image size is dynamic throw std::runtime_error(string_format("%s: image_size (%d) cannot be negative\n", __func__, hparams.image_size)); } - if (hparams.image_size > 65536) { - throw std::runtime_error(string_format("%s: image_size (%d) is too large (max 65536)\n", __func__, hparams.image_size)); + if (hparams.image_size > 8192) { + // cap prevents int32 overflow in n_patches = (image_size/patch_size)^2 + throw std::runtime_error(string_format("%s: image_size (%d) is too large (max 8192)\n", __func__, hparams.image_size)); } if (hparams.patch_size <= 0 || hparams.patch_size >= 65536) { throw std::runtime_error(string_format("%s: patch_size (%d) must be positive and less than 65536\n", __func__, hparams.patch_size)); @@ -2024,9 +2033,12 @@ struct clip_model_loader { if (hparams.image_max_pixels < hparams.image_min_pixels) { throw std::runtime_error(string_format("%s: image_max_pixels (%d) is less than image_min_pixels (%d)\n", __func__, hparams.image_max_pixels, hparams.image_min_pixels)); } - if (hparams.n_merge < 0 || hparams.n_merge >= 65536) { + if (hparams.n_merge <= 0 || hparams.n_merge >= 65536) { throw std::runtime_error(string_format("%s: n_merge (%d) must be greater than 0 and less than 65536\n", __func__, hparams.n_merge)); } + if (hparams.attn_window_size > 4096) { + throw std::runtime_error(string_format("%s: attn_window_size (%d) is too large (max 4096)\n", __func__, hparams.attn_window_size)); + } } LOG_INF("%s: projector: %s\n", __func__, proj_type.c_str()); @@ -3824,6 +3836,9 @@ struct clip_model_loader { } return; } + if (gguf_get_kv_type(ctx_gguf.get(), i) != GGUF_TYPE_ARRAY) { + throw std::runtime_error(string_format("%s: key '%s' is not an array\n", __func__, key.c_str())); + } const auto type = gguf_get_arr_type(ctx_gguf.get(), i); if (type != GGUF_TYPE_FLOAT32) { throw std::runtime_error(string_format("%s: array '%s' has type %d, expected %d (GGUF_TYPE_FLOAT32)\n", __func__, key.c_str(), type, GGUF_TYPE_FLOAT32)); @@ -3858,6 +3873,9 @@ struct clip_model_loader { } return; } + if (gguf_get_kv_type(ctx_gguf.get(), i) != GGUF_TYPE_ARRAY) { + throw std::runtime_error(string_format("%s: key '%s' is not an array\n", __func__, key.c_str())); + } const auto type = gguf_get_arr_type(ctx_gguf.get(), i); if (type != GGUF_TYPE_INT32) { throw std::runtime_error(string_format("%s: array '%s' has type %d, expected %d (GGUF_TYPE_INT32)\n", __func__, key.c_str(), type, GGUF_TYPE_INT32)); @@ -4307,18 +4325,20 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_GRANITE4_VISION: { // Per-tile output token count: each projector block outputs - // query_side^2 tokens per window × n^2 windows. - // For 384×384 input: n = 24/8 = 3, query_side = 4 → 144. + // query_side^2 tokens per window x n^2 windows. + // For 384x384 input: n = 24/8 = 3, query_side = 4 -> 144. const int window_side = ctx->model.hparams.downsample_window_side; const int query_side = ctx->model.hparams.downsample_query_side; const int side = img->nx() / params.patch_size; const int n = side / window_side; - n_patches = (query_side * n) * (query_side * n); - if (img->add_newline) { - // For single-tile case: append 1 newline row. - // For multi-tile rowwise: handled by caller, but here we - // report the per-tile count including one trailing newline. - n_patches += 1; + const int out_side = query_side * n; + n_patches = out_side * out_side; + if (img->anyres.is_tiled()) { + // overview tile, then the unpadded tile grid with one newline per row + int off_x, off_y, w, h; + clip_anyres_unpad(img->anyres.grid_x * out_side, img->anyres.grid_y * out_side, + img->anyres.orig_nx, img->anyres.orig_ny, off_x, off_y, w, h); + n_patches += h * (w + 1); } } break; default: @@ -5498,13 +5518,13 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { const int context_size = ctx->model.hparams.audio_chunk_size; const int max_pos_emb = ctx->model.hparams.audio_max_pos_emb; - std::vector dists(context_size * context_size); + std::vector dists((size_t) context_size * (size_t) context_size); for (int i = 0; i < context_size; i++) { for (int j = 0; j < context_size; j++) { int d = i - j; if (d < -context_size) d = -context_size; if (d > context_size) d = context_size; - dists[i * context_size + j] = d + max_pos_emb; + dists[(size_t) i * (size_t) context_size + (size_t) j] = d + max_pos_emb; } } set_input_i32("attn_dists", dists); @@ -5513,13 +5533,13 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { const int remainder = n_frames % context_size; if (remainder > 0) { const int num_blocks = (n_frames + context_size - 1) / context_size; - std::vector mask(context_size * context_size * num_blocks, 0.0f); + std::vector mask((size_t) context_size * (size_t) context_size * (size_t) num_blocks, 0.0f); const float neg_inf = -INFINITY; - const int last_block_offset = (num_blocks - 1) * context_size * context_size; + const size_t last_block_offset = (size_t) (num_blocks - 1) * (size_t) context_size * (size_t) context_size; for (int q = 0; q < context_size; q++) { for (int k = 0; k < context_size; k++) { if (q >= remainder || k >= remainder) { - mask[last_block_offset + q * context_size + k] = neg_inf; + mask[last_block_offset + (size_t) q * (size_t) context_size + (size_t) k] = neg_inf; } } } @@ -5583,10 +5603,18 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { return idx; }; + // the same permutation is applied to every tile of the stacked image auto upload = [&](const std::string & name, const std::vector & idx) { ggml_tensor * t = ggml_graph_get_tensor(gf, name.c_str()); GGML_ASSERT(t); - ggml_backend_tensor_set(t, idx.data(), 0, idx.size() * sizeof(int32_t)); + GGML_ASSERT(ggml_nelements(t) % (int64_t) idx.size() == 0); + const int n_rep = ggml_nelements(t) / idx.size(); + std::vector buf; + buf.reserve(idx.size() * n_rep); + for (int i = 0; i < n_rep; ++i) { + buf.insert(buf.end(), idx.begin(), idx.end()); + } + ggml_backend_tensor_set(t, buf.data(), 0, ggml_nbytes(t)); }; // Stage 1b only uses block 0's permutations; future stages diff --git a/tools/mtmd/models/granite4-vision.cpp b/tools/mtmd/models/granite4-vision.cpp index 1b252543c..a75f1cee9 100644 --- a/tools/mtmd/models/granite4-vision.cpp +++ b/tools/mtmd/models/granite4-vision.cpp @@ -14,18 +14,39 @@ * Stage 1a: SigLIP vision tower (N layers, post-norm) * Stage 1b: WindowQFormer blocks (deepstack + spatial) * Stage 1c: Concatenate and pack outputs - * Stage 1d: Append newline tokens if add_newline is set + * Stage 1d: Assemble the anyres tiles into one token sequence */ // --------------------------------------------------------------------------- // Member method implementations // --------------------------------------------------------------------------- +// split the stacked tiles into the batch axis, then run the usual patch embedding +ggml_tensor * clip_graph_granite4_vision::build_tile_inp() { + ggml_tensor * inp_raw = build_inp_raw(); + + if (n_tiles > 1) { + const int px = img.nx(); + inp_raw = ggml_reshape_4d(ctx0, inp_raw, px * px, n_tiles, 3, 1); + inp_raw = ggml_cont(ctx0, ggml_permute(ctx0, inp_raw, 0, 2, 1, 3)); + inp_raw = ggml_reshape_4d(ctx0, inp_raw, px, px, 3, n_tiles); + } + + ggml_tensor * inp = ggml_conv_2d(ctx0, model.patch_embeddings_0, inp_raw, patch_size, patch_size, 0, 0, 1, 1); + inp = ggml_reshape_3d(ctx0, inp, tile_side * tile_side, n_embd, n_tiles); + inp = ggml_cont(ctx0, ggml_transpose(ctx0, inp)); + if (model.patch_bias) { + inp = ggml_add(ctx0, inp, model.patch_bias); + } + return inp; +} + ggml_tensor * clip_graph_granite4_vision::gather( ggml_tensor * src, const std::string & name, int idx_len) { - ggml_tensor * idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, idx_len); + // one index row per tile, all rows hold the same permutation + ggml_tensor * idx = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, idx_len, n_tiles); ggml_set_name(idx, name.c_str()); ggml_set_input(idx); return ggml_get_rows(ctx0, src, idx); @@ -36,12 +57,15 @@ ggml_tensor * clip_graph_granite4_vision::interp_down( int side, int new_side) { const int n_embd = src->ne[0]; - ggml_tensor * t = ggml_reshape_4d(ctx0, src, n_embd, side, side, 1); + ggml_tensor * t = ggml_reshape_4d(ctx0, src, n_embd, side, side, n_tiles); t = ggml_cont(ctx0, ggml_permute(ctx0, t, 2, 0, 1, 3)); + // fold the tile axis into the channel axis, ggml_pool_2d only pools the first two axes + t = ggml_reshape_3d(ctx0, t, side, side, n_embd * n_tiles); const int kernel = side / new_side; t = ggml_pool_2d(ctx0, t, GGML_OP_POOL_AVG, kernel, kernel, kernel, kernel, 0, 0); + t = ggml_reshape_4d(ctx0, t, new_side, new_side, n_embd, n_tiles); t = ggml_cont(ctx0, ggml_permute(ctx0, t, 1, 2, 0, 3)); - return ggml_reshape_2d(ctx0, t, n_embd, new_side * new_side); + return ggml_reshape_3d(ctx0, t, n_embd, new_side * new_side, n_tiles); } // --------------------------------------------------------------------------- @@ -63,6 +87,7 @@ ggml_tensor * clip_graph_granite4_vision::build_block( const int n = image_side / window_side; const int new_side = n * query_side; const int n_windows = n * n; + const int n_win_all = n_windows * n_tiles; // windows of every tile, batched together const int enc_len = window_side * window_side; const int query_len = query_side * query_side; @@ -82,7 +107,7 @@ ggml_tensor * clip_graph_granite4_vision::build_block( ggml_tensor * enc_flat = gather(x, "g4v_blk" + std::to_string(bid) + "_win_idx", image_side * image_side); - enc = ggml_reshape_3d(ctx0, enc_flat, n_embd, enc_len, n_windows); + enc = ggml_reshape_3d(ctx0, enc_flat, n_embd, enc_len, n_win_all); } cbx(enc, "enc"); @@ -104,7 +129,7 @@ ggml_tensor * clip_graph_granite4_vision::build_block( ggml_tensor * dw_flat = gather(d, "g4v_blk" + std::to_string(bid) + "_qwin_idx", new_side * new_side); - ggml_tensor * dw = ggml_reshape_3d(ctx0, dw_flat, n_embd, query_len, n_windows); + ggml_tensor * dw = ggml_reshape_3d(ctx0, dw_flat, n_embd, query_len, n_win_all); q_in = ggml_add(ctx0, dw, blk.qf_proj_query); } cbx(q_in, "query_embeds"); @@ -140,12 +165,12 @@ ggml_tensor * clip_graph_granite4_vision::build_block( ggml_tensor * K = linear(q, pl.k_w, pl.k_b); ggml_tensor * V = linear(q, pl.v_w, pl.v_b); - Q = ggml_reshape_4d(ctx0, Q, d_h, n_head, nq, n_windows); - K = ggml_reshape_4d(ctx0, K, d_h, n_head, nq, n_windows); - V = ggml_reshape_4d(ctx0, V, d_h, n_head, nq, n_windows); + Q = ggml_reshape_4d(ctx0, Q, d_h, n_head, nq, n_win_all); + K = ggml_reshape_4d(ctx0, K, d_h, n_head, nq, n_win_all); + V = ggml_reshape_4d(ctx0, V, d_h, n_head, nq, n_win_all); sa_out = build_attn(pl.o_w, pl.o_b, Q, K, V, nullptr, scale, bid); - sa_out = ggml_reshape_3d(ctx0, sa_out, n_embd, nq, n_windows); + sa_out = ggml_reshape_3d(ctx0, sa_out, n_embd, nq, n_win_all); sa_out = ggml_add(ctx0, sa_out, q); sa_out = build_norm(sa_out, pl.ln_1_w, pl.ln_1_b, @@ -166,13 +191,13 @@ ggml_tensor * clip_graph_granite4_vision::build_block( ggml_tensor * K = linear(e_in, pl.cross_attn_k_w, pl.cross_attn_k_b); ggml_tensor * V = linear(e_in, pl.cross_attn_v_w, pl.cross_attn_v_b); - Q = ggml_reshape_4d(ctx0, Q, d_h, n_head, nq, n_windows); - K = ggml_reshape_4d(ctx0, K, d_h, n_head, nkv, n_windows); - V = ggml_reshape_4d(ctx0, V, d_h, n_head, nkv, n_windows); + Q = ggml_reshape_4d(ctx0, Q, d_h, n_head, nq, n_win_all); + K = ggml_reshape_4d(ctx0, K, d_h, n_head, nkv, n_win_all); + V = ggml_reshape_4d(ctx0, V, d_h, n_head, nkv, n_win_all); ca_out = build_attn(pl.cross_attn_o_w, pl.cross_attn_o_b, Q, K, V, nullptr, scale, bid); - ca_out = ggml_reshape_3d(ctx0, ca_out, n_embd, nq, n_windows); + ca_out = ggml_reshape_3d(ctx0, ca_out, n_embd, nq, n_win_all); ca_out = ggml_add(ctx0, ca_out, sa_out); ca_out = build_norm(ca_out, pl.cross_attn_norm_w, pl.cross_attn_norm_b, @@ -183,13 +208,13 @@ ggml_tensor * clip_graph_granite4_vision::build_block( // 6c. FFN ggml_tensor * ffn; { - ggml_tensor * t = ggml_reshape_2d(ctx0, ca_out, n_embd, query_len * n_windows); + ggml_tensor * t = ggml_reshape_2d(ctx0, ca_out, n_embd, query_len * n_win_all); t = build_mm(pl.ff_up_w, t); if (pl.ff_up_b) t = ggml_add(ctx0, t, pl.ff_up_b); t = ggml_gelu_erf(ctx0, t); t = build_mm(pl.ff_down_w, t); if (pl.ff_down_b) t = ggml_add(ctx0, t, pl.ff_down_b); - t = ggml_reshape_3d(ctx0, t, n_embd, query_len, n_windows); + t = ggml_reshape_3d(ctx0, t, n_embd, query_len, n_win_all); ffn = ggml_add(ctx0, t, ca_out); ffn = build_norm(ffn, pl.ln_2_w, pl.ln_2_b, NORM_TYPE_NORMAL, qformer_eps, bid); } @@ -198,7 +223,7 @@ ggml_tensor * clip_graph_granite4_vision::build_block( // 7. _unwin back to raster ggml_tensor * unwinned; { - ggml_tensor * flat = ggml_reshape_2d(ctx0, ffn, n_embd, query_len * n_windows); + ggml_tensor * flat = ggml_reshape_3d(ctx0, ffn, n_embd, query_len * n_windows, n_tiles); unwinned = gather(flat, "g4v_blk" + std::to_string(bid) + "_unwin_idx", new_side * new_side); @@ -244,13 +269,42 @@ ggml_tensor * clip_graph_granite4_vision::build_newline_row(ggml_context * ctx0) return ggml_reshape_2d(ctx0, nl_row_2d, n_mmproj_embd, 1); } -// Append a single newline row at the end of the tile output. -ggml_tensor * clip_graph_granite4_vision::append_rowwise_newlines(ggml_context * ctx0, ggml_tensor * tile_output) { - // For the single-tile case, append one newline row at the end. - // For the multi-tile rowwise case, this will be called per-tile - // (though currently only the single-tile path uses it). - ggml_tensor * nl_row = build_newline_row(ctx0); - return ggml_concat(ctx0, tile_output, nl_row, 1); +// Assemble [overview, tile(0,0), tile(0,1), ...] into one token sequence: +// the overview tokens first, then the tile grid read in raster order with one newline per row. +// ref: https://github.com/huggingface/transformers/blob/v5.0.0/src/transformers/models/llava_next/modeling_llava_next.py#L266 +ggml_tensor * clip_graph_granite4_vision::build_anyres_assembly(ggml_tensor * cur, int out_side) { + const int n_dim = cur->ne[0]; + const int grid_x = anyres.grid_x; + const int grid_y = anyres.grid_y; + const int cur_w = grid_x * out_side; + const int cur_h = grid_y * out_side; + GGML_ASSERT(cur->ne[1] == out_side * out_side); + GGML_ASSERT(cur->ne[2] == 1 + grid_x * grid_y); + + ggml_tensor * base = ggml_view_2d(ctx0, cur, n_dim, out_side * out_side, cur->nb[1], 0); + + ggml_tensor * tiles = ggml_view_3d(ctx0, cur, n_dim, out_side * out_side, grid_x * grid_y, + cur->nb[1], cur->nb[2], cur->nb[2]); + + // (n_dim*out_side, out_side, grid_x, grid_y) -> interleave the tiles of a grid row + tiles = ggml_reshape_4d(ctx0, tiles, n_dim * out_side, out_side, grid_x, grid_y); + tiles = ggml_cont(ctx0, ggml_permute(ctx0, tiles, 0, 2, 1, 3)); + tiles = ggml_reshape_3d(ctx0, tiles, n_dim, cur_w, cur_h); + + // drop the tokens that only cover the padding added when resizing to the grid + int off_x, off_y, w, h; + clip_anyres_unpad(cur_w, cur_h, anyres.orig_nx, anyres.orig_ny, off_x, off_y, w, h); + if (w != cur_w || h != cur_h) { + tiles = ggml_cont(ctx0, ggml_view_3d(ctx0, tiles, n_dim, w, h, + tiles->nb[1], tiles->nb[2], + off_x * tiles->nb[1] + off_y * tiles->nb[2])); + } + + ggml_tensor * nl = ggml_repeat_4d(ctx0, build_newline_row(ctx0), n_dim, 1, h, 1); + tiles = ggml_concat(ctx0, tiles, nl, 1); + tiles = ggml_reshape_2d(ctx0, tiles, n_dim, (w + 1) * h); + + return ggml_concat(ctx0, base, tiles, 1); } ggml_cgraph * clip_graph_granite4_vision::build() { @@ -260,10 +314,12 @@ ggml_cgraph * clip_graph_granite4_vision::build() { GGML_ASSERT(!model.qf_proj_blocks.empty()); // --- Stage 1a: SigLIP encoder producing intermediate hidden states --- - ggml_tensor * inp = build_inp(); + ggml_tensor * inp = build_tile_inp(); inp = ggml_add(ctx0, inp, model.position_embeddings); cb(inp, "pos_embed", -1); + const int tile_n_patches = tile_side * tile_side; + ggml_tensor * inpL = inp; std::vector layer_outs(n_layer, nullptr); @@ -281,12 +337,13 @@ ggml_cgraph * clip_graph_granite4_vision::build() { ggml_tensor * Vcur = build_mm(layer.v_w, cur); if (layer.v_b) Vcur = ggml_add(ctx0, Vcur, layer.v_b); - Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_patches); - Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_patches); - Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_patches); + Qcur = ggml_reshape_4d(ctx0, Qcur, d_head, n_head, tile_n_patches, n_tiles); + Kcur = ggml_reshape_4d(ctx0, Kcur, d_head, n_head, tile_n_patches, n_tiles); + Vcur = ggml_reshape_4d(ctx0, Vcur, d_head, n_head, tile_n_patches, n_tiles); cur = build_attn(layer.o_w, layer.o_b, Qcur, Kcur, Vcur, nullptr, kq_scale, il); + cur = ggml_reshape_3d(ctx0, cur, n_embd, tile_n_patches, n_tiles); cur = ggml_add(ctx0, cur, inpL); inpL = cur; @@ -318,7 +375,7 @@ ggml_cgraph * clip_graph_granite4_vision::build() { ggml_tensor * stream = build_block( blk, h, bid, hparams.proj_spatial_offsets[bid], - n_patches_x, + tile_side, hparams.downsample_window_side, hparams.downsample_query_side, qformer_eps); @@ -326,10 +383,11 @@ ggml_cgraph * clip_graph_granite4_vision::build() { mmproj = mmproj ? ggml_concat(ctx0, mmproj, stream, 0) : stream; } - // --- Stage 1d: Append newline tokens if add_newline is set --- - if (add_newline) { - mmproj = append_rowwise_newlines(ctx0, mmproj); - ggml_set_name(mmproj, "g4v_mmproj_out_nl"); + // --- Stage 1d: assemble the tiles and weave in the newline tokens --- + if (anyres.is_tiled()) { + const int out_side = tile_side / hparams.downsample_window_side * hparams.downsample_query_side; + mmproj = build_anyres_assembly(mmproj, out_side); + ggml_set_name(mmproj, "g4v_mmproj_out_anyres"); } else { ggml_set_name(mmproj, "g4v_mmproj_out"); } diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index ed8c1ea51..3631d849b 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -402,16 +402,19 @@ struct clip_graph_exaone4_5 : clip_graph { struct clip_graph_granite4_vision : clip_graph { clip_graph_granite4_vision(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img), - add_newline(img.add_newline) {} + anyres(img.anyres), + n_tiles(img.ny() / img.nx()), + tile_side(img.nx() / patch_size) {} ggml_cgraph * build() override; private: - // The graph is per-tile since only batch-size 1 is supported in clip. As - // such, this value is set at construct time based on the tile that will be - // encoded, then used during build to determine how to handle newlines. - const bool add_newline; + // the input image is a stack of tiles on the Y axis: [overview, tile(0,0), tile(0,1), ...] + const clip_image_f32::anyres_info anyres; + const int n_tiles; + const int tile_side; // patches per tile side + ggml_tensor * build_tile_inp(); ggml_tensor * gather(ggml_tensor * src, const std::string & name, int idx_len); ggml_tensor * interp_down(ggml_tensor * src, int side, int new_side); ggml_tensor * build_block(const qf_block & blk, ggml_tensor * h, int bid, @@ -419,7 +422,7 @@ private: int query_side, float qformer_eps); ggml_tensor * build_newline_row(ggml_context * ctx0); - ggml_tensor * append_rowwise_newlines(ggml_context * ctx0, ggml_tensor * tile_output); + ggml_tensor * build_anyres_assembly(ggml_tensor * cur, int out_side); }; struct clip_graph_muse_glimmer : clip_graph { diff --git a/tools/mtmd/mtmd-helper-common.h b/tools/mtmd/mtmd-helper-common.h index 968b4df9c..f907346c7 100644 --- a/tools/mtmd/mtmd-helper-common.h +++ b/tools/mtmd/mtmd-helper-common.h @@ -82,7 +82,7 @@ struct decode_embd_batch { llama_batch batch; decode_embd_batch(float * embd, int32_t n_tokens, int n_pos_per_embd, int n_mmproj_embd) : n_pos_per_embd(n_pos_per_embd), n_mmproj_embd(n_mmproj_embd) { GGML_ASSERT(n_tokens > 0 && n_pos_per_embd > 0 && n_mmproj_embd > 0); - pos .resize(n_tokens * n_pos_per_embd); + pos .resize((size_t) n_tokens * (size_t) n_pos_per_embd); n_seq_id.resize(n_tokens); seq_ids .resize(n_tokens + 1); logits .resize(n_tokens); @@ -115,10 +115,12 @@ struct decode_embd_batch { GGML_ASSERT(!rel_pos.empty() && (int32_t)rel_pos.size() == batch.n_tokens); seq_id_0[0] = seq_id; for (int32_t i = 0; i < batch.n_tokens; i++) { - pos[i ] = rel_pos[i].t; - pos[i + batch.n_tokens ] = rel_pos[i].y; - pos[i + batch.n_tokens * 2] = rel_pos[i].x; - pos[i + batch.n_tokens * 3] = rel_pos[i].z; + const size_t idx = (size_t) i; + const size_t n_tokens = (size_t) batch.n_tokens; + pos[idx ] = rel_pos[i].t; + pos[idx + n_tokens ] = rel_pos[i].y; + pos[idx + n_tokens * 2 ] = rel_pos[i].x; + pos[idx + n_tokens * 3 ] = rel_pos[i].z; } for (int i = 0; i < batch.n_tokens; i++) { batch.n_seq_id[i] = 1; @@ -132,10 +134,12 @@ struct decode_embd_batch { GGML_ASSERT(n_pos_per_embd == 4); seq_id_0[0] = seq_id; for (int i = 0; i < batch.n_tokens; i++) { - pos[i ] = pos_0 + i; - pos[i + batch.n_tokens ] = pos_0 + i; - pos[i + batch.n_tokens * 2] = pos_0 + i; - pos[i + batch.n_tokens * 3] = pos_0 + i; + const size_t idx = (size_t) i; + const size_t n_tokens = (size_t) batch.n_tokens; + pos[idx ] = pos_0 + i; + pos[idx + n_tokens ] = pos_0 + i; + pos[idx + n_tokens * 2 ] = pos_0 + i; + pos[idx + n_tokens * 3 ] = pos_0 + i; } for (int i = 0; i < batch.n_tokens; i++) { batch.n_seq_id[i] = 1; @@ -148,7 +152,7 @@ struct decode_embd_batch { GGML_ASSERT(offset >= 0 && n_tokens > 0 && offset + n_tokens <= batch.n_tokens); llama_pos * pos_ptr; pos_view.clear(); - pos_view.reserve(n_tokens * n_pos_per_embd); + pos_view.reserve((size_t) n_tokens * (size_t) n_pos_per_embd); if (n_pos_per_embd > 1) { // mrope // for example, with layout of src: 1234...1234...1234...1234... @@ -157,7 +161,7 @@ struct decode_embd_batch { // assume n_tokens is less than or equal to batch.n_tokens // batch.n_tokens is number of **total** tokens // n_tokens is number of viewed token - size_t src_idx = i * batch.n_tokens + offset; + size_t src_idx = (size_t) i * (size_t) batch.n_tokens + (size_t) offset; pos_view.insert(pos_view.end(), pos.data() + src_idx, pos.data() + src_idx + n_tokens); diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index 813fe493f..e795b47e0 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -1317,7 +1317,7 @@ void mtmd_image_preprocessor_step3vl::img_u8_resize_bilinear_to_f32( const float scale_x = static_cast(src_size.width) / target_width; const float scale_y = static_cast(src_size.height) / target_height; - std::vector local_buf(3 * target_width * target_height); + std::vector local_buf((size_t) 3 * (size_t) target_width * (size_t) target_height); for (int y = 0; y < target_height; ++y) { const float src_y = (static_cast(y) + 0.5f) * scale_y - 0.5f; @@ -1338,7 +1338,7 @@ void mtmd_image_preprocessor_step3vl::img_u8_resize_bilinear_to_f32( const auto p10 = src.get_pixel(x0, y1); const auto p11 = src.get_pixel(x1, y1); - const size_t idx_dst = 3 * (y * target_width + x); + const size_t idx_dst = (size_t) 3 * ((size_t) y * (size_t) target_width + (size_t) x); for (int c = 0; c < 3; ++c) { const float v00 = (static_cast(p00[c]) / 255.0f - mean[c]) / std[c]; const float v01 = (static_cast(p01[c]) / 255.0f - mean[c]) / std[c]; @@ -1602,17 +1602,50 @@ mtmd_image_preproc_out mtmd_image_preprocessor_youtuvl::preprocess(const clip_im } mtmd_image_preproc_out mtmd_image_preprocessor_granite::preprocess(const clip_image_u8 & img) { - auto output = mtmd_image_preprocessor_llava_uhd::preprocess(img); - if (output.entries.size() == 0) { - // Single-tile (overview only): append one newline row. - output.overview.add_newline = true; - } else { - // Multi-tile: overview gets no newline, grid tiles get one. - output.overview.add_newline = false; - for (size_t i = 0; i < output.entries.size(); ++i) { - output.entries[i].add_newline = true; + GGML_ASSERT(!hparams.image_res_candidates.empty()); + + const clip_image_size orig_size = img.get_size(); + const int tile_size = hparams.image_size; + + // llava-next always encodes an overview plus a grid of tiles, even for small images + const clip_image_size refined_size = select_best_resolution(orig_size, hparams.image_res_candidates); + const int grid_x = refined_size.width / tile_size; + const int grid_y = refined_size.height / tile_size; + + clip_image_u8 overview; + img_tool::resize(img, overview, {tile_size, tile_size}, hparams.image_resize_algo_ov, + hparams.image_pad_ov, hparams.image_pad_color_ov); + + clip_image_u8 refined; + img_tool::resize(img, refined, refined_size, hparams.image_resize_algo_rf, + hparams.image_pad_rf, hparams.image_pad_color_rf); + + // stack the overview and the tiles on the Y axis, so the whole grid goes through one graph + clip_image_u8 stacked; + stacked.set_size({tile_size, tile_size * (1 + grid_x * grid_y)}, false); + auto copy_tile = [&](const clip_image_u8 & src, int src_x, int src_y, int dst_idx) { + for (int py = 0; py < tile_size; py++) { + for (int px = 0; px < tile_size; px++) { + stacked.set_pixel(px, dst_idx * tile_size + py, src.get_pixel(src_x + px, src_y + py)); + } + } + }; + copy_tile(overview, 0, 0, 0); + for (int ty = 0; ty < grid_y; ty++) { + for (int tx = 0; tx < grid_x; tx++) { + copy_tile(refined, tx * tile_size, ty * tile_size, 1 + ty * grid_x + tx); } } + + LOG_DBG("%s: grid size: %d x %d (%d tiles) + overview\n", __func__, grid_x, grid_y, grid_x * grid_y); + + mtmd_image_preproc_out output; + output.append(hparams, stacked, true); + auto & entry = output.entries.back(); + entry.anyres.grid_x = grid_x; + entry.anyres.grid_y = grid_y; + entry.anyres.orig_nx = orig_size.width; + entry.anyres.orig_ny = orig_size.height; return output; } diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index 0669aa112..238332934 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -85,9 +85,6 @@ struct mtmd_image_preprocessor_llava_uhd : mtmd_image_preprocessor { protected: clip_image_size get_best_resize(const clip_image_size & original_size, int scale_resolution, int patch_size, bool allow_upscale = false); -private: - clip_image_size resize_maintain_aspect_ratio(const clip_image_size & orig, const clip_image_size & target_max); - /** * Selects the best resolution from a list of possible resolutions based on the original size. * @@ -104,6 +101,9 @@ private: * @return The best fit resolution */ clip_image_size select_best_resolution(const clip_image_size & original_size, const std::vector & possible_resolutions); + +private: + clip_image_size resize_maintain_aspect_ratio(const clip_image_size & orig, const clip_image_size & target_max); int ensure_divide(int length, int patch_size); clip_image_size get_refine_size(const clip_image_size & original_size, const clip_image_size & grid, int scale_resolution, int patch_size, bool allow_upscale = false); clip_image_size get_best_grid(const int max_slice_nums, const int multiple, const float log_ratio); @@ -225,7 +225,7 @@ struct mtmd_image_preprocessor_youtuvl : mtmd_image_preprocessor { mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; }; -// similar to llava_uhd, but has add_newline +// llava-next "anyres": stacks the overview and all tiles into one image, assembled by clip in a single graph struct mtmd_image_preprocessor_granite : mtmd_image_preprocessor_llava_uhd { mtmd_image_preprocessor_granite(const clip_ctx * ctx) : mtmd_image_preprocessor_llava_uhd(ctx) {} mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 4b9c45d62..6596b2785 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -891,10 +891,10 @@ struct mtmd_context { } break; case PROJECTOR_TYPE_GRANITE4_VISION: { - img_beg = ""; - img_end = ""; + // ... (image embeddings) \n ... + img_beg = ""; + img_end = "\n"; image_preproc = std::make_unique(ctx_v); - ov_img_first = true; } break; default: throw std::runtime_error(string_format("%s: unexpected vision projector type %d\n", __func__, proj)); diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 5ff7685bb..7ed486528 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -1292,12 +1292,15 @@ json oaicompat_chat_params_parse( throw std::invalid_argument("invalid type for \"enable_thinking\" (expected boolean, got string)"); } - // Parse also the OAI "reasoning_effort": "none" specific value + // Parse the OAI "reasoning_effort" field; "none" disables reasoning. if (body.contains("reasoning_effort")) { auto reasoning_effort = json_value(body, "reasoning_effort", std::string("")); if (reasoning_effort == "none") { inputs.enable_thinking = false; - } // other reasoning_effort values are model-specific and not yet handled + inputs.chat_template_kwargs.erase("reasoning_effort"); + } else if (!reasoning_effort.empty()) { + inputs.chat_template_kwargs["reasoning_effort"] = json(reasoning_effort).dump(); + } } inputs.force_pure_content = opt.force_pure_content; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index f02a1da68..a7d382886 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -688,97 +688,99 @@ struct server_slot { other.prompt = prompt.clone(); other.init_sampler(); } - - // returns 0 on success - // caller need to update prompt.tokens after a successful call to keep track of the processing progress - int process_mtmd_chunk(size_t idx, size_t & n_tokens_out) { - GGML_ASSERT(mctx); - const auto & input_tokens = task->tokens; - const auto & chunk = input_tokens.find_chunk(idx); - int32_t res = 0; - - auto try_decode = [&]() -> int32_t { - if (mbatch) { - float * embd = mtmd_batch_get_output_embd(mbatch.get(), chunk.get()); - if (embd) { - void * cb_data = spec; - static auto cb = [](llama_batch batch, void * user_data) { - common_speculative * spec = static_cast(user_data); - if (!common_speculative_process(spec, batch)) { - return 1; - } - return 0; - }; - - llama_pos new_n_past; // unused for now - res = mtmd_helper_decode_image_chunk( - mctx, - ctx_tgt, - chunk.get(), - embd, - prompt.tokens.pos_next(), - id, - llama_n_batch(ctx_tgt), - &new_n_past, - cb, - cb_data - ); - if (res != 0) { - SLT_ERR(*this, "failed to decode mtmd chunk, idx = %zu, res = %d\n", idx, res); - return -1; - } - n_tokens_out = mtmd_input_chunk_get_n_tokens(chunk.get()); - return 0; // success - } - } - return 1; // (non-error) need to create & encode batch - }; - - // if the batch is already exist, try searching & encode - res = try_decode(); - if (res == 0) { - return 0; - } - if (res < 0) { - // fatal error - return res; - } - - // otherwise, the batch is either uninitialized or is used up - // we need to create & encode a new batch - mbatch.reset(mtmd_batch_init(mctx)); - res = mtmd_batch_add_chunk(mbatch.get(), chunk.get()); - GGML_ASSERT(res == 0); // we should never have an empty batch - - // try batching as much as possible - int n_added = 1; - size_t idx_cur = idx; - while (res == 0) { - auto [next_chunk, next_idx] = input_tokens.find_next_media_chunk(idx_cur); - if (next_chunk == nullptr) { - break; - } - res = mtmd_batch_add_chunk(mbatch.get(), next_chunk->get()); - n_added += (res == 0 ? 1 : 0); - idx_cur = next_idx; - SLT_DBG(*this, "try adding media chunk idx = %zu to batch, res = %d\n", next_idx, res); - // if res != 0, batch is full or chunk is not compatible -> this loop breaks - } - - // TODO @ngxson : move this log line to debug when it become more stable - SLT_TRC(*this, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added); - - res = mtmd_batch_encode(mbatch.get()); - if (res != 0) { - SLT_ERR(*this, "failed to encode mtmd batch for chunk idx = %zu, res = %d\n", idx, res); - return -1; - } - - return try_decode(); - } }; +// returns 0 on success +// caller need to update prompt.tokens after a successful call to keep track of the processing progress +// note: this is not a member of server_slot because we want to run it inside yield_to_queue +// slot is passed as const to avoid accidental modification of the slot state +// some pointers are allowed to be used, they are not used by to_json() +static int process_mtmd_chunk(const server_slot & slot, mtmd::batch_ptr & mbatch, size_t idx, size_t & n_tokens_out) { + GGML_ASSERT(slot.mctx); + const auto & mctx = slot.mctx; + const auto & input_tokens = slot.task->tokens; + const auto & chunk = input_tokens.find_chunk(idx); + int32_t res = 0; + auto try_decode = [&]() -> int32_t { + if (mbatch) { + float * embd = mtmd_batch_get_output_embd(mbatch.get(), chunk.get()); + if (embd) { + void * cb_data = slot.spec; + static auto cb = [](llama_batch batch, void * user_data) { + common_speculative * spec = static_cast(user_data); + if (!common_speculative_process(spec, batch)) { + return 1; + } + return 0; + }; + + llama_pos new_n_past; // unused for now + res = mtmd_helper_decode_image_chunk( + mctx, + slot.ctx_tgt, + chunk.get(), + embd, + slot.prompt.tokens.pos_next(), + slot.id, + llama_n_batch(slot.ctx_tgt), + &new_n_past, + cb, + cb_data + ); + if (res != 0) { + SLT_ERR(slot, "failed to decode mtmd chunk, idx = %zu, res = %d\n", idx, res); + return -1; + } + n_tokens_out = mtmd_input_chunk_get_n_tokens(chunk.get()); + return 0; // success + } + } + return 1; // (non-error) need to create & encode batch + }; + + // if the batch is already exist, try searching & encode + res = try_decode(); + if (res == 0) { + return 0; + } + if (res < 0) { + // fatal error + return res; + } + + // otherwise, the batch is either uninitialized or is used up + // we need to create & encode a new batch + mbatch.reset(mtmd_batch_init(mctx)); + res = mtmd_batch_add_chunk(mbatch.get(), chunk.get()); + GGML_ASSERT(res == 0); // we should never have an empty batch + + // try batching as much as possible + int n_added = 1; + size_t idx_cur = idx; + while (res == 0) { + auto [next_chunk, next_idx] = input_tokens.find_next_media_chunk(idx_cur); + if (next_chunk == nullptr) { + break; + } + res = mtmd_batch_add_chunk(mbatch.get(), next_chunk->get()); + n_added += (res == 0 ? 1 : 0); + idx_cur = next_idx; + SLT_DBG(slot, "try adding media chunk idx = %zu to batch, res = %d\n", next_idx, res); + // if res != 0, batch is full or chunk is not compatible -> this loop breaks + } + + // TODO @ngxson : move this log line to debug when it become more stable + SLT_TRC(slot, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added); + + res = mtmd_batch_encode(mbatch.get()); + if (res != 0) { + SLT_ERR(slot, "failed to encode mtmd batch for chunk idx = %zu, res = %d\n", idx, res); + return -1; + } + + return try_decode(); +} // // server_context_impl (private implementation) @@ -1354,8 +1356,8 @@ private: GGML_ASSERT(!sleeping); // wiring up server queues - queue_tasks.on_new_task([this](server_task && task) { - process_single_task(std::move(task)); + queue_tasks.on_new_task([this](server_task && task, bool is_yielding) { + return process_single_task(std::move(task), is_yielding); }); queue_tasks.on_update_slots([this]() { update_slots(); @@ -2286,7 +2288,14 @@ private: cur.pos_max, cur.n_tokens, (float) cur.size() / 1024 / 1024); } - void process_single_task(server_task && task) { + // returns false to decline the task, it is offered again after the decode is done + bool process_single_task(server_task && task, bool is_yielding) { + // while yielding, an encode / decode is running and only accessing metrics is safe + if (is_yielding && task.type != SERVER_TASK_TYPE_METRICS) { + SRV_DBG("decoding, decline task, id_task = %d\n", task.id); + return false; + } + switch (task.type) { case SERVER_TASK_TYPE_COMPLETION: case SERVER_TASK_TYPE_INFILL: @@ -2620,6 +2629,8 @@ private: queue_results.send(std::move(res)); } break; } + + return true; } void iterate(std::vector & slots, std::function callback) { @@ -3382,8 +3393,13 @@ private: // so the timing is queued and flushed on the next sync metrics_pre_decode(); + // encode on the worker thread, so we can still handle metrics tasks size_t n_tokens_out = 0; - int32_t res = slot.process_mtmd_chunk(cur_token_idx, n_tokens_out); + int32_t res = 0; + queue_tasks.yield_to_queue([&]() { + res = process_mtmd_chunk(slot, slot.mbatch, cur_token_idx, n_tokens_out); + }); + if (res != 0) { SLT_ERR(slot, "failed to process mtmd chunk, res = %d\n", res); send_error(slot, "failed to process mtmd chunk", ERROR_TYPE_SERVER); @@ -3557,7 +3573,20 @@ private: } } - const int ret = llama_decode(ctx_tgt, batch_view); + bool has_output = false; + for (int i = off; i < off + batch_view.n_tokens; ++i) { + has_output |= batch.tokens[i].output; + } + + // decode on the worker thread, so we can still handle metrics tasks while waiting + // note: the sync is done here too, so that the wait also happens off the main thread + int ret = 0; + queue_tasks.yield_to_queue([&]() { + ret = llama_decode(ctx_tgt, batch_view); + if (ret == 0 && has_output) { + llama_synchronize(ctx_tgt); + } + }); if (ret != 0) { { @@ -3609,7 +3638,7 @@ private: return false; // retry with the updated n_batch } else { // success, apply batch metrics - metrics_post_decode(off, batch_view.n_tokens); + metrics_post_decode(off, batch_view.n_tokens, has_output); } // TODO: avoid restoring the draft context and re-evaluating the drafted tokens when not needed [TAG_SPEC_AVOID_DRAFT_REEVAL] @@ -3922,7 +3951,8 @@ private: n_prompt_queued = 0; } - void metrics_post_decode(int32_t off, int32_t n_tokens) { + // has_output is computed by the caller, which also already synchronized the context if it is set + void metrics_post_decode(int32_t off, int32_t n_tokens, bool has_output) { metrics.n_decode++; for (const auto & slot : slots) { if (slot.is_processing()) { @@ -3935,13 +3965,10 @@ private: // note: a slot can be released before we get here, which clears its stats // the tokens were still computed, counted in the global metrics, not in slot uint64_t n_prompt_tokens = 0; - bool has_output = false; for (int i = off; i < off + n_tokens; ++i) { const auto & t = batch.tokens[i]; - has_output |= t.output; - if (!t.is_prompt) { continue; // generated tokens are handled after sampling } @@ -3957,14 +3984,12 @@ private: metrics_queue_prompt(n_prompt_tokens); if (has_output) { - // sync if we have at least one output in batch - // so that we can calculate the timings correctly - llama_synchronize(ctx_tgt); + // the context is already synchronized, so the timings are correct metrics_flush_prompt(); } // advance the prompt timing of the slots that had tokens in this batch - // note: a second pass, it must run after the sync above to reflect the compute + // note: a second pass, it must run after the sync to reflect the compute const int64_t t_now = ggml_time_us(); for (int i = off; i < off + n_tokens; ++i) { const auto & t = batch.tokens[i]; diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 5d37c3453..d22527ffe 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -4,6 +4,7 @@ #include "log.h" #include +#include #define QUE_INF(fmt, ...) LOG_INF("que %12.*s: " fmt, 12, __func__, __VA_ARGS__) #define QUE_WRN(fmt, ...) LOG_WRN("que %12.*s: " fmt, 12, __func__, __VA_ARGS__) @@ -122,10 +123,135 @@ void server_queue::terminate() { condition_tasks.notify_all(); } +bool server_queue::process_new_tasks(bool is_yielding) { + while (true) { + std::unique_lock lock(mutex_tasks); + if (!running) { + QUE_DBG("%s", "terminate\n"); + return true; + } + if (queue_tasks.empty()) { + return false; + } + server_task task = std::move(queue_tasks.front()); + queue_tasks.pop_front(); + lock.unlock(); + + QUE_DBG("processing task, id = %d\n", task.id); + if (!callback_new_task(std::move(task), is_yielding)) { + // set it aside, do not put it back in the queue, else we offer it again in a loop + GGML_ASSERT(is_yielding && "a task can only be declined while yielding"); + QUE_DBG("task declined, id = %d\n", task.id); + lock.lock(); + queue_tasks_unhandled.push_back(std::move(task)); + } + } +} + +void server_queue::worker_loop() { + while (true) { + std::function work; + { + std::unique_lock lock(mutex_tasks); + worker.cv.wait(lock, [&]{ + return worker.stop || worker.work != nullptr; + }); + if (worker.stop) { + return; + } + work = std::move(worker.work); + worker.work = nullptr; + } + + // note: do not hold any lock here, work() may post new tasks + std::exception_ptr exception; + try { + work(); + } catch (...) { + exception = std::current_exception(); + } + + // signal completion to yield_to_queue() + std::unique_lock lock(mutex_tasks); + worker.exception = std::move(exception); + worker.busy = false; + condition_tasks.notify_all(); + } +} + +void server_queue::worker_stop() { + if (!worker.thread.joinable()) { + return; + } + { + std::unique_lock lock(mutex_tasks); + worker.stop = true; + } + worker.cv.notify_one(); + worker.thread.join(); +} + +void server_queue::yield_to_queue(std::function && work) { + GGML_ASSERT(worker.thread.joinable() && "yield_to_queue() requires start_loop() to be running"); + + QUE_DBG("%s", "yielding to queue\n"); + + { + std::unique_lock lock(mutex_tasks); + GGML_ASSERT(!worker.busy && "yield_to_queue() cannot be nested"); + worker.busy = true; + worker.work = std::move(work); + } + worker.cv.notify_one(); + + while (true) { + // note: on terminate this is a no-op, but we still wait for the work to finish + process_new_tasks(true); + + std::unique_lock lock(mutex_tasks); + // declined tasks are moved to queue_tasks_unhandled, so a non-empty queue always has something new + condition_tasks.wait(lock, [&]{ + return !worker.busy || (running && !queue_tasks.empty()); + }); + if (!worker.busy) { + break; + } + } + + std::exception_ptr exception; + { + std::unique_lock lock(mutex_tasks); + + // put the declined tasks back, keeping their order + while (!queue_tasks_unhandled.empty()) { + queue_tasks.push_front(std::move(queue_tasks_unhandled.back())); + queue_tasks_unhandled.pop_back(); + } + + // make sure to avoid idle timeout here + time_last_task = ggml_time_ms(); + + // the worker is idle now, take the exception it may have left behind + std::swap(exception, worker.exception); + } + + QUE_DBG("%s", "done yielding to queue\n"); + + // note: rethrow only after the declined tasks are back in the queue, so they are not lost + if (exception) { + std::rethrow_exception(exception); + } +} + void server_queue::start_loop(int64_t idle_sleep_ms) { running = true; time_last_task = ggml_time_ms(); + // spawn the worker thread used by yield_to_queue() + GGML_ASSERT(!worker.thread.joinable() && "start_loop() is already running"); + worker.stop = false; + worker.thread = std::thread([this]() { worker_loop(); }); + constexpr auto max_wait_time = std::chrono::seconds(1); auto should_sleep = [&]() -> bool { // caller must hold mutex_tasks @@ -138,24 +264,10 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { while (true) { QUE_DBG("%s", "processing new tasks\n"); - - while (true) { - std::unique_lock lock(mutex_tasks); - if (!running) { - QUE_DBG("%s", "terminate\n"); - return; - } - if (queue_tasks.empty()) { - lock.unlock(); - break; - } - server_task task = std::move(queue_tasks.front()); - queue_tasks.pop_front(); - lock.unlock(); - - QUE_DBG("processing task, id = %d\n", task.id); - callback_new_task(std::move(task)); + if (process_new_tasks(false)) { + break; // terminate } + // all tasks in the current loop is processed, slots data is now ready QUE_DBG("%s", "update slots\n"); @@ -206,6 +318,8 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { } } } + + worker_stop(); } void server_queue::cleanup_pending_task(int id_target) { @@ -214,11 +328,15 @@ void server_queue::cleanup_pending_task(int id_target) { return task.id == id_target; }; queue_tasks.erase( - std::remove_if(queue_tasks.begin(), queue_tasks.end(), rm_func), + std::remove_if(queue_tasks.begin(), queue_tasks.end(), rm_func), queue_tasks.end()); queue_tasks_deferred.erase( - std::remove_if(queue_tasks_deferred.begin(), queue_tasks_deferred.end(), rm_func), + std::remove_if(queue_tasks_deferred.begin(), queue_tasks_deferred.end(), rm_func), queue_tasks_deferred.end()); + // a task declined while yielding is not in queue_tasks yet, but it can still be cancelled + queue_tasks_unhandled.erase( + std::remove_if(queue_tasks_unhandled.begin(), queue_tasks_unhandled.end(), rm_func), + queue_tasks_unhandled.end()); } // diff --git a/tools/server/server-queue.h b/tools/server/server-queue.h index 0b674d6ff..69a3e554f 100644 --- a/tools/server/server-queue.h +++ b/tools/server/server-queue.h @@ -4,7 +4,9 @@ #include #include +#include #include +#include #include #include @@ -21,16 +23,32 @@ private: // queues std::deque queue_tasks; std::deque queue_tasks_deferred; + // tasks declined while yielding, put back in queue_tasks once the yield is done + // note: kept as a member so that cleanup_pending_task() can also reach them + std::deque queue_tasks_unhandled; std::mutex mutex_tasks; std::condition_variable condition_tasks; + // used by yield_to_queue, all fields are guarded by mutex_tasks + struct worker_t { + std::thread thread; + std::condition_variable cv; // the worker sleeps on this until there is work + std::function work; // pending work, picked up by the thread + std::exception_ptr exception; // exception thrown by work(), if any + bool stop = false; + bool busy = false; + }; + worker_t worker; + // callback functions - std::function callback_new_task; - std::function callback_update_slots; - std::function callback_sleeping_state; + std::function callback_new_task; + std::function callback_update_slots; + std::function callback_sleeping_state; public: + ~server_queue() { worker_stop(); } + // Add a new task to the end of the queue int post(server_task && task, bool front = false); @@ -75,6 +93,15 @@ public: */ void start_loop(int64_t idle_sleep_ms = -1); + // run work() on a separate thread, while the current thread calls process_new_tasks + // returns once work() is done (may throw exceptions) + // must be called from start_loop() thread (ideally inside callback_update_slots) + // use case: return metrics while encode/decode is running + // ref: https://github.com/ggml-org/llama.cpp/pull/27041 + // + // tasks declined by callback_new_task are put back in the queue once this returns + void yield_to_queue(std::function && work); + // for metrics size_t queue_tasks_deferred_size() { std::unique_lock lock(mutex_tasks); @@ -86,7 +113,10 @@ public: // // Register function to process a new task - void on_new_task(std::function callback) { + // the second argument tells whether the queue is currently yielding (see yield_to_queue) + // only then may the callback return false to decline the task, and it must leave it + // untouched, so that it can be put back in the queue later + void on_new_task(std::function callback) { callback_new_task = std::move(callback); } @@ -112,6 +142,15 @@ public: private: void cleanup_pending_task(int id_target); + + // process all pending tasks in the queue + // returns true if the queue is terminated, false if there is no more task to process + // while yielding, declined tasks are moved to queue_tasks_unhandled + bool process_new_tasks(bool is_yielding); + + // for worker_t + void worker_loop(); + void worker_stop(); }; // struct for managing server responses diff --git a/vendor/cpp-httplib/CMakeLists.txt b/vendor/cpp-httplib/CMakeLists.txt index 49c419094..6a6eefed1 100644 --- a/vendor/cpp-httplib/CMakeLists.txt +++ b/vendor/cpp-httplib/CMakeLists.txt @@ -41,7 +41,7 @@ if (LLAMA_BUILD_BORINGSSL) set(FIPS OFF CACHE BOOL "Enable FIPS (BoringSSL)") set(BORINGSSL_GIT "https://boringssl.googlesource.com/boringssl" CACHE STRING "BoringSSL git repository") - set(BORINGSSL_VERSION "0.20260803.0" CACHE STRING "BoringSSL version") + set(BORINGSSL_VERSION "0.20260813.0" CACHE STRING "BoringSSL version") message(STATUS "Fetching BoringSSL version ${BORINGSSL_VERSION}") diff --git a/vendor/cpp-httplib/httplib.cpp b/vendor/cpp-httplib/httplib.cpp index 3b687ff4f..81cdcfe3a 100644 --- a/vendor/cpp-httplib/httplib.cpp +++ b/vendor/cpp-httplib/httplib.cpp @@ -7146,6 +7146,10 @@ PathParamsMatcher::PathParamsMatcher(const std::string &pattern) bool PathParamsMatcher::match(Request &request) const { request.matches = std::smatch(); request.path_params.clear(); + + // A pattern without parameters is just a literal path to compare against + if (param_names_.empty()) { return request.path == pattern(); } + request.path_params.reserve(param_names_.size()); // One past the position at which the path matched the pattern last time @@ -7188,6 +7192,11 @@ bool PathParamsMatcher::match(Request &request) const { bool RegexMatcher::match(Request &request) const { request.path_params.clear(); + // See CPPHTTPLIB_REGEX_ROUTE_PATH_MAX_LENGTH: an overlong path is treated as + // a non-match rather than risking a stack overflow in std::regex_match. + if (request.path.length() > CPPHTTPLIB_REGEX_ROUTE_PATH_MAX_LENGTH) { + return false; + } return std::regex_match(request.path, request.matches, regex_); } @@ -7613,11 +7622,21 @@ Server::~Server() = default; std::unique_ptr Server::make_matcher(const std::string &pattern) { + // Path params take precedence, so "/users/:id/(.*)" keeps being matched as + // a path params pattern if (pattern.find("/:") != std::string::npos) { return detail::make_unique(pattern); - } else { - return detail::make_unique(pattern); } + + // A pattern with no regex metacharacter only has to be compared literally, + // which is what PathParamsMatcher already does when it captures no + // parameter, so std::regex is only worth building for the patterns that + // actually need it + if (pattern.find_first_of(".^$|()[]{}*+?\\") == std::string::npos) { + return detail::make_unique(pattern); + } + + return detail::make_unique(pattern); } Server &Server::Get(const std::string &pattern, Handler handler) { @@ -8259,15 +8278,12 @@ bool Server::read_content_core( } } if (has_data) { - auto result = - detail::read_content_without_length(strm, payload_max_length_, out); - if (result == detail::ReadContentResult::PayloadTooLarge) { - res.status = StatusCode::PayloadTooLarge_413; - return false; - } else if (result != detail::ReadContentResult::Success) { - return false; - } - return true; + // Route through the same decompressing reader used by the + // length-framed and chunked paths below, so payload_max_length_ is + // enforced on the decompressed size here too instead of only on the + // compressed wire bytes. + return detail::read_content(strm, req, payload_max_length_, res.status, + nullptr, out, true); } } return true; diff --git a/vendor/cpp-httplib/httplib.h b/vendor/cpp-httplib/httplib.h index f82e58e49..6fc86c7c7 100644 --- a/vendor/cpp-httplib/httplib.h +++ b/vendor/cpp-httplib/httplib.h @@ -8,8 +8,8 @@ #ifndef CPPHTTPLIB_HTTPLIB_H #define CPPHTTPLIB_HTTPLIB_H -#define CPPHTTPLIB_VERSION "0.53.0" -#define CPPHTTPLIB_VERSION_NUM "0x003500" +#define CPPHTTPLIB_VERSION "0.53.1" +#define CPPHTTPLIB_VERSION_NUM "0x003501" #ifdef _WIN32 #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00 @@ -138,6 +138,18 @@ #define CPPHTTPLIB_RANGE_MAX_COUNT 1024 #endif +// std::regex_match's backtracking implementation (most acutely on libstdc++) +// recurses roughly once per matched character for quantified patterns such +// as "(.*)", so a long enough path can exhaust the calling thread's stack; on +// a default ~8MB thread stack that has been observed to take on the order of +// a couple thousand characters for a simple pattern. 256 leaves a wide safety +// margin below that (well under the 8192-byte request URI limit) while still +// fitting any realistic route segment; raise it if a route legitimately needs +// longer paths. Regex routes are never applied to paths longer than this. +#ifndef CPPHTTPLIB_REGEX_ROUTE_PATH_MAX_LENGTH +#define CPPHTTPLIB_REGEX_ROUTE_PATH_MAX_LENGTH 256 +#endif + #ifndef CPPHTTPLIB_TCP_NODELAY #define CPPHTTPLIB_TCP_NODELAY false #endif @@ -839,6 +851,15 @@ inline bool parse_url(const std::string &url, UrlComponents &uc) { } pos = close + 1; + + // The IPv6 literal is the whole host, so ']' must be followed by a port, + // path, query or fragment delimiter (or the end of input). Otherwise the + // trailing bytes would be folded into the path while the connection + // still targets the bracketed address. + if (pos < url.size()) { + auto c = url[pos]; + if (c != ':' && c != '/' && c != '?' && c != '#') { return false; } + } } else { auto end = url.find_first_of(":/?#", pos); if (end == std::string::npos) { end = url.size(); }