From ce16fff2ad1d1ddfb39365fedfb4fcfd6db178f1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 18 Jul 2026 04:55:29 +0000 Subject: [PATCH 1/2] Add TML Inkling architecture Hybrid attention model: 55 sliding-window plus 11 global layers, banded content-dependent relative position bias instead of RoPE, per-layer short convolution state, fine-grained MoE (256 experts top-6 plus 2 shared), attention log-scaling past 128K, 1M context. Includes the GGML_OP_FLASH_ATTN_EXT_BANDED operator (CPU and CUDA, fused into the MMA flash attention kernel with an fp16 accumulator overflow guard), HF to GGUF conversion, chat template with typed content block parsing (interleaved thinking, narration and tool calls), mmproj vision and audio support, and backend op tests at production shapes. --- common/chat-peg-parser.cpp | 26 +- common/chat-peg-parser.h | 12 +- common/chat.cpp | 127 +++++ conversion/__init__.py | 2 + conversion/base.py | 31 ++ conversion/inkling.py | 350 +++++++++++++ ggml/include/ggml-rpc.h | 4 +- ggml/include/ggml.h | 20 +- ggml/src/ggml-backend-meta.cpp | 13 + ggml/src/ggml-cpu/ggml-cpu.c | 3 + ggml/src/ggml-cpu/ops.cpp | 44 +- ggml/src/ggml-cuda/argsort.cu | 28 +- ggml/src/ggml-cuda/fattn-banded.cu | 247 +++++++++ ggml/src/ggml-cuda/fattn-banded.cuh | 7 + ggml/src/ggml-cuda/fattn-common.cuh | 17 +- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 112 ++++- ggml/src/ggml-cuda/fattn.cu | 10 +- ggml/src/ggml-cuda/ggml-cuda.cu | 68 ++- ggml/src/ggml-cuda/mmf.cuh | 8 +- ggml/src/ggml-cuda/mmq.cuh | 12 +- ggml/src/ggml-cuda/mmvf.cu | 4 +- ggml/src/ggml-cuda/mmvq.cu | 4 +- ggml/src/ggml-cuda/pad.cu | 2 +- ggml/src/ggml-cuda/ssm-conv.cu | 28 +- ggml/src/ggml-rpc/ggml-rpc.cpp | 1 + ggml/src/ggml.c | 59 ++- gguf-py/gguf/constants.py | 50 ++ gguf-py/gguf/tensor_mapping.py | 49 ++ models/templates/Inkling.jinja | 514 +++++++++++++++++++ src/llama-arch.cpp | 35 ++ src/llama-arch.h | 24 + src/llama-graph.cpp | 8 +- src/llama-graph.h | 3 +- src/llama-hparams.cpp | 5 + src/llama-hparams.h | 11 + src/llama-kv-cache.cpp | 97 ++++ src/llama-kv-cache.h | 13 +- src/llama-model-saver.cpp | 1 + src/llama-model.cpp | 6 +- src/llama-model.h | 9 + src/llama-quant.cpp | 10 + src/llama-vocab.cpp | 10 + src/llama-vocab.h | 1 + src/models/inkling.cpp | 666 +++++++++++++++++++++++++ src/models/models.h | 13 + tests/CMakeLists.txt | 8 + tests/test-backend-ops.cpp | 114 ++++- tests/test-chat.cpp | 79 +++ tests/test-flash-attn-bias.cpp | 525 +++++++++++++++++++ tests/test-flash-attn-generic-hash.cpp | 142 ++++++ tests/test-llama-archs.cpp | 3 + tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-impl.h | 9 + tools/mtmd/clip-model.h | 13 +- tools/mtmd/clip.cpp | 62 ++- tools/mtmd/models/inkling.cpp | 106 ++++ tools/mtmd/models/models.h | 11 + tools/mtmd/mtmd-audio.cpp | 71 +++ tools/mtmd/mtmd-audio.h | 10 + tools/mtmd/mtmd-image.cpp | 195 +++++++- tools/mtmd/mtmd-image.h | 20 + tools/mtmd/mtmd.cpp | 37 +- 62 files changed, 4048 insertions(+), 122 deletions(-) create mode 100644 conversion/inkling.py create mode 100644 ggml/src/ggml-cuda/fattn-banded.cu create mode 100644 ggml/src/ggml-cuda/fattn-banded.cuh create mode 100644 models/templates/Inkling.jinja create mode 100644 src/models/inkling.cpp create mode 100644 tests/test-flash-attn-bias.cpp create mode 100644 tests/test-flash-attn-generic-hash.cpp create mode 100644 tools/mtmd/models/inkling.cpp diff --git a/common/chat-peg-parser.cpp b/common/chat-peg-parser.cpp index a309f02765b7..52829b7c4b20 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 b3ffd7de2dd8..a2b7cbd4d303 100644 --- a/common/chat-peg-parser.h +++ b/common/chat-peg-parser.h @@ -121,7 +121,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 @@ -145,13 +146,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, @@ -159,7 +162,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 22d2ee4a2a11..f6cf981e75f7 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -2175,6 +2175,125 @@ static common_chat_params common_chat_params_init_cohere2moe(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_tag = 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) { @@ -2592,6 +2711,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 465c010f667a..73da6ca5cc1d 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -109,6 +109,7 @@ "HunYuanVLForConditionalGeneration": "hunyuan", "HYV3ForCausalLM": "hunyuan", "IQuestCoderForCausalLM": "llama", + "InklingForConditionalGeneration": "inkling", "InternLM2ForCausalLM": "internlm", "InternLM3ForCausalLM": "internlm", "JAISLMHeadModel": "jais", @@ -273,6 +274,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 1b85ef0a1b89..b780901887c3 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 00ffacf29921..5d51a956d4f2 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -392,6 +392,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: @@ -544,7 +552,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 4aaa4f367afc..0dbb38a26f2a 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; @@ -5110,6 +5152,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 ad3d6a5a6832..d7f22f1e759c 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -434,7 +434,7 @@ template static __device__ __forceinline_ continue; } - 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]; } } } @@ -475,7 +475,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma( continue; } - 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]; } } } @@ -944,7 +944,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; if (ids_dst) { col_low = expert_bounds[zt + 0]; @@ -1017,7 +1017,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; if (ids_dst) { col_low = expert_bounds[zt + 0]; @@ -1091,7 +1091,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; if (ids_dst) { col_low = expert_bounds[zt + 0]; @@ -1219,7 +1219,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 a7d1fe7d94be..9c3d172a8086 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 63ac2ed1f5cf..71809cf21f94 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -540,6 +540,7 @@ class MODEL_ARCH(IntEnum): KIMI_LINEAR = auto() TALKIE = auto() MELLUM = auto() + INKLING = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -763,6 +764,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() @@ -1122,6 +1131,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.KIMI_LINEAR: "kimi-linear", MODEL_ARCH.TALKIE: "talkie", MODEL_ARCH.MELLUM: "mellum", + MODEL_ARCH.INKLING: "inkling", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -1343,6 +1353,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", @@ -4044,6 +4061,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, @@ -4665,6 +4714,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 9efb36f8a447..3fbeb2e53743 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 @@ -2415,6 +2417,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 72968607db80..a01bad68df92 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -141,6 +141,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, { LLM_ARCH_TALKIE, "talkie" }, { LLM_ARCH_MELLUM, "mellum" }, + { LLM_ARCH_INKLING, "inkling" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; @@ -311,6 +312,18 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" }, { 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" }, @@ -583,9 +596,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" }, @@ -780,6 +803,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}}, @@ -821,6 +847,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}}, @@ -946,6 +979,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; @@ -1001,6 +1035,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_MINIMAX_M2: 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 b74d53af4a32..9d2eb4a92028 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -146,6 +146,7 @@ enum llm_arch { LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, LLM_ARCH_DFLASH, + LLM_ARCH_INKLING, LLM_ARCH_UNKNOWN, }; @@ -358,6 +359,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, @@ -425,6 +437,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, @@ -586,6 +602,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 c8ecb0a2854c..9a390c7f26f3 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 9d0683d2fec4..d0cf45f64711 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -181,6 +181,11 @@ uint32_t llama_hparams::n_embd_v_gqa_max() 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 8be5f28f39e6..38b8640326cb 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -79,6 +79,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 e70583e64152..812f8293b271 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1240,6 +1240,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); @@ -1779,6 +1824,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)); @@ -2567,6 +2645,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(); } @@ -2627,6 +2720,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 531d99dbdec1..08248f17c73f 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; @@ -214,6 +217,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; @@ -363,6 +370,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; @@ -395,6 +403,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; @@ -402,8 +411,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 a3928523ba8d..0d39901abf71 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -28,6 +28,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) { case LLM_ARCH_MIMO2: case LLM_ARCH_STEP35: case LLM_ARCH_MELLUM: + case LLM_ARCH_INKLING: return false; default: return true; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 4c10e4126f62..b7c36ca95b36 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -304,6 +304,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) + "'"); } @@ -2102,7 +2104,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) { @@ -2444,6 +2447,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 45b054cedf1d..57c7caf37ddb 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -525,6 +525,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 b66759b27760..94b30f9ac57c 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -322,6 +322,16 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param quantize &= name.find("ssm_conv1d") == std::string::npos; quantize &= name.find("shortconv.conv.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; quantize &= name.find("time_mix_w0.weight") == std::string::npos; diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index fdd447147d43..3197843e002d 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 @@ -2286,6 +2292,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 707cd4bac4bf..be5d77b53ad7 100644 --- a/src/llama-vocab.h +++ b/src/llama-vocab.h @@ -64,6 +64,7 @@ enum llama_vocab_pre_type { LLAMA_VOCAB_PRE_TYPE_WHITESPACE = 53, LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI = 54, LLAMA_VOCAB_PRE_TYPE_MELLUM2 = 55, + LLAMA_VOCAB_PRE_TYPE_INKLING = 56, }; 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 a86ae05aae6b..d13c04c57fca 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1835,6 +1835,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 e3553c5e9ca1..243d2aaeef73 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -6912,6 +6912,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; @@ -7491,6 +7586,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]; @@ -7514,7 +7610,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; } @@ -9542,6 +9638,22 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(64, 128, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q1_0)); test_cases.emplace_back(new test_flash_attn_ext(128, 64, 4, {1, 1}, 64, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q1_0, 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})); test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 93685ec8ff42..1b8ed9b1de02 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -2909,6 +2909,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 57d33a627d5d..9ada66a242fa 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -418,6 +418,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 some models are segfaulting with WebGPU: #ifdef GGML_USE_WEBGPU diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index ea684d9f156d..12ff725e37ca 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -31,6 +31,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 5b413681f040..c874a2c7c99c 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -258,6 +258,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" @@ -320,6 +327,7 @@ struct clip_ctx; enum projector_type { + PROJECTOR_TYPE_INKLING, PROJECTOR_TYPE_MLP, PROJECTOR_TYPE_MLP_NORM, PROJECTOR_TYPE_LDP, @@ -375,6 +383,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 6d4336c4010b..9b5cb09a8f6f 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 @@ -347,6 +347,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; @@ -361,6 +366,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 b8866506493e..dbd07081bf73 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -869,6 +869,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: @@ -1277,6 +1281,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: @@ -1856,7 +1875,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; @@ -1942,6 +1962,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: { @@ -2906,6 +2943,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; } @@ -3303,6 +3345,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: @@ -3675,6 +3727,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: @@ -4551,6 +4607,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 5f1493fa603e..e50785e525bf 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 b72fd067a508..7d6cc83961d8 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 ad96bd847cfc..29613d11146c 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 e10ccf186b14..3e81e44143fa 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: @@ -663,6 +669,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: { @@ -1179,23 +1192,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; From d8f87500a3dfab0933cf49bfa5d447e83662d3f1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 26 Jul 2026 11:55:01 +0000 Subject: [PATCH 2/2] inkling : adapt to the renamed thinking_end_tags chat API Upstream 910196f6b renamed common_chat_params::thinking_end_tag to thinking_end_tags and changed it from std::string to a vector, so data.thinking_end_tag = END_MESSAGE no longer compiles. --- common/chat.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/common/chat.cpp b/common/chat.cpp index 5dd8fd8979bd..c589675e52ec 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -2291,7 +2291,7 @@ static common_chat_params common_chat_params_init_inkling(const common_chat_temp data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; data.supports_thinking = true; data.thinking_start_tag = THINK; - data.thinking_end_tag = END_MESSAGE; + 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,