diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp index f786f5ff2314..fc300d89a839 100644 --- a/common/chat-peg-parser.cpp +++ b/common/chat-peg-parser.cpp @@ -594,7 +594,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key( const std::string & args_key, const std::string & effective_args_key, const std::string & call_id_key, - const std::string & gen_call_id_key) { + const std::string & gen_call_id_key, + bool require_object_args) { auto tool_choices = choice(); @@ -631,10 +632,10 @@ common_peg_parser common_chat_peg_builder::build_json_tools_function_is_key( // Arguments — either wrapped in args_key or parsed directly common_peg_parser args_parser = eps(); if (args_key.empty()) { - args_parser = tool_args(schema(json(), "tool-" + name + "-schema", params)); + args_parser = tool_args(schema(require_object_args ? json_object() : json(), "tool-" + name + "-schema", params)); } else { args_parser = literal("\"" + effective_args_key + "\"") + space() + literal(":") + space() + - tool_args(schema(json(), "tool-" + name + "-schema", params)); + tool_args(schema(require_object_args ? json_object() : json(), "tool-" + name + "-schema", params)); } inner_fields.push_back(args_parser); @@ -673,7 +674,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys( const std::string & effective_name_key, const std::string & effective_args_key, const std::string & call_id_key, - const std::string & gen_call_id_key) { + const std::string & gen_call_id_key, + bool require_object_args) { auto tool_choices = choice(); @@ -695,7 +697,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_nested_keys( auto nested_name = literal("\"" + nested_name_field + "\"") + space() + literal(":") + space() + atomic(literal("\"") + tool_name(literal(name)) + literal("\"")); auto nested_args = literal("\"" + nested_args_field + "\"") + space() + literal(":") + space() + - tool_args(schema(json(), "tool-" + name + "-schema", params)); + tool_args(schema(require_object_args ? json_object() : json(), "tool-" + name + "-schema", params)); auto nested_object = literal("{") + space() + nested_name + space() + literal(",") + space() + @@ -747,7 +749,8 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys( const std::string & call_id_key, const std::string & gen_call_id_key, const std::vector & parameters_order, - bool accept_openai_wrapper) { + bool accept_openai_wrapper, + bool require_object_args) { auto tool_choices = choice(); auto name_key_parser = literal("\"" + effective_name_key + "\""); @@ -764,7 +767,7 @@ common_peg_parser common_chat_peg_builder::build_json_tools_flat_keys( auto tool_name_ = name_key_parser + space() + literal(":") + space() + atomic(literal("\"") + tool_name(literal(name)) + literal("\"")); auto tool_args_ = args_key_parser + space() + literal(":") + space() + - tool_args(schema(json(), "tool-" + name + "-schema", params)); + tool_args(schema(require_object_args ? json_object() : json(), "tool-" + name + "-schema", params)); // Build ID parsers if keys are provided common_peg_parser id_parser = eps(); @@ -879,7 +882,8 @@ common_peg_parser common_chat_peg_builder::standard_json_tools( const std::string & call_id_key, const std::string & gen_call_id_key, const std::vector & parameters_order, - bool accept_openai_wrapper) { + bool accept_openai_wrapper, + bool require_object_args) { if (!tools.is_array() || tools.empty()) { return eps(); } @@ -890,14 +894,14 @@ common_peg_parser common_chat_peg_builder::standard_json_tools( // Dispatch to the appropriate builder based on the JSON layout mode common_peg_parser tool_choices = eps(); if (function_is_key) { - tool_choices = build_json_tools_function_is_key(tools, args_key, effective_args_key, call_id_key, gen_call_id_key); + tool_choices = build_json_tools_function_is_key(tools, args_key, effective_args_key, call_id_key, gen_call_id_key, require_object_args); } else { auto name_spec = parse_key_spec(effective_name_key); auto args_spec = parse_key_spec(effective_args_key); if (!name_spec.first.empty() || !args_spec.first.empty()) { - tool_choices = build_json_tools_nested_keys(tools, effective_name_key, effective_args_key, call_id_key, gen_call_id_key); + tool_choices = build_json_tools_nested_keys(tools, effective_name_key, effective_args_key, call_id_key, gen_call_id_key, require_object_args); } else { - tool_choices = build_json_tools_flat_keys(tools, effective_name_key, effective_args_key, call_id_key, gen_call_id_key, parameters_order, accept_openai_wrapper); + tool_choices = build_json_tools_flat_keys(tools, effective_name_key, effective_args_key, call_id_key, gen_call_id_key, parameters_order, accept_openai_wrapper, require_object_args); } } diff --git a/common/chat-peg-parser.h b/common/chat-peg-parser.h index cd14f2c11750..2540155f2742 100644 --- a/common/chat-peg-parser.h +++ b/common/chat-peg-parser.h @@ -133,7 +133,8 @@ class common_chat_peg_builder : public common_peg_parser_builder { const std::string & call_id_key = "", const std::string & gen_call_id_key = "", const std::vector & parameters_order = {}, - bool accept_openai_wrapper = false); + bool accept_openai_wrapper = false, + bool require_object_args = false); // Legacy-compatible helper for building XML/tagged style tool calls // Used by tests and manual parsers @@ -157,13 +158,15 @@ class common_chat_peg_builder : public common_peg_parser_builder { const std::string & args_key, const std::string & effective_args_key, const std::string & call_id_key, - const std::string & gen_call_id_key); + const std::string & gen_call_id_key, + bool require_object_args); common_peg_parser build_json_tools_nested_keys(const nlohmann::ordered_json & tools, const std::string & effective_name_key, const std::string & effective_args_key, const std::string & call_id_key, - const std::string & gen_call_id_key); + const std::string & gen_call_id_key, + bool require_object_args); common_peg_parser build_json_tools_flat_keys(const nlohmann::ordered_json & tools, const std::string & effective_name_key, @@ -171,7 +174,8 @@ class common_chat_peg_builder : public common_peg_parser_builder { const std::string & call_id_key, const std::string & gen_call_id_key, const std::vector & parameters_order, - bool accept_openai_wrapper); + bool accept_openai_wrapper, + bool require_object_args); }; inline common_peg_arena build_chat_peg_parser( diff --git a/common/chat.cpp b/common/chat.cpp index 7740f35c0edc..50910242aedf 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -2530,6 +2530,125 @@ static common_chat_params common_chat_params_init_minimax_m3(const common_chat_t return data; } +// Inkling / TML typed-content-block parser: <|end_message|> separates blocks within a turn, +// <|content_model_end_sampling|> is the sole end-of-generation token (mirrors sglang TmlDetector). +static common_chat_params common_chat_params_init_inkling(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + const std::string MSG_MODEL = "<|message_model|>"; + const std::string MSG_USER = "<|message_user|>"; + const std::string MSG_SYSTEM = "<|message_system|>"; + const std::string MSG_TOOL = "<|message_tool|>"; + const std::string THINK = "<|content_thinking|>"; + const std::string TEXT = "<|content_text|>"; + const std::string END_MESSAGE = "<|end_message|>"; + const std::string END_SAMPLING = "<|content_model_end_sampling|>"; + const std::string INVOKE_TOOL = "<|content_invoke_tool_json|>"; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + data.thinking_start_tag = THINK; + data.thinking_end_tags = {END_MESSAGE}; + data.preserved_tokens = { + MSG_MODEL, MSG_USER, MSG_SYSTEM, MSG_TOOL, + THINK, TEXT, END_MESSAGE, END_SAMPLING, INVOKE_TOOL, + }; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + data.message_delimiters = { + { COMMON_CHAT_ROLE_ASSISTANT, MSG_MODEL }, + { COMMON_CHAT_ROLE_USER, MSG_USER }, + { COMMON_CHAT_ROLE_SYSTEM, MSG_SYSTEM }, + { COMMON_CHAT_ROLE_TOOL, MSG_TOOL }, + }; + + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + + if (inputs.has_continuation()) { + const auto & msg = inputs.continue_msg; + + data.generation_prompt = MSG_MODEL + THINK + msg.reasoning_content; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += END_MESSAGE + TEXT + msg.render_content(); + } + + data.prompt += data.generation_prompt; + } + + auto parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto generation_prompt = p.literal(MSG_MODEL); + auto end = p.end(); + + // thinking block; may also reappear mid-turn (after content), so it is both an optional + // prefix and a choice inside the block loops. With reasoning_format=NONE keep it + // (markers included) inline as content + common_peg_parser reasoning_block = p.eps(); + if (extract_reasoning) { + reasoning_block = p.literal(THINK) + + p.reasoning(p.until_one_of({ END_MESSAGE, TEXT, END_SAMPLING })) + + p.optional(p.literal(END_MESSAGE)); + } else { + reasoning_block = p.content(p.literal(THINK) + + p.until_one_of({ END_MESSAGE, TEXT, END_SAMPLING }) + + p.optional(p.literal(END_MESSAGE))); + } + auto reasoning = p.optional(reasoning_block); + + // TML re-emits <|message_model|> before each content block; a turn may contain several + // text blocks (one per content part), so the block repeats and bodies concatenate. + // THINK stops the content scan so a mid-turn thinking block is never leaked as text + auto text_block = p.optional(p.literal(MSG_MODEL)) + + p.optional(p.literal(TEXT)) + + p.content(p.until_one_of({ THINK, END_MESSAGE, END_SAMPLING })) + + p.optional(p.literal(END_MESSAGE)); + auto text_content = p.one_or_more(p.choice({ reasoning_block, text_block })); + + if (!has_tools || inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_NONE) { + return generation_prompt + reasoning + text_content + + p.optional(p.literal(END_SAMPLING)) + end; + } + + // each call is its own block (role opener + bare name echo + JSON section); + // force_tool_calls=true makes the JSON section required so a pure-text answer fails the + // block cleanly; parallel calls are separate blocks, hence repeat + parallel=false + auto tool_section = p.standard_json_tools( + INVOKE_TOOL, END_MESSAGE, inputs.tools, /* parallel_tool_calls = */ false, + /* force_tool_calls = */ true, + /* name_key = */ "name", + /* args_key = */ "args", + /* array_wrapped = */ false, + /* function_is_key = */ false, + /* call_id_key = */ "", + /* gen_call_id_key = */ "", + /* parameters_order = */ {}, + /* accept_openai_wrapper = */ false, + /* require_object_args = */ true); + // the name-echo scan must stop at any block marker: a greedy until(INVOKE_TOOL) returns + // NEED_MORE_INPUT mid-stream, which choice() treats as a match and shadows the text branch + auto tool_block = p.optional(p.literal(MSG_MODEL)) + + p.until_one_of({ INVOKE_TOOL, TEXT, THINK, END_MESSAGE, END_SAMPLING }) + + tool_section; + auto tool_calls = inputs.parallel_tool_calls ? p.one_or_more(tool_block) : tool_block; + // turns may interleave narration, thinking and calls; parse block-by-block (tool block + // first) since a whole-body choice would let the text branch swallow tool blocks into + // visible content + auto mixed_body = p.one_or_more(p.choice({ tool_block, reasoning_block, text_block })); + auto body = inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED + ? tool_calls + : mixed_body; + + return generation_prompt + reasoning + body + + p.optional(p.literal(END_SAMPLING)) + end; + }); + + data.parser = parser.save(); + + return data; +} + namespace workaround { static void map_developer_role_to_system(json & messages) { @@ -2947,6 +3066,14 @@ std::optional common_chat_try_specialized_template( return common_chat_params_init_cohere2moe(tmpl, params); } + // Inkling / TML: this marker combination is unique to the template + if (src.find("<|content_thinking|>") != std::string::npos && + src.find("<|content_text|>") != std::string::npos && + src.find("<|message_model|>") != std::string::npos) { + LOG_DBG("Using specialized template: Inkling\n"); + return common_chat_params_init_inkling(tmpl, params); + } + if (is_lfm2_template(src)) { LOG_DBG("Using specialized template: LFM2\n"); return common_chat_params_init_lfm2(tmpl, params, /* tool_list_tokens = */ true); diff --git a/conversion/__init__.py b/conversion/__init__.py index 1a47b851a0e6..abdf76cdba30 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -111,6 +111,7 @@ "HunYuanVLForConditionalGeneration": "hunyuan", "HYV3ForCausalLM": "hunyuan", "IQuestCoderForCausalLM": "llama", + "InklingForConditionalGeneration": "inkling", "InternLM2ForCausalLM": "internlm", "InternLM3ForCausalLM": "internlm", "JAISLMHeadModel": "jais", @@ -279,6 +280,7 @@ "GraniteSpeechPlusForConditionalGeneration": "granite", "HunYuanVLForConditionalGeneration": "hunyuan", "Idefics3ForConditionalGeneration": "smolvlm", + "InklingForConditionalGeneration": "inkling", "InternVisionModel": "internvl", "JanusForConditionalGeneration": "januspro", "KimiK25ForConditionalGeneration": "kimivl", diff --git a/conversion/base.py b/conversion/base.py index a7cd3fd904aa..720cfdca3a21 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -268,8 +268,17 @@ def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Call data_gen = lambda data=data_torch: LazyTorchTensor.from_eager(data) # noqa: E731 else: data_gen = lambda data=data_torch: data # noqa: E731 + # the index maps each tensor to one shard; a duplicate would silently overwrite it + if weight_map and name in weight_map and weight_map[name] != part_name: + raise ValueError( + f"tensor '{name}' found in '{part_name}' but the index assigns " + f"it to '{weight_map[name]}'; refusing to load a wrong-shard copy") if titem := self.filter_tensors((name, data_gen)): tname, tgen = titem + if tname in tensors: + raise ValueError( + f"duplicate tensor '{tname}' found in multiple model parts; " + f"refusing to silently overwrite") tensors[tname] = tgen # verify tensor name presence and identify potentially missing files @@ -537,6 +546,20 @@ def dequant_packed(w: Tensor, scale: Tensor, shape_tensor: Tensor, zero_point: T else: raise NotImplementedError(f"Quant format {quant_format!r} for method {quant_method!r} is not yet supported") elif quant_method == "modelopt": + # Stacked-expert NVFP4 checkpoints (e.g. Inkling-NVFP4) store experts as + # w13_weight/w2_weight with .scale/.scale2/.input_amax/.original_shape + # auxiliaries; the main tensors do not end in .weight, so the NVFP4 path + # below would silently skip them. Reject them with a clear error. + stacked_expert_aux = [ + n for n in self.model_tensors + if n.endswith((".scale2", ".input_amax", ".original_shape")) + ] + if stacked_expert_aux: + raise NotImplementedError( + "This checkpoint stores quantized experts in the stacked ModelOpt NVFP4 layout " + f"({len(stacked_expert_aux)} auxiliary tensors like {stacked_expert_aux[0]!r}), " + "which is not supported yet. Convert from the unquantized (BF16) checkpoint instead." + ) # Mixed-precision ModelOpt models: NVFP4 tensors are handled by # _generate_nvfp4_tensors; FP8 tensors have 1D weight_scale and # are dequantized here. k/v scale tensors are unused. @@ -1023,6 +1046,14 @@ def write_vocab(self): def write(self): self.prepare_tensors() + # zero tensors means the shards were never discovered, yet a metadata-only GGUF logs success + n_written = sum(len(shard) for shard in self.gguf_writer.tensors) + if n_written == 0: + raise ValueError( + "no tensors were written: the model shards could not be found. " + "Check that the safetensors filenames are discoverable (they must " + "start with 'model') or that model.safetensors exists." + ) self.prepare_metadata(vocab_only=False) self.gguf_writer.write_header_to_file(path=self.fname_out) self.gguf_writer.write_kv_data_to_file() diff --git a/conversion/inkling.py b/conversion/inkling.py new file mode 100644 index 000000000000..90ef65c3bdcc --- /dev/null +++ b/conversion/inkling.py @@ -0,0 +1,350 @@ +from __future__ import annotations + +from typing import Callable, Iterable, TYPE_CHECKING + +if TYPE_CHECKING: + from torch import Tensor + +from .base import MmprojModel, ModelBase, TextModel, gguf, logger + + +@ModelBase.register("InklingForConditionalGeneration") +class InklingModel(TextModel): + model_arch = gguf.MODEL_ARCH.INKLING + undo_permute = False + + _SKIP_PREFIXES = ("model.visual.", "model.audio.", "model.mtp.") + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # explicit raises (not assert, stripped by python -O) guard the single supported variant + hp = self.hparams + + # normalize keys renamed by HF-port re-saved configs back to checkpoint names + if "dense_intermediate_size" not in hp and "moe_intermediate_size" in hp: + hp["dense_intermediate_size"] = hp["intermediate_size"] + hp["intermediate_size"] = hp["moe_intermediate_size"] + if "sconv_kernel_size" not in hp and "conv_kernel_size" in hp: + hp["sconv_kernel_size"] = hp["conv_kernel_size"] + if "dense_mlp_idx" not in hp and hp.get("mlp_layer_types"): + types = hp["mlp_layer_types"] + hp["dense_mlp_idx"] = next((i for i, t in enumerate(types) if t != "dense"), len(types)) + + if hp.get("gate_activation", "sigmoid") != "sigmoid": + raise NotImplementedError( + f"unsupported gate_activation {hp.get('gate_activation')!r}; only 'sigmoid' is implemented" + ) + for flag, want in ( + ("norm_after_topk", True), + ("shared_expert_sink", True), + ("use_sconv", True), + ("use_embed_norm", True), + ("use_gate_bias", True), + ("use_global_scale", True), + ): + if hp.get(flag, want) is not want: + raise NotImplementedError(f"unsupported {flag}={hp.get(flag)!r}; only {want} is implemented") + if hp.get("q_bias", False) is not False or hp.get("o_bias", False) is not False: + raise NotImplementedError("attention q_bias / o_bias are not supported") + if hp.get("final_logit_softcapping") not in (None, 0, 0.0): + raise NotImplementedError( + f"final_logit_softcapping={hp.get('final_logit_softcapping')!r} is not supported" + ) + if hp["swa_head_dim"] != hp["head_dim"]: + raise ValueError(f"swa_head_dim {hp['swa_head_dim']} must equal head_dim {hp['head_dim']}") + if hp["swa_num_attention_heads"] != hp["num_attention_heads"]: + raise ValueError( + f"swa_num_attention_heads {hp['swa_num_attention_heads']} must equal " + f"num_attention_heads {hp['num_attention_heads']}" + ) + + # context length comes from model_max_length per the design contract + if (mml := hp.get("model_max_length")) is not None: + self.hparams["max_position_embeddings"] = mml + + # checked by the base find_hparam list before the MoE "intermediate_size" + self.hparams["prefix_dense_intermediate_size"] = hp["dense_intermediate_size"] + + self._local_layer_flags = self._get_local_layer_flags() + self.hparams["num_key_value_heads"] = [ + hp["swa_num_key_value_heads"] if is_local else hp["num_key_value_heads"] + for is_local in self._local_layer_flags + ] + + def _get_local_layer_flags(self) -> list[bool]: + # local_layer_ids is authoritative; a round-tripped layer_types may be stale and must not override it + n_layer = self.hparams["num_hidden_layers"] + local_ids = self.hparams.get("local_layer_ids") + if local_ids is None: + # default: global at id % 6 == 5; omitted/null must not collapse to all-global (explicit [] does) + local_ids = [i for i in range(n_layer) if i % 6 != 5] + local_ids = set(local_ids) + return [i in local_ids for i in range(n_layer)] + + def get_vocab_base(self) -> tuple[list[str], list[int], str]: + tokens, toktypes, tokpre = super().get_vocab_base() + # dedicated pre-type: o200k-family regex that keeps combining marks attached to base letters + tokpre = "inkling" + import gguf as _gguf + n_vocab = self.hparams["vocab_size"] + n_unpadded = self.hparams.get("unpadded_vocab_size") or n_vocab + if len(tokens) != n_vocab: + raise ValueError(f"Inkling tokenizer produced {len(tokens)} entries, expected {n_vocab}") + # force-CONTROL special ids from added_tokens_decoder, else the trailing-60 convention + try: + import json as _json + import pathlib as _pl + tc = _json.loads((_pl.Path(self.dir_model) / "tokenizer_config.json").read_text()) + special_ids = sorted(int(i) for i, d in tc.get("added_tokens_decoder", {}).items() if d.get("special")) + except Exception: + special_ids = list(range(n_unpadded - 60, n_unpadded)) + for tid in special_ids: + if 0 <= tid < n_vocab: + toktypes[tid] = _gguf.TokenType.CONTROL + if any(t != _gguf.TokenType.UNUSED for t in toktypes[n_unpadded:]): + raise ValueError("real tokens found at/above unpadded_vocab_size; padded-vocab mask would hide them") + return tokens, toktypes, tokpre + + def set_vocab(self): + self._set_vocab_gpt2() + eos_id = int(self.hparams.get("eos_token_id", 200006)) + if eos_id < 199998: + # HF-port configs re-save generic bos/eos defaults; the real EOS lives at 199998+ + eos_id = 200006 + # 200006 is the SOLE end-of-generation token; <|end_message|> (200010) is an + # intra-turn block separator and must NOT be registered eot/eog + self.gguf_writer.add_eos_token_id(eos_id) + # no BOS is ever prepended; pin bos to EOS so a stale base-tokenizer bos id never surfaces + self.gguf_writer.add_bos_token_id(eos_id) + self.gguf_writer.add_add_bos_token(False) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hp = self.hparams + + self.gguf_writer.add_vocab_size(hp["vocab_size"]) + self.gguf_writer.add_expert_feed_forward_length(hp["intermediate_size"]) + self.gguf_writer.add_expert_shared_count(hp["n_shared_experts"]) + self.gguf_writer.add_expert_weights_scale(hp["route_scale"]) + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) + + # sliding_window_size is canonical; explicit is-None fallback so a serialized 0 cannot bypass the mismatch check + canonical_window = hp["sliding_window_size"] + sliding_window = hp.get("sliding_window") + if sliding_window is None: + sliding_window = canonical_window + elif sliding_window != canonical_window: + raise ValueError( + f"sliding_window {sliding_window} disagrees with sliding_window_size " + f"{canonical_window!r}" + ) + if sliding_window <= 0: + raise ValueError(f"sliding_window must be positive, got {sliding_window}") + self.gguf_writer.add_sliding_window(sliding_window) + # true = local (swa) layer + self.gguf_writer.add_sliding_window_pattern(self._local_layer_flags) + + # no RoPE (arch-determined NONE); custom inkling.* keys per INKLING_DESIGN.md + arch = gguf.MODEL_ARCH_NAMES[self.model_arch] + self.gguf_writer.add_uint32(f"{arch}.d_rel", hp["d_rel"]) + self.gguf_writer.add_uint32(f"{arch}.rel_extent", hp["rel_extent"]) + self.gguf_writer.add_uint32(f"{arch}.rel_extent_swa", sliding_window) + self.gguf_writer.add_uint32(f"{arch}.shortconv_kernel", hp["sconv_kernel_size"]) + self.gguf_writer.add_uint32(f"{arch}.dense_block_count", hp["dense_mlp_idx"]) + self.gguf_writer.add_float32(f"{arch}.logit_scale_denom", hp["logits_mup_width_multiplier"]) + self.gguf_writer.add_uint32(f"{arch}.log_scaling_n_floor", int(hp.get("log_scaling_n_floor") or 0)) + self.gguf_writer.add_float32(f"{arch}.log_scaling_alpha", hp.get("log_scaling_alpha", 0.0)) + self.gguf_writer.add_uint32(f"{arch}.unpadded_vocab_size", hp["unpadded_vocab_size"]) + + logger.info(f"gguf: (inkling) swa pattern (true=local) = {self._local_layer_flags}") + logger.info(f"gguf: (inkling) unpadded_vocab_size = {hp['unpadded_vocab_size']}") + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + + if name.startswith(cls._SKIP_PREFIXES): + return None + + name = name.replace("model.llm.", "model.") + # parameter has no ".weight"-style suffix in the checkpoint + name = name.replace("rel_logits_proj.proj", "rel_logits_proj.weight") + + return super().filter_tensors((name, gen)) + + @staticmethod + def _deinterleave_w13(w13: Tensor) -> tuple[Tensor, Tensor]: + # interleaved SwiGLU along the output rows: silu(z[..., ::2]) * z[..., 1::2] + gate = w13[..., 0::2, :].contiguous() + up = w13[..., 1::2, :].contiguous() + return gate, up + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # short convs: [C, 1, K] -> [C, K] (same layout as LFM2 shortconv.conv) + if name.endswith("_sconv.weight"): + data_torch = data_torch.squeeze(1) + return [(self.map_tensor_name(name), data_torch)] + + if name.endswith(".mlp.w13_dn.weight"): + assert bid is not None + gate, up = self._deinterleave_w13(data_torch) + return [ + (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE, bid), gate), + (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP, bid), up), + ] + + if name.endswith(".mlp.global_scale") or name.endswith(".mlp.gate.global_scale"): + assert bid is not None + return [(self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GSCALE, bid), data_torch.float())] + + if name.endswith(".mlp.gate.bias"): + assert bid is not None + return [(self.format_tensor_name(gguf.MODEL_TENSOR.FFN_EXP_PROBS_B, bid, ".bias"), data_torch.float())] + + if name.endswith(".mlp.experts.w13_weight"): + assert bid is not None + gate, up = self._deinterleave_w13(data_torch) + return [ + (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid), gate), + (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_EXP, bid), up), + ] + if name.endswith(".mlp.experts.w2_weight"): + assert bid is not None + return [(self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid), data_torch)] + + # shared experts stored stacked for mul_mat_id + if name.endswith(".mlp.shared_experts.shared_w13_weight"): + assert bid is not None + gate, up = self._deinterleave_w13(data_torch) + return [ + (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_SHEXP, bid), gate), + (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_SHEXP, bid), up), + ] + if name.endswith(".mlp.shared_experts.shared_w2_weight"): + assert bid is not None + return [(self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, bid), data_torch)] + + return [(self.map_tensor_name(name), data_torch)] + + def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int): + # used in fp32 rel-bias math; keep full precision + if new_name.endswith("attn_rel_proj.weight"): + return gguf.GGMLQuantizationType.F32 + # ggml_ssm_conv kernels are F32-only + if ".shortconv_" in new_name: + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + +@ModelBase.register("InklingForConditionalGeneration") +class InklingMmprojModel(MmprojModel): + """Export Inkling's hMLP and dMel towers as one mtmd projector.""" + + has_vision_encoder = True + has_audio_encoder = True + + _IMAGE_MEAN = [0.48145466, 0.4578275, 0.40821073] + _IMAGE_STD = [0.26862954, 0.2613026, 0.2757771] + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + assert self.hparams_vision is not None + hp = self.hparams_vision + expected = { + "vision_encoder_type": "hmlp", + "patch_size": 40, + "temporal_patch_size": 2, + "n_channels": 3, + "n_layers": 4, + "decoder_dmodel": 6144, + "use_vision_norm": True, + } + for key, want in expected.items(): + got = hp.get(key, want) + if got != want: + raise NotImplementedError( + f"Inkling mmproj requires vision_config.{key}={want!r}, got {got!r}" + ) + + assert self.hparams_audio is not None + ahp = self.hparams_audio + audio_expected = { + "audio_mode": "dmel", + "decoder_dmodel": 6144, + "n_mel_bins": 80, + "mel_vocab_size": 16, + "use_audio_norm": True, + } + for key, want in audio_expected.items(): + got = ahp.get(key, want) + if got != want: + raise NotImplementedError( + f"Inkling mmproj requires audio_config.{key}={want!r}, got {got!r}" + ) + + def set_gguf_parameters(self): + hp = self.hparams_vision + assert hp is not None + + self.gguf_writer.add_file_type(self.ftype) + self.gguf_writer.add_clip_has_vision_encoder(True) + self.gguf_writer.add_clip_has_audio_encoder(True) + self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.INKLING) + self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.INKLING) + self.gguf_writer.add_vision_projection_dim(hp["decoder_dmodel"]) + + # clip.cpp requires these common fields even though hMLP is not a ViT. + self.gguf_writer.add_vision_image_size(hp["patch_size"]) + self.gguf_writer.add_vision_patch_size(hp["patch_size"]) + self.gguf_writer.add_vision_embedding_length(hp["n_channels"]) + self.gguf_writer.add_vision_feed_forward_length(0) + self.gguf_writer.add_vision_block_count(hp["n_layers"]) + self.gguf_writer.add_vision_head_count(1) + self.gguf_writer.add_vision_attention_layernorm_eps(1e-6) + self.gguf_writer.add_vision_image_mean(self._IMAGE_MEAN) + self.gguf_writer.add_vision_image_std(self._IMAGE_STD) + + ahp = self.hparams_audio + assert ahp is not None + self.gguf_writer.add_audio_projection_dim(ahp["decoder_dmodel"]) + self.gguf_writer.add_audio_embedding_length(ahp["decoder_dmodel"]) + self.gguf_writer.add_audio_feed_forward_length(0) + self.gguf_writer.add_audio_block_count(0) + self.gguf_writer.add_audio_head_count(1) + self.gguf_writer.add_audio_attention_layernorm_eps(1e-6) + self.gguf_writer.add_audio_num_mel_bins(ahp["n_mel_bins"]) + + @classmethod + def filter_tensors(cls, item): + name, gen = item + if not name.startswith(("model.visual.", "visual.", "model.audio.", "audio.")): + return None + return name, gen + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): + del bid + if name.startswith(("model.audio.", "audio.")): + prefix = "model.audio." if name.startswith("model.audio.") else "audio." + local = name.removeprefix(prefix) + if local == "encoder.weight": + yield "a.dmel.embedding.weight", data_torch + return + if local == "final_norm.weight": + yield "a.dmel.final_norm.weight", data_torch + return + raise ValueError(f"unexpected Inkling audio tensor {name!r}") + + prefix = "model.visual." if name.startswith("model.visual.") else "visual." + local = name.removeprefix(prefix) + if local == "final_norm.weight": + yield "v.hmlp.final_norm.weight", data_torch + return + + parts = local.split(".") + if len(parts) == 3 and parts[0] == "layers" and parts[2] == "weight": + kind, sep, layer_s = parts[1].partition("_") + if sep and kind in ("linear", "norm") and layer_s.isdigit(): + yield f"v.hmlp.{int(layer_s)}.{kind}.weight", data_torch + return + + raise ValueError(f"unexpected Inkling vision tensor {name!r}") diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 16ca33947a2e..7ca7f5fbf54e 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 4 #define RPC_PROTO_MINOR_VERSION 0 -#define RPC_PROTO_PATCH_VERSION 3 +#define RPC_PROTO_PATCH_VERSION 4 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 35f0c44ec421..6e51f5734ee5 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -435,8 +435,9 @@ extern "C" { // precision enum ggml_prec { - GGML_PREC_DEFAULT = 0, // stored as ggml_tensor.op_params, 0 by default - GGML_PREC_F32 = 10, + GGML_PREC_DEFAULT = 0, // stored as ggml_tensor.op_params, 0 by default + GGML_PREC_F32 = 10, + GGML_PREC_F32_PEDANTIC = 11, }; // op hint @@ -558,6 +559,7 @@ extern "C" { GGML_OP_FILL, GGML_OP_FLASH_ATTN_EXT, + GGML_OP_FLASH_ATTN_EXT_BANDED, GGML_OP_FLASH_ATTN_BACK, GGML_OP_SSM_CONV, GGML_OP_SSM_SCAN, @@ -1432,6 +1434,7 @@ extern "C" { // change the precision of a matrix multiplication // set to GGML_PREC_F32 for higher precision (useful for phi-2) + // or GGML_PREC_F32_PEDANTIC to require true F32 arithmetic GGML_API void ggml_mul_mat_set_prec( struct ggml_tensor * a, enum ggml_prec prec); @@ -2426,6 +2429,19 @@ extern "C" { float max_bias, float logit_softcap); + // flash attention with an additive banded relative-position bias, applied after scale, no dense bias tensor: + // rel_logits: [rel_extent, n_head, n_batch, ne3]; rel_dist = q_idx + (n_kv - n_batch) - kv_idx + // score += rel_logits[rel_dist, head, q_idx, batch] iff 0 <= rel_dist < rel_extent + GGML_API struct ggml_tensor * ggml_flash_attn_ext_banded( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * mask, + struct ggml_tensor * rel_logits, + float scale, + int64_t rel_extent); + GGML_API void ggml_flash_attn_ext_set_prec( struct ggml_tensor * a, enum ggml_prec prec); diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index a5a3a58ad054..b397fb6ba1b8 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -753,6 +753,16 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; }; + auto handle_flash_attn_ext_banded = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + GGML_ASSERT( src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_2); + GGML_ASSERT( src_ss[1].axis == GGML_BACKEND_SPLIT_AXIS_2); + GGML_ASSERT( src_ss[2].axis == GGML_BACKEND_SPLIT_AXIS_2); + GGML_ASSERT(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + // rel_logits is [E, H, Q, B], so its head shard is axis 1. + GGML_ASSERT( src_ss[5].axis == GGML_BACKEND_SPLIT_AXIS_1); + return {GGML_BACKEND_SPLIT_AXIS_1, {0}, {1}, 1}; + }; + auto handle_ssm_conv = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { if (src_ss[0].axis == src_ss[1].axis) { if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) { @@ -964,6 +974,9 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( case GGML_OP_FLASH_ATTN_EXT: { split_state = handle_flash_attn_ext(src_ss); } break; + case GGML_OP_FLASH_ATTN_EXT_BANDED: { + split_state = handle_flash_attn_ext_banded(src_ss); + } break; case GGML_OP_FLASH_ATTN_BACK: { split_state = handle_generic(src_ss, /*scalar_only =*/ true); } break; diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 491316f74912..1c335b07b89e 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -1998,6 +1998,7 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm ggml_compute_forward_fill(params, tensor); } break; case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_FLASH_ATTN_EXT_BANDED: { ggml_compute_forward_flash_attn_ext(params, tensor); } break; @@ -2396,6 +2397,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_ARGSORT: case GGML_OP_TOP_K: case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_FLASH_ATTN_EXT_BANDED: case GGML_OP_FLASH_ATTN_BACK: case GGML_OP_SSM_CONV: case GGML_OP_SSM_SCAN: @@ -2944,6 +2946,7 @@ struct ggml_cplan ggml_graph_plan( cur += sizeof(int32_t)*node->src[0]->ne[0]*n_tasks; } break; case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_FLASH_ATTN_EXT_BANDED: { const int64_t neq2 = node->src[0]->ne[2]; // number of query heads const int64_t DK = node->src[1]->ne[0]; diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 42ec809ce521..ba8e8c69bd8a 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -8332,10 +8332,11 @@ template struct cmp_argsort { const float * data; bool operator()(int32_t a, int32_t b) const { + // ties must resolve to the lower id (MoE routers); std::sort is unstable if constexpr (order == GGML_SORT_ORDER_ASC) { - return data[a] < data[b]; + return data[a] < data[b] || (data[a] == data[b] && a < b); } else { - return data[a] > data[b]; + return data[a] > data[b] || (data[a] == data[b] && a < b); } } }; @@ -8404,7 +8405,8 @@ void ggml_compute_forward_argsort( struct cmp_top_k { const float * data; bool operator()(int32_t a, int32_t b) const { - return data[a] > data[b]; + // ties must resolve to the lower id so the selected set matches the CUDA backend + return data[a] > data[b] || (data[a] == data[b] && a < b); } }; @@ -8465,6 +8467,30 @@ void ggml_compute_forward_top_k( } } +static inline float ggml_flash_attn_ext_banded_load( + const ggml_tensor * rel, + int64_t iq1, + int64_t iq2, + int64_t iq3, + int64_t rel_idx) { + const char * ptr = (const char *) rel->data + + (size_t) rel_idx * rel->nb[0] + + (size_t) iq2 * rel->nb[1] + + (size_t) iq1 * rel->nb[2] + + (size_t) (iq3 % rel->ne[3]) * rel->nb[3]; + + switch (rel->type) { + case GGML_TYPE_F32: + return *(const float *) ptr; + case GGML_TYPE_F16: + return GGML_CPU_FP16_TO_FP32(*(const ggml_fp16_t *) ptr); + case GGML_TYPE_BF16: + return GGML_BF16_TO_FP32(*(const ggml_bf16_t *) ptr); + default: + GGML_ABORT("banded flash attention: unsupported rel_logits type"); + } +} + static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( const ggml_compute_params * params, ggml_tensor * dst, @@ -8478,6 +8504,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( const ggml_tensor * v = dst->src[2]; const ggml_tensor * mask = dst->src[3]; const ggml_tensor * sinks = dst->src[4]; + const ggml_tensor * rel = dst->src[5]; GGML_TENSOR_LOCALS(int64_t, neq, q, ne) GGML_TENSOR_LOCALS(size_t, nbq, q, nb) @@ -8606,6 +8633,14 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( s = logit_softcap*tanhf(s); } + if (rel) { + // the offset aligns a short decode Q block to the tail of K (FA4 seqlen_k - seqlen_q convention) + const int64_t rel_dist = iq1 + (nek1 - neq1) - ic; + if (rel_dist >= 0 && rel_dist < rel->ne[0]) { + s += ggml_flash_attn_ext_banded_load(rel, iq1, iq2, iq3, rel_dist); + } + } + s += mv; // apply mask const float Mold = M; @@ -9070,6 +9105,7 @@ static void ggml_compute_forward_flash_attn_ext_f16( const ggml_tensor * q = dst->src[0]; const ggml_tensor * k = dst->src[1]; const ggml_tensor * v = dst->src[2]; + const ggml_tensor * rel = dst->src[5]; GGML_TENSOR_LOCALS(int64_t, neq, q, ne) GGML_TENSOR_LOCALS(size_t, nbq, q, nb) @@ -9169,7 +9205,7 @@ static void ggml_compute_forward_flash_attn_ext_f16( const int64_t dr = (nr + nchunk - 1) / nchunk; static constexpr int64_t Q_TILE_SZ = ggml_fa_tile_config::Q; - bool use_tiled = !use_ref && + bool use_tiled = !use_ref && rel == nullptr && (q->type == GGML_TYPE_F32 && kv_is_f32_or_f16 && k->type == v->type && diff --git a/ggml/src/ggml-cuda/argsort.cu b/ggml/src/ggml-cuda/argsort.cu index 26af90025972..0352e6baca06 100644 --- a/ggml/src/ggml-cuda/argsort.cu +++ b/ggml/src/ggml-cuda/argsort.cu @@ -163,6 +163,22 @@ static inline __device__ void ggml_cuda_swap(T & a, T & b) { b = tmp; } +// true if ia sorts after ib; padded indices sink to the end, ties break to the lower index (matches CPU cmp_argsort) +template +static inline __device__ bool argsort_ranks_after(const float * x_row, int ia, int ib, int ncols) { + const bool a_pad = ia >= ncols; + const bool b_pad = ib >= ncols; + if (a_pad || b_pad) { + return a_pad && (!b_pad || ia > ib); + } + const float xa = x_row[ia]; + const float xb = x_row[ib]; + if (xa != xb) { + return order == GGML_SORT_ORDER_ASC ? (xa > xb) : (xa < xb); + } + return ia > ib; +} + template static __global__ void k_argsort_f32_i32(const float * x, int * dst, const int ncols, int ncols_pad) { // bitonic sort @@ -186,19 +202,11 @@ static __global__ void k_argsort_f32_i32(const float * x, int * dst, const int n int ixj = col ^ j; if (ixj > col) { if ((col & k) == 0) { - if (dst_row[col] >= ncols || - (dst_row[ixj] < ncols && (order == GGML_SORT_ORDER_ASC ? - x_row[dst_row[col]] > x_row[dst_row[ixj]] : - x_row[dst_row[col]] < x_row[dst_row[ixj]])) - ) { + if (argsort_ranks_after(x_row, dst_row[col], dst_row[ixj], ncols)) { ggml_cuda_swap(dst_row[col], dst_row[ixj]); } } else { - if (dst_row[ixj] >= ncols || - (dst_row[col] < ncols && (order == GGML_SORT_ORDER_ASC ? - x_row[dst_row[col]] < x_row[dst_row[ixj]] : - x_row[dst_row[col]] > x_row[dst_row[ixj]])) - ) { + if (argsort_ranks_after(x_row, dst_row[ixj], dst_row[col], ncols)) { ggml_cuda_swap(dst_row[col], dst_row[ixj]); } } diff --git a/ggml/src/ggml-cuda/fattn-banded.cu b/ggml/src/ggml-cuda/fattn-banded.cu new file mode 100644 index 000000000000..98458ae86860 --- /dev/null +++ b/ggml/src/ggml-cuda/fattn-banded.cu @@ -0,0 +1,247 @@ +#include "common.cuh" +#include "fattn-banded.cuh" +#include "fattn.cuh" + +#include + +static __device__ __forceinline__ float fattn_banded_load( + const char * ptr, const int type) { + switch (type) { + case GGML_TYPE_F32: + return *(const float *) ptr; + case GGML_TYPE_F16: + return __half2float(*(const half *) ptr); + case GGML_TYPE_BF16: { + // Read BF16 as raw bits so this kernel needs no native BF16 support; all math stays FP32. + const uint32_t bits = uint32_t(*(const uint16_t *) ptr) << 16; + return __uint_as_float(bits); + } + default: + return 0.0f; + } +} + +template +static __global__ void flash_attn_ext_banded_f32( + const char * __restrict__ q, + const char * __restrict__ k, + const char * __restrict__ v, + const char * __restrict__ mask, + const char * __restrict__ rel, + float * __restrict__ dst, + float scale, + int type_k, + int type_v, + int type_rel, + int64_t n_q, + int64_t n_kv, + int64_t n_head_q, + int64_t n_head_kv, + int64_t n_batch, + int64_t rel_extent, + int64_t mask_ne2, + int64_t mask_ne3, + uint64_t q_nb1, + uint64_t q_nb2, + uint64_t q_nb3, + uint64_t k_nb0, + uint64_t k_nb1, + uint64_t k_nb2, + uint64_t k_nb3, + uint64_t v_nb0, + uint64_t v_nb1, + uint64_t v_nb2, + uint64_t v_nb3, + uint64_t m_nb1, + uint64_t m_nb2, + uint64_t m_nb3, + uint64_t r_nb0, + uint64_t r_nb1, + uint64_t r_nb2, + uint64_t r_nb3, + int64_t rel_ne3) { + constexpr int values_per_lane = D / WARP_SIZE; + static_assert(D == 64 || D == 128, "banded FA supports head dimensions 64 and 128"); + static_assert(D % WARP_SIZE == 0, "head dimension must be divisible by warp size"); + + const int lane = threadIdx.x % WARP_SIZE; + const int warp = threadIdx.x / WARP_SIZE; + const int64_t iq = int64_t(blockIdx.x) * WARPS_PER_BLOCK + warp; + const int64_t ih = blockIdx.y; + const int64_t ib = blockIdx.z; + + if (iq >= n_q || ih >= n_head_q || ib >= n_batch) { + return; + } + + const int64_t ih_kv = ih / (n_head_q / n_head_kv); + const char * q_row = q + uint64_t(iq)*q_nb1 + uint64_t(ih)*q_nb2 + uint64_t(ib)*q_nb3; + + float q_reg[values_per_lane]; + float out[values_per_lane]; +#pragma unroll + for (int j = 0; j < values_per_lane; ++j) { + const int d = lane + j*WARP_SIZE; + q_reg[j] = *(const float *)(q_row + uint64_t(d)*sizeof(float)); + out[j] = 0.0f; + } + + float row_max = -INFINITY; + float row_sum = 0.0f; + const int64_t q_offset = n_kv - n_q; + + for (int64_t ik = 0; ik < n_kv; ++ik) { + const char * k_row = k + uint64_t(ik)*k_nb1 + uint64_t(ih_kv)*k_nb2 + uint64_t(ib)*k_nb3; + float dot = 0.0f; +#pragma unroll + for (int j = 0; j < values_per_lane; ++j) { + const int d = lane + j*WARP_SIZE; + dot += q_reg[j] * fattn_banded_load(k_row + uint64_t(d)*k_nb0, type_k); + } + dot = warp_reduce_sum(dot); + + float score = dot * scale; + if (lane == 0) { + const int64_t rel_dist = iq + q_offset - ik; + if (rel_dist >= 0 && rel_dist < rel_extent) { + const char * rel_value = rel + + uint64_t(rel_dist)*r_nb0 + uint64_t(ih)*r_nb1 + + uint64_t(iq)*r_nb2 + uint64_t(ib % rel_ne3)*r_nb3; + score += fattn_banded_load(rel_value, type_rel); + } + if (mask) { + const char * mask_value = mask + uint64_t(ik)*sizeof(half) + + uint64_t(iq)*m_nb1 + uint64_t(ih % mask_ne2)*m_nb2 + + uint64_t(ib % mask_ne3)*m_nb3; + score += __half2float(*(const half *) mask_value); + } + } + score = __shfl_sync(0xffffffff, score, 0, WARP_SIZE); + + if (score == -INFINITY) { + continue; + } + + float old_scale = 1.0f; + float value_scale = 1.0f; + if (score > row_max) { + old_scale = expf(row_max - score); + row_max = score; + } else { + value_scale = expf(score - row_max); + } + + const char * v_row = v + uint64_t(ik)*v_nb1 + uint64_t(ih_kv)*v_nb2 + uint64_t(ib)*v_nb3; +#pragma unroll + for (int j = 0; j < values_per_lane; ++j) { + const int d = lane + j*WARP_SIZE; + const float vv = fattn_banded_load(v_row + uint64_t(d)*v_nb0, type_v); + out[j] = out[j]*old_scale + vv*value_scale; + } + row_sum = row_sum*old_scale + value_scale; + } + + const float inv_sum = row_sum == 0.0f ? 0.0f : 1.0f/row_sum; + float * dst_row = dst + ((ib*n_q + iq)*n_head_q + ih)*D; +#pragma unroll + for (int j = 0; j < values_per_lane; ++j) { + const int d = lane + j*WARP_SIZE; + dst_row[d] = out[j]*inv_sum; + } +} + +static bool fattn_banded_type_supported(ggml_type type) { + return type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || type == GGML_TYPE_BF16; +} + +bool ggml_cuda_flash_attn_ext_banded_supported(int device, const ggml_tensor * dst) { + GGML_UNUSED(device); +#if defined(GGML_USE_MUSA) + GGML_UNUSED(dst); + return false; +#else + if (dst->op != GGML_OP_FLASH_ATTN_EXT_BANDED) { + return false; + } + + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * m = dst->src[3]; + const ggml_tensor * rel = dst->src[5]; + if (!q || !k || !v || !rel || q->type != GGML_TYPE_F32) { + return false; + } + if (!fattn_banded_type_supported(k->type) || + !fattn_banded_type_supported(v->type) || + !fattn_banded_type_supported(rel->type)) { + return false; + } + if ((q->ne[0] != 64 && q->ne[0] != 128) || v->ne[0] != q->ne[0] || k->ne[0] != q->ne[0]) { + return false; + } + if (q->ne[2] % k->ne[2] != 0 || q->ne[2] % v->ne[2] != 0 || k->ne[2] != v->ne[2]) { + return false; + } + if (q->ne[3] != k->ne[3] || q->ne[3] != v->ne[3]) { + return false; + } + if (q->nb[0] != sizeof(float) || k->nb[0] != ggml_type_size(k->type) || + v->nb[0] != ggml_type_size(v->type) || rel->nb[0] != ggml_type_size(rel->type)) { + return false; + } + if (rel->ne[0] <= 0 || rel->ne[1] != q->ne[2] || rel->ne[2] != q->ne[1] || + (rel->ne[3] != 1 && rel->ne[3] != q->ne[3])) { + return false; + } + return !m || (m->type == GGML_TYPE_F16 && ggml_is_contiguous(m) && + q->ne[2] % m->ne[2] == 0 && q->ne[3] % m->ne[3] == 0); +#endif +} + +void ggml_cuda_flash_attn_ext_banded(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + GGML_ASSERT(ggml_cuda_flash_attn_ext_banded_supported(ctx.device, dst)); + + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * m = dst->src[3]; + const ggml_tensor * rel = dst->src[5]; + + // route F16/BF16 K/V to the MMA kernel; keep this FP32 kernel for mixed types and strided rel + if (k->type != GGML_TYPE_F32 && v->type != GGML_TYPE_F32 && + rel->type == GGML_TYPE_F32 && ggml_is_contiguous(rel) && + // MMA ABI indexes rel by Q's batch: a singleton rel batch must take the stride-aware fallback + rel->ne[3] == q->ne[3] && rel->ne[0] <= (1 << 20)) { + ggml_cuda_flash_attn_ext(ctx, dst); + return; + } + + float scale; + memcpy(&scale, dst->op_params, sizeof(scale)); + // the tensor extent (not op_params) is authoritative after graph cloning + const int64_t rel_extent = rel->ne[0]; + + constexpr int warps_per_block = 4; + const dim3 blocks((q->ne[1] + warps_per_block - 1) / warps_per_block, q->ne[2], q->ne[3]); + const dim3 threads(warps_per_block * WARP_SIZE, 1, 1); + cudaStream_t stream = ctx.stream(); + +#define LAUNCH_BANDED(D) \ + flash_attn_ext_banded_f32<<>>( \ + (const char *) q->data, (const char *) k->data, (const char *) v->data, \ + m ? (const char *) m->data : nullptr, (const char *) rel->data, (float *) dst->data, \ + scale, k->type, v->type, rel->type, q->ne[1], k->ne[1], q->ne[2], k->ne[2], q->ne[3], \ + rel_extent, m ? m->ne[2] : 1, m ? m->ne[3] : 1, \ + q->nb[1], q->nb[2], q->nb[3], k->nb[0], k->nb[1], k->nb[2], k->nb[3], \ + v->nb[0], v->nb[1], v->nb[2], v->nb[3], \ + m ? m->nb[1] : 0, m ? m->nb[2] : 0, m ? m->nb[3] : 0, \ + rel->nb[0], rel->nb[1], rel->nb[2], rel->nb[3], rel->ne[3]) + + if (q->ne[0] == 64) { + LAUNCH_BANDED(64); + } else { + LAUNCH_BANDED(128); + } +#undef LAUNCH_BANDED +} diff --git a/ggml/src/ggml-cuda/fattn-banded.cuh b/ggml/src/ggml-cuda/fattn-banded.cuh new file mode 100644 index 000000000000..bb56b8986720 --- /dev/null +++ b/ggml/src/ggml-cuda/fattn-banded.cuh @@ -0,0 +1,7 @@ +#pragma once + +#include "common.cuh" + +void ggml_cuda_flash_attn_ext_banded(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +bool ggml_cuda_flash_attn_ext_banded_supported(int device, const ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index e67cc7fdf784..b07a0122db09 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -52,7 +52,7 @@ struct ggml_cuda_flash_attn_ext_f16_extra_data { static inline ggml_cuda_flash_attn_ext_f16_extra_data ggml_cuda_flash_attn_ext_get_f16_extra_data( const ggml_tensor * dst, const bool need_f16_K, const bool need_f16_V) { - GGML_ASSERT(dst->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(dst->op == GGML_OP_FLASH_ATTN_EXT || dst->op == GGML_OP_FLASH_ATTN_EXT_BANDED); const ggml_tensor * K = dst->src[1]; const ggml_tensor * V = dst->src[2]; @@ -983,7 +983,8 @@ void launch_fattn( const bool V_is_K_view = V->view_src && (V->view_src == K || (V->view_src == K->view_src && V->view_offs == K->view_offs)); const ggml_tensor * mask = dst->src[3]; - const ggml_tensor * sinks = dst->src[4]; + const ggml_tensor * rel = dst->op == GGML_OP_FLASH_ATTN_EXT_BANDED ? dst->src[5] : nullptr; + const ggml_tensor * sinks = rel ? rel : dst->src[4]; ggml_tensor * KQV = dst; @@ -1192,6 +1193,13 @@ void launch_fattn( memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float)); memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float)); + // banded op reuses the MMA ABI: negative max_bias tags the branch, sinks_ptr carries rel_logits, -max_bias is E + if (rel) { + GGML_ASSERT(rel->type == GGML_TYPE_F32 && ggml_is_contiguous(rel)); + GGML_ASSERT(rel->ne[0] <= (1 << 20)); // exactly representable in float + max_bias = -float(rel->ne[0]); + } + if (logit_softcap != 0.0f) { scale /= logit_softcap; } @@ -1199,8 +1207,9 @@ void launch_fattn( const uint32_t n_head = Q->ne[2]; const uint32_t n_head_log2 = 1u << uint32_t(floorf(log2f(float(n_head)))); - const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); - const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + // m0/m1 are unused for banded bias; the negative tag would otherwise blow up the exponent + const float m0 = rel ? 1.0f : powf(2.0f, -(max_bias ) / n_head_log2); + const float m1 = rel ? 1.0f : powf(2.0f, -(max_bias / 2.0f) / n_head_log2); // TODO other tensor dimensions after removal of WMMA kernel: const uint3 ne01 = init_fastdiv_values(Q->ne[1]); diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 7f4cfd5511ff..f6db6a79b24a 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -535,13 +535,17 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( const half2 * const __restrict__ K_h2, const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, + const float * const __restrict__ rel_f, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, const float scale, const float slope, const float logit_softcap, + const int rel_extent, + const int head_q0, const uint3 ne01, const int ne02, + const int ne11, const int stride_K, const int stride_V, const int stride_mask, @@ -688,10 +692,42 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( #pragma unroll for (int col = 0; col < cols_per_thread; ++col) { KQ_max_new[col] = KQ_max[col]; + // The fp16 VKQ accumulator holds the still-unnormalized sum of softmax weights times V. + // With large-magnitude V (~1e3) it overflows after a few thousand KV positions even though + // all inputs are finite. Gated on rel_f so only the banded op (whose V is that large) pays + // for it. When the per-thread partial row sum of weights gets large, force + // the running maximum up: the regular rescale below then shrinks the accumulator, row sum, + // and all further weights by 2^-8, which the final normalization cancels exactly. The bump + // reaches all threads sharing the column via the KQ_max_new warp reduction, and re-triggers + // only after the row sum regrows 256x, so the number of bumps is logarithmic in n_kv. + if (rel_f && KQ_rowsum[col] > 4.0f) { + KQ_max_new[col] += 8.0f*0.6931f; + } } float KQ_rowsum_add[cols_per_thread] = {0.0f}; if constexpr (cols_per_warp == 8) { + if (rel_f) { +#pragma unroll + for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::I) { + const int i0 = i00 + (threadIdx.y % np)*T_C_KQ::I; +#pragma unroll + for (int l = 0; l < T_C_KQ::ne; ++l) { + const int i = i0 + T_C_KQ::get_i(l); + const int jc = (threadIdx.y / np)*T_C_KQ::J + T_C_KQ::get_j(l); + const int j = jc / ncols2; + const int c = jc % ncols2; + const int64_t q_idx = int64_t(jt)*ncols1 + j; + const int64_t kv_idx = int64_t(k_VKQ_0) + i; + const int64_t dist = q_idx + (int64_t(ne11) - ne01.z) - kv_idx; + if (q_idx < ne01.z && head_q0 + c < ne02 && kv_idx < ne11 && + dist >= 0 && dist < rel_extent) { + KQ_C[i00/(np*T_C_KQ::I)].x[l] += + rel_f[(q_idx*ne02 + head_q0 + c)*rel_extent + dist]; + } + } + } + } if (ncols2 > 1 || mask_h) { #pragma unroll for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::I) { @@ -754,6 +790,27 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } } } else { // not Turing mma or T_B_KQ::I > 8 + if (rel_f) { +#pragma unroll + for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::J) { + const int i0 = i00 + (threadIdx.y % np)*T_C_KQ::J; +#pragma unroll + for (int l = 0; l < T_C_KQ::ne; ++l) { + const int i = i0 + T_C_KQ::get_j(l); + const int jc = (threadIdx.y / np)*cols_per_warp + T_C_KQ::get_i(l); + const int j = jc / ncols2; + const int c = jc % ncols2; + const int64_t q_idx = int64_t(jt)*ncols1 + j; + const int64_t kv_idx = int64_t(k_VKQ_0) + i; + const int64_t dist = q_idx + (int64_t(ne11) - ne01.z) - kv_idx; + if (q_idx < ne01.z && head_q0 + c < ne02 && kv_idx < ne11 && + dist >= 0 && dist < rel_extent) { + KQ_C[i00/(np*T_C_KQ::J)].x[l] += + rel_f[(q_idx*ne02 + head_q0 + c)*rel_extent + dist]; + } + } + } + } if (ncols2 > 1 || mask_h) { #pragma unroll for (int i00 = 0; i00 < nbatch_fa; i00 += np*T_C_KQ::J) { @@ -1015,8 +1072,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, - scale, slope, logit_softcap, ne01, ne02, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, rel_f, dstk, dstk_fixup, + scale, slope, logit_softcap, rel_extent, head_q0, ne01, ne02, ne11, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, kb0); @@ -1120,11 +1177,14 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const half2 * const __restrict__ V_h2, const half * const __restrict__ mask_h, const float * const __restrict__ sinks_f, + const float * const __restrict__ rel_f, float2 * const __restrict__ dstk, float2 * const __restrict__ dstk_fixup, const float scale, const float slope, const float logit_softcap, + const int rel_extent, + const int head_q0, const uint3 ne01, const int ne02, const int gqa_ratio, @@ -1278,8 +1338,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( flash_attn_ext_f16_iter - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, - ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + (Q_f2, K_h2, V_h2, mask_h, rel_f, dstk, dstk_fixup, scale, slope, logit_softcap, + rel_extent, head_q0, ne01, ne02, ne11, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } constexpr bool last_iter = true; @@ -1287,8 +1347,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( flash_attn_ext_f16_iter - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, - ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + (Q_f2, K_h2, V_h2, mask_h, rel_f, dstk, dstk_fixup, scale, slope, logit_softcap, + rel_extent, head_q0, ne01, ne02, ne11, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } else { constexpr bool oob_check = false; @@ -1298,8 +1358,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( flash_attn_ext_f16_iter - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, - ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + (Q_f2, K_h2, V_h2, mask_h, rel_f, dstk, dstk_fixup, scale, slope, logit_softcap, + rel_extent, head_q0, ne01, ne02, ne11, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } constexpr bool last_iter = true; @@ -1307,8 +1367,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( flash_attn_ext_f16_iter - (Q_f2, K_h2, V_h2, mask_h, dstk, dstk_fixup, scale, slope, logit_softcap, - ne01, ne02, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, + (Q_f2, K_h2, V_h2, mask_h, rel_f, dstk, dstk_fixup, scale, slope, logit_softcap, + rel_extent, head_q0, ne01, ne02, ne11, stride_K, stride_V, stride_mask, tile_Q, tile_K, tile_V, tile_mask, Q_B, VKQ_C, KQ_max, KQ_rowsum, jt, kb0, k_VKQ_sup); } @@ -1692,8 +1752,8 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( } } #else - GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dstk_fixup, - scale, slope, logit_softcap, ne01, ne02, gqa_ratio, + GGML_UNUSED_VARS(Q_f2, K_h2, V_h2, mask_h, sinks_f, rel_f, dstk, dstk_fixup, + scale, slope, logit_softcap, rel_extent, head_q0, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, kb0_start, kb0_stop); NO_DEVICE_CODE; @@ -1734,6 +1794,8 @@ static __global__ void flash_attn_ext_f16( const int * GGML_CUDA_RESTRICT KV_max = KV_max_ptr; float * GGML_CUDA_RESTRICT dst = dst_ptr; float2 * GGML_CUDA_RESTRICT dst_meta = dst_meta_ptr; + const bool banded_bias = max_bias < 0.0f; + const int rel_extent = banded_bias ? int(-max_bias) : 0; // Skip unused kernel variants for faster compilation: if (use_logit_softcap && !(DKQ == 128 || DKQ == 256 || DKQ == 512)) { @@ -1819,9 +1881,12 @@ static __global__ void flash_attn_ext_f16( float2 * dstk = ((float2 *) dst) + (sequence*ne01.z*ne02 + zt_Q) * (DV/2); const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); - const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; + const float * sinks_f = sinks && !banded_bias ? (const float *) sinks + zt_Q : nullptr; + const float * rel_f = banded_bias ? + (const float *) sinks + int64_t(sequence)*ne01.z*ne02*rel_extent : nullptr; - const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; + const float slope = banded_bias ? 1.0f : + (ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f); if (KV_max) { kb0_stop = min(kb0_stop, KV_max[sequence*iter_j + jt] / nbatch_fa); @@ -1830,13 +1895,13 @@ static __global__ void flash_attn_ext_f16( if (kb0_start == 0) { constexpr bool needs_fixup = false; // CUDA block is working on an entire tile. flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, - ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); + (Q_f2, K_h2, V_h2, mask_h, sinks_f, rel_f, dstk, dst_meta, scale, slope, logit_softcap, + rel_extent, zt_Q, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } else { constexpr bool needs_fixup = true; // CUDA block is missing the beginning of a tile. flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, - ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); + (Q_f2, K_h2, V_h2, mask_h, sinks_f, rel_f, dstk, dst_meta, scale, slope, logit_softcap, + rel_extent, zt_Q, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } kbc += iter_k; @@ -1865,9 +1930,12 @@ static __global__ void flash_attn_ext_f16( float2 * dstk = ((float2 *) dst) + (sequence*ne01.z*ne02 + zt_Q) * (DV/2); const half2 * V_h2 = V_is_K_view ? K_h2 : (const half2 *) (V + nb23*sequence + nb22*z_KV); - const float * sinks_f = sinks ? (const float *) sinks + zt_Q : nullptr; + const float * sinks_f = sinks && !banded_bias ? (const float *) sinks + zt_Q : nullptr; + const float * rel_f = banded_bias ? + (const float *) sinks + int64_t(sequence)*ne01.z*ne02*rel_extent : nullptr; - const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f; + const float slope = banded_bias ? 1.0f : + (ncols2 == 1 ? get_alibi_slope(max_bias, zt_Q, n_head_log2, m0, m1) : 1.0f); if (KV_max) { kb0_stop = min(kb0_stop, KV_max[sequence*iter_j + jt] / nbatch_fa); @@ -1876,8 +1944,8 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. constexpr bool needs_fixup = false; flash_attn_ext_f16_process_tile - (Q_f2, K_h2, V_h2, mask_h, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, - ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); + (Q_f2, K_h2, V_h2, mask_h, sinks_f, rel_f, dstk, dst_meta, scale, slope, logit_softcap, + rel_extent, zt_Q, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); #else GGML_UNUSED_VARS(Q_ptr, K_ptr, V_ptr, mask_ptr, sinks_ptr, KV_max_ptr, dst_ptr, dst_meta_ptr, scale, max_bias, m0, m1, n_head_log2, logit_softcap, diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ab7a3b297c07..43849b9cec22 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -390,6 +390,14 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const const int cc = ggml_cuda_info().devices[device].cc; + // banded bias lives in the F16 MMA loop; F32 K/V keeps the dedicated FP32 warp kernel + if (dst->op == GGML_OP_FLASH_ATTN_EXT_BANDED && + dst->src[5]->type == GGML_TYPE_F32 && K->type != GGML_TYPE_F32 && V->type != GGML_TYPE_F32 && + dst->src[5]->ne[3] == Q->ne[3] && + turing_mma_available(cc) && (Q->ne[0] == 64 || Q->ne[0] == 128) && V->ne[0] == Q->ne[0]) { + return BEST_FATTN_KERNEL_MMA_F16; + } + switch (K->ne[0]) { case 40: case 64: @@ -534,7 +542,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const } size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * dst) { - GGML_ASSERT(dst->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(dst->op == GGML_OP_FLASH_ATTN_EXT || dst->op == GGML_OP_FLASH_ATTN_EXT_BANDED); const ggml_tensor * K = dst->src[1]; const ggml_tensor * V = dst->src[2]; diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index e73a7b8906ce..5373b64d9332 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -25,6 +25,7 @@ #include "ggml-cuda/diagmask.cuh" #include "ggml-cuda/diag.cuh" #include "ggml-cuda/fattn.cuh" +#include "ggml-cuda/fattn-banded.cuh" #include "ggml-cuda/fwht.cuh" #include "ggml-cuda/getrows.cuh" #include "ggml-cuda/im2col.cuh" @@ -906,7 +907,7 @@ static size_t ggml_backend_cuda_buffer_type_get_alignment(ggml_backend_buffer_ty static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *) buft->context; - size_t size = tensor->op == GGML_OP_FLASH_ATTN_EXT + size_t size = (tensor->op == GGML_OP_FLASH_ATTN_EXT || tensor->op == GGML_OP_FLASH_ATTN_EXT_BANDED) ? ggml_cuda_flash_attn_ext_get_alloc_size(buft_ctx->device, tensor) : ggml_nbytes(tensor); int64_t ne0 = tensor->ne[0]; @@ -1496,6 +1497,9 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const size_t nbd2 = dst->nb[2]; size_t nbd3 = dst->nb[3]; + const bool f32_pedantic = compute_type == GGML_TYPE_F32 && + src0->type == GGML_TYPE_F32 && ggml_prec(dst->op_params[0]) == GGML_PREC_F32_PEDANTIC; + cublasComputeType_t cu_compute_type = traits::compute_type; cudaDataType_t cu_data_type = traits::data_type; cudaDataType_t cu_data_type_a = traits::data_type; @@ -1527,6 +1531,18 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const } } + if (f32_pedantic) { +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11020 + cu_compute_type = CUBLAS_COMPUTE_32F_PEDANTIC; +#else + // no pedantic compute enum here; ordinary F32 is the strongest available contract + cu_compute_type = CUBLAS_COMPUTE_32F; +#endif + } + + const auto cu_gemm_algo = f32_pedantic ? + CUBLAS_GEMM_DEFAULT : CUBLAS_GEMM_DEFAULT_TENSOR_OP; + GGML_ASSERT(ne12 % ne02 == 0); GGML_ASSERT(ne13 % ne03 == 0); @@ -1538,12 +1554,23 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const // However, for some old NVIDIA and AMD GPUs the strided/Ex GEMM is much slower, // probably because the internal kernel selection logic is suboptimal. if (compute_type == GGML_TYPE_F32 && ne12 == 1 && ne13 == 1) { - CUBLAS_CHECK( - cublasSgemm(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, - ne01, ne11, ne10, - (const float *) alpha, (const float *) src0_ptr, s01, - (const float *) src1_ptr, s11, - (const float *) beta, (float *) dst_ptr, ne0)); + if (f32_pedantic) { + CUBLAS_CHECK( + cublasGemmEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, + ne01, ne11, ne10, + alpha, src0_ptr, CUDA_R_32F, s01, + src1_ptr, CUDA_R_32F, s11, + beta, dst_ptr, CUDA_R_32F, ne0, + cu_compute_type, + cu_gemm_algo)); + } else { + CUBLAS_CHECK( + cublasSgemm(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, + ne01, ne11, ne10, + (const float *) alpha, (const float *) src0_ptr, s01, + (const float *) src1_ptr, s11, + (const float *) beta, (float *) dst_ptr, ne0)); + } } else if (ne12 == 1 && ne13 == 1) { CUBLAS_CHECK( cublasGemmEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, @@ -1552,7 +1579,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const src1_ptr, cu_data_type_b, s11, beta, dst_ptr, cu_data_type, ne0, cu_compute_type, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + cu_gemm_algo)); } else if (r2 == 1 && r3 == 1 && is_src0_cont_2 && is_src1_cont_2) { // with a [0, 2, 1, 3] perm. and ne02==1 the matrix strides need to be determined from dim 3: const int64_t sma = ne02 == 1 ? s03 : s02; @@ -1568,7 +1595,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const beta, dst_ptr, cu_data_type, ne0, ne1*ne0, // strideC ne12*ne13, cu_compute_type, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + cu_gemm_algo)); } else { // use cublasGemmBatchedEx const int64_t ne23 = ne12*ne13; @@ -1606,7 +1633,7 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const beta, ( void **) (ptrs_dst.get() + 0*ne23), cu_data_type, ne0, ne23, cu_compute_type, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + cu_gemm_algo)); } // Convert output back to F32 if needed @@ -1644,6 +1671,11 @@ static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml } } + // a scoped pedantic request overrides the process-wide compute type, for F32 weights only + if (src0->type == GGML_TYPE_F32 && ggml_prec(dst->op_params[0]) == GGML_PREC_F32_PEDANTIC) { + compute_type = GGML_TYPE_F32; + } + switch (compute_type) { case GGML_TYPE_F32: ggml_cuda_mul_mat_cublas_impl(ctx, src0, src1, dst); @@ -1829,6 +1861,8 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor const int cc = ggml_cuda_info().devices[ctx.device].cc; const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; + const bool f32_pedantic = src0->type == GGML_TYPE_F32 && + ggml_prec(dst->op_params[0]) == GGML_PREC_F32_PEDANTIC; if (ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11)) { // The custom F16 vector kernel can be used over batched cuBLAS GEMM. @@ -1836,7 +1870,8 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); return; } - if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) { + if (!f32_pedantic && ggml_cuda_should_use_mmf( + src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) { ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); return; } @@ -1862,6 +1897,8 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * GGML_TENSOR_BINARY_OP_LOCALS const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const bool f32_pedantic = src0->type == GGML_TYPE_F32 && + ggml_prec(dst->op_params[0]) == GGML_PREC_F32_PEDANTIC; // [TAG_MUL_MAT_ID_CUDA_GRAPHS] if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { @@ -1886,7 +1923,8 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * return; } - if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { + if (!f32_pedantic && ggml_cuda_should_use_mmf( + src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst); return; } @@ -1994,6 +2032,7 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst_slice.nb[2] = dst_slice.ne[1] * dst_slice.nb[1]; dst_slice.nb[3] = dst_slice.ne[2] * dst_slice.nb[2]; dst_slice.data = dst_data_cur; + dst_slice.op_params[0] = dst->op_params[0]; ggml_cuda_mul_mat(ctx, &src0_slice, &src1_slice, &dst_slice); CUDA_CHECK(cudaGetLastError()); @@ -2305,6 +2344,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_FLASH_ATTN_EXT: ggml_cuda_flash_attn_ext(ctx, dst); break; + case GGML_OP_FLASH_ATTN_EXT_BANDED: + ggml_cuda_flash_attn_ext_banded(ctx, dst); + break; case GGML_OP_CROSS_ENTROPY_LOSS: ggml_cuda_cross_entropy_loss(ctx, dst); break; @@ -5142,6 +5184,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g op->type == GGML_TYPE_F32; case GGML_OP_FLASH_ATTN_EXT: return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); + case GGML_OP_FLASH_ATTN_EXT_BANDED: + return ggml_cuda_flash_attn_ext_banded_supported(dev_ctx->device, op); case GGML_OP_CROSS_ENTROPY_LOSS: case GGML_OP_CROSS_ENTROPY_LOSS_BACK: case GGML_OP_OPT_STEP_ADAMW: diff --git a/ggml/src/ggml-cuda/mmf.cuh b/ggml/src/ggml-cuda/mmf.cuh index d55cc1ec7b52..8b2be232ec9b 100644 --- a/ggml/src/ggml-cuda/mmf.cuh +++ b/ggml/src/ggml-cuda/mmf.cuh @@ -110,9 +110,9 @@ static __global__ void mul_mat_f( const int sample_x = sample_dst / sample_ratio; const int sample_y = sample_dst; - x += int64_t(sample_x) *stride_sample_x + channel_x *stride_channel_x + row0*stride_row ; - y += int64_t(sample_y) *stride_sample_y + (has_ids ? 0 : channel_y *stride_channel_y); - dst += int64_t(sample_dst)*stride_sample_dst + (has_ids ? 0 : channel_dst*stride_channel_dst); + x += int64_t(sample_x) *stride_sample_x + int64_t(channel_x)*stride_channel_x + int64_t(row0)*stride_row; + y += int64_t(sample_y) *stride_sample_y + (has_ids ? 0 : int64_t(channel_y) *stride_channel_y); + dst += int64_t(sample_dst)*stride_sample_dst + (has_ids ? 0 : int64_t(channel_dst)*stride_channel_dst); if constexpr (has_ids) { constexpr int y_stride_scale = std::is_same_v ? 1 : 2; @@ -362,7 +362,7 @@ static __global__ void mul_mat_f_ids( const int sample_x = sample_dst / sample_ratio; const int sample_y = sample_dst; - x += int64_t(sample_x) *stride_sample_x + channel_x *stride_channel_x + row0*stride_row; + x += int64_t(sample_x) *stride_sample_x + int64_t(channel_x)*stride_channel_x + int64_t(row0)*stride_row; y += int64_t(sample_y) *stride_sample_y; dst += int64_t(sample_dst)*stride_sample_dst; diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 71e3b2647a8e..70371e45753e 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -438,12 +438,12 @@ template static __device__ __forceinline_ if constexpr (type == GGML_TYPE_NVFP4) { if (y_scale_used) { - dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; + dst[(int64_t) ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; } else { - dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; + dst[(int64_t) ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; } } else { - dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; + dst[(int64_t) ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size]; GGML_UNUSED(y_scale_used); } } @@ -491,12 +491,12 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma( if constexpr (type == GGML_TYPE_NVFP4) { if (y_scale_used) { - dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l]; + dst[(int64_t) ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l]; } else { - dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; + dst[(int64_t) ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; } } else { - dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; + dst[(int64_t) ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; GGML_UNUSED(y_scale_used); } } @@ -969,7 +969,7 @@ static __global__ void mul_mat_q( int col_high = ncols_dst; int col_diff = ncols_dst; int offset_y = wt*stride_sample_y + zt*stride_channel_y; - int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int64_t offset_dst = (int64_t) wt*stride_sample_dst + (int64_t) zt*stride_channel_dst + (int64_t) jt*J*stride_col_dst; int offset_y_scale; if constexpr (type == GGML_TYPE_NVFP4) { offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; @@ -1057,7 +1057,7 @@ static __global__ void mul_mat_q( int col_high = ncols_dst; int col_diff = ncols_dst; int offset_y = wt*stride_sample_y + zt*stride_channel_y; - int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int64_t offset_dst = (int64_t) wt*stride_sample_dst + (int64_t) zt*stride_channel_dst + (int64_t) jt*J*stride_col_dst; int offset_y_scale; if constexpr (type == GGML_TYPE_NVFP4) { offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; @@ -1146,7 +1146,7 @@ static __global__ void mul_mat_q( int col_high = ncols_dst; int col_diff = ncols_dst; int offset_y = wt*stride_sample_y + zt*stride_channel_y; - int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst; + int64_t offset_dst = (int64_t) wt*stride_sample_dst + (int64_t) zt*stride_channel_dst + (int64_t) jt*J*stride_col_dst; int offset_y_scale; if constexpr (type == GGML_TYPE_NVFP4) { offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y; @@ -1289,7 +1289,7 @@ static __global__ void mul_mat_q_stream_k_fixup( const int it = tmp2.x; if (!ids_dst) { - const int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst + it*I; + const int64_t offset_dst = (int64_t) wt*stride_sample_dst + (int64_t) zt*stride_channel_dst + (int64_t) jt*J*stride_col_dst + (int64_t) it*I; dst += offset_dst; const int i_max = nrows_x - it*I - 1; diff --git a/ggml/src/ggml-cuda/mmvf.cu b/ggml/src/ggml-cuda/mmvf.cu index d7dbc8b99282..bbf862b3f5d7 100644 --- a/ggml/src/ggml-cuda/mmvf.cu +++ b/ggml/src/ggml-cuda/mmvf.cu @@ -44,7 +44,7 @@ static __global__ void mul_mat_vec_f( constexpr int warp_size = ggml_cuda_get_physical_warp_size(); - x += int64_t(sample_x) *stride_sample_x + channel_x *stride_channel_x + row*stride_row; + x += int64_t(sample_x) *stride_sample_x + int64_t(channel_x)*stride_channel_x + int64_t(row)*stride_row; y += int64_t(sample_y) *stride_sample_y + channel_y *stride_channel_y; dst += int64_t(sample_dst)*stride_sample_dst + channel_dst*stride_channel_dst; if constexpr (is_multi_token_id) { @@ -81,7 +81,7 @@ static __global__ void mul_mat_vec_f( } if (use_gate) { - gate_x += int64_t(sample_x) *stride_sample_x + channel_x *stride_channel_x + row*stride_row; + gate_x += int64_t(sample_x) *stride_sample_x + int64_t(channel_x)*stride_channel_x + int64_t(row)*stride_row; } if constexpr (has_fusion) { diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index e18ada5377d5..f62c4ebb2389 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -587,7 +587,7 @@ static __global__ void mul_mat_vec_q( float tmp_gate[ncols_dst][rows_per_cuda_block] = {{0.0f}}; const block_q8_1 * y = ((const block_q8_1 *) vy) + sample_y*stride_sample_y + channel_y*stride_channel_y; - const int kbx_offset = sample_x*stride_sample_x + channel_x*stride_channel_x + row0*stride_row_x; + const int64_t kbx_offset = int64_t(sample_x)*stride_sample_x + int64_t(channel_x)*stride_channel_x + int64_t(row0)*stride_row_x; for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { const int kby = kbx * (qk/QK8_1); // y block index that aligns with kbx @@ -739,7 +739,7 @@ static __global__ void mul_mat_vec_q_moe( const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y); const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y; - const int kbx_offset = channel_x*stride_channel_x + row0*stride_row_x; + const int64_t kbx_offset = int64_t(channel_x)*stride_channel_x + int64_t(row0)*stride_row_x; // partial sum for each thread float tmp[c_rows_per_block] = {0.0f}; diff --git a/ggml/src/ggml-cuda/pad.cu b/ggml/src/ggml-cuda/pad.cu index 31cd00f77816..013bd8386e93 100644 --- a/ggml/src/ggml-cuda/pad.cu +++ b/ggml/src/ggml-cuda/pad.cu @@ -25,7 +25,7 @@ static __global__ void pad_f32(const float * src, size_t s00, size_t s01, size_t return; } - const int64_t dst_idx = i3 * (ne0 * ne1 * ne2) + i2 * (ne0 * ne1) + i1 * ne0 + i0; + const int64_t dst_idx = (int64_t) i3 * ne0 * ne1 * ne2 + (int64_t) i2 * ne0 * ne1 + (int64_t) i1 * ne0 + i0; if (!circular) { if ((i0 >= lp0 && i0 < ne0 - rp0) && (i1 >= lp1 && i1 < ne1 - rp1) && (i2 >= lp2 && i2 < ne2 - rp2) && diff --git a/ggml/src/ggml-cuda/ssm-conv.cu b/ggml/src/ggml-cuda/ssm-conv.cu index 1463169cf78b..1e09177a4172 100644 --- a/ggml/src/ggml-cuda/ssm-conv.cu +++ b/ggml/src/ggml-cuda/ssm-conv.cu @@ -5,8 +5,8 @@ template static __global__ void ssm_conv_f32(const float * src0_ptr, const float * src1_ptr, const float * bias_ptr, - const int src0_nb0, const int src0_nb1, const int src0_nb2, const int src1_nb1, - float * dst_ptr, const int dst_nb0, const int dst_nb1, const int dst_nb2, + const int64_t src0_nb0, const int64_t src0_nb1, const int64_t src0_nb2, const int64_t src1_nb1, + float * dst_ptr, const int64_t dst_nb0, const int64_t dst_nb1, const int64_t dst_nb2, const int64_t n_t) { ggml_cuda_pdl_lc(); const float * GGML_CUDA_RESTRICT src0 = src0_ptr; @@ -22,9 +22,9 @@ static __global__ void ssm_conv_f32(const float * src0_ptr, const float * src1_p const float * w_block = (const float *) ((const char *) src1 + bidy * split_d_inner * src1_nb1); float * y_block = (float *) ((char *) dst + bidx * dst_nb2 + bidy * split_d_inner * dst_nb0); - const int stride_x = src0_nb1 / sizeof(float); - const int stride_w = src1_nb1 / sizeof(float); - const int stride_y = dst_nb1 / sizeof(float); + const int64_t stride_x = src0_nb1 / sizeof(float); + const int64_t stride_w = src1_nb1 / sizeof(float); + const int64_t stride_y = dst_nb1 / sizeof(float); float x[d_conv] = { 0.0f }; float w[d_conv] = { 0.0f }; @@ -60,9 +60,9 @@ static __global__ void ssm_conv_f32(const float * src0_ptr, const float * src1_p template static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0, const float * __restrict__ src1, const float * __restrict__ bias, - const int src0_nb0, const int src0_nb1, const int src0_nb2, - const int src1_nb1, float * __restrict__ dst, const int dst_nb0, - const int dst_nb1, const int dst_nb2, const int64_t n_t) { + const int64_t src0_nb0, const int64_t src0_nb1, const int64_t src0_nb2, + const int64_t src1_nb1, float * __restrict__ dst, const int64_t dst_nb0, + const int64_t dst_nb1, const int64_t dst_nb2, const int64_t n_t) { const int tid = threadIdx.x; const int bidx = blockIdx.x; const int bidy = blockIdx.y; @@ -74,9 +74,9 @@ static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0, float * y_block = (float *) ((char *) dst + bidx * dst_nb2 + bidz * split_n_t * dst_nb1 + bidy * split_d_inner * dst_nb0); - const int stride_x = src0_nb1 / sizeof(float); - const int stride_w = src1_nb1 / sizeof(float); - const int stride_y = dst_nb1 / sizeof(float); + const int64_t stride_x = src0_nb1 / sizeof(float); + const int64_t stride_w = src1_nb1 / sizeof(float); + const int64_t stride_y = dst_nb1 / sizeof(float); const int64_t local_n_t = min(split_n_t, n_t - bidz * split_n_t); const int n_cols = d_conv - 1 + split_n_t; @@ -124,9 +124,9 @@ static __global__ void ssm_conv_long_token_f32(const float * __restrict__ src0, } template -static void ssm_conv_f32_cuda(const float * src0, const float * src1, const float * bias, const int src0_nb0, const int src0_nb1, - const int src0_nb2, const int src1_nb1, float * dst, const int dst_nb0, const int dst_nb1, - const int dst_nb2, const int64_t nc, const int64_t nr, const int64_t n_t, +static void ssm_conv_f32_cuda(const float * src0, const float * src1, const float * bias, const int64_t src0_nb0, const int64_t src0_nb1, + const int64_t src0_nb2, const int64_t src1_nb1, float * dst, const int64_t dst_nb0, const int64_t dst_nb1, + const int64_t dst_nb2, const int64_t nc, const int64_t nr, const int64_t n_t, const int64_t n_s, cudaStream_t stream) { const int threads = 128; GGML_ASSERT(nr % threads == 0); diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index d38057721834..1166178b7249 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -600,6 +600,7 @@ static size_t ggml_backend_rpc_buffer_type_get_alloc_size(ggml_backend_buffer_ty // ops that require additional memory for fleeting data on certain backends // ref: https://github.com/ggml-org/llama.cpp/pull/15966 rpc_get |= tensor->op == GGML_OP_FLASH_ATTN_EXT; + rpc_get |= tensor->op == GGML_OP_FLASH_ATTN_EXT_BANDED; rpc_get |= tensor->op == GGML_OP_MUL_MAT_ID; if (rpc_get) { diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 59191c663eb0..a03f29bb5425 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1067,6 +1067,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "FILL", "FLASH_ATTN_EXT", + "FLASH_ATTN_EXT_BANDED", "FLASH_ATTN_BACK", "SSM_CONV", "SSM_SCAN", @@ -1100,7 +1101,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1182,6 +1183,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "fill(x, c)", "flash_attn_ext(x)", + "flash_attn_ext_banded(x)", "flash_attn_back(x)", "ssm_conv(x)", "ssm_scan(x)", @@ -1215,7 +1217,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); +static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -5443,11 +5445,60 @@ struct ggml_tensor * ggml_flash_attn_ext( return result; } +struct ggml_tensor * ggml_flash_attn_ext_banded( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * mask, + struct ggml_tensor * rel_logits, + float scale, + int64_t rel_extent) { + GGML_ASSERT(ggml_can_mul_mat(k, q)); + GGML_ASSERT(q->type == GGML_TYPE_F32); + GGML_ASSERT(q->ne[3] == k->ne[3]); + GGML_ASSERT(q->ne[3] == v->ne[3]); + GGML_ASSERT(q->ne[2] % k->ne[2] == 0); + GGML_ASSERT(q->ne[2] % v->ne[2] == 0); + + GGML_ASSERT(rel_logits != NULL); + GGML_ASSERT(rel_logits->type == GGML_TYPE_F32 || + rel_logits->type == GGML_TYPE_F16 || + rel_logits->type == GGML_TYPE_BF16); + GGML_ASSERT(rel_extent > 0); + GGML_ASSERT(rel_logits->ne[0] == rel_extent); + GGML_ASSERT(rel_logits->ne[1] == q->ne[2]); + GGML_ASSERT(rel_logits->ne[2] == q->ne[1]); + GGML_ASSERT(rel_logits->ne[3] == 1 || rel_logits->ne[3] == q->ne[3]); + + if (mask) { + GGML_ASSERT(mask->type == GGML_TYPE_F16); + GGML_ASSERT(ggml_is_contiguous(mask)); + GGML_ASSERT(q->ne[2] % mask->ne[2] == 0); + GGML_ASSERT(q->ne[3] % mask->ne[3] == 0); + } + + int64_t ne[4] = { v->ne[0], q->ne[2], q->ne[1], q->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + float params[] = { scale, 0.0f, 0.0f }; + ggml_set_op_params(result, params, sizeof(params)); + memcpy(result->op_params + 16, &rel_extent, sizeof(rel_extent)); + + result->op = GGML_OP_FLASH_ATTN_EXT_BANDED; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = mask; + result->src[5] = rel_logits; + + return result; +} void ggml_flash_attn_ext_set_prec( struct ggml_tensor * a, enum ggml_prec prec) { - GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT || a->op == GGML_OP_FLASH_ATTN_EXT_BANDED); const int32_t prec_i32 = (int32_t) prec; @@ -5456,7 +5507,7 @@ void ggml_flash_attn_ext_set_prec( enum ggml_prec ggml_flash_attn_ext_get_prec( const struct ggml_tensor * a) { - GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT || a->op == GGML_OP_FLASH_ATTN_EXT_BANDED); const int32_t prec_i32 = ggml_get_op_params_i32(a, 3); diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 124ea28b0616..643063e662e0 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -556,6 +556,7 @@ class MODEL_ARCH(IntEnum): TALKIE = auto() MELLUM = auto() NANBEIGE = auto() + INKLING = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -779,6 +780,14 @@ class MODEL_TENSOR(IntEnum): SHORTCONV_CONV = auto() SHORTCONV_INPROJ = auto() SHORTCONV_OUTPROJ = auto() + # inkling + ATTN_R = auto() + ATTN_REL_PROJ = auto() + SHORTCONV_K = auto() + SHORTCONV_V = auto() + SHORTCONV_ATTN = auto() + SHORTCONV_MLP = auto() + FFN_GSCALE = auto() VISEXP_ATTN_QKV = auto() VISEXP_ATTN_OUT = auto() VISEXP_GATE = auto() @@ -1168,6 +1177,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.TALKIE: "talkie", MODEL_ARCH.MELLUM: "mellum", MODEL_ARCH.NANBEIGE: "nanbeige", + MODEL_ARCH.INKLING: "inkling", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -1389,6 +1399,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.SHORTCONV_CONV: "blk.{bid}.shortconv.conv", MODEL_TENSOR.SHORTCONV_INPROJ: "blk.{bid}.shortconv.in_proj", MODEL_TENSOR.SHORTCONV_OUTPROJ: "blk.{bid}.shortconv.out_proj", + MODEL_TENSOR.ATTN_R: "blk.{bid}.attn_r", # inkling + MODEL_TENSOR.ATTN_REL_PROJ: "blk.{bid}.attn_rel_proj", # inkling + MODEL_TENSOR.SHORTCONV_K: "blk.{bid}.shortconv_k", # inkling + MODEL_TENSOR.SHORTCONV_V: "blk.{bid}.shortconv_v", # inkling + MODEL_TENSOR.SHORTCONV_ATTN: "blk.{bid}.shortconv_attn", # inkling + MODEL_TENSOR.SHORTCONV_MLP: "blk.{bid}.shortconv_mlp", # inkling + MODEL_TENSOR.FFN_GSCALE: "blk.{bid}.ffn_gscale", # inkling MODEL_TENSOR.VISEXP_ATTN_QKV: "blk.{bid}.vis_attn_qkv", MODEL_TENSOR.VISEXP_ATTN_OUT: "blk.{bid}.vis_attn_output", MODEL_TENSOR.VISEXP_GATE: "blk.{bid}.vis_gate", @@ -4161,6 +4178,38 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_UP_EXP, MODEL_TENSOR.FFN_EXP_PROBS_B, ], + MODEL_ARCH.INKLING: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.TOKEN_EMBD_NORM, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_R, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_REL_PROJ, + MODEL_TENSOR.SHORTCONV_K, + MODEL_TENSOR.SHORTCONV_V, + MODEL_TENSOR.SHORTCONV_ATTN, + MODEL_TENSOR.SHORTCONV_MLP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_GSCALE, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + ], MODEL_ARCH.SMALLTHINKER: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, @@ -4834,6 +4883,7 @@ def get_type(val: Any) -> GGUFValueType: class VisionProjectorType: + INKLING = "inkling" GEMMA3 = "gemma3" GEMMA3NV = "gemma3nv" GEMMA3NA = "gemma3na" diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 1e991b873cea..5a0175a56c8e 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -63,6 +63,7 @@ class TensorNameMap: "model.layers.0.pre_norm", # rwkv7 "backbone.norm", # wavtokenizer "model.embedding_norm", # lfm2 + "model.embed_norm", # inkling ), # Position embeddings @@ -86,6 +87,7 @@ class TensorNameMap: "lm_head", # llama4 "model.transformer.ff_out", # llada "head.decoder", # modern-bert + "model.unembed", # inkling ), MODEL_TENSOR.DENSE_2_OUT: ( "dense_2_out", # embeddinggemma @@ -2548,6 +2550,53 @@ class TensorNameMap: # architecture-specific block mappings arch_block_mappings_cfg: dict[MODEL_ARCH, dict[MODEL_TENSOR, tuple[str, ...]]] = { + MODEL_ARCH.INKLING: { + MODEL_TENSOR.ATTN_NORM: ( + "model.layers.{bid}.attn_norm", + ), + MODEL_TENSOR.ATTN_Q: ( + "model.layers.{bid}.attn.wq_du", + ), + MODEL_TENSOR.ATTN_K: ( + "model.layers.{bid}.attn.wk_dv", + ), + MODEL_TENSOR.ATTN_V: ( + "model.layers.{bid}.attn.wv_dv", + ), + MODEL_TENSOR.ATTN_R: ( + "model.layers.{bid}.attn.wr_du", + ), + MODEL_TENSOR.ATTN_OUT: ( + "model.layers.{bid}.attn.wo_ud", + ), + MODEL_TENSOR.ATTN_Q_NORM: ( + "model.layers.{bid}.attn.q_norm", + ), + MODEL_TENSOR.ATTN_K_NORM: ( + "model.layers.{bid}.attn.k_norm", + ), + MODEL_TENSOR.ATTN_REL_PROJ: ( + "model.layers.{bid}.attn.rel_logits_proj", + ), + MODEL_TENSOR.SHORTCONV_K: ( + "model.layers.{bid}.attn.k_sconv", + ), + MODEL_TENSOR.SHORTCONV_V: ( + "model.layers.{bid}.attn.v_sconv", + ), + MODEL_TENSOR.SHORTCONV_ATTN: ( + "model.layers.{bid}.attn_sconv", + ), + MODEL_TENSOR.SHORTCONV_MLP: ( + "model.layers.{bid}.mlp_sconv", + ), + MODEL_TENSOR.FFN_NORM: ( + "model.layers.{bid}.mlp_norm", + ), + MODEL_TENSOR.FFN_DOWN: ( + "model.layers.{bid}.mlp.w2_md", + ), + }, MODEL_ARCH.ARCTIC: { MODEL_TENSOR.FFN_NORM: ( "model.layers.{bid}.residual_layernorm", diff --git a/models/templates/Inkling.jinja b/models/templates/Inkling.jinja new file mode 100644 index 000000000000..512d92685eac --- /dev/null +++ b/models/templates/Inkling.jinja @@ -0,0 +1,514 @@ +{#- Keep Python's floating type spelling when the jinja engine compacts 1.0 to 1. -#} +{%- macro json_scalar(value) -%} + {%- set serialized = value | tojson(ensure_ascii=false, separators=(',', ':')) -%} + {%- if value is float and '.' not in serialized and 'e' not in (serialized | lower) -%} + {{- serialized -}}{{- '.0' -}} + {%- else -%} + {{- serialized -}} + {%- endif -%} +{%- endmacro -%} + +{%- macro canonical_json(value) -%} + {%- if value is mapping -%} + {{- '{' -}} + {%- for key, item in value | dictsort(case_sensitive=true) -%} + {{- (key | string) | tojson(ensure_ascii=false, separators=(',', ':')) -}}{{- ':' -}} + {{- canonical_json(item) -}} + {%- if not loop.last -%}{{- ',' -}}{%- endif -%} + {%- endfor -%} + {{- '}' -}} + {%- elif value is sequence and value is not string -%} + {{- '[' -}} + {%- for item in value -%} + {{- canonical_json(item) -}} + {%- if not loop.last -%}{{- ',' -}}{%- endif -%} + {%- endfor -%} + {{- ']' -}} + {%- else -%} + {{- json_scalar(value) -}} + {%- endif -%} +{%- endmacro -%} + +{#- Advance across insignificant JSON whitespace. -#} +{%- macro json_skip_ws(source, state) -%} + {%- set scan = namespace(done=false) -%} + {%- for ignored in range(source | length) -%} + {%- if not scan.done and state.pos < source | length -%} + {%- set ch = source[state.pos] -%} + {%- if ch == ' ' or ch == '\t' or ch == '\r' or ch == '\n' -%} + {%- set state.pos = state.pos + 1 -%} + {%- else -%} + {%- set scan.done = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} +{%- endmacro -%} + +{#- + Decode a JSON string under Transformers, then re-encode it like json.dumps. + llama.cpp converts string arguments to objects after capability detection; + the fallback keeps engine-only parsing safe without runtime-only filters. +-#} +{%- macro json_string(source, state) -%} + {%- if lipsum is defined -%} + {%- set out = namespace(decoded='', done=false, code=0, low=0) -%} + {%- set hex_values = {'0': 0, '1': 1, '2': 2, '3': 3, '4': 4, '5': 5, '6': 6, '7': 7, '8': 8, '9': 9, 'a': 10, 'b': 11, 'c': 12, 'd': 13, 'e': 14, 'f': 15} -%} + {%- set state.pos = state.pos + 1 -%} + {%- for ignored in range(source | length) -%} + {%- if not out.done and state.pos < source | length -%} + {%- set ch = source[state.pos] -%} + {%- set state.pos = state.pos + 1 -%} + {%- if ch == '"' -%} + {%- set out.done = true -%} + {%- elif ch != '\\' -%} + {%- set out.decoded = out.decoded + ch -%} + {%- elif state.pos < source | length -%} + {%- set escape = source[state.pos] -%} + {%- set state.pos = state.pos + 1 -%} + {%- if escape == '"' -%}{%- set out.decoded = out.decoded + '"' -%} + {%- elif escape == '\\' -%}{%- set out.decoded = out.decoded + '\\' -%} + {%- elif escape == '/' -%}{%- set out.decoded = out.decoded + '/' -%} + {%- elif escape == 'b' -%}{%- set out.decoded = out.decoded + '\b' -%} + {%- elif escape == 'f' -%}{%- set out.decoded = out.decoded + '\f' -%} + {%- elif escape == 'n' -%}{%- set out.decoded = out.decoded + '\n' -%} + {%- elif escape == 'r' -%}{%- set out.decoded = out.decoded + '\r' -%} + {%- elif escape == 't' -%}{%- set out.decoded = out.decoded + '\t' -%} + {%- elif escape == 'u' -%} + {%- set out.code = 0 -%} + {%- for offset in range(4) -%} + {%- set digit = source[state.pos + offset] | lower -%} + {%- set out.code = out.code * 16 + hex_values[digit] -%} + {%- endfor -%} + {%- set state.pos = state.pos + 4 -%} + {%- if out.code >= 55296 and out.code <= 56319 and source[state.pos:state.pos + 2] == '\\u' -%} + {%- set out.low = 0 -%} + {%- for offset in range(4) -%} + {%- set digit = source[state.pos + 2 + offset] | lower -%} + {%- set out.low = out.low * 16 + hex_values[digit] -%} + {%- endfor -%} + {%- if out.low >= 56320 and out.low <= 57343 -%} + {%- set out.code = 65536 + (out.code - 55296) * 1024 + out.low - 56320 -%} + {%- set state.pos = state.pos + 6 -%} + {%- endif -%} + {%- endif -%} + {%- set out.decoded = out.decoded + ('%c' % out.code) -%} + {%- else -%} + {%- set out.decoded = out.decoded + escape -%} + {%- endif -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + {%- set state.string_value = out.decoded -%} + {{- out.decoded | tojson(ensure_ascii=false, separators=(',', ':')) -}} + {%- else -%} + {%- set out = namespace(start=state.pos, escaped=false, done=false) -%} + {%- set state.pos = state.pos + 1 -%} + {%- for ignored in range(source | length) -%} + {%- if not out.done and state.pos < source | length -%} + {%- set ch = source[state.pos] -%} + {%- set state.pos = state.pos + 1 -%} + {%- if out.escaped -%} + {%- set out.escaped = false -%} + {%- elif ch == '\\' -%} + {%- set out.escaped = true -%} + {%- elif ch == '"' -%} + {%- set out.done = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + {%- set state.string_value = source[out.start:state.pos] -%} + {{- source[out.start:state.pos] -}} + {%- endif -%} +{%- endmacro -%} + +{#- + Move the cursor across one complete JSON value without rendering it. Object + parsing uses the raw slice to sort members before recursively rendering them. +-#} +{%- macro json_scan_value(source, state) -%} + {{- json_skip_ws(source, state) -}} + {%- set scan = namespace(depth=0, quoted=false, escaped=false, done=false) -%} + {%- for ignored in range(source | length) -%} + {%- if not scan.done and state.pos < source | length -%} + {%- set ch = source[state.pos] -%} + {%- if scan.quoted -%} + {%- set state.pos = state.pos + 1 -%} + {%- if ch == '"' and not scan.escaped -%} + {%- set scan.quoted = false -%} + {%- elif ch == '\\' and not scan.escaped -%} + {%- set scan.escaped = true -%} + {%- else -%} + {%- set scan.escaped = false -%} + {%- endif -%} + {%- elif ch == '"' -%} + {%- set scan.quoted = true -%} + {%- set state.pos = state.pos + 1 -%} + {%- elif ch == '{' or ch == '[' -%} + {%- set scan.depth = scan.depth + 1 -%} + {%- set state.pos = state.pos + 1 -%} + {%- elif ch == '}' or ch == ']' -%} + {%- if scan.depth > 0 -%} + {%- set scan.depth = scan.depth - 1 -%} + {%- set state.pos = state.pos + 1 -%} + {%- else -%} + {%- set scan.done = true -%} + {%- endif -%} + {%- elif ch == ',' and scan.depth == 0 -%} + {%- set scan.done = true -%} + {%- else -%} + {%- set state.pos = state.pos + 1 -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} +{%- endmacro -%} + +{#- Parse and canonically render one JSON value from source at state.pos. -#} +{%- macro canonical_json_text_value(source, state) -%} + {{- json_skip_ws(source, state) -}} + {%- if state.pos >= source | length -%} + {{- '{}' -}} + {%- elif source[state.pos] == '"' -%} + {{- json_string(source, state) -}} + {%- elif source[state.pos] == '{' -%} + {%- set state.pos = state.pos + 1 -%} + {{- json_skip_ws(source, state) -}} + {%- set object_state = namespace(pairs=[], done=false) -%} + {%- if state.pos < source | length and source[state.pos] == '}' -%} + {%- set state.pos = state.pos + 1 -%} + {%- set object_state.done = true -%} + {%- endif -%} + {%- for ignored in range(source | length) -%} + {%- if not object_state.done -%} + {{- json_skip_ws(source, state) -}} + {%- set key = json_string(source, state) -%} + {%- set sort_key = state.string_value -%} + {{- json_skip_ws(source, state) -}} + {%- if state.pos < source | length and source[state.pos] == ':' -%} + {%- set state.pos = state.pos + 1 -%} + {%- endif -%} + {{- json_skip_ws(source, state) -}} + {%- set value_start = state.pos -%} + {{- json_scan_value(source, state) -}} + {#- json.loads keeps the final member when a key is duplicated. -#} + {%- set unique = namespace(pairs=[]) -%} + {%- for previous in object_state.pairs -%} + {%- if previous[0] != sort_key -%} + {%- set unique.pairs = unique.pairs + [previous] -%} + {%- endif -%} + {%- endfor -%} + {%- set object_state.pairs = unique.pairs + [[sort_key, key, source[value_start:state.pos]]] -%} + {{- json_skip_ws(source, state) -}} + {%- if state.pos < source | length and source[state.pos] == ',' -%} + {%- set state.pos = state.pos + 1 -%} + {%- else -%} + {%- if state.pos < source | length and source[state.pos] == '}' -%} + {%- set state.pos = state.pos + 1 -%} + {%- endif -%} + {%- set object_state.done = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + {{- '{' -}} + {%- for pair in object_state.pairs | sort(case_sensitive=true, attribute=0) -%} + {{- pair[1] -}}{{- ':' -}} + {%- set child_state = namespace(pos=0, string_value='') -%} + {{- canonical_json_text_value(pair[2], child_state) -}} + {%- if not loop.last -%}{{- ',' -}}{%- endif -%} + {%- endfor -%} + {{- '}' -}} + {%- elif source[state.pos] == '[' -%} + {%- set state.pos = state.pos + 1 -%} + {{- '[' -}} + {{- json_skip_ws(source, state) -}} + {%- set array_state = namespace(done=false, first=true) -%} + {%- if state.pos < source | length and source[state.pos] == ']' -%} + {%- set state.pos = state.pos + 1 -%} + {%- set array_state.done = true -%} + {%- endif -%} + {%- for ignored in range(source | length) -%} + {%- if not array_state.done -%} + {%- if not array_state.first -%}{{- ',' -}}{%- endif -%} + {%- set array_state.first = false -%} + {{- canonical_json_text_value(source, state) -}} + {{- json_skip_ws(source, state) -}} + {%- if state.pos < source | length and source[state.pos] == ',' -%} + {%- set state.pos = state.pos + 1 -%} + {%- else -%} + {%- if state.pos < source | length and source[state.pos] == ']' -%} + {%- set state.pos = state.pos + 1 -%} + {%- endif -%} + {%- set array_state.done = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + {{- ']' -}} + {%- else -%} + {%- set scalar = namespace(start=state.pos, done=false) -%} + {%- for ignored in range(source | length) -%} + {%- if not scalar.done and state.pos < source | length -%} + {%- set ch = source[state.pos] -%} + {%- if ch == ',' or ch == '}' or ch == ']' or ch == ' ' or ch == '\t' or ch == '\r' or ch == '\n' -%} + {%- set scalar.done = true -%} + {%- else -%} + {%- set state.pos = state.pos + 1 -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + {%- set token = source[scalar.start:state.pos] -%} + {%- if token == 'true' or token == 'false' or token == 'null' -%} + {{- token -}} + {%- elif '.' in token or 'e' in token or 'E' in token -%} + {{- json_scalar(token | float) -}} + {%- else -%} + {{- token | int | tojson(ensure_ascii=false, separators=(',', ':')) -}} + {%- endif -%} + {%- endif -%} +{%- endmacro -%} + +{%- macro canonical_json_text(source) -%} + {%- set text = source | trim -%} + {%- if not text -%} + {{- '{}' -}} + {%- else -%} + {%- set state = namespace(pos=0, string_value='') -%} + {{- canonical_json_text_value(text, state) -}} + {%- endif -%} +{%- endmacro -%} + +{#- + OpenAI clients use either an argument mapping or a JSON-encoded object. + Non-object/empty oddities degrade to {} instead of raising. Object parsing + recursively sorts keys, keeps the last duplicate key (json.loads behavior), + decodes JSON escapes under Transformers, and preserves array order. +-#} +{%- macro canonical_arguments(arguments) -%} + {%- if arguments is mapping -%} + {{- canonical_json(arguments) -}} + {%- elif arguments is string -%} + {%- set source = arguments | trim -%} + {%- if source and source[0] == '{' -%} + {{- canonical_json_text(source) -}} + {%- else -%} + {{- '{}' -}} + {%- endif -%} + {%- else -%} + {{- '{}' -}} + {%- endif -%} +{%- endmacro -%} + +{#- + Match Python f"{float(effort):.2f}" followed by trailing-zero removal. + llama.cpp's jinja engine has no round filter; the bit table handles binary64 midpoint equality. +-#} +{%- macro reasoning_effort_text(effort) -%} + {%- set eff = effort -%} + {%- if eff is string -%} + {%- set e = eff | trim | lower -%} + {%- if e == 'none' -%}{%- set eff = 0.0 -%} + {%- elif e == 'minimal' -%}{%- set eff = 0.1 -%} + {%- elif e == 'low' -%}{%- set eff = 0.2 -%} + {%- elif e == 'medium' -%}{%- set eff = 0.7 -%} + {%- elif e == 'high' -%}{%- set eff = 0.9 -%} + {%- elif e == 'xhigh' -%}{%- set eff = 0.99 -%} + {%- elif e == 'max' -%}{%- set eff = 0.99 -%} + {%- else -%}{%- set eff = e | float(-1.0) -%} + {%- endif -%} + {%- endif -%} + {%- set value = eff | float -%} + {%- if value < 0 or value > 0.99 -%} + {{- raise_exception('Invalid reasoning_effort: ' + (effort | string) + '; expected none/minimal/low/medium/high/xhigh/max or a number in [0.0, 0.99]') -}} + {%- endif -%} + {%- if value == value | int and value >= 0 and value <= 1 -%} + {{- value | int -}} + {%- else -%} + {%- set midpoint_rounds_up = '1011011011010100100100100111000111000111100011100011111100000001111110000001111110000001111111000000' -%} + {%- set rounded = namespace(hundredths=0) -%} + {%- for lower_hundredth in range(100) -%} + {%- set boundary = (lower_hundredth + 0.5) / 100 -%} + {%- if value > boundary or (value == boundary and midpoint_rounds_up[lower_hundredth] == '1') -%} + {%- set rounded.hundredths = lower_hundredth + 1 -%} + {%- endif -%} + {%- endfor -%} + {%- if rounded.hundredths == 100 -%} + {{- '1' -}} + {%- elif rounded.hundredths == 0 -%} + {{- '0' -}} + {%- elif rounded.hundredths % 10 == 0 -%} + {{- '0.' -}}{{- (rounded.hundredths / 10) | int -}} + {%- elif rounded.hundredths < 10 -%} + {{- '0.0' -}}{{- rounded.hundredths -}} + {%- else -%} + {{- '0.' -}}{{- rounded.hundredths -}} + {%- endif -%} + {%- endif -%} +{%- endmacro -%} + +{%- macro role_token(role) -%} + {%- if role == 'user' -%}{{- '<|message_user|>' -}} + {%- elif role == 'assistant' -%}{{- '<|message_model|>' -}} + {%- elif role == 'system' or role == 'developer' -%}{{- '<|message_system|>' -}} + {%- elif role == 'tool' -%}{{- '<|message_tool|>' -}} + {%- endif -%} +{%- endmacro -%} + +{%- macro emit_message(role, kind, content='', author_name='') -%} + {{- role_token(role) -}} + {%- if author_name -%}{{- author_name -}}{%- endif -%} + {%- if kind == 'text' -%} + {{- '<|content_text|>' -}}{{- content -}} + {%- elif kind == 'thinking' -%} + {{- '<|content_thinking|>' -}}{{- content -}} + {%- elif kind == 'xml' -%} + {{- '<|content_xml|>' -}}{{- content -}} + {%- elif kind == 'invoke_tool_json' -%} + {{- '<|content_invoke_tool_json|>' -}}{{- content -}} + {%- elif kind == 'image' -%} + {{- '<|content_image|><|image|>' -}} + {%- elif kind == 'audio' -%} + {{- '<|content_audio_input|><|audio|><|audio_end|>' -}} + {%- endif -%} + {{- '<|end_message|>' -}} +{%- endmacro -%} + +{%- set effort_value = reasoning_effort if (reasoning_effort is defined and reasoning_effort is not none) else 0.9 -%} +{%- set eff_ns = namespace(emitted=false) -%} +{%- set first_ns = namespace(idx=-1) -%} +{%- for m in messages -%} + {%- if first_ns.idx == -1 and m.get('role') not in ['system', 'developer'] -%} + {%- set first_ns.idx = loop.index0 -%} + {%- endif -%} +{%- endfor -%} + +{%- if tools is defined and tools -%} + {{- '<|message_system|>tool_declare<|content_xml|>[' -}} + {%- for tool in tools -%} + {%- set function = tool.get('function', {}) if tool.get('function', {}) is mapping else {} -%} + {%- set description = function.get('description') or '' -%} + {%- set parameters = function.get('parameters') or {} -%} + {%- set tool_type = tool.get('type', 'function') -%} + {{- '{"description":' -}}{{- canonical_json(description) -}} + {{- ',"name":' -}}{{- canonical_json(function.get('name')) -}} + {{- ',"parameters":' -}}{{- canonical_json(parameters) -}} + {{- ',"type":' -}}{{- canonical_json(tool_type) -}}{{- '}' -}} + {%- if not loop.last -%}{{- ',' -}}{%- endif -%} + {%- endfor -%} + {{- ']<|end_message|>' -}} +{%- endif -%} + +{#- Last-user boundary used only by the opt-in preserve_thinking=false mode. -#} +{%- set thinking_state = namespace(last_user_index=messages | length - 1, found_user=false) -%} +{%- for index in range(messages | length - 1, -1, -1) -%} + {%- if not thinking_state.found_user and messages[index].get('role') == 'user' -%} + {%- set thinking_state.last_user_index = index -%} + {%- set thinking_state.found_user = true -%} + {%- endif -%} +{%- endfor -%} + +{%- for message in messages -%} + {%- set message_index = loop.index0 -%} + {%- set role = message.get('role') -%} + {%- if not eff_ns.emitted and loop.index0 == first_ns.idx -%} + {{- emit_message('system', 'text', 'Thinking effort level: ' + reasoning_effort_text(effort_value)) -}} + {%- set eff_ns.emitted = true -%} + {%- endif -%} + {%- if role == 'tool' -%} + {%- set resolved = namespace(name=message.get('name') or '') -%} + {%- if not resolved.name and message.get('tool_call_id') -%} + {%- for prior in messages[:message_index] -%} + {%- if prior.get('role') == 'assistant' -%} + {%- for call in prior.get('tool_calls') or [] -%} + {%- if call.get('id') and (call.get('id') | string) == message.get('tool_call_id') -%} + {%- set prior_function = call.get('function', {}) if call.get('function', {}) is mapping else {} -%} + {%- set resolved.name = prior_function.get('name') or '' -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- set tool_content = message.get('content', '') -%} + {%- if tool_content is none -%} + {%- set tool_content = '' -%} + {%- elif tool_content is mapping or (tool_content is sequence and tool_content is not string) -%} + {%- set tool_content = canonical_json(tool_content) -%} + {%- elif tool_content is not string -%} + {%- set tool_content = canonical_json(tool_content) -%} + {%- endif -%} + {{- emit_message('tool', 'text', tool_content, resolved.name | string) -}} + {%- elif role == 'user' or role == 'assistant' or role == 'system' or role == 'developer' -%} + {%- set turn_out -%} + {%- if role == 'assistant' and message.get('reasoning_content') is string and message.get('reasoning_content') and + ((preserve_thinking is not defined) or preserve_thinking is not false or message_index > thinking_state.last_user_index) -%} + {{- emit_message('assistant', 'thinking', message.get('reasoning_content')) -}} + {%- endif -%} + + {%- set content = message.get('content', '') -%} + {#- Makes llama.cpp retain typed arrays rather than flattening them. -#} + {%- set content_probe = content[0] if content is sequence and content | length > 0 else none -%} + {%- if content is string -%} + {%- if '<__media_' in content -%} + {#- Flattened media markers: each part becomes its own message block. The runtime may + randomize the marker suffix, so split on the stable prefix and re-emit the exact + marker text; the runtime then expands it into the typed content sentinel plus the + media embedding rows. -#} + {%- for segment in content.split('<__media_') -%} + {%- if loop.first -%} + {%- if segment -%}{{- emit_message(role, 'text', segment) -}}{%- endif -%} + {%- else -%} + {%- set mparts = segment.split('>') -%} + {%- set rest = mparts[1:] | join('>') -%} + {{- role_token(role) -}}{{- '<__media_' + mparts[0] + '>' -}}{{- '<|end_message|>' -}} + {%- if rest -%}{{- emit_message(role, 'text', rest) -}}{%- endif -%} + {%- endif -%} + {%- endfor -%} + {%- elif content -%} + {{- emit_message(role, 'text', content) -}} + {%- endif -%} + {%- elif content is sequence -%} + {%- for part in content -%} + {%- if part is string -%} + {{- emit_message(role, 'text', part) -}} + {%- elif part is mapping -%} + {%- set part_type = part.get('type') -%} + {%- if part_type is none or part_type == 'text' or part_type == 'input_text' -%} + {%- set part_text = part.get('text', '') -%} + {{- emit_message(role, 'text', part_text if part_text is string else '') -}} + {%- elif part_type == 'image' or part_type == 'input_image' or part_type == 'image_url' -%} + {{- emit_message(role, 'image') -}} + {%- elif part_type == 'audio' or part_type == 'input_audio' or part_type == 'audio_url' -%} + {{- emit_message(role, 'audio') -}} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + + {%- if role == 'assistant' -%} + {%- for call in message.get('tool_calls') or [] -%} + {%- set function = call.get('function', {}) if call.get('function', {}) is mapping else {} -%} + {%- if function.get('name') is string -%} + {%- set raw_arguments = function.get('arguments') or {} -%} + {%- set arguments_json = canonical_arguments(raw_arguments) -%} + {%- set invocation = '{"name":' + canonical_json(function.get('name')) + ',"args":' + arguments_json + '}' -%} + {{- emit_message('assistant', 'invoke_tool_json', invocation, function.get('name')) -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endset -%} + {{- turn_out -}} + {#- Close each historical model turn, but never emit a bare terminator for an + assistant message that rendered no blocks. -#} + {%- if role == 'assistant' and turn_out -%} + {{- '<|content_model_end_sampling|>' -}} + {%- endif -%} + {%- else -%} + {{- raise_exception('Unknown message role: ' + (role | string)) -}} + {%- endif -%} +{%- endfor -%} + +{%- if not eff_ns.emitted -%} + {{- emit_message('system', 'text', 'Thinking effort level: ' + reasoning_effort_text(effort_value)) -}} +{%- endif -%} + +{%- if add_generation_prompt is defined and add_generation_prompt -%} + {{- '<|message_model|>' -}} +{%- endif -%} +{#- Unsloth translation to jinja from TML's parser #} diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index e81ff647eee4..99fad9738a06 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -144,6 +144,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_TALKIE, "talkie" }, { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_NANBEIGE, "nanbeige" }, + { LLM_ARCH_INKLING, "inkling" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; @@ -320,6 +321,18 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" }, { LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" }, + + // inkling (private arch) + { LLM_KV_INKLING_D_REL, "%s.d_rel" }, + { LLM_KV_INKLING_REL_EXTENT, "%s.rel_extent" }, + { LLM_KV_INKLING_REL_EXTENT_SWA, "%s.rel_extent_swa" }, + { LLM_KV_INKLING_SHORTCONV_KERNEL, "%s.shortconv_kernel" }, + { LLM_KV_INKLING_DENSE_BLOCK_COUNT, "%s.dense_block_count" }, + { LLM_KV_INKLING_LOGIT_SCALE_DENOM, "%s.logit_scale_denom" }, + { LLM_KV_INKLING_LOG_SCALING_N_FLOOR, "%s.log_scaling_n_floor" }, + { LLM_KV_INKLING_LOG_SCALING_ALPHA, "%s.log_scaling_alpha" }, + { LLM_KV_INKLING_UNPADDED_VOCAB_SIZE, "%s.unpadded_vocab_size" }, + // sentence-transformers dense modules feature dims { LLM_KV_DENSE_2_FEAT_IN, "%s.dense_2_feat_in" }, { LLM_KV_DENSE_2_FEAT_OUT, "%s.dense_2_feat_out" }, @@ -592,9 +605,19 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_SHORTCONV_CONV, "blk.%d.shortconv.conv" }, { LLM_TENSOR_SHORTCONV_INPROJ, "blk.%d.shortconv.in_proj" }, { LLM_TENSOR_SHORTCONV_OUTPROJ, "blk.%d.shortconv.out_proj" }, + { LLM_TENSOR_ATTN_R, "blk.%d.attn_r" }, + { LLM_TENSOR_ATTN_REL_PROJ, "blk.%d.attn_rel_proj" }, + { LLM_TENSOR_SHORTCONV_K, "blk.%d.shortconv_k" }, + { LLM_TENSOR_SHORTCONV_V, "blk.%d.shortconv_v" }, + { LLM_TENSOR_SHORTCONV_ATTN, "blk.%d.shortconv_attn" }, + { LLM_TENSOR_SHORTCONV_MLP, "blk.%d.shortconv_mlp" }, + { LLM_TENSOR_FFN_GSCALE, "blk.%d.ffn_gscale" }, { LLM_TENSOR_FFN_GATE_CHEXPS, "blk.%d.ffn_gate_chexps" }, { LLM_TENSOR_FFN_DOWN_CHEXPS, "blk.%d.ffn_down_chexps" }, { LLM_TENSOR_FFN_UP_CHEXPS, "blk.%d.ffn_up_chexps" }, + { LLM_TENSOR_FFN_GATE_SHEXPS, "blk.%d.ffn_gate_shexp" }, + { LLM_TENSOR_FFN_DOWN_SHEXPS, "blk.%d.ffn_down_shexp" }, + { LLM_TENSOR_FFN_UP_SHEXPS, "blk.%d.ffn_up_shexp" }, { LLM_TENSOR_VISEXP_ATTN_QKV, "blk.%d.vis_attn_qkv" }, { LLM_TENSOR_VISEXP_ATTN_OUT, "blk.%d.vis_attn_output" }, { LLM_TENSOR_VISEXP_FFN_GATE, "blk.%d.vis_gate" }, @@ -795,6 +818,9 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_FFN_UP_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_GATE_UP_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_DOWN_CHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_DOWN_SHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_GATE_SHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_UP_SHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_GATE_CHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_UP_CHEXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_EXP_PROBS_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, @@ -836,6 +862,13 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_SHORTCONV_CONV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, {LLM_TENSOR_SHORTCONV_INPROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_SHORTCONV_OUTPROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_R, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_REL_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_SHORTCONV_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, + {LLM_TENSOR_SHORTCONV_V, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, + {LLM_TENSOR_SHORTCONV_ATTN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, + {LLM_TENSOR_SHORTCONV_MLP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, + {LLM_TENSOR_FFN_GSCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_VISEXP_ATTN_QKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_VISEXP_ATTN_OUT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_VISEXP_FFN_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, @@ -968,6 +1001,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_INKLING: return true; default: return false; @@ -1024,6 +1058,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_INKLING: return false; default: return true; diff --git a/src/llama-arch.h b/src/llama-arch.h index cbc97085ea79..6a283eb6b5a1 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -149,6 +149,7 @@ enum llm_arch { LLM_ARCH_MINIMAX_M3, LLM_ARCH_DFLASH, LLM_ARCH_NANBEIGE, + LLM_ARCH_INKLING, LLM_ARCH_UNKNOWN, }; @@ -367,6 +368,17 @@ enum llm_kv { LLM_KV_SHORTCONV_L_CACHE, + // inkling (private arch) + LLM_KV_INKLING_D_REL, + LLM_KV_INKLING_REL_EXTENT, + LLM_KV_INKLING_REL_EXTENT_SWA, + LLM_KV_INKLING_SHORTCONV_KERNEL, + LLM_KV_INKLING_DENSE_BLOCK_COUNT, + LLM_KV_INKLING_LOGIT_SCALE_DENOM, + LLM_KV_INKLING_LOG_SCALING_N_FLOOR, + LLM_KV_INKLING_LOG_SCALING_ALPHA, + LLM_KV_INKLING_UNPADDED_VOCAB_SIZE, + LLM_KV_XIELU_ALPHA_N, LLM_KV_XIELU_ALPHA_P, LLM_KV_XIELU_BETA, @@ -434,6 +446,10 @@ enum llm_tensor { LLM_TENSOR_FFN_DOWN_CHEXPS, LLM_TENSOR_FFN_GATE_CHEXPS, LLM_TENSOR_FFN_UP_CHEXPS, + // Inkling: same GGUF names as *_SHEXP but registered as MUL_MAT_ID (3D shared-expert bank via ggml_mul_mat_id) + LLM_TENSOR_FFN_DOWN_SHEXPS, + LLM_TENSOR_FFN_GATE_SHEXPS, + LLM_TENSOR_FFN_UP_SHEXPS, LLM_TENSOR_FFN_EXP_PROBS_B, LLM_TENSOR_FFN_LATENT_DOWN, LLM_TENSOR_FFN_LATENT_UP, @@ -595,6 +611,14 @@ enum llm_tensor { LLM_TENSOR_SHORTCONV_CONV, LLM_TENSOR_SHORTCONV_INPROJ, LLM_TENSOR_SHORTCONV_OUTPROJ, + // inkling (private arch) + LLM_TENSOR_ATTN_R, + LLM_TENSOR_ATTN_REL_PROJ, + LLM_TENSOR_SHORTCONV_K, + LLM_TENSOR_SHORTCONV_V, + LLM_TENSOR_SHORTCONV_ATTN, + LLM_TENSOR_SHORTCONV_MLP, + LLM_TENSOR_FFN_GSCALE, LLM_TENSOR_VISEXP_ATTN_QKV, LLM_TENSOR_VISEXP_ATTN_OUT, LLM_TENSOR_VISEXP_FFN_GATE, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 6d1c8f4e42a8..980d393db4eb 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1382,9 +1382,15 @@ ggml_tensor * llm_graph_context::build_cvec( ggml_tensor * llm_graph_context::build_lora_mm( ggml_tensor * w, ggml_tensor * cur, - ggml_tensor * w_s) const { + ggml_tensor * w_s, + enum ggml_prec prec) const { ggml_tensor * res = ggml_mul_mat(ctx0, w, cur); + if (prec != GGML_PREC_DEFAULT) { + // Set precision on the base MUL_MAT before an optional scale/LoRA attachment changes the root op. + ggml_mul_mat_set_prec(res, prec); + } + if (w_s) { res = ggml_mul(ctx0, res, w_s); } diff --git a/src/llama-graph.h b/src/llama-graph.h index 7ed490ce6728..d3270a1a3469 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -954,7 +954,8 @@ struct llm_graph_context { ggml_tensor * build_lora_mm( ggml_tensor * w, ggml_tensor * cur, - ggml_tensor * w_s = nullptr) const; + ggml_tensor * w_s = nullptr, + enum ggml_prec prec = GGML_PREC_DEFAULT) const; // do mat_mul_id, while optionally apply lora and per-expert scale ggml_tensor * build_lora_mm_id( diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 50af97f358c3..36568582e32b 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -191,6 +191,11 @@ uint32_t llama_hparams::n_embd_k_idx(uint32_t il) const { } uint32_t llama_hparams::n_embd_r() const { + if (n_embd_r_impl != 0) { + // explicit override (e.g. inkling: 4 packed shortconv streams per layer) + return n_embd_r_impl; + } + if (wkv_head_size != 0) { // for RWKV models return token_shift_count * n_embd; diff --git a/src/llama-hparams.h b/src/llama-hparams.h index fc770bf003e6..c31955700919 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -80,6 +80,17 @@ struct llama_hparams { uint32_t n_shortconv_l_cache = 0; + // explicit override for the rolling state size per layer (see n_embd_r()) + uint32_t n_embd_r_impl = 0; + + // inkling (private arch) + uint32_t inkling_d_rel = 0; + uint32_t inkling_rel_extent = 0; // global (non-SWA) layers + uint32_t inkling_rel_extent_swa = 0; // local (SWA) layers + uint32_t inkling_log_n_floor = 0; // 0 = log-N scaling disabled + float inkling_log_alpha = 0.0f; + uint32_t inkling_unpadded_n_vocab = 0; // 0 = no padded-vocab masking + std::array n_head_arr; std::array n_head_kv_arr; std::array n_ff_arr; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 44cb1668dacf..76ba5dbabd68 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1353,6 +1353,51 @@ uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { return result; } +uint32_t llama_kv_cache::get_n_kv_pos_contiguous(const slot_info & sinfo, const llama_ubatch & ubatch) const { + if (sinfo.n_stream() != 1 || ubatch.n_seqs_unq != 1 || ubatch.n_tokens == 0 || + ubatch.pos == nullptr || ubatch.n_seq_id == nullptr || + ubatch.seq_id == nullptr || ubatch.seq_id[0] == nullptr) { + return 0; + } + + const llama_seq_id seq_id = ubatch.seq_id[0][0]; + if (seq_id < 0 || (size_t) seq_id >= seq_to_stream.size()) { + return 0; + } + + const uint32_t stream = seq_to_stream[seq_id]; + if (sinfo.strm[0] < 0 || (uint32_t) sinfo.strm[0] != stream || stream >= v_cells.size()) { + return 0; + } + + const auto & cells = v_cells[stream]; + const llama_pos pos_max = cells.seq_pos_max(seq_id); + + if (pos_max < 0 || pos_max >= (llama_pos) cells.size()) { + return 0; + } + + // the banded op aligns Q to the tail of K: the ubatch must be that monotonic tail, else dense bias + if ((uint32_t) pos_max + 1 < ubatch.n_tokens) { + return 0; + } + const llama_pos pos_start = pos_max + 1 - ubatch.n_tokens; + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + if (ubatch.pos[i] != pos_start + (llama_pos) i || + ubatch.n_seq_id[i] < 1 || ubatch.seq_id[i] == nullptr || ubatch.seq_id[i][0] != seq_id) { + return 0; + } + } + + for (llama_pos pos = 0; pos <= pos_max; ++pos) { + if (cells.is_empty(pos) || cells.pos_get(pos) != pos || !cells.seq_has(pos, seq_id)) { + return 0; + } + } + + return pos_max + 1; +} + ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { const int32_t ikv = map_layer_ids.at(il); @@ -1931,6 +1976,39 @@ void llama_kv_cache::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch } } +void llama_kv_cache::set_input_pos_rel_flat(ggml_tensor * dst, const llama_ubatch * ubatch, uint32_t extent) const { + const int64_t n_tokens = ubatch->n_tokens; + + GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); + + int32_t * data = (int32_t *) dst->data; + + const int64_t n_kv = dst->ne[0]; + GGML_ASSERT(dst->ne[1] == n_tokens); + + // [n_kv, n_tokens] in GLOBAL token order (stream-major, same as the KQ mask) + for (int64_t i = 0; i < n_tokens; ++i) { + const llama_seq_id seq_id = ubatch->seq_id[i][0]; + + const auto & cells = v_cells[seq_to_stream[seq_id]]; + + const llama_pos p1 = ubatch->pos[i]; + + for (int64_t j = 0; j < n_kv; ++j) { + // use the ACTUAL absolute position in the KV cell; physical slot order is not monotonic + int32_t rel = (int32_t) extent; // zero-bias column + if (!cells.is_empty(j)) { + const llama_pos d = p1 - cells.pos_get(j); + if (d >= 0 && d < (llama_pos) extent) { + rel = (int32_t) d; + } + } + + data[i*n_kv + j] = (int32_t) (i*(extent + 1)) + rel; + } + } +} + void llama_kv_cache::set_input_k_rot(ggml_tensor * dst) const { GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); @@ -2828,6 +2906,21 @@ uint32_t llama_kv_cache_context::get_n_kv() const { return n_kv; } +uint32_t llama_kv_cache_context::get_n_kv_pos_contiguous() const { + // Full-cache and update contexts do not carry a concrete ubatch/slot pair. + if (ubatches.empty() || sinfos.empty() || i_cur >= ubatches.size() || i_cur >= sinfos.size()) { + // reserve context: report the whole cache as position-contiguous so the worst-case graph + // is the banded path; reserving the dense fallback is unallocatable at large n_ctx + if (kv != nullptr && lctx == nullptr && kv->get_n_stream() == 1) { + return n_kv; + } + return 0; + } + + const uint32_t result = kv->get_n_kv_pos_contiguous(sinfos[i_cur], ubatches[i_cur]); + return result <= (uint32_t) n_kv ? result : 0; +} + ggml_type llama_kv_cache_context::type_k() const { return kv->type_k(); } @@ -2896,6 +2989,10 @@ void llama_kv_cache_context::set_input_pos_bucket(ggml_tensor * dst, const llama kv->set_input_pos_bucket(dst, ubatch); } +void llama_kv_cache_context::set_input_pos_rel_flat(ggml_tensor * dst, const llama_ubatch * ubatch, uint32_t extent) const { + kv->set_input_pos_rel_flat(dst, ubatch, extent); +} + void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const { kv->set_input_k_rot(dst); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index d5a92f4405b5..0ea921d8d4b9 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -170,6 +170,9 @@ class llama_kv_cache : public llama_memory_i { uint32_t get_n_kv(const slot_info & sinfo) const; + // active cell count when position p lives in physical cell p; 0 for any non-contiguous layout + uint32_t get_n_kv_pos_contiguous(const slot_info & sinfo, const llama_ubatch & ubatch) const; + // get views of the current state of the cache ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; @@ -216,6 +219,10 @@ class llama_kv_cache : public llama_memory_i { void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const; + // inkling: fill dst I32 [n_kv, n_tokens] with flat rel-bias gather indices, + // idx(i, j) = i*(extent + 1) + rel; empty/out-of-band cells map to the zero-bias column `extent` + void set_input_pos_rel_flat(ggml_tensor * dst, const llama_ubatch * ubatch, uint32_t extent) const; + void set_input_k_rot(ggml_tensor * dst) const; void set_input_v_rot(ggml_tensor * dst) const; @@ -371,6 +378,7 @@ class llama_kv_cache_context : public llama_memory_context_i { // uint32_t get_n_kv() const; + uint32_t get_n_kv_pos_contiguous() const; ggml_type type_k() const; ggml_type type_v() const; @@ -405,6 +413,7 @@ class llama_kv_cache_context : public llama_memory_context_i { void set_input_k_shift (ggml_tensor * dst) const; void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const; + void set_input_pos_rel_flat(ggml_tensor * dst, const llama_ubatch * ubatch, uint32_t extent) const; // inkling void set_input_k_rot(ggml_tensor * dst) const; void set_input_v_rot(ggml_tensor * dst) const; @@ -412,8 +421,8 @@ class llama_kv_cache_context : public llama_memory_context_i { private: llama_memory_status status; - llama_kv_cache * kv; - llama_context * lctx; + llama_kv_cache * kv = nullptr; + llama_context * lctx = nullptr; // // update context diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 3812c594e795..264b0b5b1411 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -29,6 +29,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) { case LLM_ARCH_STEP35: case LLM_ARCH_MELLUM: case LLM_ARCH_LAGUNA: + case LLM_ARCH_INKLING: return false; default: return true; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 7a70585fa4d8..cc39004dbd4d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -311,6 +311,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_kimi_linear(params); case LLM_ARCH_STEP35: return new llama_model_step35(params); + case LLM_ARCH_INKLING: + return new llama_model_inkling(params); default: throw std::runtime_error(std::string("unsupported model architecture: '") + llm_arch_name(arch) + "'"); } @@ -2114,7 +2116,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; - if (arch == LLM_ARCH_FALCON_H1) { + if (arch == LLM_ARCH_FALCON_H1 || arch == LLM_ARCH_INKLING) { + // all layers have both an attention KV cache and a recurrent (conv) state filter_attn = [&](uint32_t) { return true; }; filter_recr = [&](uint32_t) { return true; }; } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { @@ -2454,6 +2457,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_INKLING: return LLAMA_ROPE_TYPE_NONE; // use what we call a normal RoPE, operating on pairs of consecutive head values diff --git a/src/llama-model.h b/src/llama-model.h index 056a6efa59e8..4184111212f3 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -533,6 +533,15 @@ struct llama_layer { struct llama_layer_shortconv shortconv; struct llama_layer_nextn nextn; + + // inkling (private arch) + struct ggml_tensor * wr = nullptr; // attn_r [n_embd, n_head*d_rel] + struct ggml_tensor * attn_rel_proj = nullptr; // [rel_extent, d_rel] (checkpoint [d_rel, E] orientation) + struct ggml_tensor * shortconv_k = nullptr; // [K, kvw] + struct ggml_tensor * shortconv_v = nullptr; // [K, kvw] + struct ggml_tensor * shortconv_attn = nullptr; // [K, n_embd] + struct ggml_tensor * shortconv_mlp = nullptr; // [K, n_embd] + struct ggml_tensor * ffn_gscale = nullptr; // F32 [1] }; struct llama_device { diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 92ebc11b99f3..38f97b0cc7f1 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -329,6 +329,15 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param // do not quantize MiniMax's indexer projection weights, they are tiny quantize &= name.find("indexer.k_proj.weight") == std::string::npos; quantize &= name.find("indexer.q_proj.weight") == std::string::npos; + // keep Inkling's shortconv kernels and rel-proj table unquantized; arch-gated so the + // name substrings cannot hit another architecture + if (arch == LLM_ARCH_INKLING) { + quantize &= name.find("shortconv_k.weight") == std::string::npos; + quantize &= name.find("shortconv_v.weight") == std::string::npos; + quantize &= name.find("shortconv_attn.weight") == std::string::npos; + quantize &= name.find("shortconv_mlp.weight") == std::string::npos; + quantize &= name.find("attn_rel_proj.weight") == std::string::npos; + } // do not quantize RWKV's small yet 2D weights quantize &= name.find("time_mix_first.weight") == std::string::npos; diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 9164a4dd888d..74942e8126d9 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -433,6 +433,12 @@ struct llm_tokenizer_bpe : llm_tokenizer { "[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))*((?=[\\p{L}])([^A-Z]))+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}])([^a-z]))+((?=[\\p{L}])([^A-Z]))*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", }; break; + case LLAMA_VOCAB_PRE_TYPE_INKLING: + // o200k-family with \p{M} in the letter classes; own pre-type so GPT4O / MINIMAX_M2 stay unchanged + regex_exprs = { + "[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))*((?=[\\p{L}\\p{M}])([^A-Z]))+(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))+((?=[\\p{L}\\p{M}])([^A-Z]))*(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])?|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+", + }; + break; case LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI: // Same lookaheads as GPT4O but with \p{M} added so combining marks // (diacritics) attach to their base letters. Avoids excessive @@ -2295,6 +2301,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "talkie") { pre_type = LLAMA_VOCAB_PRE_TYPE_GPT4O; clean_spaces = false; + } else if ( + tokenizer_pre == "inkling") { + pre_type = LLAMA_VOCAB_PRE_TYPE_INKLING; + clean_spaces = false; } else if ( tokenizer_pre == "granite-embed-multi-97m") { pre_type = LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI; diff --git a/src/llama-vocab.h b/src/llama-vocab.h index b7c28926338b..65e43f671897 100644 --- a/src/llama-vocab.h +++ b/src/llama-vocab.h @@ -65,6 +65,7 @@ enum llama_vocab_pre_type { LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI = 54, LLAMA_VOCAB_PRE_TYPE_MELLUM2 = 55, LLAMA_VOCAB_PRE_TYPE_LAGUNA = 56, + LLAMA_VOCAB_PRE_TYPE_INKLING = 57, }; struct LLM_KV; diff --git a/src/models/inkling.cpp b/src/models/inkling.cpp new file mode 100644 index 000000000000..f83930c6976a --- /dev/null +++ b/src/models/inkling.cpp @@ -0,0 +1,666 @@ +// Inkling (PRIVATE arch): hybrid iSWA attention + per-layer packed shortconv state; see INKLING_DESIGN.md. + +#include "models.h" + +#include "../llama-kv-cache-iswa.h" +#include "../llama-kv-cache.h" +#include "../llama-memory-hybrid-iswa.h" +#include "../llama-memory-recurrent.h" + +#include + +void llama_model_inkling::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_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); + + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; // visible iff pos_q - pos_k < n_swa (includes self) + ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl, hparams.n_layer()); + + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + hparams.is_recr_impl[il] = 1; + } + + ml.get_key(LLM_KV_INKLING_D_REL, hparams.inkling_d_rel); + ml.get_key(LLM_KV_INKLING_REL_EXTENT, hparams.inkling_rel_extent); + ml.get_key(LLM_KV_INKLING_REL_EXTENT_SWA, hparams.inkling_rel_extent_swa); + ml.get_key(LLM_KV_INKLING_SHORTCONV_KERNEL, hparams.n_shortconv_l_cache); + ml.get_key(LLM_KV_INKLING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); + + float logit_scale_denom = 0.0f; + ml.get_key(LLM_KV_INKLING_LOGIT_SCALE_DENOM, logit_scale_denom); + GGML_ASSERT(logit_scale_denom != 0.0f); + hparams.f_logit_scale = 1.0f / logit_scale_denom; + + ml.get_key(LLM_KV_INKLING_LOG_SCALING_N_FLOOR, hparams.inkling_log_n_floor, false); + ml.get_key(LLM_KV_INKLING_LOG_SCALING_ALPHA, hparams.inkling_log_alpha, false); + ml.get_key(LLM_KV_INKLING_UNPADDED_VOCAB_SIZE, hparams.inkling_unpadded_n_vocab, false); + + GGML_ASSERT(hparams.n_shortconv_l_cache > 1); + GGML_ASSERT(hparams.inkling_d_rel > 0); + GGML_ASSERT(hparams.inkling_rel_extent > 0 && hparams.inkling_rel_extent_swa > 0); + + // uniform state per cell: 4 packed streams [k | v | attn | mlp] of last K-1 columns, k/v sized for the widest layer + const uint32_t d_conv = hparams.n_shortconv_l_cache - 1; + hparams.n_embd_r_impl = d_conv * (hparams.n_embd_k_gqa_max() + hparams.n_embd_v_gqa_max() + 2*hparams.n_embd); + + type = LLM_TYPE_UNKNOWN; +} + +void llama_model_inkling::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t head_dim = hparams.n_embd_head_k(); + const int64_t d_rel = hparams.inkling_d_rel; + const int64_t K = hparams.n_shortconv_l_cache; + const int64_t n_ff_exp = hparams.n_ff_exp; + const int64_t n_shexp = hparams.n_expert_shared; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + tok_norm = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD_NORM, "weight", 0), {n_embd}, 0); // bid 0: compute on the first layer's device + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + const int64_t n_head_kv_i = hparams.n_head_kv(i); + const int64_t kvw = n_head_kv_i * head_dim; + const int64_t rel_extent = hparams.is_swa(i) ? hparams.inkling_rel_extent_swa : hparams.inkling_rel_extent; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head*head_dim}, 0); + layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, kvw}, 0); + layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, kvw}, 0); + layer.wr = create_tensor(tn(LLM_TENSOR_ATTN_R, "weight", i), {n_embd, n_head*d_rel}, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head*head_dim, n_embd}, 0); + + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {head_dim}, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {head_dim}, 0); + + // stored in checkpoint orientation [d_rel, E] -> gguf ne = [E, d_rel] + layer.attn_rel_proj = create_tensor(tn(LLM_TENSOR_ATTN_REL_PROJ, "weight", i), {rel_extent, d_rel}, 0); + + layer.shortconv_k = create_tensor(tn(LLM_TENSOR_SHORTCONV_K, "weight", i), {K, kvw}, 0); + layer.shortconv_v = create_tensor(tn(LLM_TENSOR_SHORTCONV_V, "weight", i), {K, kvw}, 0); + layer.shortconv_attn = create_tensor(tn(LLM_TENSOR_SHORTCONV_ATTN, "weight", i), {K, n_embd}, 0); + layer.shortconv_mlp = create_tensor(tn(LLM_TENSOR_SHORTCONV_MLP, "weight", i), {K, n_embd}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_gscale = create_tensor(tn(LLM_TENSOR_FFN_GSCALE, "weight", i), {1}, 0); + + if (i < (int) hparams.n_layer_dense_lead) { + const int64_t n_ff_i = hparams.n_ff(i); + + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff_i}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff_i}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff_i, n_embd}, 0); + } else { + GGML_ASSERT(n_expert > 0 && n_expert_used > 0 && n_shexp > 0); + + // gate holds n_expert + n_shexp rows (incl. shared-expert sink logits) + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert + n_shexp}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); + + // shared experts stacked as an n_shexp bank, registered MUL_MAT_ID so the loader picks a mul_mat_id-capable buffer + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXPS, "weight", i), {n_embd, n_ff_exp, n_shexp}, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXPS, "weight", i), {n_embd, n_ff_exp, n_shexp}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXPS, "weight", i), {n_ff_exp, n_embd, n_shexp}, 0); + } + } +} + +class llm_graph_input_inkling : public llm_graph_input_i { +public: + llm_graph_input_inkling( + const llama_hparams & hparams, + const llama_memory_hybrid_iswa_context * mctx) : + hparams(hparams), + mctx(mctx) {} + virtual ~llm_graph_input_inkling() = default; + + void set_input(const llama_ubatch * ubatch) override { + if (tau) { + GGML_ASSERT(ggml_backend_buffer_is_host(tau->buffer)); + float * data = (float *) tau->data; + + const float n_floor = (float) hparams.inkling_log_n_floor; + const float alpha = hparams.inkling_log_alpha; + + for (int64_t i = 0; i < (int64_t) ubatch->n_tokens; ++i) { + const float eff = (float) (ubatch->pos[i] + 1) / n_floor; + data[i] = 1.0f + alpha*logf(std::max(eff, 1.0f)); + } + } + + if (rel_idx) { + mctx->get_attn()->get_base()->set_input_pos_rel_flat(rel_idx, ubatch, hparams.inkling_rel_extent); + } + + if (rel_idx_swa) { + mctx->get_attn()->get_swa()->set_input_pos_rel_flat(rel_idx_swa, ubatch, hparams.inkling_rel_extent_swa); + } + + if (vocab_mask) { + GGML_ASSERT(ggml_backend_buffer_is_host(vocab_mask->buffer)); + float * data = (float *) vocab_mask->data; + + const int64_t n_vocab = vocab_mask->ne[0]; + const int64_t n_unpadded = hparams.inkling_unpadded_n_vocab; + + for (int64_t id = 0; id < n_vocab; ++id) { + data[id] = id < n_unpadded ? 0.0f : -INFINITY; + } + } + + if (shexp_idx) { + GGML_ASSERT(ggml_backend_buffer_is_host(shexp_idx->buffer)); + int32_t * data = (int32_t *) shexp_idx->data; + + const int64_t n_shexp = shexp_idx->ne[0]; + const int64_t n_tokens = shexp_idx->ne[1]; + + for (int64_t j = 0; j < n_tokens; ++j) { + for (int64_t s = 0; s < n_shexp; ++s) { + data[j*n_shexp + s] = (int32_t) s; + } + } + } + } + + ggml_tensor * tau = nullptr; // F32 [1, 1, n_tokens] + ggml_tensor * rel_idx = nullptr; // I32 [n_kv_base, n_tokens] + ggml_tensor * rel_idx_swa = nullptr; // I32 [n_kv_swa, n_tokens] + ggml_tensor * vocab_mask = nullptr; // F32 [n_vocab] + ggml_tensor * shexp_idx = nullptr; // I32 [n_shexp, n_tokens], constant 0..n_shexp-1 + + const llama_hparams hparams; + + const llama_memory_hybrid_iswa_context * mctx; +}; + +std::unique_ptr llama_model_inkling::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_inkling::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params) { + + const int64_t head_dim = hparams.n_embd_head_k(); + const int64_t d_rel = hparams.inkling_d_rel; + const int64_t d_conv = hparams.n_shortconv_l_cache - 1; + const int64_t n_embd_r = hparams.n_embd_r(); + const int64_t kw_max = hparams.n_embd_k_gqa_max(); + const int64_t vw_max = hparams.n_embd_v_gqa_max(); + + // packed conv-state stream offsets within one cell: [k | v | attn | mlp] + const int64_t off_k = 0; + const int64_t off_v = d_conv*kw_max; + const int64_t off_attn = d_conv*(kw_max + vw_max); + const int64_t off_mlp = d_conv*(kw_max + vw_max + n_embd); + + const auto * mctx_hyb = static_cast(mctx); + const auto * mctx_recr = mctx_hyb->get_recr(); + const auto * mctx_attn = mctx_hyb->get_attn(); + + const uint32_t kv_head = mctx_recr->get_head(); + + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + const int64_t n_seqs = ubatch.n_seqs; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + const uint32_t n_kv_flash_base = cparams.flash_attn ? mctx_attn->get_base()->get_n_kv_pos_contiguous() : 0; + const uint32_t n_kv_flash_swa = cparams.flash_attn ? mctx_attn->get_swa ()->get_n_kv_pos_contiguous() : 0; + + bool has_global = false; + bool needs_rel_idx_local = false; + bool needs_rel_idx_global = false; + + const auto banded_cache_type_supported = [](ggml_type type) { + return type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || type == GGML_TYPE_BF16; + }; + + const auto use_banded_flash = [&](int il) { + const auto * cache = hparams.is_swa(il) ? mctx_attn->get_swa() : mctx_attn->get_base(); + const uint32_t n_kv_flash = hparams.is_swa(il) ? n_kv_flash_swa : n_kv_flash_base; + + // get_n_kv_pos_contiguous() is 0 for multi-sequence ubatches; the reserve context reports full n_kv + return cparams.flash_attn && + n_kv_flash > 0 && + (head_dim == 64 || head_dim == 128) && + hparams.n_embd_head_v(il) == head_dim && + hparams.n_head(il) % hparams.n_head_kv(il) == 0 && + banded_cache_type_supported(cache->type_k()) && + banded_cache_type_supported(cache->type_v()); + }; + + for (int il = 0; il < n_layer; ++il) { + if (hparams.is_swa(il)) { + needs_rel_idx_local |= !use_banded_flash(il); + } else { + has_global = true; + needs_rel_idx_global |= !use_banded_flash(il); + } + } + + auto inp = std::make_unique(hparams, mctx_hyb); + + if (hparams.inkling_log_n_floor > 0 && has_global) { + inp->tau = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tokens); + ggml_set_input(inp->tau); + ggml_set_name(inp->tau, "inkling_tau"); + } + + if (needs_rel_idx_global) { + const int64_t n_kv = mctx_attn->get_base()->get_n_kv(); + inp->rel_idx = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_tokens); + ggml_set_input(inp->rel_idx); + ggml_set_name(inp->rel_idx, "inkling_rel_idx"); + } + + if (needs_rel_idx_local) { + const int64_t n_kv_swa = mctx_attn->get_swa()->get_n_kv(); + inp->rel_idx_swa = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv_swa, n_tokens); + ggml_set_input(inp->rel_idx_swa); + ggml_set_name(inp->rel_idx_swa, "inkling_rel_idx_swa"); + } + + const int64_t n_vocab = model.vocab.n_tokens(); + if (!cparams.embeddings && hparams.inkling_unpadded_n_vocab > 0 && (int64_t) hparams.inkling_unpadded_n_vocab < n_vocab) { + inp->vocab_mask = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_vocab); + ggml_set_input(inp->vocab_mask); + ggml_set_name(inp->vocab_mask, "inkling_vocab_mask"); + } + + // shared experts go through mul_mat_id: 2D views into a repacked/quantized 3D bank are invalid + if (hparams.n_expert_shared > 0 && (uint32_t) n_layer > hparams.n_layer_dense_lead) { + inp->shexp_idx = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, hparams.n_expert_shared, n_tokens); + ggml_set_input(inp->shexp_idx); + ggml_set_name(inp->shexp_idx, "inkling_shexp_idx"); + } + + ggml_tensor * tau = inp->tau; + ggml_tensor * rel_idx = inp->rel_idx; + ggml_tensor * rel_idx_swa = inp->rel_idx_swa; + ggml_tensor * vocab_mask = inp->vocab_mask; + ggml_tensor * shexp_idx = inp->shexp_idx; + + res->add_input(std::move(inp)); + + auto * inp_hybrid = build_inp_mem_hybrid_iswa(); + + // shared by the 4 stream sub-views; build_rs must run exactly once per layer (it zero-inits fresh states) + ggml_tensor * conv_rs_cur = nullptr; + + // sconv(x) = x + causal_depthwise_conv1d(x); rolling state = last K-1 inputs + auto build_sconv = [&](ggml_tensor * x2d, ggml_tensor * kernel, int64_t off, int il) -> ggml_tensor * { + const int64_t w = x2d->ne[0]; + + ggml_tensor * x3 = ggml_reshape_3d(ctx0, x2d, w, n_seq_tokens, n_seqs); + ggml_tensor * xt = ggml_transpose(ctx0, x3); // time-major for the conv + + ggml_tensor * conv_state = mctx_recr->get_r_l(il); + ggml_tensor * conv_rs = conv_rs_cur; // {n_embd_r, n_seqs} + GGML_ASSERT(conv_rs != nullptr); + + const size_t sz = ggml_element_size(conv_rs); + + // this stream's slice of the packed state + ggml_tensor * state = ggml_view_3d(ctx0, conv_rs, d_conv, w, n_seqs, + d_conv*sz, conv_rs->nb[1], off*sz); + + ggml_tensor * sx = ggml_concat(ctx0, state, xt, 0); // {d_conv + n_seq_tokens, w, n_seqs} + + // write the last d_conv time columns back into the cache + ggml_tensor * new_state = ggml_view_3d(ctx0, sx, d_conv, w, n_seqs, + sx->nb[1], sx->nb[2], (sx->ne[0] - d_conv)*sx->nb[0]); + ggml_tensor * state_dst = ggml_view_3d(ctx0, conv_state, d_conv, w, n_seqs, + d_conv*sz, n_embd_r*sz, (kv_head*n_embd_r + off)*sz); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, new_state, state_dst)); + + ggml_tensor * conv_out = ggml_ssm_conv(ctx0, sx, kernel); // {w, n_seq_tokens, n_seqs} + + ggml_tensor * y = ggml_add(ctx0, x3, conv_out); // built-in residual, no activation + + return ggml_reshape_2d(ctx0, y, w, n_seq_tokens*n_seqs); + }; + + auto build_attn_block = [&](ggml_tensor * cur, int il) -> ggml_tensor * { + const auto & layer = model.layers[il]; + + const bool is_swa = hparams.is_swa(il); + const int64_t n_head_kv = hparams.n_head_kv(il); + const int64_t rel_extent = is_swa ? hparams.inkling_rel_extent_swa : hparams.inkling_rel_extent; + + ggml_tensor * q = build_lora_mm(layer.wq, cur); + ggml_tensor * k = build_lora_mm(layer.wk, cur); + ggml_tensor * v = build_lora_mm(layer.wv, cur); + ggml_tensor * r = build_lora_mm(layer.wr, cur); + cb(q, "inkling_attn_q", il); + cb(k, "inkling_attn_k", il); + cb(v, "inkling_attn_v", il); + cb(r, "inkling_attn_r", il); + + // k/v short convs on the flat projections, before the head reshape + k = build_sconv(k, layer.shortconv_k, off_k, il); + v = build_sconv(v, layer.shortconv_v, off_v, il); + cb(k, "inkling_attn_k_sconv", il); + cb(v, "inkling_attn_v_sconv", il); + + q = ggml_reshape_3d(ctx0, q, head_dim, n_head, n_tokens); + k = ggml_reshape_3d(ctx0, k, head_dim, n_head_kv, n_tokens); + v = ggml_reshape_3d(ctx0, v, head_dim, n_head_kv, n_tokens); + + q = build_norm(q, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); + k = build_norm(k, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + cb(q, "inkling_attn_q_norm", il); + cb(k, "inkling_attn_k_norm", il); + + // log-N tau on global layers only, after q_norm + if (tau && !is_swa) { + q = ggml_mul(ctx0, q, tau); + } + + // relative position bias + ggml_tensor * r2 = ggml_reshape_2d(ctx0, r, d_rel, n_head*n_tokens); + + // proj stored [E, d_rel]; transpose so ggml_mul_mat contracts over d_rel + ggml_tensor * proj = ggml_cont(ctx0, ggml_transpose(ctx0, layer.attn_rel_proj)); // {d_rel, E} + + ggml_tensor * rel = ggml_mul_mat(ctx0, proj, r2); // {E, n_head*n_tokens} + ggml_mul_mat_set_prec(rel, GGML_PREC_F32_PEDANTIC); + rel = ggml_reshape_3d(ctx0, rel, rel_extent, n_head, n_tokens); + + if (tau && !is_swa) { + rel = ggml_mul(ctx0, rel, tau); + } + cb(rel, "inkling_rel_logits", il); + + auto * inp_attn = inp_hybrid->get_attn(); + const int64_t n_stream = (is_swa ? inp_attn->get_kq_mask_swa() : inp_attn->get_kq_mask())->ne[3]; + GGML_ASSERT(n_tokens % n_stream == 0); + + if (use_banded_flash(il)) { + GGML_ASSERT(q->type == GGML_TYPE_F32); + auto * k_rot = is_swa ? inp_attn->self_k_rot_swa : inp_attn->self_k_rot; + auto * v_rot = is_swa ? inp_attn->self_v_rot_swa : inp_attn->self_v_rot; + + if (k_rot) { + q = llama_mul_mat_hadamard(ctx0, q, k_rot); + k = llama_mul_mat_hadamard(ctx0, k, k_rot); + } + if (v_rot) { + v = llama_mul_mat_hadamard(ctx0, v, v_rot); + } + + ggml_build_forward_expand(gf, q); + ggml_build_forward_expand(gf, k); + ggml_build_forward_expand(gf, v); + + const auto * cache = is_swa ? inp_attn->mctx->get_swa() : inp_attn->mctx->get_base(); + const auto & k_idxs = is_swa ? inp_attn->get_k_idxs_swa() : inp_attn->get_k_idxs(); + const auto & v_idxs = is_swa ? inp_attn->get_v_idxs_swa() : inp_attn->get_v_idxs(); + + ggml_build_forward_expand(gf, cache->cpy_k(ctx0, k, k_idxs, il)); + ggml_build_forward_expand(gf, cache->cpy_v(ctx0, v, v_idxs, il)); + + ggml_tensor * q_fa = ggml_view_4d(ctx0, q, + q->ne[0], q->ne[1], q->ne[2]/n_stream, n_stream, + q->nb[1], q->nb[2], q->nb[3]/n_stream, 0); + ggml_tensor * k_fa = cache->get_k(ctx0, il); + ggml_tensor * v_fa = cache->get_v(ctx0, il); + + const int64_t n_kv_flash = is_swa ? n_kv_flash_swa : n_kv_flash_base; + GGML_ASSERT(n_stream == 1 && n_kv_flash <= k_fa->ne[2]); + + k_fa = ggml_view_4d(ctx0, k_fa, + k_fa->ne[0], k_fa->ne[1], n_kv_flash, k_fa->ne[3], + k_fa->nb[1], k_fa->nb[2], k_fa->nb[3], 0); + + const bool v_trans = v_fa->nb[1] > v_fa->nb[2]; + if (v_trans) { + GGML_ASSERT(n_kv_flash <= v_fa->ne[0]); + v_fa = ggml_view_4d(ctx0, v_fa, + n_kv_flash, v_fa->ne[1], v_fa->ne[2], v_fa->ne[3], + v_fa->nb[1], v_fa->nb[2], v_fa->nb[3], 0); + } else { + GGML_ASSERT(n_kv_flash <= v_fa->ne[2]); + v_fa = ggml_view_4d(ctx0, v_fa, + v_fa->ne[0], v_fa->ne[1], n_kv_flash, v_fa->ne[3], + v_fa->nb[1], v_fa->nb[2], v_fa->nb[3], 0); + } + + q_fa = ggml_permute(ctx0, q_fa, 0, 2, 1, 3); + k_fa = ggml_permute(ctx0, k_fa, 0, 2, 1, 3); + v_fa = ggml_permute(ctx0, v_fa, 0, 2, 1, 3); + + if (v_trans) { + v_fa = ggml_transpose(ctx0, v_fa); + } + if (k_fa->type == GGML_TYPE_F32) { + k_fa = ggml_cast(ctx0, k_fa, GGML_TYPE_F16); + } + if (v_fa->type == GGML_TYPE_F32) { + v_fa = ggml_cast(ctx0, v_fa, GGML_TYPE_F16); + } + + ggml_tensor * rel_fa = ggml_reshape_4d(ctx0, rel, + rel_extent, n_head, n_tokens/n_stream, n_stream); + ggml_tensor * mask = is_swa ? inp_attn->get_kq_mask_swa() : inp_attn->get_kq_mask(); + mask = ggml_cont(ctx0, ggml_view_4d(ctx0, mask, + n_kv_flash, mask->ne[1], mask->ne[2], mask->ne[3], + mask->nb[1], mask->nb[2], mask->nb[3], 0)); + + cur = ggml_flash_attn_ext_banded(ctx0, q_fa, k_fa, v_fa, mask, rel_fa, + 1.0f/float(head_dim), rel_extent); + ggml_flash_attn_ext_set_prec(cur, GGML_PREC_F32); + res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); + + cur = ggml_reshape_2d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2]*cur->ne[3]); + ggml_build_forward_expand(gf, cur); + cb(cur, "kqv_out", il); + + if (v_rot) { + cur = llama_mul_mat_hadamard(ctx0, cur, v_rot); + } + cur = build_lora_mm(layer.wo, cur); + } else { + // soft_max_ext scales kq + kq_b jointly: fold 1/head_dim into q to keep the bias unscaled + q = ggml_scale(ctx0, q, 1.0f/float(head_dim)); + + // zero column at index E is gathered by out-of-band / empty-cell indices + rel = ggml_pad(ctx0, rel, 1, 0, 0, 0); // {E+1, n_head, n_tokens} + rel = ggml_cont(ctx0, ggml_permute(ctx0, rel, 1, 0, 2, 3)); // {n_head, E+1, n_tokens} + rel = ggml_reshape_2d(ctx0, rel, n_head, (rel_extent + 1)*n_tokens); + + ggml_tensor * idx = is_swa ? rel_idx_swa : rel_idx; // {n_kv, n_tokens} + GGML_ASSERT(idx != nullptr); + const int64_t n_kv = idx->ne[0]; + + ggml_tensor * idx1 = ggml_reshape_1d(ctx0, idx, n_kv*n_tokens); + + ggml_tensor * kq_b = ggml_get_rows(ctx0, rel, idx1); // {n_head, n_kv*n_tokens} + kq_b = ggml_reshape_3d(ctx0, kq_b, n_head, n_kv, n_tokens); + kq_b = ggml_cont(ctx0, ggml_permute(ctx0, kq_b, 2, 0, 1, 3)); // {n_kv, n_tokens, n_head} + cb(kq_b, "inkling_kq_b", il); + + // streamed kq is [n_kv, n_tokens/n_stream, n_head, n_stream], tokens stream-major: view kq_b to match (same trick as the KQ mask) + if (n_stream > 1) { + kq_b = ggml_view_4d(ctx0, kq_b, n_kv, n_tokens/n_stream, n_head, n_stream, + kq_b->nb[1], + kq_b->nb[2], + (n_tokens/n_stream)*kq_b->nb[1], + 0); + } + + cur = build_attn(inp_attn, + layer.wo, NULL, NULL, + q, k, v, kq_b, nullptr, nullptr, 1.0f, il); + } + cb(cur, "inkling_attn_o", il); + + return cur; + }; + + auto build_dense_ffn = [&](ggml_tensor * cur, int il) -> ggml_tensor * { + cur = build_ffn(cur, + model.layers[il].ffn_up, NULL, NULL, + model.layers[il].ffn_gate, NULL, NULL, + model.layers[il].ffn_down, NULL, NULL, + NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); + cur = ggml_mul(ctx0, cur, model.layers[il].ffn_gscale); + cb(cur, "inkling_dense_ffn_out", il); + return cur; + }; + + // custom MoE routing (not expressible via build_moe_ffn): select by top-k(sigmoid(logits) + bias), weight by softmax(logsigmoid(raw logits)) * scales + auto build_moe = [&](ggml_tensor * cur, int il) -> ggml_tensor * { + const auto & layer = model.layers[il]; + + const int64_t n_shexp = hparams.n_expert_shared; + + ggml_tensor * logits = build_lora_mm( + layer.ffn_gate_inp, cur, nullptr, GGML_PREC_F32_PEDANTIC); // {n_expert + n_shexp, n_tokens} + cb(logits, "inkling_moe_logits", il); + + const size_t lsz = ggml_element_size(logits); + + ggml_tensor * routed = ggml_cont(ctx0, ggml_view_2d(ctx0, logits, n_expert, n_tokens, logits->nb[1], 0)); + ggml_tensor * shared_logits = ggml_view_2d(ctx0, logits, n_shexp, n_tokens, logits->nb[1], n_expert*lsz); + + // bias affects selection only, not the weights + ggml_tensor * scores = ggml_sigmoid(ctx0, routed); + scores = ggml_add(ctx0, scores, layer.ffn_exp_probs_b); + cb(scores, "inkling_moe_scores", il); + + ggml_tensor * selected = ggml_argsort_top_k(ctx0, scores, n_expert_used); // I32 {n_expert_used, n_tokens} + cb(selected, "inkling_moe_topk", il); + + // weights use the raw top-k logits, not the biased scores + ggml_tensor * routed3 = ggml_reshape_3d(ctx0, routed, 1, n_expert, n_tokens); + ggml_tensor * topk_logits = ggml_get_rows(ctx0, routed3, selected); // {1, n_expert_used, n_tokens} + topk_logits = ggml_reshape_2d(ctx0, topk_logits, n_expert_used, n_tokens); + + ggml_tensor * all_logits = ggml_concat(ctx0, topk_logits, shared_logits, 0); // {n_expert_used + n_shexp, n_tokens} + + // logsigmoid(x) = -softplus(-x) + ggml_tensor * w = ggml_neg(ctx0, ggml_softplus(ctx0, ggml_neg(ctx0, all_logits))); + w = ggml_soft_max(ctx0, w); + w = ggml_scale(ctx0, w, hparams.expert_weights_scale); + w = ggml_mul(ctx0, w, layer.ffn_gscale); // gate global_scale (F32 [1]) + cb(w, "inkling_moe_weights", il); + + const size_t wsz = ggml_element_size(w); + + ggml_tensor * weights = ggml_cont(ctx0, ggml_view_2d(ctx0, w, n_expert_used, n_tokens, w->nb[1], 0)); + weights = ggml_reshape_3d(ctx0, weights, 1, n_expert_used, n_tokens); + + ggml_tensor * xr = ggml_reshape_3d(ctx0, cur, n_embd, 1, n_tokens); + ggml_tensor * gate = build_lora_mm_id(layer.ffn_gate_exps, xr, selected); // {n_ff_exp, n_expert_used, n_tokens} + ggml_tensor * up = build_lora_mm_id(layer.ffn_up_exps, xr, selected); + ggml_tensor * h = ggml_swiglu_split(ctx0, gate, up); + + ggml_tensor * experts = build_lora_mm_id(layer.ffn_down_exps, h, selected); // {n_embd, n_expert_used, n_tokens} + experts = ggml_mul(ctx0, experts, weights); + + ggml_tensor * moe_out = nullptr; + for (int64_t i = 0; i < n_expert_used; ++i) { + ggml_tensor * e = ggml_view_2d(ctx0, experts, n_embd, n_tokens, experts->nb[2], i*experts->nb[1]); + moe_out = moe_out ? ggml_add(ctx0, moe_out, e) : e; + } + + // shared experts: mul_mat_id with constant ids (never 2D-view a quantized/repacked weight) + GGML_ASSERT(shexp_idx != nullptr); + ggml_tensor * gs = build_lora_mm_id(layer.ffn_gate_shexp, xr, shexp_idx); // {n_ff_exp, n_shexp, n_tokens} + ggml_tensor * us = build_lora_mm_id(layer.ffn_up_shexp, xr, shexp_idx); + ggml_tensor * hs = ggml_swiglu_split(ctx0, gs, us); + + // gammas (last n_shexp weight rows) must scale hs BEFORE the down-proj to match reference rounding in bf16/quant + ggml_tensor * gammas = ggml_cont(ctx0, ggml_view_2d(ctx0, w, n_shexp, n_tokens, w->nb[1], n_expert_used*wsz)); + hs = ggml_mul(ctx0, hs, ggml_reshape_3d(ctx0, gammas, 1, n_shexp, n_tokens)); + ggml_tensor * ds = build_lora_mm_id(layer.ffn_down_shexp, hs, shexp_idx); // {n_embd, n_shexp, n_tokens} + + for (int64_t s = 0; s < n_shexp; ++s) { + ggml_tensor * e = ggml_view_2d(ctx0, ds, n_embd, n_tokens, ds->nb[2], s*ds->nb[1]); + moe_out = ggml_add(ctx0, moe_out, e); + } + cb(moe_out, "inkling_moe_out", il); + + return moe_out; + }; + + ggml_tensor * cur = build_inp_embd(model.tok_embd); + // mtmd embd rows arrive pre-normalized; embed_norm applies to text token lookups only + if (ubatch.token) { + cur = build_norm(cur, model.tok_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "inkling_embd_norm", -1); + } else { + cb(cur, "inkling_mm_embd", -1); + } + + ggml_build_forward_expand(gf, cur); + + for (int il = 0; il < n_layer; ++il) { + conv_rs_cur = build_rs(inp_hybrid->get_recr(), mctx_recr->get_r_l(il), n_embd_r, n_seqs); + + // h += attn_sconv(attn(attn_norm(h))) + ggml_tensor * attn_in = build_norm(cur, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(attn_in, "inkling_attn_norm", il); + ggml_tensor * attn_out = build_attn_block(attn_in, il); + attn_out = build_sconv(attn_out, model.layers[il].shortconv_attn, off_attn, il); + cb(attn_out, "inkling_attn_sconv", il); + + cur = ggml_add(ctx0, cur, attn_out); + + // h += mlp_sconv(mlp(mlp_norm(h))) + ggml_tensor * ffn_in = build_norm(cur, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); + cb(ffn_in, "inkling_ffn_norm", il); + ggml_tensor * ffn_out = il < (int) hparams.n_layer_dense_lead ? + build_dense_ffn(ffn_in, il) : build_moe(ffn_in, il); + ffn_out = build_sconv(ffn_out, model.layers[il].shortconv_mlp, off_mlp, il); + cb(ffn_out, "inkling_ffn_sconv", il); + + cur = ggml_add(ctx0, cur, ffn_out); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + } + + // conv states need every layer to see ALL tokens, so trim outputs only after the full stack + ggml_tensor * inp_out_ids = build_inp_out_ids(); + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + if (!cparams.embeddings) { + cur = ggml_scale(ctx0, cur, hparams.f_logit_scale); + cur = build_lora_mm( + model.output, cur, nullptr, + model.output->type == GGML_TYPE_F32 ? GGML_PREC_F32_PEDANTIC : GGML_PREC_DEFAULT); + + // padded vocab rows get -inf so samplers never emit a padded id + if (vocab_mask) { + cur = ggml_add(ctx0, cur, vocab_mask); + } + cb(cur, "result_output", -1); + res->t_logits = cur; + } + + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index 92ebfafa1e29..cadea985dcbd 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1866,6 +1866,19 @@ struct llama_model_lfm2moe : public llama_model_base { }; +struct llama_model_inkling : public llama_model_base { + llama_model_inkling(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_smallthinker : public llama_model_base { llama_model_smallthinker(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 7a93b19a0765..8fda15806815 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -265,6 +265,14 @@ if (NOT LLAMA_SANITIZE_ADDRESS AND NOT GGML_SCHED_NO_REALLOC) endif() llama_build_and_test(test-backend-ops.cpp) +add_executable(test-flash-attn-bias test-flash-attn-bias.cpp) +target_link_libraries(test-flash-attn-bias PRIVATE ggml) +add_test(NAME test-flash-attn-bias COMMAND test-flash-attn-bias) + +add_executable(test-flash-attn-generic-hash test-flash-attn-generic-hash.cpp) +target_link_libraries(test-flash-attn-generic-hash PRIVATE ggml) +add_test(NAME test-flash-attn-generic-hash COMMAND test-flash-attn-generic-hash) + llama_build_and_test(test-model-load-cancel.cpp LABEL "model") llama_build_and_test(test-autorelease.cpp LABEL "model") llama_build_and_test(test-backend-sampler.cpp LABEL "model") diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index a5b660f47a0a..3dbac4eb3e75 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -6924,6 +6924,101 @@ struct test_flash_attn_ext : public test_case { } }; +// GGML_OP_FLASH_ATTN_EXT_BANDED +struct test_flash_attn_ext_banded : public test_case { + const int64_t d; + const int64_t n_head; + const int64_t n_head_kv; + const int64_t n_q; + const int64_t n_kv; + const int64_t rel_extent; + const int mask_kind; // 0: none, 1: causal, 2: causal sliding window with dist < rel_extent + const ggml_type kv_type; + const ggml_type rel_type; + const bool strided; + + std::string vars() override { + return VARS_TO_STR10(d, n_head, n_head_kv, n_q, n_kv, rel_extent, mask_kind, kv_type, rel_type, strided); + } + + double max_nmse_err() override { + if (kv_type == GGML_TYPE_F32 && rel_type == GGML_TYPE_F32) { + return 2e-6; + } + // fp16 VKQ accumulation error grows with the KV length (plus periodic accumulator + // rescales guarding against fp16 overflow) + return n_kv > 8192 ? 1e-3 : 5e-4; + } + + uint64_t op_flops(ggml_tensor * t) override { + GGML_UNUSED(t); + return 4*n_head*n_q*n_kv*d; + } + + test_flash_attn_ext_banded( + int64_t d, int64_t n_head, int64_t n_head_kv, + int64_t n_q, int64_t n_kv, int64_t rel_extent, + int mask_kind, ggml_type kv_type, ggml_type rel_type, bool strided = false) + : d(d), n_head(n_head), n_head_kv(n_head_kv), n_q(n_q), n_kv(n_kv), + rel_extent(rel_extent), mask_kind(mask_kind), kv_type(kv_type), rel_type(rel_type), strided(strided) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d, n_q, n_head, 1); + ggml_tensor * k; + ggml_tensor * v; + ggml_tensor * r; + if (strided) { + // gaps between rows/heads force the kernels to use the 64-bit byte strides + ggml_tensor * kb = ggml_new_tensor_4d(ctx, kv_type, 2*d, n_kv, n_head_kv, 1); + ggml_tensor * vb = ggml_new_tensor_4d(ctx, kv_type, 2*d, n_kv, n_head_kv, 1); + ggml_tensor * rb = ggml_new_tensor_4d(ctx, rel_type, 2*rel_extent, n_head, n_q, 1); + k = ggml_view_4d(ctx, kb, d, n_kv, n_head_kv, 1, kb->nb[1], kb->nb[2], kb->nb[3], 0); + v = ggml_view_4d(ctx, vb, d, n_kv, n_head_kv, 1, vb->nb[1], vb->nb[2], vb->nb[3], 0); + r = ggml_view_4d(ctx, rb, rel_extent, n_head, n_q, 1, rb->nb[1], rb->nb[2], rb->nb[3], 0); + } else { + k = ggml_new_tensor_4d(ctx, kv_type, d, n_kv, n_head_kv, 1); + v = ggml_new_tensor_4d(ctx, kv_type, d, n_kv, n_head_kv, 1); + r = ggml_new_tensor_4d(ctx, rel_type, rel_extent, n_head, n_q, 1); + } + ggml_set_name(q, "q"); + ggml_set_name(k, "k"); + ggml_set_name(v, "v"); + ggml_set_name(r, "rel_logits"); + + ggml_tensor * m = nullptr; + if (mask_kind != 0) { + m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_q, 1, 1); + ggml_set_name(m, "m"); + } + + ggml_tensor * out = ggml_flash_attn_ext_banded(ctx, q, k, v, m, r, 1.0f/float(d), rel_extent); + ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "m") == 0) { + std::vector data(n_q*n_kv); + for (int64_t iq = 0; iq < n_q; ++iq) { + for (int64_t ik = 0; ik < n_kv; ++ik) { + const int64_t rel_dist = iq + (n_kv - n_q) - ik; + const bool visible = rel_dist >= 0 && (mask_kind == 1 || rel_dist < rel_extent); + data[iq*n_kv + ik] = ggml_fp32_to_fp16(visible ? 0.0f : -INFINITY); + } + } + ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(data[0])); + } else if (strcmp(t->name, "rel_logits") == 0) { + // A larger range makes rel_dist = E versus E-1 mistakes immediately visible. + init_tensor_uniform(t, -1.0f, 1.0f); + } else { + init_tensor_uniform(t, -0.25f, 0.25f); + } + } + } +}; + // GGML_OP_CROSS_ENTROPY_LOSS struct test_cross_entropy_loss : public test_case { const ggml_type type; @@ -7503,6 +7598,7 @@ struct test_generic_op : public test_case { case GGML_OP_RWKV_WKV7: return 5e-3; case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_FLASH_ATTN_EXT_BANDED: { // Scale error with kv length to account for accumulating floating point error const int64_t kv = sources[1].ne[1]; @@ -7526,7 +7622,7 @@ struct test_generic_op : public test_case { } // FLASH_ATTN_EXT: src[3] is the KQ mask - if (op == GGML_OP_FLASH_ATTN_EXT && i == 3) { + if ((op == GGML_OP_FLASH_ATTN_EXT || op == GGML_OP_FLASH_ATTN_EXT_BANDED) && i == 3) { init_tensor_kq_mask(t); continue; } @@ -9568,6 +9664,21 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(128, 128, 8, {4, 1}, kv, 512, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); } + // banded score-bias coverage: band edges, masks, decode offset, GQA, head sizes, table types + test_cases.emplace_back(new test_flash_attn_ext_banded( 64, 2, 1, 8, 8, 8, 1, GGML_TYPE_F32, GGML_TYPE_F32)); + test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 2, 16, 16, 8, 1, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext_banded( 64, 8, 2, 1, 64, 8, 1, GGML_TYPE_F32, GGML_TYPE_F32)); + test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 2, 64, 64, 8, 2, GGML_TYPE_BF16, GGML_TYPE_BF16)); + test_cases.emplace_back(new test_flash_attn_ext_banded( 64, 8, 1, 64, 64, 512, 1, GGML_TYPE_F16, GGML_TYPE_F32)); + test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 2, 17, 33, 8, 1, GGML_TYPE_F16, GGML_TYPE_F16, true)); + // production-scale n_kv straddling the observed ~16.4-16.9K garbage threshold + test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 8192, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); + test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 16384, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); + test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 16403, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); + test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 16896, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); + test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 17024, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); + test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 1, 17024, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); + test_cases.emplace_back(new test_flash_attn_ext_banded(128, 8, 1, 512, 32768, 1024, 1, GGML_TYPE_F16, GGML_TYPE_F32)); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 01b07953a627..d0f985049c89 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -2996,6 +2996,85 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .run(); } + { + // Inkling / TML typed content blocks: <|end_message|> separates blocks, <|content_model_end_sampling|> ends the turn. + auto tst = peg_tester("models/templates/Inkling.jinja", detailed_debug); + + // reasoning + visible answer + tst.test("<|content_thinking|>I'm\nthinking<|end_message|>" + "<|message_model|><|content_text|>Hello, world!\nWhat's up?<|end_message|>" + "<|content_model_end_sampling|>") + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .expect(message_assist_thoughts) + .run(); + + // Visible answer only (reasoning_effort=0 -> no thinking block). + tst.test("<|content_text|>Hello, world!\nWhat's up?<|end_message|>" + "<|content_model_end_sampling|>") + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .expect(message_assist) + .run(); + + // Empty thinking block, then the answer. + tst.test("<|content_thinking|><|end_message|>" + "<|message_model|><|content_text|>Hello, world!\nWhat's up?<|end_message|>" + "<|content_model_end_sampling|>") + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .expect(message_assist) + .run(); + + // single tool call with reasoning; the bare tool-name echo and role opener are dropped + tst.test("<|content_thinking|>I'm\nthinking<|end_message|>" + "<|message_model|>special_function<|content_invoke_tool_json|>" + "{\"name\":\"special_function\",\"args\":{\"arg1\":1}}<|end_message|>" + "<|content_model_end_sampling|>") + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ special_function_tool }) + .expect(message_with_reasoning_and_tool_call("I'm\nthinking", "special_function", "{\"arg1\": 1}")) + .run(); + + // Tool call, no reasoning. + tst.test("<|message_model|>special_function<|content_invoke_tool_json|>" + "{\"name\":\"special_function\",\"args\":{\"arg1\":1}}<|end_message|>" + "<|content_model_end_sampling|>") + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ special_function_tool }) + .expect(message_assist_call) + .run(); + + // regression: the tool branch must not swallow a pure-text answer + tst.test("<|content_thinking|>I'm\nthinking<|end_message|>" + "<|message_model|><|content_text|>Hello, world!\nWhat's up?<|end_message|>" + "<|content_model_end_sampling|>") + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ special_function_tool }) + .expect(message_assist_thoughts) + .run(); + + // tools available, content only, no thinking (reasoning_effort=0 leak scenario) + tst.test("<|content_text|>Hello, world!\nWhat's up?<|end_message|>" + "<|content_model_end_sampling|>") + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .tools({ special_function_tool }) + .expect(message_assist) + .run(); + + // parallel tool calls are separate marker-wrapped blocks + tst.test("<|message_model|>special_function<|content_invoke_tool_json|>" + "{\"name\":\"special_function\",\"args\":{\"arg1\":1}}<|end_message|>" + "<|message_model|>python<|content_invoke_tool_json|>" + "{\"name\":\"python\",\"args\":{\"code\":\"print('hey')\"}}<|end_message|>" + "<|content_model_end_sampling|>") + .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK) + .parallel_tool_calls(true) + .tools({ special_function_tool, python_tool }) + .expect_tool_calls({ + { "special_function", R"({"arg1": 1})", "" }, + { "python", "{\"code\": \"print('hey')\"}", "" }, + }) + .run(); + } + { // Google Gemma 2 2B - does not support tool calling auto tst = peg_tester("models/templates/google-gemma-2-2b-it.jinja"); diff --git a/tests/test-flash-attn-bias.cpp b/tests/test-flash-attn-bias.cpp new file mode 100644 index 000000000000..94232b32dbcb --- /dev/null +++ b/tests/test-flash-attn-bias.cpp @@ -0,0 +1,525 @@ +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cpp.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct bias_test_config { + const char * name; + int64_t d; + int64_t nq; + int64_t nkv; + int64_t hq; + int64_t hkv; + int64_t extent; + ggml_type type; + bool use_mask; + bool sliding; + int64_t n_batch = 1; + int64_t rel_batch = 1; + bool strided_rel = false; +}; + +struct test_data { + std::vector q; + std::vector k; + std::vector v; + std::vector rel; + std::vector mask; + std::vector k_typed; + std::vector v_typed; + std::vector rel_typed; + std::vector k_rounded; + std::vector v_rounded; + std::vector rel_rounded; + std::vector dense_bias; +}; + +struct run_result { + std::vector output; + size_t allocated_bytes; + double ms; +}; + +static std::vector convert_type(ggml_type type, const std::vector & src, std::vector & rounded) { + rounded.resize(src.size()); + if (type == GGML_TYPE_F32) { + rounded = src; + std::vector bytes(src.size()*sizeof(float)); + memcpy(bytes.data(), src.data(), bytes.size()); + return bytes; + } + if (type == GGML_TYPE_F16) { + std::vector tmp(src.size()); + ggml_fp32_to_fp16_row(src.data(), tmp.data(), src.size()); + ggml_fp16_to_fp32_row(tmp.data(), rounded.data(), src.size()); + std::vector bytes(tmp.size()*sizeof(tmp[0])); + memcpy(bytes.data(), tmp.data(), bytes.size()); + return bytes; + } + GGML_ASSERT(type == GGML_TYPE_BF16); + std::vector tmp(src.size()); + ggml_fp32_to_bf16_row_ref(src.data(), tmp.data(), src.size()); + ggml_bf16_to_fp32_row(tmp.data(), rounded.data(), src.size()); + std::vector bytes(tmp.size()*sizeof(tmp[0])); + memcpy(bytes.data(), tmp.data(), bytes.size()); + return bytes; +} + +static test_data make_data(const bias_test_config & c) { + test_data data; + data.q.resize(c.d*c.nq*c.hq*c.n_batch); + data.k.resize(c.d*c.nkv*c.hkv*c.n_batch); + data.v.resize(c.d*c.nkv*c.hkv*c.n_batch); + data.rel.resize(c.extent*c.hq*c.nq*c.rel_batch); + data.mask.resize(c.nkv*c.nq); + + for (size_t i = 0; i < data.q.size(); ++i) { + data.q[i] = 0.20f*std::sin(float(i)*0.017f + 0.13f); + } + for (size_t i = 0; i < data.k.size(); ++i) { + data.k[i] = 0.25f*std::cos(float(i)*0.013f - 0.29f); + data.v[i] = 0.30f*std::sin(float(i)*0.019f + 0.71f); + } + for (int64_t ib = 0; ib < c.rel_batch; ++ib) { + for (int64_t iq = 0; iq < c.nq; ++iq) { + for (int64_t ih = 0; ih < c.hq; ++ih) { + for (int64_t ie = 0; ie < c.extent; ++ie) { + const size_t idx = ((ib*c.nq + iq)*c.hq + ih)*c.extent + ie; + data.rel[idx] = 0.75f*std::sin(float(idx)*0.007f + float(ie)*0.021f + 0.31f); + } + } + } + } + for (int64_t iq = 0; iq < c.nq; ++iq) { + for (int64_t ik = 0; ik < c.nkv; ++ik) { + const int64_t dist = iq + (c.nkv - c.nq) - ik; + const bool visible = !c.use_mask || (dist >= 0 && (!c.sliding || dist < c.extent)); + data.mask[iq*c.nkv + ik] = ggml_fp32_to_fp16(visible ? 0.0f : -INFINITY); + } + } + + data.k_typed = convert_type(c.type, data.k, data.k_rounded); + data.v_typed = convert_type(c.type, data.v, data.v_rounded); + data.rel_typed = convert_type(GGML_TYPE_F32, data.rel, data.rel_rounded); + + data.dense_bias.assign(c.nkv*c.nq*c.hq*c.n_batch, 0.0f); + for (int64_t ib = 0; ib < c.n_batch; ++ib) { + const int64_t irb = ib % c.rel_batch; + for (int64_t ih = 0; ih < c.hq; ++ih) { + for (int64_t iq = 0; iq < c.nq; ++iq) { + for (int64_t ik = 0; ik < c.nkv; ++ik) { + const int64_t dist = iq + (c.nkv - c.nq) - ik; + if (dist >= 0 && dist < c.extent) { + data.dense_bias[((ib*c.hq + ih)*c.nq + iq)*c.nkv + ik] = + data.rel_rounded[((irb*c.nq + iq)*c.hq + ih)*c.extent + dist]; + } + } + } + } + } + return data; +} + +static run_result run_graph( + ggml_backend_t backend, + const bias_test_config & c, + const test_data & data, + bool dense, + int repeats) { + ggml_init_params params = { + /* .mem_size = */ ggml_tensor_overhead()*64 + ggml_graph_overhead_custom(64, false), + /* .mem_base = */ nullptr, + /* .no_alloc = */ true, + }; + ggml_context_ptr ctx(ggml_init(params)); + GGML_ASSERT(ctx); + + ggml_tensor * q = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, c.d, c.nq, c.hq, c.n_batch); + ggml_tensor * k = ggml_new_tensor_4d(ctx.get(), c.type, c.d, c.nkv, c.hkv, c.n_batch); + ggml_tensor * v = ggml_new_tensor_4d(ctx.get(), c.type, c.d, c.nkv, c.hkv, c.n_batch); + ggml_tensor * r_storage = nullptr; + ggml_tensor * r; + if (c.strided_rel) { + r_storage = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, + 2*c.extent, c.hq, c.nq, c.rel_batch); + r = ggml_view_4d(ctx.get(), r_storage, c.extent, c.hq, c.nq, c.rel_batch, + r_storage->nb[1], r_storage->nb[2], r_storage->nb[3], 0); + } else { + r = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, + c.extent, c.hq, c.nq, c.rel_batch); + } + ggml_tensor * m = c.use_mask ? ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, c.nkv, c.nq, 1, 1) : nullptr; + ggml_set_name(q, "q"); + ggml_set_name(k, "k"); + ggml_set_name(v, "v"); + ggml_set_name(r, "rel_logits"); + if (m) { + ggml_set_name(m, "mask"); + } + + ggml_tensor * out; + ggml_tensor * bias = nullptr; + if (!dense) { + out = ggml_flash_attn_ext_banded(ctx.get(), q, k, v, m, r, 1.0f/float(c.d), c.extent); + } else { + bias = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, c.nkv, c.nq, c.hq, c.n_batch); + ggml_set_name(bias, "dense_bias"); + ggml_tensor * scores = ggml_mul_mat(ctx.get(), k, q); + ggml_mul_mat_set_prec(scores, GGML_PREC_F32); + scores = ggml_scale(ctx.get(), scores, 1.0f/float(c.d)); + scores = ggml_add(ctx.get(), scores, bias); + scores = ggml_soft_max_ext(ctx.get(), scores, m, 1.0f, 0.0f); + ggml_tensor * vt = ggml_cont(ctx.get(), ggml_transpose(ctx.get(), v)); + out = ggml_mul_mat(ctx.get(), vt, scores); + ggml_mul_mat_set_prec(out, GGML_PREC_F32); + out = ggml_cont(ctx.get(), ggml_permute(ctx.get(), out, 0, 2, 1, 3)); + } + ggml_set_name(out, dense ? "out_dense" : "out_flash"); + + GGML_ASSERT(ggml_backend_supports_op(backend, out)); + ggml_backend_buffer_ptr buffer(ggml_backend_alloc_ctx_tensors(ctx.get(), backend)); + GGML_ASSERT(buffer); + + ggml_backend_tensor_set(q, data.q.data(), 0, data.q.size()*sizeof(float)); + ggml_backend_tensor_set(k, data.k_typed.data(), 0, data.k_typed.size()); + ggml_backend_tensor_set(v, data.v_typed.data(), 0, data.v_typed.size()); + if (r_storage) { + std::vector physical(2*c.extent*c.hq*c.nq*c.rel_batch, 0.0f); + for (int64_t ib = 0; ib < c.rel_batch; ++ib) { + for (int64_t iq = 0; iq < c.nq; ++iq) { + for (int64_t ih = 0; ih < c.hq; ++ih) { + const size_t logical = ((ib*c.nq + iq)*c.hq + ih)*c.extent; + const size_t storage = ((ib*c.nq + iq)*c.hq + ih)*(2*c.extent); + memcpy(physical.data() + storage, data.rel_rounded.data() + logical, + c.extent*sizeof(float)); + } + } + } + ggml_backend_tensor_set(r_storage, physical.data(), 0, physical.size()*sizeof(float)); + } else { + ggml_backend_tensor_set(r, data.rel_typed.data(), 0, data.rel_typed.size()); + } + if (m) { + ggml_backend_tensor_set(m, data.mask.data(), 0, data.mask.size()*sizeof(data.mask[0])); + } + if (bias) { + ggml_backend_tensor_set(bias, data.dense_bias.data(), 0, data.dense_bias.size()*sizeof(float)); + } + + ggml_cgraph * graph = ggml_new_graph_custom(ctx.get(), 64, false); + ggml_build_forward_expand(graph, out); + GGML_ASSERT(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + + const int64_t start = ggml_time_us(); + for (int i = 0; i < repeats; ++i) { + GGML_ASSERT(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + } + ggml_backend_synchronize(backend); + const int64_t elapsed = ggml_time_us() - start; + + run_result result; + result.output.resize(ggml_nelements(out)); + ggml_backend_tensor_get(out, result.output.data(), 0, result.output.size()*sizeof(float)); + result.allocated_bytes = ggml_backend_buffer_get_size(buffer.get()); + result.ms = double(elapsed)/1000.0/repeats; + return result; +} + +static std::vector naive_materialized(const bias_test_config & c, const test_data & data) { + const int64_t nrows = c.n_batch*c.hq*c.nq; + std::vector scores(nrows*c.nkv); + std::vector output(c.d*c.hq*c.nq*c.n_batch, 0.0f); + std::atomic next_row(0); + const unsigned nt = std::max(1u, std::thread::hardware_concurrency()); + std::vector workers; + workers.reserve(nt); + + for (unsigned it = 0; it < nt; ++it) { + workers.emplace_back([&]() { + while (true) { + const int64_t row = next_row.fetch_add(1); + if (row >= nrows) { + break; + } + const int64_t ib = row / (c.hq*c.nq); + const int64_t ih = (row / c.nq) % c.hq; + const int64_t iq = row % c.nq; + const int64_t ihkv = ih / (c.hq/c.hkv); + float row_max = -INFINITY; + for (int64_t ik = 0; ik < c.nkv; ++ik) { + float dot = 0.0f; + for (int64_t id = 0; id < c.d; ++id) { + dot += data.q[((ib*c.hq + ih)*c.nq + iq)*c.d + id] * + data.k_rounded[((ib*c.hkv + ihkv)*c.nkv + ik)*c.d + id]; + } + const float mask = ggml_fp16_to_fp32(data.mask[iq*c.nkv + ik]); + const float score = dot/float(c.d) + data.dense_bias[row*c.nkv + ik] + mask; + scores[row*c.nkv + ik] = score; + row_max = std::max(row_max, score); + } + float sum = 0.0f; + for (int64_t ik = 0; ik < c.nkv; ++ik) { + const float p = std::exp(scores[row*c.nkv + ik] - row_max); + scores[row*c.nkv + ik] = p; + sum += p; + } + for (int64_t ik = 0; ik < c.nkv; ++ik) { + const float p = scores[row*c.nkv + ik]/sum; + for (int64_t id = 0; id < c.d; ++id) { + output[((ib*c.nq + iq)*c.hq + ih)*c.d + id] += + p*data.v_rounded[((ib*c.hkv + ihkv)*c.nkv + ik)*c.d + id]; + } + } + } + }); + } + for (auto & worker : workers) { + worker.join(); + } + return output; +} + +static void error_stats(const std::vector & got, const std::vector & ref, + double & max_abs, double & mean_abs, double & max_rel, double & mean_rel, double & rmse) { + double sq = 0.0; + double abs_sum = 0.0; + double rel_sum = 0.0; + max_abs = 0.0; + max_rel = 0.0; + for (size_t i = 0; i < got.size(); ++i) { + const double ae = std::abs(double(got[i]) - ref[i]); + const double re = ae/std::max(1e-5, std::abs(double(ref[i]))); + max_abs = std::max(max_abs, ae); + max_rel = std::max(max_rel, re); + abs_sum += ae; + rel_sum += re; + sq += ae*ae; + } + mean_abs = abs_sum/got.size(); + mean_rel = rel_sum/got.size(); + rmse = std::sqrt(sq/got.size()); +} + +static void overflow_arithmetic_self_test() { + // mirrors the scalar-path offset math; exact offset checked at 128 bits, lands past 2^31 + const uint64_t nb0 = sizeof(float); + const uint64_t nb1 = 1024*nb0; + const uint64_t nb2 = 64*nb1; + const uint64_t nb3 = 131072*nb2; + const uint64_t dist = 1023, head = 63, query = 131071, batch = 3; + const uint64_t offset = dist*nb0 + head*nb1 + query*nb2 + batch*nb3; + __extension__ typedef unsigned __int128 uint128_t; + const uint128_t exact = uint128_t(dist)*nb0 + uint128_t(head)*nb1 + + uint128_t(query)*nb2 + uint128_t(batch)*nb3; + GGML_ASSERT(exact <= UINT64_MAX && offset == (uint64_t) exact && offset > INT32_MAX); + + const int64_t nq = int64_t(1) << 40; + const int64_t nkv = nq + 8192; + const int64_t iq = nq - 1; + const int64_t ik = nkv - 1024; + const int64_t rel_dist = iq + (nkv - nq) - ik; + GGML_ASSERT(rel_dist == 1023); + printf("overflow_check offset=%llu (>INT32_MAX) large_T=%lld rel_dist=%lld PASS\n", + (unsigned long long) offset, (long long) nq, (long long) rel_dist); +} + +static bool overflow_kernel_test(ggml_backend_t backend, const char * backend_kind) { + // rel-logits row for query 1 sits beyond INT32_MAX; only two small logical rows are touched + const bias_test_config c = { + "overflow_kernel_stride", 64, 2, 2, 2, 1, 8, GGML_TYPE_F32, true, false, + }; + test_data data = make_data(c); + const uint64_t rel_nb2 = (UINT64_C(1) << 31) + 4096; + const size_t rel_row_bytes = c.extent*c.hq*sizeof(float); + const uint64_t storage_bytes = rel_nb2 + rel_row_bytes; + + ggml_init_params params = { + /* .mem_size = */ ggml_tensor_overhead()*32 + ggml_graph_overhead_custom(32, false), + /* .mem_base = */ nullptr, + /* .no_alloc = */ true, + }; + ggml_context_ptr ctx(ggml_init(params)); + GGML_ASSERT(ctx); + ggml_tensor * q = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, c.d, c.nq, c.hq, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, c.d, c.nkv, c.hkv, 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F32, c.d, c.nkv, c.hkv, 1); + ggml_tensor * m = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, c.nkv, c.nq, 1, 1); + ggml_tensor * r_storage = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_F32, + (storage_bytes + sizeof(float) - 1)/sizeof(float)); + ggml_tensor * r = ggml_view_4d(ctx.get(), r_storage, c.extent, c.hq, c.nq, 1, + c.extent*sizeof(float), rel_nb2, rel_nb2*c.nq, 0); + ggml_tensor * out = ggml_flash_attn_ext_banded( + ctx.get(), q, k, v, m, r, 1.0f/float(c.d), c.extent); + ggml_backend_buffer_ptr buffer(ggml_backend_alloc_ctx_tensors(ctx.get(), backend)); + GGML_ASSERT(buffer); + + ggml_backend_tensor_set(q, data.q.data(), 0, data.q.size()*sizeof(float)); + ggml_backend_tensor_set(k, data.k_typed.data(), 0, data.k_typed.size()); + ggml_backend_tensor_set(v, data.v_typed.data(), 0, data.v_typed.size()); + ggml_backend_tensor_set(m, data.mask.data(), 0, data.mask.size()*sizeof(data.mask[0])); + ggml_backend_tensor_set(r_storage, data.rel_typed.data(), 0, rel_row_bytes); + ggml_backend_tensor_set(r_storage, data.rel_typed.data() + rel_row_bytes, rel_nb2, rel_row_bytes); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx.get(), 32, false); + ggml_build_forward_expand(graph, out); + GGML_ASSERT(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + std::vector got(ggml_nelements(out)); + ggml_backend_tensor_get(out, got.data(), 0, got.size()*sizeof(float)); + const std::vector ref = naive_materialized(c, data); + double max_abs, mean_abs, max_rel, mean_rel, rmse; + error_stats(got, ref, max_abs, mean_abs, max_rel, mean_rel, rmse); + const bool pass = max_abs <= 2e-5; + printf("overflow_kernel backend=%s rel_query_stride=%llu allocated_bytes=%zu " + "naive_max_abs=%.9g naive_mean_abs=%.9g naive_max_rel=%.9g naive_mean_rel=%.9g naive_rmse=%.9g %s\n", + backend_kind, (unsigned long long) rel_nb2, ggml_backend_buffer_get_size(buffer.get()), + max_abs, mean_abs, max_rel, mean_rel, rmse, pass ? "PASS" : "FAIL"); + return pass; +} + +int main(int argc, char ** argv) { + std::string backend_kind = argc > 1 ? argv[1] : "cpu"; + std::string suite = argc > 2 ? argv[2] : "small"; + int repeats = argc > 3 ? std::max(1, atoi(argv[3])) : 1; + + overflow_arithmetic_self_test(); + ggml_backend_load_all(); + ggml_backend_dev_t chosen = nullptr; + const enum ggml_backend_dev_type wanted = backend_kind == "cuda" ? + GGML_BACKEND_DEVICE_TYPE_GPU : GGML_BACKEND_DEVICE_TYPE_CPU; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) == wanted) { + chosen = dev; + break; + } + } + GGML_ASSERT(chosen); + ggml_backend_ptr backend(ggml_backend_dev_init(chosen, nullptr)); + GGML_ASSERT(backend); + + if (suite == "overflow") { + return overflow_kernel_test(backend.get(), backend_kind.c_str()) ? 0 : 1; + } + + if (suite == "perf") { + const std::vector perf_cases = { + {"prefill_t1024", 64, 1024, 1024, 8, 2, 512, GGML_TYPE_F16, true, false}, + {"prefill_t2048", 64, 2048, 2048, 8, 2, 512, GGML_TYPE_F16, true, false}, + {"prefill_t4096", 64, 4096, 4096, 8, 2, 512, GGML_TYPE_F16, true, false}, + {"decode_8k", 128, 1, 8192, 8, 1, 1024, GGML_TYPE_F16, true, false}, + {"heads64_gqa", 64, 1024, 1024, 64, 8, 512, GGML_TYPE_F16, true, false}, + }; + printf("backend=%s suite=perf device=%s repeats=%d\n", backend_kind.c_str(), + ggml_backend_dev_description(chosen), repeats); + for (const bias_test_config & c : perf_cases) { + test_data data = make_data(c); + const run_result flash = run_graph(backend.get(), c, data, false, repeats); + const run_result dense = run_graph(backend.get(), c, data, true, repeats); + printf("%s type=%s D=%lld nq=%lld nkv=%lld hq=%lld hkv=%lld E=%lld " + "flash_ms=%.6f dense_ms=%.6f speedup=%.6f flash_bytes=%zu dense_bytes=%zu memory_ratio=%.6f\n", + c.name, ggml_type_name(c.type), (long long)c.d, (long long)c.nq, + (long long)c.nkv, (long long)c.hq, (long long)c.hkv, (long long)c.extent, + flash.ms, dense.ms, dense.ms/flash.ms, flash.allocated_bytes, + dense.allocated_bytes, double(dense.allocated_bytes)/flash.allocated_bytes); + } + return 0; + } + + if (suite == "memory") { + const std::vector memory_cases = { + {"memory_t1024", 64, 1024, 1024, 8, 2, 512, GGML_TYPE_F16, false, false}, + {"memory_t2048", 64, 2048, 2048, 8, 2, 512, GGML_TYPE_F16, false, false}, + {"memory_t4096", 64, 4096, 4096, 8, 2, 512, GGML_TYPE_F16, false, false}, + }; + printf("backend=%s suite=memory device=%s\n", backend_kind.c_str(), + ggml_backend_dev_description(chosen)); + for (const bias_test_config & c : memory_cases) { + test_data data = make_data(c); + const run_result flash = run_graph(backend.get(), c, data, false, 1); + const run_result dense = run_graph(backend.get(), c, data, true, 1); + double max_abs, mean_abs, max_rel, mean_rel, rmse; + error_stats(flash.output, dense.output, max_abs, mean_abs, max_rel, mean_rel, rmse); + printf("%s T=%lld flash_bytes=%zu dense_bytes=%zu memory_ratio=%.6f " + "dense_max_abs=%.9g dense_mean_abs=%.9g PASS\n", + c.name, (long long)c.nq, flash.allocated_bytes, dense.allocated_bytes, + double(dense.allocated_bytes)/flash.allocated_bytes, max_abs, mean_abs); + } + return 0; + } + + std::vector configs; + if (suite == "small") { + configs = { + {"edge_e8_f32", 64, 16, 16, 2, 1, 8, GGML_TYPE_F32, true, false}, + {"gqa_d128_f16", 128, 16, 16, 8, 2, 8, GGML_TYPE_F16, true, false}, + {"sliding_bf16", 64, 64, 64, 8, 2, 8, GGML_TYPE_BF16, true, true }, + {"decode_offset", 64, 1, 513, 8, 2, 8, GGML_TYPE_F32, true, false}, + {"extent_512_edge", 64, 64, 64, 8, 1, 512, GGML_TYPE_F16, true, false}, + {"strided_rel_f16", 64, 17, 33, 8, 2, 8, GGML_TYPE_F16, true, false, 1, 1, true}, + {"batch_distinct", 64, 16, 16, 8, 2, 8, GGML_TYPE_F16, true, false, 2, 2, false}, + {"batch_broadcast", 64, 16, 16, 8, 2, 8, GGML_TYPE_F16, true, false, 2, 1, false}, + {"heads64_gqa4", 64, 16, 16, 64, 16, 8, GGML_TYPE_F16, true, false}, + {"heads64_gqa8", 64, 16, 16, 64, 8, 8, GGML_TYPE_F16, true, false}, + }; + } else if (suite == "medium") { + configs = { + {"medium_f32_e512", 64, 1024, 1024, 8, 2, 512, GGML_TYPE_F32, true, false}, + {"medium_f16_e1024", 128, 1024, 1024, 8, 2, 1024, GGML_TYPE_F16, true, false}, + {"medium_bf16_local", 64, 2048, 2048, 8, 2, 512, GGML_TYPE_BF16, true, true }, + {"decode_8k_e1024", 128, 1, 8192, 8, 1, 1024, GGML_TYPE_F16, true, false}, + }; + } else if (suite == "hard") { + configs = { + {"heads64_gqa", 64, 1024, 1024, 64, 8, 512, GGML_TYPE_F16, true, false}, + }; + } else { + fprintf(stderr, "unknown suite: %s (expected small, medium, hard, perf, memory, or overflow)\n", + suite.c_str()); + return 2; + } + + bool ok = true; + printf("backend=%s suite=%s device=%s\n", backend_kind.c_str(), suite.c_str(), ggml_backend_dev_description(chosen)); + for (const bias_test_config & c : configs) { + test_data data = make_data(c); + run_result flash = run_graph(backend.get(), c, data, false, repeats); + run_result dense = run_graph(backend.get(), c, data, true, repeats); + // always compare to an independently materialized oracle (O(T^2), test-only) + const std::vector naive = naive_materialized(c, data); + + double abs_dense, mean_abs_dense, rel_dense, mean_rel_dense, rmse_dense; + error_stats(flash.output, dense.output, abs_dense, mean_abs_dense, rel_dense, mean_rel_dense, rmse_dense); + double abs_naive = 0.0, mean_abs_naive = 0.0, rel_naive = 0.0, mean_rel_naive = 0.0, rmse_naive = 0.0; + error_stats(flash.output, naive, abs_naive, mean_abs_naive, rel_naive, mean_rel_naive, rmse_naive); + const double tol = c.type == GGML_TYPE_F32 ? 2e-5 : 2e-3; + const bool pass = abs_naive <= tol; + ok = ok && pass; + printf("%s type=%s D=%lld nq=%lld nkv=%lld hq=%lld hkv=%lld E=%lld mask=%d sliding=%d " + "batch=%lld rel_batch=%lld strided_rel=%d " + "dense_max_abs=%.9g dense_mean_abs=%.9g dense_max_rel=%.9g dense_mean_rel=%.9g dense_rmse=%.9g " + "naive_max_abs=%.9g naive_mean_abs=%.9g naive_max_rel=%.9g naive_mean_rel=%.9g naive_rmse=%.9g " + "flash_ms=%.4f dense_ms=%.4f speedup=%.4f " + "flash_bytes=%zu dense_bytes=%zu memory_ratio=%.4f %s\n", + c.name, ggml_type_name(c.type), (long long)c.d, (long long)c.nq, (long long)c.nkv, + (long long)c.hq, (long long)c.hkv, (long long)c.extent, c.use_mask, c.sliding, + (long long)c.n_batch, (long long)c.rel_batch, c.strided_rel, + abs_dense, mean_abs_dense, rel_dense, mean_rel_dense, rmse_dense, + abs_naive, mean_abs_naive, rel_naive, mean_rel_naive, rmse_naive, + flash.ms, dense.ms, dense.ms/flash.ms, + flash.allocated_bytes, dense.allocated_bytes, double(dense.allocated_bytes)/flash.allocated_bytes, + pass ? "PASS" : "FAIL"); + } + return ok ? 0 : 1; +} diff --git a/tests/test-flash-attn-generic-hash.cpp b/tests/test-flash-attn-generic-hash.cpp new file mode 100644 index 000000000000..4e3aca44ffab --- /dev/null +++ b/tests/test-flash-attn-generic-hash.cpp @@ -0,0 +1,142 @@ +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include + +// Deterministic FLASH_ATTN_EXT probe; banded-API-free so the same source builds on base and final trees. +int main(int argc, char ** argv) { + const std::string backend_kind = argc > 1 ? argv[1] : "cpu"; + const std::string output_path = argc > 2 ? argv[2] : ""; + + ggml_backend_load_all(); + const enum ggml_backend_dev_type wanted = backend_kind == "cuda" ? + GGML_BACKEND_DEVICE_TYPE_GPU : GGML_BACKEND_DEVICE_TYPE_CPU; + ggml_backend_dev_t chosen = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_type(dev) == wanted) { + chosen = dev; + break; + } + } + if (!chosen) { + fprintf(stderr, "requested backend is unavailable: %s\n", backend_kind.c_str()); + return 2; + } + ggml_backend_t backend = ggml_backend_dev_init(chosen, nullptr); + if (!backend) { + return 2; + } + + constexpr int64_t d = 64; + constexpr int64_t nq = 33; + constexpr int64_t nkv = 47; + constexpr int64_t hq = 8; + constexpr int64_t hkv = 2; + ggml_init_params params = { + /* .mem_size = */ ggml_tensor_overhead()*16 + ggml_graph_overhead_custom(16, false), + /* .mem_base = */ nullptr, + /* .no_alloc = */ true, + }; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + ggml_backend_free(backend); + return 2; + } + + ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d, nq, hq, 1); + ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, nkv, hkv, 1); + ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, d, nkv, hkv, 1); + ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, nkv, nq, 1, 1); + ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/float(d), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); + if (!ggml_backend_supports_op(backend, out)) { + fprintf(stderr, "ordinary flash attention is unsupported on %s\n", + ggml_backend_dev_description(chosen)); + ggml_free(ctx); + ggml_backend_free(backend); + return 2; + } + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buffer) { + ggml_free(ctx); + ggml_backend_free(backend); + return 2; + } + + std::vector q_data(ggml_nelements(q)); + std::vector k_f32(ggml_nelements(k)); + std::vector v_f32(ggml_nelements(v)); + std::vector k_data(k_f32.size()); + std::vector v_data(v_f32.size()); + std::vector mask(ggml_nelements(m)); + for (size_t i = 0; i < q_data.size(); ++i) { + q_data[i] = 0.20f*std::sin(0.017f*float(i) + 0.11f); + } + for (size_t i = 0; i < k_f32.size(); ++i) { + k_f32[i] = 0.23f*std::cos(0.013f*float(i) - 0.29f); + v_f32[i] = 0.31f*std::sin(0.019f*float(i) + 0.71f); + } + ggml_fp32_to_fp16_row(k_f32.data(), k_data.data(), k_data.size()); + ggml_fp32_to_fp16_row(v_f32.data(), v_data.data(), v_data.size()); + for (int64_t iq = 0; iq < nq; ++iq) { + for (int64_t ik = 0; ik < nkv; ++ik) { + const int64_t dist = iq + (nkv - nq) - ik; + mask[iq*nkv + ik] = ggml_fp32_to_fp16(dist >= 0 && dist < 29 ? 0.0f : -INFINITY); + } + } + ggml_backend_tensor_set(q, q_data.data(), 0, q_data.size()*sizeof(q_data[0])); + ggml_backend_tensor_set(k, k_data.data(), 0, k_data.size()*sizeof(k_data[0])); + ggml_backend_tensor_set(v, v_data.data(), 0, v_data.size()*sizeof(v_data[0])); + ggml_backend_tensor_set(m, mask.data(), 0, mask.size()*sizeof(mask[0])); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 16, false); + ggml_build_forward_expand(graph, out); + const ggml_status status = ggml_backend_graph_compute(backend, graph); + ggml_backend_synchronize(backend); + if (status != GGML_STATUS_SUCCESS) { + fprintf(stderr, "graph failed with status %d\n", int(status)); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); + return 1; + } + + std::vector result(ggml_nelements(out)); + ggml_backend_tensor_get(out, result.data(), 0, result.size()*sizeof(result[0])); + uint64_t fnv = UINT64_C(1469598103934665603); + const uint8_t * bytes = reinterpret_cast(result.data()); + for (size_t i = 0; i < result.size()*sizeof(result[0]); ++i) { + fnv ^= bytes[i]; + fnv *= UINT64_C(1099511628211); + } + if (!output_path.empty()) { + FILE * fp = fopen(output_path.c_str(), "wb"); + if (!fp || fwrite(result.data(), sizeof(result[0]), result.size(), fp) != result.size()) { + fprintf(stderr, "failed to write %s\n", output_path.c_str()); + if (fp) { + fclose(fp); + } + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); + return 1; + } + fclose(fp); + } + printf("generic_hash backend=%s device=%s bytes=%zu fnv1a64=%016llx\n", + backend_kind.c_str(), ggml_backend_dev_description(chosen), + result.size()*sizeof(result[0]), (unsigned long long) fnv); + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); + return 0; +} diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index a1ed2a76f879..792617731a04 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -427,6 +427,9 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_DEEPSEEK4) { return false; } + if (arch == LLM_ARCH_INKLING) { + return false; // TODO fixture params for the arch-specific hparams (d_rel, rel_extent, shortconv, logit_scale_denom) + } // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 15040e4af5f9..461c4f429cfb 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -38,6 +38,7 @@ add_library(mtmd models/granite4-vision.cpp models/hunyuanvl.cpp models/internvl.cpp + models/inkling.cpp models/kimivl.cpp models/kimik25.cpp models/nemotron-v2-vl.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 589fc724ed08..fc011f714704 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -285,6 +285,13 @@ #define TN_A_FFN_POST_NORM "%s.blk.%d.ffn_post_norm.%s" #define TN_A_FFN_POST_NORM_1 "%s.blk.%d.ffn_post_norm_1.%s" +// Inkling hMLP vision and dMel audio towers. +#define TN_INKLING_HMLP_LINEAR "v.hmlp.%d.linear.weight" +#define TN_INKLING_HMLP_NORM "v.hmlp.%d.norm.weight" +#define TN_INKLING_HMLP_FINAL_NORM "v.hmlp.final_norm.weight" +#define TN_INKLING_DMEL_EMBD "a.dmel.embedding.weight" +#define TN_INKLING_DMEL_FINAL_NORM "a.dmel.final_norm.weight" + // mobilenetv5 (gemma3n) definitions #define TN_MNV5_STEM_CONV "v.conv_stem.conv.weight" #define TN_MNV5_STEM_BIAS "v.conv_stem.conv.bias" @@ -353,6 +360,7 @@ struct clip_ctx; enum projector_type { + PROJECTOR_TYPE_INKLING, PROJECTOR_TYPE_MLP, PROJECTOR_TYPE_MLP_NORM, PROJECTOR_TYPE_LDP, @@ -411,6 +419,7 @@ enum projector_type { }; static std::map PROJECTOR_TYPE_NAMES = { + { PROJECTOR_TYPE_INKLING, "inkling" }, { PROJECTOR_TYPE_MLP, "mlp" }, { PROJECTOR_TYPE_LDP, "ldp" }, { PROJECTOR_TYPE_LDPV2, "ldpv2"}, diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 146eabce23b7..bfe021922098 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -33,7 +33,7 @@ enum resize_algo { RESIZE_ALGO_BILINEAR, // stretch to target resolution RESIZE_ALGO_BICUBIC, // center-crop when aspect ratio doesn't match RESIZE_ALGO_BICUBIC_PILLOW, - // RESIZE_ALGO_LANCZOS, // TODO + RESIZE_ALGO_LANCZOS, }; // Padding style for img_tool::resize @@ -363,6 +363,11 @@ struct qf_block { std::vector qf_proj_layers; }; +struct inkling_hmlp_layer { + ggml_tensor * linear_w = nullptr; + ggml_tensor * norm_w = nullptr; +}; + struct clip_model { clip_modality modality = CLIP_MODALITY_VISION; projector_type proj_type = PROJECTOR_TYPE_MLP; @@ -377,6 +382,12 @@ struct clip_model { ggml_tensor * norm_embd_w = nullptr; ggml_tensor * norm_embd_b = nullptr; + // Inkling towers (neither uses standard transformer blocks). + std::vector inkling_hmlp_layers; + ggml_tensor * inkling_hmlp_final_norm_w = nullptr; + ggml_tensor * inkling_dmel_embd_w = nullptr; + ggml_tensor * inkling_dmel_final_norm_w = nullptr; + // "indexed" patch embedding norms ggml_tensor * patch_norm_1_w = nullptr; ggml_tensor * patch_norm_1_b = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 11f9820edeb8..f1ec237a4bda 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -878,6 +878,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const std::unique_ptr builder; switch (ctx->proj_type()) { + case PROJECTOR_TYPE_INKLING: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_GEMMA3: case PROJECTOR_TYPE_IDEFICS3: case PROJECTOR_TYPE_LFM2: @@ -1298,6 +1302,21 @@ struct clip_model_loader { // model-specific params switch (model.proj_type) { + case PROJECTOR_TYPE_INKLING: + { + if (is_vision) { + hparams.image_resize_algo = RESIZE_ALGO_BICUBIC_PILLOW; + hparams.image_resize_pad = PAD_NONE; + hparams.warmup_image_size = hparams.patch_size; + } else { + hparams.audio_chunk_len = 0; + hparams.audio_sample_rate = 16000; + hparams.audio_n_fft = 1600; + hparams.audio_window_len = 1600; + hparams.audio_hop_len = 800; + hparams.warmup_audio_size = 4; + } + } break; case PROJECTOR_TYPE_MLP: case PROJECTOR_TYPE_MLP_NORM: case PROJECTOR_TYPE_LDP: @@ -1971,7 +1990,8 @@ struct clip_model_loader { model.position_embeddings = get_tensor(string_format(TN_POS_EMBD, prefix), false); const bool has_standard_layers = ( - model.proj_type != PROJECTOR_TYPE_GEMMA3NV); + model.proj_type != PROJECTOR_TYPE_GEMMA3NV && + model.proj_type != PROJECTOR_TYPE_INKLING); // layers const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0; @@ -2057,6 +2077,23 @@ struct clip_model_loader { switch (model.proj_type) { + case PROJECTOR_TYPE_INKLING: + { + if (model.modality == CLIP_MODALITY_VISION) { + model.inkling_hmlp_layers.resize(hparams.n_layer); + for (int il = 0; il < hparams.n_layer; ++il) { + auto & layer = model.inkling_hmlp_layers[il]; + layer.linear_w = get_tensor(string_format(TN_INKLING_HMLP_LINEAR, il)); + layer.norm_w = il + 1 < hparams.n_layer + ? get_tensor(string_format(TN_INKLING_HMLP_NORM, il)) + : nullptr; + } + model.inkling_hmlp_final_norm_w = get_tensor(TN_INKLING_HMLP_FINAL_NORM); + } else { + model.inkling_dmel_embd_w = get_tensor(TN_INKLING_DMEL_EMBD); + model.inkling_dmel_final_norm_w = get_tensor(TN_INKLING_DMEL_FINAL_NORM); + } + } break; case PROJECTOR_TYPE_MLP: case PROJECTOR_TYPE_MLP_NORM: { @@ -3144,6 +3181,11 @@ struct clip_model_loader { LOG_INF("%s: warmup with audio size = %d\n", __func__, hparams.warmup_audio_size); } batch.entries.push_back(img); + // One logical Inkling vision patch always contains two adjacent temporal slices (including warmup reserve). + if (ctx_clip.model.modality == CLIP_MODALITY_VISION && + ctx_clip.model.proj_type == PROJECTOR_TYPE_INKLING) { + batch.entries.push_back(img); + } return batch; } @@ -3541,6 +3583,16 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { projector_type proj = ctx->proj_type(); switch (proj) { + case PROJECTOR_TYPE_INKLING: + { + if (ctx->model.modality == CLIP_MODALITY_AUDIO) { + // One dMel row is one soft audio token. + n_patches = img->nx(); + } else { + // Batch entries are temporal slices; each pair emits one token. + n_patches = 1; + } + } break; case PROJECTOR_TYPE_MLP: case PROJECTOR_TYPE_MLP_NORM: case PROJECTOR_TYPE_JANUS_PRO: @@ -3927,6 +3979,10 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 // set input per projector switch (ctx->model.proj_type) { + case PROJECTOR_TYPE_INKLING: + { + // inp_raw is the only graph input. + } break; case PROJECTOR_TYPE_MINICPMV: { // inspired from siglip: @@ -4955,6 +5011,10 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 int clip_n_mmproj_embd(const struct clip_ctx * ctx) { switch (ctx->model.proj_type) { + case PROJECTOR_TYPE_INKLING: + return ctx->model.modality == CLIP_MODALITY_AUDIO + ? ctx->model.inkling_dmel_final_norm_w->ne[0] + : ctx->model.inkling_hmlp_final_norm_w->ne[0]; case PROJECTOR_TYPE_LDP: return ctx->model.mm_model_block_1_block_2_1_b->ne[0]; case PROJECTOR_TYPE_LDPV2: diff --git a/tools/mtmd/models/inkling.cpp b/tools/mtmd/models/inkling.cpp new file mode 100644 index 000000000000..c072f66a2c40 --- /dev/null +++ b/tools/mtmd/models/inkling.cpp @@ -0,0 +1,106 @@ +#include "models.h" + +ggml_tensor * clip_graph_inkling::build_mm(ggml_tensor * w, ggml_tensor * x) const { + ggml_tensor * cur = ggml_mul_mat(ctx0, w, x); + ggml_mul_mat_set_prec(cur, GGML_PREC_F32); + return cur; +} + +// fold square neighborhoods from W/H into channels; folded order is [h_fold, w_fold, C] +static ggml_tensor * inkling_fold_spatial( + ggml_context * ctx0, + ggml_tensor * cur, + int scale) { + GGML_ASSERT(scale > 0); + GGML_ASSERT(cur->ne[1] % scale == 0); + GGML_ASSERT(cur->ne[2] % scale == 0); + + const int64_t c = cur->ne[0]; + const int64_t w = cur->ne[1]; + const int64_t h = cur->ne[2]; + const int64_t b = cur->ne[3]; + + cur = ggml_reshape_4d(ctx0, cur, c * scale, w / scale, h, b); + cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 0, 2, 1, 3)); + cur = ggml_reshape_4d(ctx0, cur, c * scale * scale, h / scale, w / scale, b); + cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 0, 2, 1, 3)); + return cur; +} + +ggml_cgraph * clip_graph_inkling::build() { + return model.modality == CLIP_MODALITY_AUDIO ? build_audio() : build_vision(); +} + +ggml_cgraph * clip_graph_inkling::build_vision() { + static constexpr int temporal_patch_size = 2; + static constexpr int spatial_folds[] = {5, 2, 4}; + + GGML_ASSERT(img.nx() == 40 && img.ny() == 40); + GGML_ASSERT(n_batch > 0 && n_batch % temporal_patch_size == 0); + GGML_ASSERT(model.inkling_hmlp_layers.size() == 4); + GGML_ASSERT(model.inkling_hmlp_final_norm_w); + + // Raw input is [W,H,RGB,temporal*patch]. Put RGB on ne[0]. + ggml_tensor * cur = build_inp_raw(3); + cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 2, 0, 3)); + + for (int il = 0; il < 3; ++il) { + cur = inkling_fold_spatial(ctx0, cur, spatial_folds[il]); + const int64_t w = cur->ne[1]; + const int64_t h = cur->ne[2]; + const int64_t b = cur->ne[3]; + + cur = ggml_reshape_2d(ctx0, cur, cur->ne[0], w * h * b); + cur = build_mm(model.inkling_hmlp_layers[il].linear_w, cur); + cur = build_norm(cur, model.inkling_hmlp_layers[il].norm_w, + nullptr, NORM_TYPE_RMS, eps, il); + cur = ggml_gelu_erf(ctx0, cur); + cur = ggml_reshape_4d(ctx0, cur, cur->ne[0], w, h, b); + } + + GGML_ASSERT(cur->ne[1] == 1 && cur->ne[2] == 1); + const int64_t n_patches = n_batch / temporal_patch_size; + cur = ggml_reshape_2d(ctx0, cur, cur->ne[0] * temporal_patch_size, n_patches); + cur = build_mm(model.inkling_hmlp_layers[3].linear_w, cur); + cur = build_norm(cur, model.inkling_hmlp_final_norm_w, + nullptr, NORM_TYPE_RMS, eps, 3); + + // Batched mtmd convention: one token in ne[1], patch count in ne[2]. + cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], 1, n_patches); + ggml_build_forward_expand(gf, cur); + return gf; +} + +ggml_cgraph * clip_graph_inkling::build_audio() { + static constexpr int n_mels = 80; + static constexpr int mel_vocab_size = 16; + static constexpr int n_embd = 6144; + + GGML_ASSERT(img.ny() == n_mels); + GGML_ASSERT(model.inkling_dmel_embd_w); + GGML_ASSERT(model.inkling_dmel_final_norm_w); + GGML_ASSERT(model.inkling_dmel_embd_w->ne[0] == n_embd); + GGML_ASSERT(model.inkling_dmel_embd_w->ne[1] == n_mels * mel_vocab_size); + + const int64_t n_tokens = img.nx(); + // mtmd audio storage is mel-major and represented as [token, mel]. + ggml_tensor * bins = build_inp_raw(1); + bins = ggml_cont(ctx0, ggml_transpose(ctx0, bins)); + ggml_tensor * offsets = ggml_arange(ctx0, 0, n_mels, 1); + offsets = ggml_scale(ctx0, offsets, mel_vocab_size); + offsets = ggml_reshape_2d(ctx0, offsets, n_mels, 1); + ggml_tensor * indices = ggml_cast(ctx0, ggml_add(ctx0, bins, offsets), GGML_TYPE_I32); + + indices = ggml_reshape_1d(ctx0, indices, n_mels * n_tokens); + ggml_tensor * cur = ggml_get_rows(ctx0, model.inkling_dmel_embd_w, indices); + cur = ggml_reshape_3d(ctx0, cur, n_embd, n_mels, n_tokens); + cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 0, 2, 3)); + cur = ggml_sum_rows(ctx0, cur); + cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 0, 2, 3)); + cur = ggml_reshape_2d(ctx0, cur, n_embd, n_tokens); + cur = build_norm(cur, model.inkling_dmel_final_norm_w, + nullptr, NORM_TYPE_RMS, eps, -1); + + ggml_build_forward_expand(gf, cur); + return gf; +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index e54366a086f4..6735c9cdebf9 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -12,6 +12,17 @@ struct clip_graph_siglip : clip_graph { ggml_cgraph * build() override; }; +struct clip_graph_inkling : clip_graph { + clip_graph_inkling(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; + ggml_tensor * build_mm(ggml_tensor * w, ggml_tensor * x) const override; + bool support_batch() const override { return true; } + +private: + ggml_cgraph * build_vision(); + ggml_cgraph * build_audio(); +}; + struct clip_graph_gemma4v : clip_graph { clip_graph_gemma4v(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index fea03557d05c..13e31ff54e9a 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -283,6 +283,7 @@ struct filter_params { bool norm_per_feature = false; bool use_magnitude = false; // |X| instead of |X|^2 float mel_floor = 5.960464477539063e-08f; + float power_floor = 0.0f; // clamp |X|^2 to this before sqrt (0 = disabled) }; static void log_mel_spectrogram_worker_thread(int ith, @@ -327,6 +328,7 @@ static void log_mel_spectrogram_worker_thread(int ith, // Calculate modulus^2 (power) or modulus (magnitude) for (int j = 0; j < n_fft_bins; j++) { float power = (fft_out[2 * j + 0] * fft_out[2 * j + 0] + fft_out[2 * j + 1] * fft_out[2 * j + 1]); + power = std::max(power, params.power_floor); fft_out[j] = params.use_magnitude ? sqrtf(power) : power; } @@ -537,6 +539,75 @@ static bool log_mel_spectrogram( return true; } +void mtmd_audio_preprocessor_inkling::initialize() { + GGML_ASSERT(hparams.n_mel_bins == 80); + GGML_ASSERT(hparams.audio_n_fft == 1600); + cache.fill_sin_cos_table(hparams.audio_n_fft); + cache.fill_hann_window(hparams.audio_window_len, true); + cache.fill_mel_filterbank_matrix( + hparams.n_mel_bins, hparams.audio_n_fft, hparams.audio_sample_rate, + 0.0f, -1.0f, true, 1.0f, false); +} + +bool mtmd_audio_preprocessor_inkling::preprocess( + const float * samples, + size_t n_samples, + std::vector & output) { + if (n_samples == 0) { + return false; + } + GGML_ASSERT(hparams.audio_hop_len == 800); + GGML_ASSERT(hparams.audio_window_len == 1600); + GGML_ASSERT(!cache.filters.data.empty()); + + // Reference behavior: left pad by n_fft-hop, then right pad to a whole hop. + const size_t left_pad = static_cast(hparams.audio_n_fft - hparams.audio_hop_len); + const size_t right_pad = + (static_cast(hparams.audio_hop_len) - n_samples % hparams.audio_hop_len) + % hparams.audio_hop_len; + std::vector padded(left_pad + n_samples + right_pad, 0.0f); + std::copy(samples, samples + n_samples, padded.begin() + left_pad); + + filter_params params; + params.n_mel = hparams.n_mel_bins; + params.n_fft_bins = 1 + hparams.audio_n_fft / 2; + params.hann_window_size = hparams.audio_window_len; + params.hop_length = hparams.audio_hop_len; + params.sample_rate = hparams.audio_sample_rate; + params.no_padding = true; + params.use_natural_log = false; + params.use_magnitude = true; + params.mel_floor = 1e-10f; + params.power_floor = 1e-10f; // reference clamps |X|^2 before sqrt + + mtmd_audio_mel dmel; + if (!log_mel_spectrogram( + padded.data(), static_cast(padded.size()), 4, + params, cache, dmel)) { + return false; + } + dmel.n_len_org = static_cast(n_samples); + + // nearest of 16 float64 centers over [-7,2]; strict '<' keeps the lower bin on midpoints (torch.argmin) + for (float & value_f32 : dmel.data) { + const double value = std::max(-7.0, std::min(2.0, static_cast(value_f32))); + int best = 0; + double best_dist = INFINITY; + for (int bin = 0; bin < 16; ++bin) { + const double center = -7.0 + 9.0 * static_cast(bin) / 15.0; + const double dist = std::abs(value - center); + if (dist < best_dist) { + best = bin; + best_dist = dist; + } + } + value_f32 = static_cast(best); + } + + output.push_back(std::move(dmel)); + return true; +} + // // mtmd_audio_preprocessor_whisper // diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index f65f282d96e2..e79519553e9b 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -69,6 +69,16 @@ struct mtmd_audio_preprocessor_whisper : mtmd_audio_preprocessor { mtmd_audio_cache cache; }; +// Inkling dMel: 100 ms Slaney-mel magnitude windows at a 50 ms hop, quantized to 16 bins. +struct mtmd_audio_preprocessor_inkling : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_inkling(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} + void initialize() override; + bool preprocess(const float * samples, size_t n_samples, std::vector & output) override; + +private: + mtmd_audio_cache cache; +}; + struct mtmd_audio_preprocessor_conformer : mtmd_audio_preprocessor { mtmd_audio_preprocessor_conformer(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} void initialize() override; diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index 36cd463b20eb..e58c79882a73 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -68,6 +68,9 @@ struct img_tool { case RESIZE_ALGO_BICUBIC_PILLOW: resize_bicubic_pillow(src, dst, target_resolution.width, target_resolution.height); break; + case RESIZE_ALGO_LANCZOS: + resize_lanczos_pillow(src, dst, target_resolution.width, target_resolution.height); + break; default: throw std::runtime_error("Unsupported resize algorithm"); } @@ -97,6 +100,9 @@ struct img_tool { case RESIZE_ALGO_BICUBIC_PILLOW: resize_bicubic_pillow(src, resized_image, new_width, new_height); break; + case RESIZE_ALGO_LANCZOS: + resize_lanczos_pillow(src, resized_image, new_width, new_height); + break; default: throw std::runtime_error("Unsupported resize algorithm"); } @@ -345,6 +351,19 @@ struct img_tool { // 2. Pre-computes normalized filter coefficients for each output pixel // 3. Applies convolution using fixed-point integer arithmetic for performance static bool resize_bicubic_pillow(const clip_image_u8 & img, clip_image_u8 & dst, int target_width, int target_height) { + return resize_pillow(img, dst, target_width, target_height, false); + } + + static bool resize_lanczos_pillow(const clip_image_u8 & img, clip_image_u8 & dst, int target_width, int target_height) { + return resize_pillow(img, dst, target_width, target_height, true); + } + + static bool resize_pillow( + const clip_image_u8 & img, + clip_image_u8 & dst, + int target_width, + int target_height, + bool use_lanczos) { // Fixed-point precision: 22 bits = 32 (int32_t) - 8 (uint8_t pixels) - 2 (headroom for accumulation) // This allows encoding fractional weights as integers: weight * 2^22 const int PRECISION_BITS = 32 - 8 - 2; @@ -352,7 +371,21 @@ struct img_tool { // Bicubic filter function with a = -0.5 (Note that GGML/PyTorch takes a = -0.75) // Returns filter weight for distance x from pixel center // Support: [-2, 2], meaning the filter influences pixels within 2 units of distance - auto bicubic_filter = [](double x) -> double { + auto resample_filter = [use_lanczos](double x) -> double { + if (use_lanczos) { + if (-3.0 <= x && x < 3.0) { + auto sinc = [](double value) { + if (value == 0.0) { + return 1.0; + } + const double pix = value * 3.141592653589793238462643383279502884; + return std::sin(pix) / pix; + }; + return sinc(x) * sinc(x / 3.0); + } + return 0.0; + } + constexpr double a = -0.5; if (x < 0.0) { x = -x; @@ -367,7 +400,7 @@ struct img_tool { }; // Filter support radius: bicubic extends 2 pixels in each direction - constexpr double filter_support = 2.0; + const double filter_support = use_lanczos ? 3.0 : 2.0; // Clipping function for 8-bit values auto clip8 = [](int val) -> uint8_t { @@ -434,7 +467,7 @@ struct img_tool { // Compute filter weights for each contributing input pixel for (x = 0; x < xmax; x++) { // Distance from input pixel center to output pixel center in input space - double w = bicubic_filter((x + xmin - center + 0.5) * ss); + double w = resample_filter((x + xmin - center + 0.5) * ss); pre_weights[xx * ksize + x] = w; ww += w; // Accumulate for normalization } @@ -463,6 +496,12 @@ struct img_tool { const double fxp_scale = std::ldexp(1.0, PRECISION_BITS); // 1.0 * 2^PRECISION_BITS for (int i = 0; i < outSize * ksize; i++) { + if (use_lanczos) { + // Pillow adds +/- 0.5 then truncates toward zero; std::round would round twice + const double rounded = pre_weights[i] * fxp_scale + (pre_weights[i] < 0 ? -0.5 : 0.5); + weights[i] = static_cast(rounded); + continue; + } double tmp_val = pre_weights[i] * fxp_scale; if (pre_weights[i] < 0) { tmp_val -= 0.5; @@ -606,6 +645,82 @@ struct img_tool { } }; +mtmd_inkling_image_preproc_out mtmd_image_preprocess_inkling( + const clip_image_u8 & img, + resize_algo algo) { + constexpr int patch_size = 40; + constexpr int temporal_patch_size = 2; + constexpr int max_upscaled_long_edge = 2048; + constexpr double rescale_image_frac = 2.0; + constexpr float image_mean[3] = { 0.48145466f, 0.4578275f, 0.40821073f }; + constexpr float image_std[3] = { 0.26862954f, 0.26130258f, 0.27577711f }; + + mtmd_inkling_image_preproc_out output; + output.source_size = img.get_size(); + GGML_ASSERT(output.source_size.width > 0 && output.source_size.height > 0); + GGML_ASSERT(!img.is_placeholder()); + + const int long_edge = std::max(output.source_size.width, output.source_size.height); + const double target_long_edge = std::min( + static_cast(long_edge) * rescale_image_frac, + static_cast(std::max(max_upscaled_long_edge, long_edge))); + const double ratio = target_long_edge / long_edge; + const auto scaled = [ratio](int size) { + // The reference explicitly requests half-up rounding for positive sizes. + return std::max(1, static_cast(std::floor(size * ratio + 0.5))); + }; + output.resized_size = { + scaled(output.source_size.width), + scaled(output.source_size.height), + }; + + clip_image_u8 resized; + img_tool::resize(img, resized, output.resized_size, algo, PAD_NONE); + output.resized_rgb = resized.get_ro_buf(); + + output.patch_rows = (output.resized_size.height + patch_size - 1) / patch_size; + // Inkling always appends a right-hand patch, even at exact multiples of 40 + output.patch_cols = output.resized_size.width / patch_size + 1; + const size_t n_patches = static_cast(output.patch_rows) * output.patch_cols; + const size_t values_per_patch = static_cast(temporal_patch_size) * + patch_size * patch_size * 3; + output.pixel_values_bthwc.resize(n_patches * values_per_patch); + + // preserve torchvision's fused (raw - mean*255)/(std*255) order (raw=-1 for padding) for ulp parity + float mean_255[3]; + float std_255[3]; + for (int c = 0; c < 3; ++c) { + mean_255[c] = image_mean[c] * 255.0f; + std_255[c] = image_std[c] * 255.0f; + } + + for (int py = 0; py < output.patch_rows; ++py) { + for (int px = 0; px < output.patch_cols; ++px) { + const size_t patch_index = static_cast(py) * output.patch_cols + px; + for (int t = 0; t < temporal_patch_size; ++t) { + for (int y = 0; y < patch_size; ++y) { + const int iy = py * patch_size + y; + for (int x = 0; x < patch_size; ++x) { + const int ix = px * patch_size + x; + const bool in_image = iy < output.resized_size.height && ix < output.resized_size.width; + const std::array rgb = in_image + ? resized.get_pixel(ix, iy) + : std::array{ 0, 0, 0 }; + for (int c = 0; c < 3; ++c) { + const float raw = in_image ? static_cast(rgb[c]) : -1.0f; + const size_t offset = (((((patch_index * temporal_patch_size + t) * + patch_size + y) * patch_size + x) * 3) + c); + output.pixel_values_bthwc[offset] = (raw - mean_255[c]) / std_255[c]; + } + } + } + } + } + } + + return output; +} + // // mtmd_image_preprocessor_llava_uhd @@ -885,6 +1000,80 @@ mtmd_image_preproc_out mtmd_image_preprocessor_fixed_size::preprocess(const clip return output; } +// Inkling hMLP patches +mtmd_image_preproc_out mtmd_image_preprocessor_inkling::preprocess(const clip_image_u8 & img) { + GGML_ASSERT(hparams.patch_size == 40); + constexpr int temporal_patch_size = 2; + constexpr float rescale_frac = 2.0f; + constexpr int rescale_max_long_edge = 2048; + + const auto original = img.get_size(); + const int long_edge = std::max(original.width, original.height); + const float target_long = std::min( + static_cast(long_edge) * rescale_frac, + static_cast(std::max(long_edge, rescale_max_long_edge))); + const float ratio = long_edge > 0 ? target_long / long_edge : 1.0f; + const clip_image_size scaled_size { + std::max(1, static_cast(std::floor(original.width * ratio + 0.5f))), + std::max(1, static_cast(std::floor(original.height * ratio + 0.5f))), + }; + + clip_image_u8 scaled; + // The reference processor resizes with Pillow Lanczos (radius 3). + img_tool::resize(img, scaled, scaled_size, RESIZE_ALGO_LANCZOS, PAD_NONE); + + const int p = hparams.patch_size; + const int nph = (scaled_size.height + p - 1) / p; + // The extra right column is intentional, including exact multiples of 40. + const int npw = scaled_size.width / p + 1; + + const float pad_raw = -1.0f / 255.0f; + // the reference materializes normalized patches as bf16; round through bf16 for bit-parity + const auto bf16_round = [](float v) { + return ggml_bf16_to_fp32(ggml_fp32_to_bf16(v)); + }; + float pad_norm[3]; + for (int c = 0; c < 3; ++c) { + pad_norm[c] = bf16_round((pad_raw - hparams.image_mean[c]) / hparams.image_std[c]); + } + + mtmd_image_preproc_out output; + output.entries.reserve(static_cast(nph) * npw * temporal_patch_size); + for (int py = 0; py < nph; ++py) { + for (int px = 0; px < npw; ++px) { + std::vector data(static_cast(p) * p * 3); + for (int y = 0; y < p; ++y) { + const int iy = py * p + y; + for (int x = 0; x < p; ++x) { + const int ix = px * p + x; + const size_t off = static_cast(y * p + x) * 3; + if (iy < scaled_size.height && ix < scaled_size.width && !scaled.is_placeholder()) { + const auto rgb = scaled.get_pixel(ix, iy); + for (int c = 0; c < 3; ++c) { + const float raw = static_cast(rgb[c]) / 255.0f; + data[off + c] = bf16_round((raw - hparams.image_mean[c]) / hparams.image_std[c]); + } + } else { + for (int c = 0; c < 3; ++c) { + data[off + c] = pad_norm[c]; + } + } + } + } + + for (int t = 0; t < temporal_patch_size; ++t) { + clip_image_f32 patch; + patch.set_size({p, p}, scaled.is_placeholder(), false); + if (!scaled.is_placeholder()) { + patch.cpy_buf(data); + } + output.entries.push_back(std::move(patch)); + } + } + } + return output; +} + // // mtmd_image_preprocessor_dyn_size // diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index 115cba51e8f4..044d72f61375 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -26,6 +26,20 @@ struct mtmd_image_preproc_out { } }; +// Inkling hMLP preprocessing; pixel_values_bthwc is row-major BTHWC [n_patches, 2, 40, 40, 3]. +struct mtmd_inkling_image_preproc_out { + clip_image_size source_size; + clip_image_size resized_size; + int patch_rows = 0; + int patch_cols = 0; + std::vector resized_rgb; + std::vector pixel_values_bthwc; +}; + +mtmd_inkling_image_preproc_out mtmd_image_preprocess_inkling( + const clip_image_u8 & img, + resize_algo algo = RESIZE_ALGO_LANCZOS); + // base class, models must inherit from this class struct mtmd_image_preprocessor { const clip_hparams & hparams; @@ -115,6 +129,12 @@ struct mtmd_image_preprocessor_fixed_size : mtmd_image_preprocessor { mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; }; +// Inkling: split an image into 40x40 hMLP patches, each duplicated across a fixed temporal dimension of two. +struct mtmd_image_preprocessor_inkling : mtmd_image_preprocessor { + mtmd_image_preprocessor_inkling(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} + mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; +}; + // resize image to multiple of patch_size*n_merge, while preserving aspect ratio // if image_resize_pad is true, the resized image will be padded, otherwise it will be either stretched or center-cropped depending on image_resize_pad // this is used by models with native support for dynamic image size, for example: Qwen-VL, Pixtral, Kimi-VL, etc diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 93ca8cbcf8ae..a70f326a9673 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -393,6 +393,12 @@ struct mtmd_context { projector_type proj = clip_get_projector_type(ctx_v); switch (proj) { + case PROJECTOR_TYPE_INKLING: + { + // renderer opens each image part with <|content_image|>; block framing comes from the template + img_beg = "<|content_image|>"; + image_preproc = std::make_unique(ctx_v); + } break; case PROJECTOR_TYPE_MLP: case PROJECTOR_TYPE_MLP_NORM: case PROJECTOR_TYPE_LDP: @@ -678,6 +684,13 @@ struct mtmd_context { // set preprocessor switch (proj) { + case PROJECTOR_TYPE_INKLING: + { + // <|content_audio_input|> ... <|audio_end|>, matching the renderer's audio framing + aud_beg = "<|content_audio_input|>"; + aud_end = "<|audio_end|>"; + audio_preproc = std::make_unique(ctx_a); + } break; case PROJECTOR_TYPE_QWEN2A: case PROJECTOR_TYPE_QWEN25O: { @@ -1204,23 +1217,35 @@ struct mtmd_tokenizer { } size_t n_tokens = 0; - for (auto & e : preproc_out.entries) { - n_tokens += clip_n_output_tokens(ctx->ctx_v, &e); - if (clip_model_n_temporal_merge(ctx->ctx_v) == 2) { - // [QWEN_VIDEO] pair input is merged to the same embd, so only count as one image - break; + if (ctx->proj_type_v() == PROJECTOR_TYPE_INKLING) { + GGML_ASSERT(preproc_out.entries.size() % 2 == 0); + n_tokens = preproc_out.entries.size() / 2; + } else { + for (auto & e : preproc_out.entries) { + n_tokens += clip_n_output_tokens(ctx->ctx_v, &e); + if (clip_model_n_temporal_merge(ctx->ctx_v) == 2) { + // [QWEN_VIDEO] pair input is merged to the same embd, so only count as one image + break; + } } } mtmd_image_tokens_ptr image_tokens(new mtmd_image_tokens); // [QWEN_VIDEO] improve this in the future - image_tokens->n_temporal_merge = clip_model_n_temporal_merge(ctx->ctx_v); + // Inkling's pair is internal to each patch; it must not merge consecutive user bitmaps as video frames. + image_tokens->n_temporal_merge = ctx->proj_type_v() == PROJECTOR_TYPE_INKLING + ? 2 + : clip_model_n_temporal_merge(ctx->ctx_v); if (mtmd_decode_use_mrope(ctx)) { // for Qwen2VL, we need this information for M-RoPE decoding positions image_tokens->nx = clip_n_output_tokens_x(ctx->ctx_v, &preproc_out.entries[0]); image_tokens->ny = clip_n_output_tokens_y(ctx->ctx_v, &preproc_out.entries[0]); + } else if (ctx->proj_type_v() == PROJECTOR_TYPE_INKLING) { + // n_tokens() multiplies 1x1 by the number of temporal groups. + image_tokens->nx = 1; + image_tokens->ny = 1; } else { // other models, we only need the total number of tokens image_tokens->nx = n_tokens;