diff --git a/common/arg.cpp b/common/arg.cpp index 05e5cd878..4ef089053 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3902,6 +3902,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex common_log_set_file(common_log_main(), value.c_str()); } ).set_env("LLAMA_ARG_LOG_FILE")); + add_opt(common_arg( + {"--log-jsonl"}, + {"--no-log-jsonl"}, + "Log as JSONL (one JSON object per line) to stdout, this also disables colored logging (default: disabled)", + [](common_params &, bool value) { + common_log_set_jsonl(common_log_main(), value); + } + ).set_env("LLAMA_ARG_LOG_JSONL")); add_opt(common_arg( {"--log-prompts-dir"}, "PATH", "Log prompts to directory (auto-created if not present; only used for debugging, default: disabled)", diff --git a/common/jinja/caps.cpp b/common/jinja/caps.cpp index 567de5ba7..78a77b2cc 100644 --- a/common/jinja/caps.cpp +++ b/common/jinja/caps.cpp @@ -120,6 +120,7 @@ caps caps_get(jinja::program & prog) { JJ_DEBUG("%s\n", ">>> Running capability check: typed content"); + bool checks_for_string = false; static const std::string content_marker = "STRING_MARKER"; // case: typed content support @@ -139,6 +140,10 @@ caps caps_get(jinja::program & prog) { [&](context &, bool success, value & messages, value &, const std::string & rendered) { auto & content = messages->at(0)->at("content"); caps_print_stats(content, "messages[0].content"); + if (has_op(content, "test_is_string")) { + // checked if content is string + checks_for_string = true; + } bool used_as_array = has_op(content, "selectattr") || has_op(content, "array_access"); if (used_as_array) { // accessed as an array @@ -154,6 +159,33 @@ caps caps_get(jinja::program & prog) { } ); + if (checks_for_string) { + caps_try_execute( + prog, + [&]() { + // messages + return json::array({ + { + {"role", "user"}, + {"content", json::array({ + })} + } + }); + }, + nullptr, // ctx_fn + nullptr, // tools_fn + [&](context &, bool success, value & messages, value &, const std::string &) { + auto & content = messages->at(0)->at("content"); + caps_print_stats(content, "messages[0].content"); + bool used_as_array = has_op(content, "selectattr") || has_op(content, "array_access"); + if (used_as_array && success) { + // accessed as an array + result.supports_typed_content = true; + } + } + ); + } + JJ_DEBUG("%s\n", ">>> Running capability check: system prompt"); // case: system prompt support diff --git a/common/jinja/runtime.cpp b/common/jinja/runtime.cpp index 9cc5cfa46..d5625184c 100644 --- a/common/jinja/runtime.cpp +++ b/common/jinja/runtime.cpp @@ -415,12 +415,18 @@ value test_expression::execute_impl(context & ctx) { throw std::runtime_error("Invalid test expression"); } - auto it = builtins.find("test_is_" + test_id); - JJ_DEBUG("Test expression %s '%s' %s (using function 'test_is_%s')", operand->type().c_str(), test_id.c_str(), negate ? "(negate)" : "", test_id.c_str()); + const std::string test_name = "test_is_" + test_id; + auto it = builtins.find(test_name); + JJ_DEBUG("Test expression %s '%s' %s (using function '%s')", operand->type().c_str(), test_id.c_str(), negate ? "(negate)" : "", test_name.c_str()); if (it == builtins.end()) { throw std::runtime_error("Unknown test '" + test_id + "'"); } + if (ctx.is_get_stats) { + value_t::stats_t::mark_used(input); + input->stats.ops.insert(test_name); + } + auto res = it->second(args); if (negate) { diff --git a/common/log.cpp b/common/log.cpp index 0f0cb7902..42951190c 100644 --- a/common/log.cpp +++ b/common/log.cpp @@ -1,5 +1,6 @@ #include "common.h" #include "log.h" +#include "json.h" #include #include @@ -66,6 +67,17 @@ static const char* g_col[] = { "", }; +static const char * level_str(enum ggml_log_level level) { + switch (level) { + case GGML_LOG_LEVEL_DEBUG: return "debug"; + case GGML_LOG_LEVEL_INFO: return "info"; + case GGML_LOG_LEVEL_WARN: return "warn"; + case GGML_LOG_LEVEL_ERROR: return "error"; + case GGML_LOG_LEVEL_CONT: return "cont"; + default: return "none"; + } +} + struct common_log_entry { enum ggml_log_level level {GGML_LOG_LEVEL_INFO}; @@ -74,6 +86,7 @@ struct common_log_entry { int64_t timestamp { 0 }; bool is_end { false }; // signals the worker thread to stop bool prefix { false }; + bool jsonl { false }; common_log_entry(size_t size = 256) : msg(size) { } @@ -88,11 +101,23 @@ struct common_log_entry { fcur = stdout; - if (level != GGML_LOG_LEVEL_NONE) { + if (level != GGML_LOG_LEVEL_NONE && !jsonl) { fcur = stderr; } } + if (jsonl) { + common_json obj = { + {"type", "log"}, + {"time", timestamp}, + {"level", level_str(level)}, + {"msg", msg.data()}, + }; + fprintf(fcur, "%s\n", obj.dump_safe().c_str()); + fflush(fcur); + return; + } + if (level != GGML_LOG_LEVEL_NONE && level != GGML_LOG_LEVEL_CONT && prefix) { if (timestamp) { // [M.s.ms.us] @@ -131,6 +156,7 @@ struct common_log { file = nullptr; prefix = false; timestamps = false; + jsonl = false; running = false; t_start = t_us(); @@ -158,6 +184,7 @@ private: bool prefix; bool timestamps; + bool jsonl; bool running; int64_t t_start; @@ -246,6 +273,7 @@ public: entry.is_end = false; entry.level = level; entry.prefix = prefix; + entry.jsonl = jsonl; entry.timestamp = 0; if (timestamps) { entry.timestamp = t_us() - t_start; @@ -360,6 +388,12 @@ public: this->timestamps = timestamps; } + + void set_jsonl(bool jsonl) { + std::lock_guard lock(mtx); + + this->jsonl = jsonl; + } }; // @@ -433,6 +467,10 @@ void common_log_set_timestamps(struct common_log * log, bool timestamps) { log->set_timestamps(timestamps); } +void common_log_set_jsonl(struct common_log * log, bool jsonl) { + log->set_jsonl(jsonl); +} + void common_log_flush(struct common_log * log) { log->pause(); log->resume(); diff --git a/common/log.h b/common/log.h index f03358252..37f4de92b 100644 --- a/common/log.h +++ b/common/log.h @@ -91,6 +91,7 @@ void common_log_set_file (struct common_log * log, const char * file); // n void common_log_set_colors (struct common_log * log, log_colors colors); // not thread-safe void common_log_set_prefix (struct common_log * log, bool prefix); // whether to output prefix to each log void common_log_set_timestamps(struct common_log * log, bool timestamps); // whether to output timestamps in the prefix +void common_log_set_jsonl (struct common_log * log, bool jsonl); // print each log as a JSON object on one line, not thread-safe void common_log_flush (struct common_log * log); // flush all pending log messages // helper macros for logging diff --git a/conversion/__init__.py b/conversion/__init__.py index 94d6a49fb..4d58bcd10 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -255,6 +255,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "SeedOssForCausalLM": "olmo", "SmallThinkerForCausalLM": "smallthinker", "SmolLM3ForCausalLM": "llama", + "Spark2_5ForCausalLM": "spark2_5", "SolarOpenForCausalLM": "glm", "StableLMEpochForCausalLM": "stablelm", "StableLmForCausalLM": "stablelm", diff --git a/conversion/base.py b/conversion/base.py index c1ecf1c65..d2d80be36 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -130,7 +130,8 @@ class ModelBase: sentence_transformers_dense_modules: bool = False, target_model_dir: Path | None = None, fuse_gate_up_exps: bool = False, - fp8_as_q8: bool = False): + fp8_as_q8: bool = False, + fuse_qkv: bool = False): if type(self) is ModelBase or \ type(self) is TextModel or \ type(self) is MmprojModel: @@ -153,6 +154,15 @@ class ModelBase: self.fuse_gate_up_exps = fuse_gate_up_exps self._gate_exp_buffer: dict[int, Tensor] = {} self._up_exp_buffer: dict[int, Tensor] = {} + self.fuse_qkv = fuse_qkv + self._q_buffer: dict[int, Tensor] = {} + self._k_buffer: dict[int, Tensor] = {} + self._v_buffer: dict[int, Tensor] = {} + self._q_bias_buffer: dict[int, Tensor] = {} + self._k_bias_buffer: dict[int, Tensor] = {} + self._v_bias_buffer: dict[int, Tensor] = {} + self._fusable_qkv_weight_layers: set[int] = set() + self._fusable_qkv_bias_layers: set[int] = set() self.hparams = ModelBase.load_hparams(self.dir_model, self.is_mistral_format) if hparams is None else hparams self.model_tensors = self.index_tensors(remote_hf_model_id=remote_hf_model_id) self.metadata_override = metadata_override @@ -617,6 +627,43 @@ class ModelBase: raise ValueError(f"Can not map tensor {name!r}") return new_name + def prepare_qkv_fusion(self) -> None: + self._fusable_qkv_weight_layers.clear() + self._fusable_qkv_bias_layers.clear() + if not self.fuse_qkv or gguf.MODEL_TENSOR.ATTN_QKV not in gguf.MODEL_TENSORS[self.model_arch]: + return + + qkv_types = { + gguf.MODEL_TENSOR.ATTN_Q, + gguf.MODEL_TENSOR.ATTN_K, + gguf.MODEL_TENSOR.ATTN_V, + } + weights: dict[int, set[gguf.MODEL_TENSOR]] = {} + biases: dict[int, set[gguf.MODEL_TENSOR]] = {} + + for name in self.model_tensors: + mapped = self.tensor_map.get_type_and_name(name, try_suffixes=(".weight", ".bias")) + if mapped is None: + continue + tensor_type, new_name = mapped + if tensor_type not in qkv_types: + continue + + bid = next((int(part) for part in new_name.split(".") if part.isdecimal()), None) + if bid is None: + continue + if new_name.endswith(".weight"): + weights.setdefault(bid, set()).add(tensor_type) + elif new_name.endswith(".bias"): + biases.setdefault(bid, set()).add(tensor_type) + + for bid, weight_types in weights.items(): + bias_types = biases.get(bid, set()) + if weight_types == qkv_types and (not bias_types or bias_types == qkv_types): + self._fusable_qkv_weight_layers.add(bid) + if bias_types: + self._fusable_qkv_bias_layers.add(bid) + def set_gguf_parameters(self): raise NotImplementedError("set_gguf_parameters() must be implemented in subclasses") @@ -645,6 +692,40 @@ class ModelBase: self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.FFN_UP_EXP, bid): return [] + # Handle Q/K/V tensor fusion if enabled + qkv_bid = next((int(part) for part in new_name.split(".") if part.isdecimal()), None) if self.fuse_qkv else None + if qkv_bid is not None: + is_bias = new_name.endswith('.bias') + suffix = '.bias' if is_bias else '.weight' + fusable_layers = self._fusable_qkv_bias_layers if is_bias else self._fusable_qkv_weight_layers + if qkv_bid not in fusable_layers: + return [(new_name, data_torch)] + + buf_q = self._q_bias_buffer if is_bias else self._q_buffer + buf_k = self._k_bias_buffer if is_bias else self._k_buffer + buf_v = self._v_bias_buffer if is_bias else self._v_buffer + + if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_Q, qkv_bid, suffix): + buf_q[qkv_bid] = data_torch + elif self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_K, qkv_bid, suffix): + buf_k[qkv_bid] = data_torch + elif self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_V, qkv_bid, suffix): + buf_v[qkv_bid] = data_torch + + if qkv_bid in buf_q and qkv_bid in buf_k and qkv_bid in buf_v: + q_data = buf_q.pop(qkv_bid) + k_data = buf_k.pop(qkv_bid) + v_data = buf_v.pop(qkv_bid) + fused_data = torch.cat([q_data, k_data, v_data], dim=0) + fused_name = self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_QKV, qkv_bid, suffix=suffix) + logger.info(f"Fused Q, K, V {suffix[1:]} into QKV for layer {qkv_bid}") + return [(fused_name, fused_data)] + + if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_Q, qkv_bid, suffix) or \ + self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_K, qkv_bid, suffix) or \ + self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.ATTN_V, qkv_bid, suffix): + return [] + return [(new_name, data_torch)] def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool: @@ -899,6 +980,8 @@ class ModelBase: self.dequant_model() + self.prepare_qkv_fusion() + # Handle empty tensor_map for models with block_count=0 (like MobileNetV5) if self.tensor_map.mapping: max_name_len = max(len(s) for _, s in self.tensor_map.mapping.values()) + len(".weight,") @@ -1027,6 +1110,13 @@ class ModelBase: self.gguf_writer.add_tensor(new_name, data, raw_dtype=data_qtype) + qkv_buffers = ( + self._q_buffer, self._k_buffer, self._v_buffer, + self._q_bias_buffer, self._k_bias_buffer, self._v_bias_buffer, + ) + if any(qkv_buffers): + raise ValueError("QKV fusion did not consume all buffered tensors") + def set_type(self): self.gguf_writer.add_type(gguf.GGUFType.MODEL) @@ -1543,6 +1633,9 @@ class TextModel(ModelBase): if chkhsh == "9e454714343b69b99b71795c1d27a68c2a1d15dab111f4d353109f966af29da7": # ref: https://huggingface.co/LiquidAI/LFM2.5-8B-A1B res = "lfm2" + if chkhsh == "0a766d034107bc736a3f2dc4968fd62e54a3570f1454443e0c5a4cc6bd7941ed": + # ref: https://huggingface.co/XHToken/Spark-X2.5-1.7B + res = "spark2_5" if chkhsh == "0ef9807a4087ebef797fc749390439009c3b9eda9ad1a097abbe738f486c01e5": # ref: https://huggingface.co/meta-llama/Meta-Llama-3-8B res = "llama-bpe" diff --git a/conversion/hy_v4.py b/conversion/hy_v4.py index f564b9ec2..358e21fe5 100644 --- a/conversion/hy_v4.py +++ b/conversion/hy_v4.py @@ -9,20 +9,6 @@ from .base import ModelBase, gguf, logger from .deepseek import DeepseekV2Model -def split_kv_b_proj(weight: torch.Tensor, n_head: int, qk_nope: int, v_head_dim: int): - """Split kv_b_proj into k_b (transposed) and v_b, matching DeepSeek MLA absorption. - - weight: [n_head*(qk_nope+v_head_dim), kv_lora_rank]. - Returns (k_b, v_b): k_b [n_head, kv_lora_rank, qk_nope], v_b [n_head, v_head_dim, kv_lora_rank]. - """ - kv_lora = weight.shape[-1] - assert weight.shape[0] == n_head * (qk_nope + v_head_dim) - kv_b = weight.view(n_head, qk_nope + v_head_dim, kv_lora) - k_b, v_b = torch.split(kv_b, [qk_nope, v_head_dim], dim=1) - k_b = k_b.transpose(1, 2).contiguous() # [n_head, kv_lora, qk_nope] - return k_b, v_b.contiguous() - - def split_gate_up(weight: torch.Tensor, moe_intermediate_size: int): """Split a fused stacked gate_up expert tensor into (gate, up). @@ -36,6 +22,7 @@ def split_gate_up(weight: torch.Tensor, moe_intermediate_size: int): @ModelBase.register("HYV4ForCausalLM") +@ModelBase.example("tencent/Hy4-preview") class HYV4Model(DeepseekV2Model): """HY_V4: DeepSeek-V3 style MLA + MoE with iHC, a gated MLA output and a learnable sink. @@ -54,6 +41,8 @@ class HYV4Model(DeepseekV2Model): model_arch = gguf.MODEL_ARCH.HY_V4 + merge_expert = False + # tensors a "full" indexer layer must carry INDEXER_SUFFIXES = frozenset({ "self_attn.indexer.wq_b.weight", @@ -186,6 +175,10 @@ class HYV4Model(DeepseekV2Model): ) def prepare_tensors(self): + # Hy4-preview for some reason has num_key_value_heads equal to 8, so override it here + # without this conversion/deepseek.py fails on assert + self.hparams["num_key_value_heads"] = self.hparams["num_attention_heads"] + # validate before the base materializes tensors, so a mismatch fails early is_full = self.indexer_is_full() if is_full is not None: @@ -227,85 +220,25 @@ class HYV4Model(DeepseekV2Model): def modify_tensors(self, data_torch: torch.Tensor, name: str, bid: int | None) -> Iterable[tuple[str, torch.Tensor]]: hparams = self.hparams - n_head = hparams["num_attention_heads"] - qk_nope = hparams["qk_nope_head_dim"] - v_head_dim = hparams["v_head_dim"] moe_inter = hparams["moe_intermediate_size"] tn = self.format_tensor_name - # ---- global (non per-layer) ---- - if name == "model.embed_tokens.weight": - return [(tn(gguf.MODEL_TENSOR.TOKEN_EMBD), data_torch)] - if name == "model.norm.weight": - return [(tn(gguf.MODEL_TENSOR.OUTPUT_NORM), data_torch)] - if name == "lm_head.weight": - return [(tn(gguf.MODEL_TENSOR.OUTPUT), data_torch)] - if name == "model.hc_head.hc_head_fn": - return [(tn(gguf.MODEL_TENSOR.HC_HEAD_FN), data_torch)] - if name == "model.hc_head.hc_head_base": - return [(tn(gguf.MODEL_TENSOR.HC_HEAD_BASE), data_torch)] - if name == "model.hc_head.hc_head_scale": - return [(tn(gguf.MODEL_TENSOR.HC_HEAD_SCALE), data_torch)] - - assert bid is not None, f"expected a per-layer tensor, got {name!r}" - - # ---- per-layer, keyed by suffix after 'model.layers.{bid}.' ---- - suffix = name.split(f"model.layers.{bid}.", 1)[-1] - - # note: q_b_proj and kv_a_proj_with_mqa are mapped straight through (no RoPE permute), - # the graph rotates consecutive pairs so the rows need no reordering - simple = { - "input_layernorm.weight": (gguf.MODEL_TENSOR.ATTN_NORM, ".weight"), - "post_attention_layernorm.weight": (gguf.MODEL_TENSOR.FFN_NORM, ".weight"), - "self_attn.q_a_proj.weight": (gguf.MODEL_TENSOR.ATTN_Q_A, ".weight"), - "self_attn.q_a_layernorm.weight": (gguf.MODEL_TENSOR.ATTN_Q_A_NORM, ".weight"), - "self_attn.q_b_proj.weight": (gguf.MODEL_TENSOR.ATTN_Q_B, ".weight"), - "self_attn.kv_a_proj_with_mqa.weight": (gguf.MODEL_TENSOR.ATTN_KV_A_MQA, ".weight"), - "self_attn.kv_a_layernorm.weight": (gguf.MODEL_TENSOR.ATTN_KV_A_NORM, ".weight"), - "self_attn.o_proj.weight": (gguf.MODEL_TENSOR.ATTN_OUT, ".weight"), - "self_attn.linear_gate.weight": (gguf.MODEL_TENSOR.ATTN_GATE, ".weight"), - "self_attn.learnable_sink_param": (gguf.MODEL_TENSOR.ATTN_SINKS, ".weight"), - "self_attn.indexer.wq_b.weight": (gguf.MODEL_TENSOR.INDEXER_ATTN_Q_B, ".weight"), - "self_attn.indexer.wk.weight": (gguf.MODEL_TENSOR.INDEXER_ATTN_K, ".weight"), - "self_attn.indexer.k_norm.weight": (gguf.MODEL_TENSOR.INDEXER_K_NORM, ".weight"), - "self_attn.indexer.k_norm.bias": (gguf.MODEL_TENSOR.INDEXER_K_NORM, ".bias"), - "self_attn.indexer.weights_proj.weight": (gguf.MODEL_TENSOR.INDEXER_PROJ, ".weight"), - "hc_attn_layer.hc_pre.hc_fn": (gguf.MODEL_TENSOR.HC_ATTN_FN, ".weight"), - "hc_attn_layer.hc_pre.hc_base": (gguf.MODEL_TENSOR.HC_ATTN_BASE, ".weight"), - "hc_attn_layer.hc_pre.hc_scale": (gguf.MODEL_TENSOR.HC_ATTN_SCALE, ".weight"), - "hc_mlp_layer.hc_pre.hc_fn": (gguf.MODEL_TENSOR.HC_FFN_FN, ".weight"), - "hc_mlp_layer.hc_pre.hc_base": (gguf.MODEL_TENSOR.HC_FFN_BASE, ".weight"), - "hc_mlp_layer.hc_pre.hc_scale": (gguf.MODEL_TENSOR.HC_FFN_SCALE, ".weight"), - "mlp.gate.weight": (gguf.MODEL_TENSOR.FFN_GATE_INP, ".weight"), - "mlp.gate.e_score_correction.bias":(gguf.MODEL_TENSOR.FFN_EXP_PROBS_B, ".bias"), - "mlp.gate_proj.weight": (gguf.MODEL_TENSOR.FFN_GATE, ".weight"), - "mlp.up_proj.weight": (gguf.MODEL_TENSOR.FFN_UP, ".weight"), - "mlp.down_proj.weight": (gguf.MODEL_TENSOR.FFN_DOWN, ".weight"), - "mlp.shared_experts.gate_proj.weight": (gguf.MODEL_TENSOR.FFN_GATE_SHEXP, ".weight"), - "mlp.shared_experts.up_proj.weight": (gguf.MODEL_TENSOR.FFN_UP_SHEXP, ".weight"), - "mlp.shared_experts.down_proj.weight": (gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, ".weight"), - } - if suffix in simple: - key, sfx = simple[suffix] - return [(tn(key, bid, sfx), data_torch)] - - # kv_b_proj: split into k_b (transposed) and v_b - if suffix == "self_attn.kv_b_proj.weight": - k_b, v_b = split_kv_b_proj(data_torch, n_head, qk_nope, v_head_dim) - return [ - (tn(gguf.MODEL_TENSOR.ATTN_K_B, bid), k_b), - (tn(gguf.MODEL_TENSOR.ATTN_V_B, bid), v_b), - ] - # fused stacked experts: split gate_up into gate/up - if suffix == "mlp.experts.gate_up_proj": + if name.endswith("mlp.experts.gate_up_proj"): gate, up = split_gate_up(data_torch, moe_inter) - return [ - (tn(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid), gate), - (tn(gguf.MODEL_TENSOR.FFN_UP_EXP, bid), up), - ] - if suffix == "mlp.experts.down_proj": - return [(tn(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid), data_torch)] + yield from super().modify_tensors(gate, tn(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid), bid) + yield from super().modify_tensors(up, tn(gguf.MODEL_TENSOR.FFN_UP_EXP, bid), bid) + return - raise ValueError(f"Unsupported HY_V4 tensor {name!r} (suffix {suffix!r})") + # add .weight suffixes + if name.endswith("mlp.experts.down_proj") or name.endswith(".self_attn.learnable_sink_param"): + name += ".weight" + + if re.search(r"\.hc_head\.hc_head_(?:fn|base|scale)$", name): + name += ".weight" + + if re.search(r"\.hc_(?:attn|mlp)_layer\.hc_pre\.hc_(?:fn|base|scale)$", name): + name += ".weight" + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/conversion/qwen.py b/conversion/qwen.py index 419611896..c7e0809f3 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -379,6 +379,13 @@ class Qwen3NextModel(_QwenMtpMixin, Qwen2MoeModel): self.gguf_writer.add_ssm_group_count(self.hparams["linear_num_key_heads"]) self.gguf_writer.add_ssm_time_step_rank(self.hparams["linear_num_value_heads"]) self.gguf_writer.add_ssm_inner_size(self.hparams["linear_value_head_dim"] * self.hparams["linear_num_value_heads"]) + if (layer_types := self.hparams.get("layer_types")) is not None: + n_layer = self.hparams["num_hidden_layers"] + if len(layer_types) != n_layer: + raise ValueError(f"layer_types has {len(layer_types)} entries, expected num_hidden_layers ({n_layer})") + recurrent = [t == "linear_attention" for t in layer_types] + recurrent += [False] * (self.block_count - n_layer) + self.gguf_writer.add_recurrent_layers(recurrent) self.gguf_writer.add_full_attention_interval(self.hparams.get("full_attention_interval", 4)) if (rope_dim := self.hparams.get("head_dim")) is None: rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"] diff --git a/conversion/spark2_5.py b/conversion/spark2_5.py new file mode 100644 index 000000000..44a0bd262 --- /dev/null +++ b/conversion/spark2_5.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from collections.abc import Iterable +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from torch import Tensor + +from .base import ModelBase, TextModel, gguf + + +@ModelBase.register("Spark2_5ForCausalLM") +@ModelBase.example("XHToken/Spark-X2.5-1.7B") +class Spark2_5Model(TextModel): + model_arch = gguf.MODEL_ARCH.SPARK2_5 + + def set_gguf_parameters(self) -> None: + super().set_gguf_parameters() + + hparams = self.hparams + layer_types = hparams["layer_types"] + if len(layer_types) != self.block_count: + raise ValueError( + f"Spark2_5 layer_types length {len(layer_types)} != num_hidden_layers {self.block_count}" + ) + if any(layer_type not in ("sliding_attention", "full_attention") for layer_type in layer_types): + raise ValueError(f"Spark2_5 has unsupported layer_types: {layer_types}") + if hparams.get("gate_attn_act_mode") != "sigmoid" or hparams.get("headwise_attn_output_gate") is not True: + raise ValueError("Spark2_5 conversion requires head-wise sigmoid attention gates") + if hparams.get("hidden_act") != "gelu": + raise ValueError(f"Spark2_5 conversion requires GELU, got {hparams.get('hidden_act')!r}") + + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + self.gguf_writer.add_sliding_window(hparams["sliding_window"]) + self.gguf_writer.add_sliding_window_pattern( + [layer_type == "sliding_attention" for layer_type in layer_types] + ) + + head_dim = hparams["head_dim"] + full_rope = self.rope_parameters["full_attention"] + swa_rope = self.rope_parameters["sliding_attention"] + self.gguf_writer.add_rope_dimension_count( + int(head_dim * float(full_rope["partial_rotary_factor"])) + ) + self.gguf_writer.add_rope_dimension_count_swa( + int(head_dim * float(swa_rope["partial_rotary_factor"])) + ) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.endswith(".self_attn.q_k_v_proj.weight"): + if bid is None: + raise ValueError(f"Spark2_5 fused QKV tensor has no block id: {name}") + yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_QKV, bid), data_torch + return + + if name.endswith(".self_attn.g_proj.weight"): + if bid is None: + raise ValueError(f"Spark2_5 attention gate tensor has no block id: {name}") + expected = self.hparams["num_attention_heads"] + if data_torch.shape[0] != expected: + raise ValueError( + f"Spark2_5 layer {bid} attention gate width {data_torch.shape[0]} != head count {expected}" + ) + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 6e74c8764..9ae159a0a 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -161,6 +161,10 @@ def parse_args() -> argparse.Namespace: help="Store tensors dequantized from FP8 as Q8_0 instead of BF16/F16.", ) + parser.add_argument( + "--fuse-qkv", action="store_true", + help="Fuse separate Q, K, V weight tensors into a single QKV tensor.", + ) parser.add_argument( "--target-model-dir", type=str, default=None, help=( @@ -294,6 +298,7 @@ def main() -> None: target_model_dir=Path(args.target_model_dir) if args.target_model_dir else None, fuse_gate_up_exps=args.fuse_gate_up_exps, fp8_as_q8=args.fp8_as_q8, + fuse_qkv=args.fuse_qkv, ) if args.vocab_only: diff --git a/convert_hf_to_gguf_update.py b/convert_hf_to_gguf_update.py index c4141afa6..6af74cd87 100755 --- a/convert_hf_to_gguf_update.py +++ b/convert_hf_to_gguf_update.py @@ -191,6 +191,7 @@ pre_computed_hashes = [ {"name": "gpt-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/evilfreelancer/ruGPT3XL", "chkhsh": "0fe1cf6eda062318a1af7270f3331a85c539a01778ff948e24388e949c5282f4"}, # lfm2 variants {"name": "lfm2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/LiquidAI/LFM2.5-8B-A1B", "chkhsh": "9e454714343b69b99b71795c1d27a68c2a1d15dab111f4d353109f966af29da7"}, + {"name": "spark2_5", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/XHToken/Spark-X2.5-1.7B", "chkhsh": "0a766d034107bc736a3f2dc4968fd62e54a3570f1454443e0c5a4cc6bd7941ed"}, ] diff --git a/examples/test-cmake/test-cmake.cpp b/examples/test-cmake/test-cmake.cpp index c5c4765b4..dc1a9ae60 100644 --- a/examples/test-cmake/test-cmake.cpp +++ b/examples/test-cmake/test-cmake.cpp @@ -2,8 +2,9 @@ #include int main(void) { - printf("[test-cmake] version: %s, build: %d (%s)\n", + printf("[test-cmake] llama.cpp version: %s, build: %d (%s)\n", llama_version(), LLAMA_BUILD_NUMBER, LLAMA_BUILD_COMMIT); + printf("[test-cmake] ggml version: %s, commit: %s\n", ggml_version(), ggml_commit()); printf("[test-cmake] Initializing backend...\n"); llama_backend_init(); printf("[test-cmake] Backend initialized.\n"); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 9ed929502..f538f1b0c 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -850,7 +850,7 @@ static void ggml_backend_sched_split_inputs_grow(struct ggml_backend_sched_split int new_cap = GGML_SCHED_MAX_SPLIT_INPUTS; if (split->inputs_capacity > 0) { new_cap = 2*split->inputs_capacity; - GGML_LOG_WARN("%s: increasing split inputs capacity from %d to %d\n", __func__, split->inputs_capacity, new_cap); + GGML_LOG_DEBUG("%s: increasing split inputs capacity from %d to %d\n", __func__, split->inputs_capacity, new_cap); } auto * pnew = (struct ggml_tensor **) realloc((void *) split->inputs, new_cap * sizeof(struct ggml_tensor *)); if (pnew == NULL) { @@ -865,7 +865,7 @@ static void ggml_backend_sched_graph_inputs_grow(ggml_backend_sched_t sched) { int new_cap = GGML_SCHED_MAX_SPLIT_INPUTS; if (sched->graph_inputs_capacity > 0) { new_cap = 2*sched->graph_inputs_capacity; - GGML_LOG_WARN("%s: increasing graph inputs capacity from %d to %d\n", __func__, sched->graph_inputs_capacity, new_cap); + GGML_LOG_DEBUG("%s: increasing graph inputs capacity from %d to %d\n", __func__, sched->graph_inputs_capacity, new_cap); } auto * pnew = (struct ggml_tensor **) realloc((void *) sched->graph_inputs, new_cap * sizeof(struct ggml_tensor *)); if (pnew == NULL) { @@ -1345,17 +1345,6 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra break; } } - // check if the split has too many inputs - // FIXME: count the number of inputs instead of only checking when full - if (split->n_inputs >= split->inputs_capacity) { - const size_t id = hash_id(src); - int src_backend_id = sched->hv_tensor_backend_ids[id]; - bool supported = ggml_backend_sched_buffer_supported(sched, src, cur_backend_id); - if (src_backend_id != cur_backend_id && tensor_id_copy(id, cur_backend_id, 0) == NULL && !supported) { - need_new_split = true; - break; - } - } } } diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index b6db16efb..c2f0b1f24 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -69,6 +69,8 @@ #define GGML_CUDA_CC_GCN4 (GGML_CUDA_CC_OFFSET_AMD + 0x803) // Tonga, Fiji, Polaris, minimum for fast fp16 #define GGML_CUDA_CC_VEGA (GGML_CUDA_CC_OFFSET_AMD + 0x900) // Vega56/64, minimum for fp16 dual issue #define GGML_CUDA_CC_VEGA20 (GGML_CUDA_CC_OFFSET_AMD + 0x906) // MI50/Radeon VII, minimum for dp4a +#define GGML_CUDA_CC_GFX909 (GGML_CUDA_CC_OFFSET_AMD + 0x909) // GCN APU +#define GGML_CUDA_CC_GFX90C (GGML_CUDA_CC_OFFSET_AMD + 0x90c) // GCN APU #define GGML_CUDA_CC_CDNA1 (GGML_CUDA_CC_OFFSET_AMD + 0x908) // MI100, minimum for MFMA, acc registers #define GGML_CUDA_CC_CDNA2 (GGML_CUDA_CC_OFFSET_AMD + 0x90a) // MI210 (gfx90a), minimum acc register renaming #define GGML_CUDA_CC_CDNA3 (GGML_CUDA_CC_OFFSET_AMD + 0x942) // MI300 @@ -89,12 +91,13 @@ #define GGML_CUDA_CC_IS_RDNA3_5(cc) (cc >= GGML_CUDA_CC_RDNA3_5 && cc < GGML_CUDA_CC_RDNA4) #define GGML_CUDA_CC_IS_RDNA3(cc) (GGML_CUDA_CC_IS_RDNA3_0(cc) || GGML_CUDA_CC_IS_RDNA3_5(cc)) #define GGML_CUDA_CC_IS_RDNA4(cc) (cc >= GGML_CUDA_CC_RDNA4) -#define GGML_CUDA_CC_IS_GCN(cc) (cc > GGML_CUDA_CC_OFFSET_AMD && cc < GGML_CUDA_CC_CDNA1) -#define GGML_CUDA_CC_IS_CDNA(cc) (cc >= GGML_CUDA_CC_CDNA1 && cc < GGML_CUDA_CC_RDNA1) -#define GGML_CUDA_CC_IS_CDNA1(cc) (cc >= GGML_CUDA_CC_CDNA1 && cc < GGML_CUDA_CC_CDNA2) -#define GGML_CUDA_CC_IS_CDNA2(cc) (cc >= GGML_CUDA_CC_CDNA2 && cc < GGML_CUDA_CC_CDNA3) -#define GGML_CUDA_CC_IS_CDNA3(cc) (cc >= GGML_CUDA_CC_CDNA3 && cc < GGML_CUDA_CC_CDNA4) -#define GGML_CUDA_CC_IS_CDNA4(cc) (cc >= GGML_CUDA_CC_CDNA4 && cc < GGML_CUDA_CC_RDNA1) +#define GGML_CUDA_CC_IS_GCN_APU(cc) ((cc) == GGML_CUDA_CC_GFX909 || (cc) == GGML_CUDA_CC_GFX90C) +#define GGML_CUDA_CC_IS_GCN(cc) ((cc > GGML_CUDA_CC_OFFSET_AMD && cc < GGML_CUDA_CC_CDNA1) || GGML_CUDA_CC_IS_GCN_APU(cc)) +#define GGML_CUDA_CC_IS_CDNA(cc) (!GGML_CUDA_CC_IS_GCN_APU(cc) && cc >= GGML_CUDA_CC_CDNA1 && cc < GGML_CUDA_CC_RDNA1) +#define GGML_CUDA_CC_IS_CDNA1(cc) (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA1 && cc < GGML_CUDA_CC_CDNA2) +#define GGML_CUDA_CC_IS_CDNA2(cc) (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA2 && cc < GGML_CUDA_CC_CDNA3) +#define GGML_CUDA_CC_IS_CDNA3(cc) (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA3 && cc < GGML_CUDA_CC_CDNA4) +#define GGML_CUDA_CC_IS_CDNA4(cc) (GGML_CUDA_CC_IS_CDNA(cc) && cc >= GGML_CUDA_CC_CDNA4 && cc < GGML_CUDA_CC_RDNA1) // Moore Threads #define MUSART_HMASK 40300 // MUSA rc4.3, min. ver. for half2 -> uint mask comparisons @@ -121,6 +124,12 @@ # define GGML_CUDA_USE_PDL #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && (CUDART_VERSION >= 12030 || (!(defined(_MSC_VER) && !defined(__clang__)) && CUDART_VERSION >= 11080)) +static __device__ __forceinline__ void ggml_cuda_syncwarp() { +#ifndef GGML_USE_HIP + __syncwarp(); +#endif // GGML_USE_HIP +} + static __device__ __forceinline__ void ggml_cuda_pdl_sync() { #if defined(GGML_CUDA_USE_PDL) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER cudaGridDependencySynchronize(); @@ -977,6 +986,7 @@ template<> struct ggml_cuda_type_traits { static constexpr int qk = 1; static constexpr int qr = 1; + static constexpr int bs = sizeof(ggml_half); }; template<> @@ -984,6 +994,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK1_0; static constexpr int qr = QR1_0; static constexpr int qi = QI1_0; + static constexpr int bs = sizeof(block_q1_0); }; template<> @@ -991,6 +1002,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK2_0; static constexpr int qr = QR2_0; static constexpr int qi = QI2_0; + static constexpr int bs = sizeof(block_q2_0); }; template<> @@ -998,6 +1010,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK4_0; static constexpr int qr = QR4_0; static constexpr int qi = QI4_0; + static constexpr int bs = sizeof(block_q4_0); }; template<> @@ -1005,6 +1018,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK4_1; static constexpr int qr = QR4_1; static constexpr int qi = QI4_1; + static constexpr int bs = sizeof(block_q4_1); }; template<> @@ -1012,6 +1026,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK5_0; static constexpr int qr = QR5_0; static constexpr int qi = QI5_0; + static constexpr int bs = sizeof(block_q5_0); }; template<> @@ -1019,6 +1034,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK5_1; static constexpr int qr = QR5_1; static constexpr int qi = QI5_1; + static constexpr int bs = sizeof(block_q5_1); }; template<> @@ -1026,6 +1042,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK8_0; static constexpr int qr = QR8_0; static constexpr int qi = QI8_0; + static constexpr int bs = sizeof(block_q8_0); }; template<> @@ -1033,6 +1050,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_MXFP4; static constexpr int qr = QR_MXFP4; static constexpr int qi = QI_MXFP4; + static constexpr int bs = sizeof(block_mxfp4); }; template<> @@ -1040,6 +1058,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_NVFP4; static constexpr int qr = QR_NVFP4; static constexpr int qi = QI_NVFP4; + static constexpr int bs = sizeof(block_nvfp4); }; template<> @@ -1047,6 +1066,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR2_K; static constexpr int qi = QI2_K; + static constexpr int bs = sizeof(block_q2_K); }; template<> @@ -1054,6 +1074,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR3_K; static constexpr int qi = QI3_K; + static constexpr int bs = sizeof(block_q3_K); }; template<> @@ -1061,6 +1082,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR4_K; static constexpr int qi = QI4_K; + static constexpr int bs = sizeof(block_q4_K); }; template<> @@ -1068,6 +1090,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR5_K; static constexpr int qi = QI5_K; + static constexpr int bs = sizeof(block_q5_K); }; template<> @@ -1075,6 +1098,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR6_K; static constexpr int qi = QI6_K; + static constexpr int bs = sizeof(block_q6_K); }; template<> @@ -1082,6 +1106,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR2_XXS; static constexpr int qi = QI2_XXS; + static constexpr int bs = sizeof(block_iq2_xxs); }; template<> @@ -1089,6 +1114,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR2_XS; static constexpr int qi = QI2_XS; + static constexpr int bs = sizeof(block_iq2_xs); }; template<> @@ -1096,6 +1122,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR2_S; static constexpr int qi = QI2_S; + static constexpr int bs = sizeof(block_iq2_s); }; template<> @@ -1103,6 +1130,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR3_XXS; static constexpr int qi = QI3_XXS; + static constexpr int bs = sizeof(block_iq3_xxs); }; template<> @@ -1110,6 +1138,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR1_S; static constexpr int qi = QI1_S; + static constexpr int bs = sizeof(block_iq1_s); }; template<> @@ -1117,6 +1146,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR1_M; static constexpr int qi = QI1_M; + static constexpr int bs = sizeof(block_iq1_m); }; template<> @@ -1124,6 +1154,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK4_NL; static constexpr int qr = QR4_NL; static constexpr int qi = QI4_NL; + static constexpr int bs = sizeof(block_iq4_nl); }; template<> @@ -1131,6 +1162,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR4_XS; static constexpr int qi = QI4_XS; + static constexpr int bs = sizeof(block_iq4_xs); }; template<> @@ -1138,6 +1170,7 @@ struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR3_S; static constexpr int qi = QI3_S; + static constexpr int bs = sizeof(block_iq3_s); }; ////////////////////// diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index 126a4c452..bc5060e81 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -1545,77 +1545,77 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( } } - if (np > 1 && threadIdx.y % np == 0) { - // Combine the meta data for parallel warps via shared memory. - // Warps with threadIdx.y % np != 0 must NOT return early. - // All threads must return simultaneously to avoid race conditions with work on the next tile. - + if (np > 1) { constexpr int nmeta = np*cols_per_warp >= warp_size ? np*cols_per_warp/warp_size : 1; + float KQ_cmn; + float KQ_cms[nmeta]; + float KQ_crs; + const int jc_meta = threadIdx.y*cols_per_warp + (np*cols_per_warp < warp_size ? threadIdx.x % (np*cols_per_warp) : threadIdx.x); float2 * const meta_ptr = ((float2 *) tile_Q) + jc_meta*(tile_stride/2) + nbatch_combine/2; - float2 meta[nmeta]; -#pragma unroll - for (int imeta = 0; imeta < nmeta; ++imeta) { - meta[imeta] = meta_ptr[imeta * warp_size * tile_stride/2]; - } - float KQ_cmn = meta[0].x; // KQ combine max new, max between all parallel warps. + if (threadIdx.y % np == 0) { + // Combine the meta data for parallel warps via shared memory. + float2 meta[nmeta]; #pragma unroll - for (int imeta = 1; imeta < nmeta; ++imeta) { - KQ_cmn = fmaxf(KQ_cmn, meta[imeta].x); - } -#pragma unroll - for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) { - if (offset < warp_size) { - KQ_cmn = fmaxf(KQ_cmn, __shfl_xor_sync(0xFFFFFFFF, KQ_cmn, offset, warp_size)); + for (int imeta = 0; imeta < nmeta; ++imeta) { + meta[imeta] = meta_ptr[imeta * warp_size * tile_stride/2]; } - } - float KQ_cms[nmeta]; // KQ combine max scale per warp. + KQ_cmn = meta[0].x; // KQ combine max new, max between all parallel warps. #pragma unroll - for (int imeta = 0; imeta < nmeta; ++imeta) { - KQ_cms[imeta] = expf(meta[imeta].x - KQ_cmn); - } + for (int imeta = 1; imeta < nmeta; ++imeta) { + KQ_cmn = fmaxf(KQ_cmn, meta[imeta].x); + } +#pragma unroll + for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) { + if (offset < warp_size) { + KQ_cmn = fmaxf(KQ_cmn, __shfl_xor_sync(0xFFFFFFFF, KQ_cmn, offset, warp_size)); + } + } - float KQ_crs = KQ_cms[0]*meta[0].y; // KQ combine rowsum, scaled sum of all parallel warps. #pragma unroll - for (int imeta = 1; imeta < nmeta; ++imeta) { - KQ_crs += KQ_cms[imeta]*meta[imeta].y; - } + for (int imeta = 0; imeta < nmeta; ++imeta) { + KQ_cms[imeta] = expf(meta[imeta].x - KQ_cmn); + } + + KQ_crs = KQ_cms[0]*meta[0].y; // KQ combine rowsum, scaled sum of all parallel warps. #pragma unroll - for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) { - if (offset < warp_size) { - KQ_crs += __shfl_xor_sync(0xFFFFFFFF, KQ_crs, offset, warp_size); + for (int imeta = 1; imeta < nmeta; ++imeta) { + KQ_crs += KQ_cms[imeta]*meta[imeta].y; + } +#pragma unroll + for (int offset = np*cols_per_warp/2; offset >= cols_per_warp; offset >>= 1) { + if (offset < warp_size) { + KQ_crs += __shfl_xor_sync(0xFFFFFFFF, KQ_crs, offset, warp_size); + } } } __syncthreads(); - // Write back combined meta data: + if (threadIdx.y % np == 0) { + // Write back combined meta data: #pragma unroll - for (int imeta = 0; imeta < nmeta; ++imeta) { - if (np*cols_per_warp >= warp_size || threadIdx.x < np*cols_per_warp) { - // Combined KQ max scale + rowsum. - meta_ptr[imeta * warp_size * tile_stride/2] = make_float2(KQ_cms[imeta], KQ_crs); + for (int imeta = 0; imeta < nmeta; ++imeta) { + if (np*cols_per_warp >= warp_size || threadIdx.x < np*cols_per_warp) { + // Combined KQ max scale + rowsum. + meta_ptr[imeta * warp_size * tile_stride/2] = make_float2(KQ_cms[imeta], KQ_crs); + } + } + + // Combined KQ max + rowsum. + static_assert(cols_per_warp <= warp_size); + if (needs_fixup && (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) { + float2 * dstk_fixup_meta = dstk_fixup + blockIdx.x*ncols; + dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs); + } + if (is_fixup && (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) { + float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols; + dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs); } } - - // Combined KQ max + rowsum. - static_assert(cols_per_warp <= warp_size); - if (needs_fixup && (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) { - float2 * dstk_fixup_meta = dstk_fixup + blockIdx.x*ncols; - dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs); - } - if (is_fixup && (cols_per_warp == warp_size || threadIdx.x < cols_per_warp)) { - float2 * dstk_fixup_meta = dstk_fixup + (gridDim.x + blockIdx.x)*ncols; - dstk_fixup_meta[(threadIdx.y/np)*cols_per_warp + threadIdx.x] = make_float2(KQ_cmn, KQ_crs); - } - } else if (np > 1) { - // Warps with threadIdx.y % np == 0 execute a __syncthreads() in the if branch. - // Therefore, all other warps also need to execute a __syncthreads(). - // Otherwise the points at which warps synchronize with each other would become misaligned. - __syncthreads(); } #pragma unroll diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index 519b36b9f..57a285565 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -317,9 +317,7 @@ static __global__ void flash_attn_ext_vec( #endif // V_DOT2_F32_F16_AVAILABLE } -#ifndef GGML_USE_HIP - __syncwarp(); -#endif // GGML_USE_HIP + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < WARP_SIZE; k0 += V_cols_per_iter) { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index b3b2857bd..9ff4d911a 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -214,6 +214,7 @@ static int ggml_cuda_parse_id(char devName[]) { } archNum += archMajor * 0x100; archNum += archMinor; + return archNum; } #endif // defined(GGML_USE_HIP) diff --git a/ggml/src/ggml-cuda/mmf.cuh b/ggml/src/ggml-cuda/mmf.cuh index d55cc1ec7..879a86527 100644 --- a/ggml/src/ggml-cuda/mmf.cuh +++ b/ggml/src/ggml-cuda/mmf.cuh @@ -143,6 +143,7 @@ static __global__ void mul_mat_f( if (threadIdx.x == 0) { slot_map[j] = -1; } + ggml_cuda_syncwarp(); if (col_base + j >= ncols_dst_total) { continue; @@ -171,10 +172,12 @@ static __global__ void mul_mat_f( tile_A A[ntA][warp_size / tile_A::J]; #pragma unroll for (int itA = 0; itA < ntA; ++itA) { + ggml_cuda_syncwarp(); #pragma unroll for (int i = 0; i < tile_A::I; ++i) { tile_xy[i*tile_k_padded + threadIdx.x] = x[(itA*tile_A::I + i)*stride_row + col]; } + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < warp_size; k0 += tile_A::J) { load_ldmatrix(A[itA][k0/tile_A::J], tile_xy + k0, tile_k_padded); @@ -183,6 +186,7 @@ static __global__ void mul_mat_f( #pragma unroll for (int itB = 0; itB < ntB; ++itB) { + ggml_cuda_syncwarp(); if constexpr (std::is_same_v) { #pragma unroll for (int j0 = 0; j0 < tile_B::I; ++j0) { @@ -212,6 +216,7 @@ static __global__ void mul_mat_f( } else { static_assert(std::is_same_v, "unsupported type"); } + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < warp_size; k0 += tile_B::J) { tile_B B; @@ -229,6 +234,8 @@ static __global__ void mul_mat_f( if (nwarps > 1) { __syncthreads(); + } else { + ggml_cuda_syncwarp(); } #pragma unroll for (int itB = 0; itB < ntB; ++itB) { @@ -245,6 +252,8 @@ static __global__ void mul_mat_f( if (nwarps > 1) { __syncthreads(); + } else { + ggml_cuda_syncwarp(); } #pragma unroll @@ -382,10 +391,12 @@ static __global__ void mul_mat_f_ids( tile_A A[ntA][warp_size / tile_A::J]; #pragma unroll for (int itA = 0; itA < ntA; ++itA) { + ggml_cuda_syncwarp(); #pragma unroll for (int i = 0; i < tile_A::I; ++i) { tile_xy[i*tile_k_padded + threadIdx.x] = x[(itA*tile_A::I + i)*stride_row + col]; } + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < warp_size; k0 += tile_A::J) { load_ldmatrix(A[itA][k0/tile_A::J], tile_xy + k0, tile_k_padded); @@ -419,6 +430,7 @@ static __global__ void mul_mat_f_ids( int next_buf = 1; #pragma unroll for (int itB = 0; itB < ntB; ++itB) { + ggml_cuda_syncwarp(); #pragma unroll for (int j0 = 0; j0 < tile_B::I; ++j0) { tile_xy[j0*tile_k_padded + threadIdx.x] = vals_buf[curr_buf][j0]; @@ -428,6 +440,7 @@ static __global__ void mul_mat_f_ids( gather_tile(itB + 1, vals_buf[next_buf]); } + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < warp_size; k0 += tile_B::J) { tile_B B; @@ -472,6 +485,7 @@ static __global__ void mul_mat_f_ids( int next_buf = 1; #pragma unroll for (int itB = 0; itB < ntB; ++itB) { + ggml_cuda_syncwarp(); #pragma unroll for (int j0 = 0; j0 < tile_B::I; ++j0) { const float2 tmp = vals_buf[curr_buf][j0]; @@ -482,6 +496,7 @@ static __global__ void mul_mat_f_ids( gather_tile(itB + 1, vals_buf[next_buf]); } + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < warp_size; k0 += tile_B::J) { tile_B B; @@ -507,6 +522,8 @@ static __global__ void mul_mat_f_ids( if (nwarps > 1) { __syncthreads(); + } else { + ggml_cuda_syncwarp(); } #pragma unroll for (int itB = 0; itB < ntB; ++itB) { @@ -523,6 +540,8 @@ static __global__ void mul_mat_f_ids( if (nwarps > 1) { __syncthreads(); + } else { + ggml_cuda_syncwarp(); } #pragma unroll diff --git a/ggml/src/ggml-cuda/mmid.cu b/ggml/src/ggml-cuda/mmid.cu index ed0851dcf..0b222e63a 100644 --- a/ggml/src/ggml-cuda/mmid.cu +++ b/ggml/src/ggml-cuda/mmid.cu @@ -101,6 +101,7 @@ static __global__ void mm_ids_helper( } } nex_prev = warp_reduce_sum(nex_prev); + ggml_cuda_syncwarp(); for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) { const mm_ids_helper_store store_it = store[itc]; diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 09976eb5d..431719a08 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -377,10 +377,10 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t return true; } - // gfx900 (Vega 10) lacks native dp4a, loses to dequant + hipBLAS + // gfx900 (Vega 10), gfx909, and gfx90c lack native dp4a, losing to dequant + hipBLAS // for dense matrices; keep MMQ only for MoE, where the // hipBLAS path is much slower. - if (cc == GGML_CUDA_CC_VEGA) { + if (cc == GGML_CUDA_CC_VEGA || GGML_CUDA_CC_IS_GCN_APU(cc)) { return n_experts > 0; } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index f65e0fbcd..6305230b1 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -6,6 +6,35 @@ #include #include +// only enabled on DGX Spark, where it is a gain on every type below. On the higher-bandwidth parts the kernel +// has little exposed latency left to hide and the extra requests cost more than they save. +// For perf data, see https://github.com/ggml-org/llama.cpp/pull/26705#issuecomment-5569335031 +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK +// returns true only for those quants that benefit from prefetch and false otherwise +static constexpr __host__ __device__ bool mmvq_should_prefetch(ggml_type type) { + switch (type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q8_0: + case GGML_TYPE_MXFP4: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + return true; + default: + return false; + } +} + +static __device__ __forceinline__ void mmvq_prefetch_l2(const void * p) { + asm volatile("prefetch.global.L2 [%0];" :: "l"(p)); +} +#endif + typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs); static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) { @@ -298,9 +327,6 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) { return ne11 <= 4; case GGML_TYPE_Q3_K: return ne11 <= 6; - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - return ne11 <= 7; default: return ne11 <= MMVQ_MAX_BATCH_SIZE; } @@ -310,8 +336,9 @@ bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11) { case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: return ne11 <= 5; + case GGML_TYPE_Q5_K: + return ne11 <= 6; case GGML_TYPE_Q6_K: return ne11 <= 7; default: @@ -675,6 +702,26 @@ static __global__ void mul_mat_vec_q( // x block quant index when casting the quants to int const int kqs = vdr * (tid % (qi/vdr)); +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK + // start the next iterations' weight loads early + if constexpr (mmvq_should_prefetch(type)) { + constexpr int pf_dist = 2; // loop iterations, not blocks + const int kbx_pf = kbx + pf_dist*blocks_per_iter; + if (kbx_pf < blocks_per_row_x) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + const size_t off = (size_t)(kbx_offset + i*stride_row_x + kbx_pf) * ggml_cuda_type_traits::bs; + mmvq_prefetch_l2((const char *) vx + off); + if constexpr (has_fusion) { + if (use_gate) { + mmvq_prefetch_l2((const char *) vgate + off); + } + } + } + } + } +#endif + #pragma unroll for (int j = 0; j < ncols_dst; ++j) { #pragma unroll diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index ec117c57d..f2a6f2009 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -936,16 +936,20 @@ static __device__ __forceinline__ float vec_dot_q4_K_q8_1( v[0] = q4[0]; v[1] = q4[4]; + // branchless so nvcc can hoist this out of the ncols_dst loop const uint16_t * scales = (const uint16_t *)bq4_K->scales; + const int j = bq8_offset/2; + const int jm = j & 1; + + const uint32_t s0 = scales[jm + 0]; + const uint32_t s2 = scales[jm + 2]; + const uint32_t s4 = scales[jm + 4]; + + const uint32_t hi = (uint32_t) -(int32_t) (j >= 2); + uint16_t aux[2]; - const int j = bq8_offset/2; - if (j < 2) { - aux[0] = scales[j+0] & 0x3f3f; - aux[1] = scales[j+2] & 0x3f3f; - } else { - aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2); - aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2); - } + aux[0] = (uint16_t) (((s0 & 0x3f3f) & ~hi) | ((((s4 >> 0) & 0x0f0f) | ((s0 & 0xc0c0) >> 2)) & hi)); + aux[1] = (uint16_t) (((s2 & 0x3f3f) & ~hi) | ((((s4 >> 4) & 0x0f0f) | ((s2 & 0xc0c0) >> 2)) & hi)); const uint8_t * sc = (const uint8_t *)aux; const uint8_t * m = sc + 2; @@ -981,16 +985,21 @@ static __device__ __forceinline__ float vec_dot_q5_K_q8_1( vh[0] = qh[0] >> bq8_offset; vh[1] = qh[4] >> bq8_offset; + // same as q4_K const uint16_t * scales = (const uint16_t *)bq5_K->scales; + const int j = bq8_offset/2; + const int jm = j & 1; + + const uint32_t s0 = scales[jm + 0]; + const uint32_t s2 = scales[jm + 2]; + const uint32_t s4 = scales[jm + 4]; + + const uint32_t hi = (uint32_t) -(int32_t) (j >= 2); + uint16_t aux[2]; - const int j = bq8_offset/2; - if (j < 2) { - aux[0] = scales[j+0] & 0x3f3f; - aux[1] = scales[j+2] & 0x3f3f; - } else { - aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2); - aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2); - } + aux[0] = (uint16_t) (((s0 & 0x3f3f) & ~hi) | ((((s4 >> 0) & 0x0f0f) | ((s0 & 0xc0c0) >> 2)) & hi)); + aux[1] = (uint16_t) (((s2 & 0x3f3f) & ~hi) | ((((s4 >> 4) & 0x0f0f) | ((s2 & 0xc0c0) >> 2)) & hi)); + const uint8_t * sc = (const uint8_t *)aux; const uint8_t * m = sc + 2; diff --git a/ggml/src/ggml-cuda/vendors/hip.h b/ggml/src/ggml-cuda/vendors/hip.h index 9aa558f3f..2fc0fe9fd 100644 --- a/ggml/src/ggml-cuda/vendors/hip.h +++ b/ggml/src/ggml-cuda/vendors/hip.h @@ -176,9 +176,9 @@ #define __CUDA_ARCH__ 1300 -#if defined(__gfx900__) || defined(__gfx906__) +#if defined(__gfx900__) || defined(__gfx906__) || defined(__gfx909__) || defined(__gfx90c__) #define GCN5 -#endif // defined(__gfx900__) || defined(__gfx906__) +#endif // defined(__gfx900__) || defined(__gfx906__) || defined(__gfx909__) || defined(__gfx90c__) #if defined(__gfx803__) #define GCN4 diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m index e1129db30..6cdc4006b 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.m +++ b/ggml/src/ggml-metal/ggml-metal-context.m @@ -111,6 +111,7 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) { id queue = ggml_metal_device_get_queue(dev); if (queue == nil) { GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); + free(res); return NULL; } diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index d23cb43d5..ddaec9fda 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1492,7 +1492,9 @@ static bool ggml_metal_supports_mul_mat_op( const struct ggml_tensor * op, bool src0_f16_has_mv, bool mm_path) { - if (!has_simdgroup_reduction || op->src[0]->type == GGML_TYPE_NVFP4) { + if (!has_simdgroup_reduction || + op->src[0]->type == GGML_TYPE_NVFP4 || + op->src[0]->type == GGML_TYPE_TQ1_0) { return false; } @@ -1893,7 +1895,8 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te }; } case GGML_OP_GET_ROWS: - return op->src[0]->type != GGML_TYPE_NVFP4; + return op->src[0]->type != GGML_TYPE_NVFP4 && + op->src[0]->type != GGML_TYPE_TQ1_0; case GGML_OP_SET_ROWS: { if (op->src[0]->type == GGML_TYPE_F16) { diff --git a/ggml/src/ggml-metal/ggml-metal-tuning.cpp b/ggml/src/ggml-metal/ggml-metal-tuning.cpp index 8cdc55a0a..2323269c4 100644 --- a/ggml/src/ggml-metal/ggml-metal-tuning.cpp +++ b/ggml/src/ggml-metal/ggml-metal-tuning.cpp @@ -1248,6 +1248,153 @@ constexpr fa_vec_entry_t fa_vec_tuned_table[] = { { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 512, 512, 2, 0 }, { 4, 4 } }, { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 512, 512, 2, 1 }, { 4, 1 } }, { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_F16, 512, 512, 2, 2 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 32, 32, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 32, 32, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 32, 32, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 32, 32, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 32, 32, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 64, 64, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 96, 96, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 96, 96, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 96, 96, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 32, 32, 1, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 32, 32, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 32, 32, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 64, 64, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 64, 64, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 64, 64, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 64, 64, 2, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 64, 64, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 96, 96, 2, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 96, 96, 3, 3 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 128, 128, 1, 4 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 128, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 192, 192, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 192, 192, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 192, 192, 2, 3 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 192, 128, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q4_1, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 32, 32, 2, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 192, 1, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 192, 1, 4 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 192, 2, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 192, 3, 2 }, { 1, 2 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 128, 2, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 128, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 192, 128, 3, 3 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 256, 256, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 256, 256, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 256, 256, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 256, 256, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_0, 576, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 32, 32, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 32, 32, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 32, 32, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 32, 32, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 32, 32, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 64, 64, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 64, 64, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 64, 64, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 64, 64, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 64, 64, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 96, 96, -1, 1 }, { 4, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 96, 96, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 96, 96, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 96, 96, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 128, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 128, 128, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 128, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 128, 128, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 192, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 192, -1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 192, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 192, 2, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 192, 3, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 128, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 128, -1, 1 }, { 4, 2 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 128, 1, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 128, 1, 2 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 128, 1, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 128, 2, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 128, 3, 1 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 192, 128, 3, 4 }, { 2, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 256, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 256, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 320, 256, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 320, 256, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 512, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 512, 512, -1, 1 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 576, 512, -1, 0 }, { 1, 4 } }, + { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q5_1, 576, 512, -1, 1 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 32, 32, -1, 1 }, { 2, 4 } }, { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 32, 32, 1, 2 }, { 1, 4 } }, { { GGML_METAL_DEVICE_M2_MAX, GGML_TYPE_Q8_0, 32, 32, 2, 2 }, { 4, 4 } }, diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 3b93b8a2c..225a69172 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -677,6 +677,11 @@ static constexpr std::initializer_list> topk_qsa_edges { { 5, 1, 4 }, // add->src[1] == reshape { 6, 0, 5 }, // top_k->src[0] == add }; +static constexpr std::initializer_list rms_norm_mul_add_mul_pattern { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD, GGML_OP_MUL }; +static constexpr std::initializer_list rms_norm_mul_add_pattern { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }; +static constexpr std::initializer_list rms_norm_mul_rope_view_set_rows_pattern { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }; +static constexpr std::initializer_list rms_norm_view_set_rows_pattern { GGML_OP_RMS_NORM, GGML_OP_VIEW, GGML_OP_SET_ROWS }; +static constexpr std::initializer_list rope_view_set_rows_pattern { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }; //node #978 ( SOFT_MAX): ffn_moe_probs-15 ( 0K) [Vulka ] use=2: ffn_moe_logits-15 ( 0K) [Vulka ] //node #979 ( RESHAPE): ffn_moe_probs-15 (re ( 0K) [Vulka ] use=1: ffn_moe_probs-15 ( 0K) [Vulka ] @@ -776,6 +781,16 @@ enum topk_moe_mode { TOPK_MOE_COUNT, }; +enum rms_norm_mode { + RMS_NORM_MUL, + RMS_NORM_MUL_ADD, + RMS_NORM_MUL_ADD_MUL, + RMS_NORM_MUL_ROPE, + RMS_NORM_MUL_ROPE_VIEW_SET_ROWS, + RMS_NORM_VIEW_SET_ROWS, + RMS_NORM_COUNT, +}; + static constexpr std::initializer_list> rope_view_set_rows_edges { { 1, 0, 0 }, // view->src[0] == rope { 2, 0, 1 }, // set_rows->src[0] == view @@ -788,6 +803,11 @@ static constexpr std::initializer_list> rms_norm_mul_rope_vie { 4, 0, 3 }, // set_rows->src[0] == view }; +static constexpr std::initializer_list> rms_norm_view_set_rows_edges { + { 1, 0, 0 }, // view->src[0] == rms_norm + { 2, 0, 1 }, // set_rows->src[0] == view +}; + static constexpr std::array lightning_indexer_k_types = { GGML_TYPE_F32, GGML_TYPE_F16, @@ -1008,6 +1028,12 @@ struct vk_device_struct { vk_pipeline pipeline_group_norm_f32; vk_pipeline pipeline_rms_norm_f32; vk_pipeline pipeline_rms_norm_mul_f32; + vk_pipeline pipeline_rms_norm_mul_add_f32; + vk_pipeline pipeline_rms_norm_mul_add_mul_f32; + vk_pipeline pipeline_rms_norm_mul_add_partials_f32; + vk_pipeline pipeline_rms_norm_mul_add_mul_partials_f32; + vk_pipeline pipeline_rms_norm_set_rows_f32_f32; + vk_pipeline pipeline_rms_norm_set_rows_f32_f16; vk_pipeline pipeline_rms_norm_partials_f32; vk_pipeline pipeline_rms_norm_mul_partials_f32; vk_pipeline pipeline_rms_norm_mul_rope_f32_f32; @@ -1090,6 +1116,9 @@ struct vk_device_struct { vk_pipeline pipeline_cumsum_multipass2_f32; vk_pipeline pipeline_argmax_f32; vk_pipeline pipeline_count_equal_i32; + vk_pipeline pipeline_dsv4_hc_comb_f32; + vk_pipeline pipeline_dsv4_hc_pre_f32; + vk_pipeline pipeline_dsv4_hc_post_f32; std::map pipeline_solve_tri_f32; vk_pipeline pipeline_im2col_f32, pipeline_im2col_f32_f16; vk_pipeline pipeline_im2col_3d_f32, pipeline_im2col_3d_f32_f16; @@ -1447,6 +1476,53 @@ struct vk_op_fwht_push_constants { float scale; }; +struct vk_op_dsv4_hc_comb_push_constants { + uint32_t n_tokens; + + uint32_t nbm0; uint32_t nbm1; + uint32_t nbs0; + uint32_t nbb0; + uint32_t nbd0; uint32_t nbd1; uint32_t nbd2; + + uint32_t m_offset; + uint32_t s_offset; + uint32_t b_offset; + uint32_t d_offset; + + float eps; + uint32_t n_iter; +}; + +struct vk_op_dsv4_hc_pre_push_constants { + uint32_t n_embd; + uint32_t n_tokens; + + uint32_t nbx0; uint32_t nbx1; uint32_t nbx2; + uint32_t nbw0; uint32_t nbw1; + uint32_t nbd0; uint32_t nbd1; + + uint32_t x_offset; + uint32_t w_offset; + uint32_t d_offset; +}; + +struct vk_op_dsv4_hc_post_push_constants { + uint32_t n_embd; + uint32_t n_tokens; + + uint32_t nbx0; uint32_t nbx1; + uint32_t nbr0; uint32_t nbr1; uint32_t nbr2; + uint32_t nbp0; uint32_t nbp1; + uint32_t nbc0; uint32_t nbc1; uint32_t nbc2; + uint32_t nbd0; uint32_t nbd1; uint32_t nbd2; + + uint32_t x_offset; + uint32_t r_offset; + uint32_t p_offset; + uint32_t c_offset; + uint32_t d_offset; +}; + struct vk_op_count_experts_push_constants { uint32_t ne00; uint32_t ne01; @@ -2473,6 +2549,7 @@ struct ggml_backend_vk_context { bool fused_topk_moe_scale {}; // QSA indexer gather+add+top_k fused into one radix-select bool fused_topk_qsa {}; + rms_norm_mode fused_rms_norm_mode {RMS_NORM_COUNT}; // for GGML_VK_PERF_LOGGER std::unique_ptr perf_logger; @@ -2494,9 +2571,38 @@ static uint64_t vk_tensor_offset(const ggml_tensor * tensor) { return (uint8_t *) tensor->data - (uint8_t *) vk_ptr_base; } -static uint32_t get_misalign_bytes(const ggml_backend_vk_context * ctx, const ggml_tensor * t) -{ - return ((vk_tensor_offset(t) + t->view_offs) & (ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1));; +static void ggml_vk_host_get(const vk_device& device, const void * ptr, vk_buffer& buf, size_t& buf_offset); + +static size_t ggml_vk_tensor_buffer_offset(const ggml_backend_vk_context * ctx, const ggml_tensor * t) { + // vk_tensor_offset() is relative to vk_ptr_base, but mapped host tensors need an offset relative to their Vulkan buffer. + if (ctx->device->uma) { + vk_buffer buf = nullptr; + size_t off = 0; + ggml_vk_host_get(ctx->device, t->data, buf, off); + if (buf) { + return off; + } + } + return (size_t)(vk_tensor_offset(t) + t->view_offs); +} + +static size_t ggml_vk_descriptor_offset(size_t tensor_offset, size_t alignment, size_t type_size) { + // Move the descriptor back until its distance to the tensor is divisible by the tensor type size. + size_t descriptor_offset = tensor_offset & ~(alignment - 1); + while ((tensor_offset - descriptor_offset) % type_size != 0) { + GGML_ASSERT(descriptor_offset >= alignment); + descriptor_offset -= alignment; + } + + return descriptor_offset; +} + +static uint32_t get_misalign_bytes(const ggml_backend_vk_context * ctx, const ggml_tensor * t) { + const size_t tensor_offset = ggml_vk_tensor_buffer_offset(ctx, t); + const size_t descriptor_offset = ggml_vk_descriptor_offset( + tensor_offset, ctx->device->properties.limits.minStorageBufferOffsetAlignment, ggml_type_size(t->type)); + GGML_ASSERT(tensor_offset - descriptor_offset <= UINT32_MAX); + return tensor_offset - descriptor_offset; } static uint32_t ggml_vk_concat_unit_size(ggml_type type) { @@ -2581,6 +2687,32 @@ template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk GGML_UNUSED(src3); } +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_comb_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) { + p.m_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + p.s_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); + p.b_offset = get_misalign_bytes(ctx, src2) / ggml_type_size(src2->type); + p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + GGML_UNUSED(src3); +} + +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_pre_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) { + p.x_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + p.w_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); + p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); + + GGML_UNUSED(src2); + GGML_UNUSED(src3); +} + +template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_post_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) { + p.x_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); + p.r_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); + p.p_offset = get_misalign_bytes(ctx, src2) / ggml_type_size(src2->type); + p.c_offset = get_misalign_bytes(ctx, src3) / ggml_type_size(src3->type); + p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); +} + struct ggml_backend_vk_buffer_context { vk_device_ref device; vk_buffer dev_buffer; @@ -4733,6 +4865,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0], matmul_q8_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K], matmul_q2_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_TQ2_0], matmul_tq2_0_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_TQ1_0], matmul_tq1_0_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K], matmul_q3_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K], matmul_q4_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K], matmul_q5_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) @@ -4774,6 +4907,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0], matmul_id_subgroup_tq1_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, 5) @@ -4847,6 +4981,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0], matmul_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ1_0], matmul_tq1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K], matmul_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); @@ -4892,6 +5027,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0], matmul_id_subgroup_tq1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); @@ -4983,6 +5119,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0], matmul_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0], matmul_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ1_0], matmul_tq1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K], matmul_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); @@ -5032,6 +5169,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0], matmul_id_subgroup_tq1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); @@ -5080,6 +5218,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_q8_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_q2_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_tq2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0], matmul_id_tq1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_q3_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_q4_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_q5_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); @@ -5160,6 +5299,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f32acc, matmul_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0].f32acc, matmul_tq2_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ1_0].f32acc, matmul_tq1_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f32acc, matmul_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f32acc, matmul_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f32acc, matmul_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); @@ -5208,6 +5348,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_subgroup_q8_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_subgroup_q2_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0].f32acc, matmul_id_subgroup_tq2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0].f32acc, matmul_id_subgroup_tq1_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_subgroup_q3_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_subgroup_q4_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_subgroup_q5_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); @@ -5238,6 +5379,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_q8_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_q2_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0].f32acc, matmul_id_tq2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0].f32acc, matmul_id_tq1_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_q3_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_q4_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_q5_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); @@ -5347,6 +5489,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f32_f32", arr_dmmv_q8_0_f32_f32_len[reduc], arr_dmmv_q8_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f32_f32", arr_dmmv_q2_k_f32_f32_len[reduc16], arr_dmmv_q2_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_TQ2_0][i], "mul_mat_vec_tq2_0_f32_f32", arr_dmmv_tq2_0_f32_f32_len[reduc16], arr_dmmv_tq2_0_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_TQ1_0][i], "mul_mat_vec_tq1_0_f32_f32", arr_dmmv_tq1_0_f32_f32_len[reduc16], arr_dmmv_tq1_0_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f32_f32", arr_dmmv_q3_k_f32_f32_len[reduc16], arr_dmmv_q3_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f32_f32", arr_dmmv_q4_k_f32_f32_len[reduc16], arr_dmmv_q4_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f32_f32", arr_dmmv_q5_k_f32_f32_len[reduc16], arr_dmmv_q5_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5375,6 +5518,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f16_f32", arr_dmmv_q8_0_f16_f32_len[reduc], arr_dmmv_q8_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f16_f32", arr_dmmv_q2_k_f16_f32_len[reduc16], arr_dmmv_q2_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_TQ2_0][i], "mul_mat_vec_tq2_0_f16_f32", arr_dmmv_tq2_0_f16_f32_len[reduc16], arr_dmmv_tq2_0_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_TQ1_0][i], "mul_mat_vec_tq1_0_f16_f32", arr_dmmv_tq1_0_f16_f32_len[reduc16], arr_dmmv_tq1_0_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f16_f32", arr_dmmv_q3_k_f16_f32_len[reduc16], arr_dmmv_q3_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f16_f32", arr_dmmv_q4_k_f16_f32_len[reduc16], arr_dmmv_q4_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f16_f32", arr_dmmv_q5_k_f16_f32_len[reduc16], arr_dmmv_q5_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5430,6 +5574,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q8_0], "mul_mat_vec_id_q8_0_f32", arr_dmmv_id_q8_0_f32_f32_len[reduc], arr_dmmv_id_q8_0_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q2_K], "mul_mat_vec_id_q2_k_f32", arr_dmmv_id_q2_k_f32_f32_len[reduc16], arr_dmmv_id_q2_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_TQ2_0], "mul_mat_vec_id_tq2_0_f32", arr_dmmv_id_tq2_0_f32_f32_len[reduc16], arr_dmmv_id_tq2_0_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_TQ1_0], "mul_mat_vec_id_tq1_0_f32", arr_dmmv_id_tq1_0_f32_f32_len[reduc16], arr_dmmv_id_tq1_0_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q3_K], "mul_mat_vec_id_q3_k_f32", arr_dmmv_id_q3_k_f32_f32_len[reduc16], arr_dmmv_id_q3_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q4_K], "mul_mat_vec_id_q4_k_f32", arr_dmmv_id_q4_k_f32_f32_len[reduc16], arr_dmmv_id_q4_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q5_K], "mul_mat_vec_id_q5_k_f32", arr_dmmv_id_q5_k_f32_f32_len[reduc16], arr_dmmv_id_q5_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5496,6 +5641,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q8_0], "dequant_q8_0_transpose", dequant_q8_0_transpose_len, dequant_q8_0_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_TQ2_0], "dequant_tq2_0", dequant_tq2_0_len, dequant_tq2_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_TQ1_0], "dequant_tq1_0", dequant_tq1_0_len, dequant_tq1_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_K], "dequant_q4_k", dequant_q4_k_len, dequant_q4_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_K], "dequant_q5_k", dequant_q5_k_len, dequant_q5_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); @@ -5525,6 +5671,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q8_0], "get_rows_q8_0", get_rows_q8_0_len, get_rows_q8_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q2_K], "get_rows_q2_k", get_rows_q2_k_len, get_rows_q2_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_TQ2_0], "get_rows_tq2_0", get_rows_tq2_0_len, get_rows_tq2_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_TQ1_0], "get_rows_tq1_0", get_rows_tq1_0_len, get_rows_tq1_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q3_K], "get_rows_q3_k", get_rows_q3_k_len, get_rows_q3_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_K], "get_rows_q4_k", get_rows_q4_k_len, get_rows_q4_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_K], "get_rows_q5_k", get_rows_q5_k_len, get_rows_q5_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); @@ -5554,6 +5701,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q8_0], "get_rows_q8_0_f32", get_rows_q8_0_f32_len, get_rows_q8_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q2_K], "get_rows_q2_k_f32", get_rows_q2_k_f32_len, get_rows_q2_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_TQ2_0], "get_rows_tq2_0_f32", get_rows_tq2_0_f32_len, get_rows_tq2_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_TQ1_0], "get_rows_tq1_0_f32", get_rows_tq1_0_f32_len, get_rows_tq1_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q3_K], "get_rows_q3_k_f32", get_rows_q3_k_f32_len, get_rows_q3_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_K], "get_rows_q4_k_f32", get_rows_q4_k_f32_len, get_rows_q4_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_K], "get_rows_q5_k_f32", get_rows_q5_k_f32_len, get_rows_q5_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); @@ -5599,6 +5747,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_rms_norm_f32, "rms_norm_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_f32, "rms_norm_mul_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_f32, "rms_norm_mul_add_f32", rms_norm_mul_add_f32_len, rms_norm_mul_add_f32_data, "main", 5, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 0}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_mul_f32, "rms_norm_mul_add_mul_f32", rms_norm_mul_add_f32_len, rms_norm_mul_add_f32_data, "main", 5, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_partials_f32, "rms_norm_mul_add_partials_f32", rms_norm_mul_add_partials_f32_len, rms_norm_mul_add_partials_f32_data, "main", 6, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 0}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_add_mul_partials_f32, "rms_norm_mul_add_mul_partials_f32", rms_norm_mul_add_partials_f32_len, rms_norm_mul_add_partials_f32_data, "main", 6, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1, 1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_set_rows_f32_f32, "rms_norm_set_rows_f32_f32", rms_norm_set_rows_f32_f32_len, rms_norm_set_rows_f32_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_set_rows_f32_f16, "rms_norm_set_rows_f32_f16", rms_norm_set_rows_f32_f16_len, rms_norm_set_rows_f32_f16_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_partials_f32, "rms_norm_partials_f32", rms_norm_partials_f32_len, rms_norm_partials_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_partials_f32, "rms_norm_mul_partials_f32", rms_norm_partials_f32_len, rms_norm_partials_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true); @@ -5905,6 +6059,16 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true); } + // comb holds a token's 4x4 matrix in one 16-lane slice of a subgroup, so it + // needs at least 16 lanes, pinned to a known size. + if (device->subgroup_basic && device->subgroup_shuffle && device->subgroup_require_full_support && device->subgroup_size >= 16) { + const uint32_t tokens_per_workgroup = 4 * (device->subgroup_size / 16); + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_comb_f32, "dsv4_hc_comb_f32", dsv4_hc_comb_f32_len, dsv4_hc_comb_f32_data, "main", 4, sizeof(vk_op_dsv4_hc_comb_push_constants), {tokens_per_workgroup, 1, 1}, { device->subgroup_size }, 1, true, true, device->subgroup_size); + } + + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_pre_f32, "dsv4_hc_pre_f32", dsv4_hc_pre_f32_len, dsv4_hc_pre_f32_data, "main", 3, sizeof(vk_op_dsv4_hc_pre_push_constants), {256, 1, 1}, { 256 }, 1); + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_f32, "dsv4_hc_post_f32", dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, { 256 }, 1); + for (auto &s : device->pipeline_solve_tri_f32) { const vk_solve_tri_pipeline_state &state = s.first; @@ -7812,6 +7976,7 @@ static vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: break; default: return nullptr; @@ -7887,6 +8052,7 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_conte case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: break; default: return nullptr; @@ -7957,6 +8123,7 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: break; default: return nullptr; @@ -8051,6 +8218,7 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_co case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: break; default: return nullptr; @@ -8124,6 +8292,7 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: break; default: return nullptr; @@ -8236,10 +8405,12 @@ static vk_subbuffer ggml_vk_tensor_subbuffer( size_t size = ggml_nbytes(tensor); - size_t misalign_bytes = offset & (ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1); + const size_t descriptor_offset = ggml_vk_descriptor_offset( + offset, ctx->device->properties.limits.minStorageBufferOffsetAlignment, ggml_type_size(tensor->type)); + const size_t misalign_bytes = offset - descriptor_offset; // The shader must support misaligned offsets when indexing into the buffer GGML_ASSERT(allow_misalign || misalign_bytes == 0); - offset &= ~misalign_bytes; + offset = descriptor_offset; size += misalign_bytes; return vk_subbuffer{buffer, offset, size}; @@ -10152,6 +10323,98 @@ static void ggml_vk_fwht(ggml_backend_vk_context * ctx, vk_context& subctx, cons ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src_buf, dst_buf }, pc, { workgroups_x, 1, 1 }); } +static uint32_t ggml_vk_nb_elem(const ggml_tensor * t, int i) { + return (uint32_t)(t->nb[i] / ggml_type_size(t->type)); +} + +static void ggml_vk_dsv4_hc_comb(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * mixes, const ggml_tensor * scale, const ggml_tensor * base, ggml_tensor * dst) { + VK_LOG_DEBUG("ggml_vk_dsv4_hc_comb(" << mixes << ", " << scale << ", " << base << ", " << dst << ")"); + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_comb_f32; + GGML_ASSERT(pipeline != nullptr); + + const uint32_t n_tokens = (uint32_t)mixes->ne[1]; + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_subbuffer mixes_buf = ggml_vk_tensor_subbuffer(ctx, mixes, true); + const vk_subbuffer scale_buf = ggml_vk_tensor_subbuffer(ctx, scale, true); + const vk_subbuffer base_buf = ggml_vk_tensor_subbuffer(ctx, base, true); + const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + + vk_op_dsv4_hc_comb_push_constants pc = { + n_tokens, + ggml_vk_nb_elem(mixes, 0), ggml_vk_nb_elem(mixes, 1), + ggml_vk_nb_elem(scale, 0), + ggml_vk_nb_elem(base, 0), + ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1), ggml_vk_nb_elem(dst, 2), + 0, 0, 0, 0, + ggml_get_op_params_f32(dst, 0), + (uint32_t)ggml_get_op_params_i32(dst, 1), + }; + init_pushconst_tensor_offsets(ctx, pc, mixes, scale, base, nullptr, dst); + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { mixes_buf, scale_buf, base_buf, dst_buf }, pc, { n_tokens, 1, 1 }); +} + +static void ggml_vk_dsv4_hc_pre(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * weights, ggml_tensor * dst) { + VK_LOG_DEBUG("ggml_vk_dsv4_hc_pre(" << x << ", " << weights << ", " << dst << ")"); + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_pre_f32; + GGML_ASSERT(pipeline != nullptr); + + const uint32_t n_embd = (uint32_t)x->ne[0]; + const uint32_t n_tokens = (uint32_t)x->ne[2]; + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_subbuffer x_buf = ggml_vk_tensor_subbuffer(ctx, x, true); + const vk_subbuffer w_buf = ggml_vk_tensor_subbuffer(ctx, weights, true); + const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + + vk_op_dsv4_hc_pre_push_constants pc = { + n_embd, n_tokens, + ggml_vk_nb_elem(x, 0), ggml_vk_nb_elem(x, 1), ggml_vk_nb_elem(x, 2), + ggml_vk_nb_elem(weights, 0), ggml_vk_nb_elem(weights, 1), + ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1), + 0, 0, 0, + }; + init_pushconst_tensor_offsets(ctx, pc, x, weights, nullptr, nullptr, dst); + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, w_buf, d_buf }, pc, { n_embd, n_tokens, 1 }); +} + +static void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * residual, const ggml_tensor * post, const ggml_tensor * comb, ggml_tensor * dst) { + VK_LOG_DEBUG("ggml_vk_dsv4_hc_post(" << x << ", " << residual << ", " << post << ", " << comb << ", " << dst << ")"); + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_post_f32; + GGML_ASSERT(pipeline != nullptr); + + const uint32_t n_embd = (uint32_t)x->ne[0]; + const uint32_t n_tokens = (uint32_t)x->ne[1]; + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_subbuffer x_buf = ggml_vk_tensor_subbuffer(ctx, x, true); + const vk_subbuffer r_buf = ggml_vk_tensor_subbuffer(ctx, residual, true); + const vk_subbuffer p_buf = ggml_vk_tensor_subbuffer(ctx, post, true); + const vk_subbuffer c_buf = ggml_vk_tensor_subbuffer(ctx, comb, true); + const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + + vk_op_dsv4_hc_post_push_constants pc = { + n_embd, n_tokens, + ggml_vk_nb_elem(x, 0), ggml_vk_nb_elem(x, 1), + ggml_vk_nb_elem(residual, 0), ggml_vk_nb_elem(residual, 1), ggml_vk_nb_elem(residual, 2), + ggml_vk_nb_elem(post, 0), ggml_vk_nb_elem(post, 1), + ggml_vk_nb_elem(comb, 0), ggml_vk_nb_elem(comb, 1), ggml_vk_nb_elem(comb, 2), + ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1), ggml_vk_nb_elem(dst, 2), + 0, 0, 0, 0, 0, + }; + init_pushconst_tensor_offsets(ctx, pc, x, residual, post, comb, dst); + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, r_buf, p_buf, c_buf, d_buf }, pc, { n_embd, n_tokens, 1 }); +} + static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx) { ggml_tensor * dst = cgraph->nodes[node_idx]; ggml_tensor * src0 = dst->src[0]; @@ -11563,10 +11826,9 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const case GGML_OP_RMS_NORM: if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { if (ctx->do_add_rms_partials) { - return ctx->num_additional_fused_ops > 0 ? ctx->device->pipeline_rms_norm_mul_partials_f32 : ctx->device->pipeline_rms_norm_partials_f32; - } else { - return ctx->num_additional_fused_ops > 0 ? ctx->device->pipeline_rms_norm_mul_f32 : ctx->device->pipeline_rms_norm_f32; + return ctx->fused_rms_norm_mode == RMS_NORM_MUL ? ctx->device->pipeline_rms_norm_mul_partials_f32 : ctx->device->pipeline_rms_norm_partials_f32; } + return ctx->fused_rms_norm_mode == RMS_NORM_MUL ? ctx->device->pipeline_rms_norm_mul_f32 : ctx->device->pipeline_rms_norm_f32; } return nullptr; case GGML_OP_RMS_NORM_BACK: @@ -12135,7 +12397,9 @@ template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk const uint32_t b_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); - GGML_ASSERT(dst->op != GGML_OP_GET_ROWS || (a_offset == 0 && b_offset == 0 && d_offset == 0)); + GGML_ASSERT(a_offset <= 0xFFFF); + GGML_ASSERT(b_offset <= 0xFF); + GGML_ASSERT(d_offset <= 0xFF); p.misalign_offsets = (a_offset << 16) | (b_offset << 8) | d_offset; @@ -13512,40 +13776,121 @@ static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor * return rope; } -static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx, float * op_params) { - ggml_tensor * dst; - const ggml_tensor * src0; - const ggml_tensor * src1; - - if (ctx->num_additional_fused_ops > 0) { - // fused rms_norm + mul - ggml_tensor *mul = cgraph->nodes[node_idx + 1]; - ggml_tensor *other_src = mul->src[0] == cgraph->nodes[node_idx + 0] ? mul->src[1] : mul->src[0]; - dst = mul; - src0 = cgraph->nodes[node_idx]->src[0]; - src1 = other_src; - } else { - dst = cgraph->nodes[node_idx]; - src0 = src1 = dst->src[0]; - } - +static vk_op_binary_push_constants ggml_vk_rms_norm_push_constants( + const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, + float eps, uint32_t num_partials) { const uint32_t src0_type_size = ggml_type_size(src0->type); const uint32_t src1_type_size = ggml_type_size(src1->type); const uint32_t dst_type_size = ggml_type_size(dst->type); - uint32_t param3 = ctx->do_add_rms_partials ? ggml_vk_rms_num_partials(ctx, dst) : 0; - - vk_op_binary_push_constants bin { + return { (uint32_t)ggml_nelements(src0), (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size, (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size, (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] / dst_type_size, (uint32_t) dst->nb[1] / dst_type_size, (uint32_t) dst->nb[2] / dst_type_size, (uint32_t) dst->nb[3] / dst_type_size, 0, - op_params[0], 0.0f, (int32_t)param3, + eps, 0.0f, (int32_t)num_partials, }; +} - // more than one fused op means rms_norm+mul+rope - if (ctx->num_additional_fused_ops > 1) { +static void ggml_vk_rms_norm_finish(ggml_backend_vk_context * ctx, const ggml_tensor * src0) { + if (ctx->do_add_rms_partials_offset_calculation) { + ctx->prealloc_size_add_rms_partials_offset += ggml_vk_rms_partials_size(ctx, src0); + ctx->do_add_rms_partials = false; + ctx->do_add_rms_partials_offset_calculation = false; + } +} + +static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx, float * op_params) { + ggml_tensor * rms = cgraph->nodes[node_idx]; + const ggml_tensor * src0 = rms->src[0]; + + if (ctx->fused_rms_norm_mode == RMS_NORM_VIEW_SET_ROWS) { + GGML_ASSERT(ctx->num_additional_fused_ops == 2); + ggml_tensor * set_rows = cgraph->nodes[node_idx + 2]; + const ggml_tensor * indices = set_rows->src[1]; + vk_op_binary_push_constants pc = ggml_vk_rms_norm_push_constants(src0, src0, set_rows, op_params[0], 0); + init_pushconst_tensor_offsets(ctx, pc, src0, src0, nullptr, nullptr, set_rows); + + vk_pipeline pipeline = set_rows->type == GGML_TYPE_F16 ? + ctx->device->pipeline_rms_norm_set_rows_f32_f16 : ctx->device->pipeline_rms_norm_set_rows_f32_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + ggml_vk_tensor_subbuffer(ctx, src0, true), + ggml_vk_tensor_subbuffer(ctx, src0, true), + ggml_vk_tensor_subbuffer(ctx, set_rows, true), + ggml_vk_tensor_subbuffer(ctx, indices), + }, pc, { (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3] }); + ggml_vk_rms_norm_finish(ctx, src0); + return; + } + + if (ctx->fused_rms_norm_mode == RMS_NORM_MUL_ADD || ctx->fused_rms_norm_mode == RMS_NORM_MUL_ADD_MUL) { + ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + ggml_tensor * add = cgraph->nodes[node_idx + 2]; + const ggml_tensor * weight = mul->src[0] == rms ? mul->src[1] : mul->src[0]; + const ggml_tensor * residual = add->src[0] == mul ? add->src[1] : add->src[0]; + const bool do_post_multiply = ctx->fused_rms_norm_mode == RMS_NORM_MUL_ADD_MUL; + GGML_ASSERT(ctx->num_additional_fused_ops == (do_post_multiply ? 3 : 2)); + ggml_tensor * dst = do_post_multiply ? cgraph->nodes[node_idx + 3] : add; + const ggml_tensor * post_scale = do_post_multiply ? + (dst->src[0] == add ? dst->src[1] : dst->src[0]) : src0; + + const uint32_t num_partials = ctx->do_add_rms_partials ? ggml_vk_rms_num_partials(ctx, dst) : 0; + vk_op_binary_push_constants pc = ggml_vk_rms_norm_push_constants(src0, weight, dst, op_params[0], num_partials); + init_pushconst_tensor_offsets(ctx, pc, src0, weight, residual, post_scale, dst); + + vk_pipeline pipeline; + if (ctx->do_add_rms_partials) { + pipeline = do_post_multiply ? + ctx->device->pipeline_rms_norm_mul_add_mul_partials_f32 : ctx->device->pipeline_rms_norm_mul_add_partials_f32; + } else { + pipeline = do_post_multiply ? + ctx->device->pipeline_rms_norm_mul_add_mul_f32 : ctx->device->pipeline_rms_norm_mul_add_f32; + } + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + if (ctx->do_add_rms_partials) { + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + ggml_vk_tensor_subbuffer(ctx, src0, true), + ggml_vk_tensor_subbuffer(ctx, weight, true), + ggml_vk_tensor_subbuffer(ctx, dst, true), + ggml_vk_subbuffer(ctx, ctx->prealloc_add_rms_partials, ctx->prealloc_size_add_rms_partials_offset), + ggml_vk_tensor_subbuffer(ctx, residual), + ggml_vk_tensor_subbuffer(ctx, post_scale), + }, pc, { (uint32_t)CEIL_DIV(src0->ne[0], 128), 1, 1 }); + } else { + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + ggml_vk_tensor_subbuffer(ctx, src0, true), + ggml_vk_tensor_subbuffer(ctx, weight, true), + ggml_vk_tensor_subbuffer(ctx, dst, true), + ggml_vk_tensor_subbuffer(ctx, residual), + ggml_vk_tensor_subbuffer(ctx, post_scale), + }, pc, { (uint32_t)src0->ne[1], (uint32_t)src0->ne[2], (uint32_t)src0->ne[3] }); + } + ggml_vk_rms_norm_finish(ctx, src0); + return; + } + + ggml_tensor * dst; + const ggml_tensor * src1; + + if (ctx->fused_rms_norm_mode != RMS_NORM_COUNT) { + ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + dst = mul; + src1 = mul->src[0] == rms ? mul->src[1] : mul->src[0]; + } else { + dst = rms; + src1 = src0; + } + + const uint32_t num_partials = ctx->do_add_rms_partials ? ggml_vk_rms_num_partials(ctx, dst) : 0; + vk_op_binary_push_constants bin = ggml_vk_rms_norm_push_constants(src0, src1, dst, op_params[0], num_partials); + + if (ctx->fused_rms_norm_mode == RMS_NORM_MUL_ROPE || + ctx->fused_rms_norm_mode == RMS_NORM_MUL_ROPE_VIEW_SET_ROWS) { static constexpr uint32_t max_tensors = 7; const ggml_tensor *tensors[max_tensors] {}; @@ -13555,7 +13900,8 @@ static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor *other_src = mul->src[0] == rms ? mul->src[1] : mul->src[0]; - bool do_set_rows = ctx->num_additional_fused_ops == 4; + bool do_set_rows = ctx->fused_rms_norm_mode == RMS_NORM_MUL_ROPE_VIEW_SET_ROWS; + GGML_ASSERT(ctx->num_additional_fused_ops == (do_set_rows ? 4 : 2)); tensors[0] = rms->src[0]; tensors[1] = other_src; @@ -13622,14 +13968,11 @@ static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_vk_subbuffer(ctx, buf[6], offset[6]), }, pc, elements); } else { + GGML_ASSERT(ctx->fused_rms_norm_mode == RMS_NORM_MUL || ctx->fused_rms_norm_mode == RMS_NORM_COUNT); ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_RMS_NORM, std::move(bin)); } - if (ctx->do_add_rms_partials_offset_calculation) { - ctx->prealloc_size_add_rms_partials_offset += ggml_vk_rms_partials_size(ctx, src0); - ctx->do_add_rms_partials = false; - ctx->do_add_rms_partials_offset_calculation = false; - } + ggml_vk_rms_norm_finish(ctx, src0); } static void ggml_vk_rms_norm_back(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { @@ -16090,6 +16433,18 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr case GGML_OP_CUMSUM: ggml_vk_cumsum(ctx, compute_ctx, src0, node); + break; + case GGML_OP_DSV4_HC_COMB: + ggml_vk_dsv4_hc_comb(ctx, compute_ctx, src0, src1, src2, node); + + break; + case GGML_OP_DSV4_HC_PRE: + ggml_vk_dsv4_hc_pre(ctx, compute_ctx, src0, src1, node); + + break; + case GGML_OP_DSV4_HC_POST: + ggml_vk_dsv4_hc_post(ctx, compute_ctx, src0, src1, src2, src3, node); + break; case GGML_OP_MEAN: ggml_vk_mean(ctx, compute_ctx, src0, node); @@ -16950,7 +17305,8 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g return false; } - if (ops.size() == 2 && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { + if ((ops.size() == 2 || ops.size() == 3 || ops.size() == 4) && + ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { // additional constraints specific to this fusion const ggml_tensor *rms_norm = cgraph->nodes[node_idx]; const ggml_tensor *mul = cgraph->nodes[node_idx + 1]; @@ -16972,6 +17328,43 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g if (!ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) { return false; } + + if (ops.size() >= 3 && ops.begin()[2] == GGML_OP_ADD) { + const ggml_tensor *add = cgraph->nodes[node_idx + 2]; + const ggml_tensor *residual = add->src[0] == mul ? add->src[1] : add->src[0]; + if (add->src[0] != mul && add->src[1] != mul) { + return false; + } + if (residual->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32 || + !ggml_are_same_shape(add, residual) || !ggml_is_contiguous(residual) || + !ggml_is_contiguous(add) || get_misalign_bytes(ctx, residual) != 0) { + return false; + } + + const ggml_tensor *dst = add; + if (ops.size() == 4) { + if (ops.begin()[3] != GGML_OP_MUL) { + return false; + } + + const ggml_tensor *post_mul = cgraph->nodes[node_idx + 3]; + const ggml_tensor *scale = post_mul->src[0] == add ? post_mul->src[1] : post_mul->src[0]; + if (post_mul->src[0] != add && post_mul->src[1] != add) { + return false; + } + // The shader reads data_e[0], so the final multiply must use a scalar. + if (scale->type != GGML_TYPE_F32 || post_mul->type != GGML_TYPE_F32 || + ggml_nelements(scale) != 1 || !ggml_is_contiguous(post_mul) || + get_misalign_bytes(ctx, scale) != 0) { + return false; + } + dst = post_mul; + } + + if (get_misalign_bytes(ctx, dst) != 0) { + return false; + } + } } auto const &mm_add_ok = [&](const ggml_tensor *mul, const ggml_tensor *add) { const ggml_tensor *bias = add->src[0] == mul ? add->src[1] : add->src[0]; @@ -17353,12 +17746,11 @@ static bool ggml_vk_can_fuse_topk_qsa(ggml_backend_vk_context * ctx, const struc static bool ggml_vk_can_fuse_rope_set_rows(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx) { - GGML_UNUSED(ctx); const ggml_tensor *rope = cgraph->nodes[node_idx + 0]; const ggml_tensor *view = cgraph->nodes[node_idx + 1]; const ggml_tensor *set_rows = cgraph->nodes[node_idx + 2]; - // ne3 not tested + // The set_rows epilogue uses one index per ne2 slice and does not encode ne3. if (rope->src[0]->ne[3] != 1) { return false; } @@ -17367,19 +17759,50 @@ static bool ggml_vk_can_fuse_rope_set_rows(ggml_backend_vk_context * ctx, const return false; } - if (set_rows->src[1]->type != GGML_TYPE_I64) { + // The shader reads each aligned I64 index as a uvec2 and uses its low 32 bits. + if (set_rows->src[1]->type != GGML_TYPE_I64 || !ggml_is_contiguous(set_rows->src[1]) || + set_rows->nb[0] != ggml_type_size(set_rows->type) || get_misalign_bytes(ctx, set_rows->src[1]) != 0) { return false; } - // The view should flatten two dims of rope into one dim + // SET_ROWS consumes one flattened [ne0*ne1] row for each ne2 slice. if (!ggml_is_contiguous(view) || - view->ne[0] != rope->ne[0] * rope->ne[1]) { + view->ne[0] != rope->ne[0] * rope->ne[1] || view->ne[1] != rope->ne[2] || + view->ne[2] != 1 || view->ne[3] != 1 || + ggml_nelements(set_rows->src[1]) != rope->ne[2]) { return false; } - // Only norm/neox/mrope shaders have the fusion code + // Only norm/neox/mrope/imrope shaders have the fusion code const int mode = ((const int32_t *) rope->op_params)[2]; - if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX && mode != GGML_ROPE_TYPE_MROPE) { + if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX && + mode != GGML_ROPE_TYPE_MROPE && mode != GGML_ROPE_TYPE_IMROPE) { + return false; + } + + return true; +} + +static bool ggml_vk_can_fuse_rms_norm_set_rows(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, + int node_idx) { + const ggml_tensor * rms = cgraph->nodes[node_idx]; + const ggml_tensor * view = cgraph->nodes[node_idx + 1]; + const ggml_tensor * set_rows = cgraph->nodes[node_idx + 2]; + + // The RMS kernel reads F32 and writes directly to the F32 or F16 SET_ROWS destination. + if (rms->src[0]->type != GGML_TYPE_F32 || rms->type != GGML_TYPE_F32 || + (set_rows->type != GGML_TYPE_F32 && set_rows->type != GGML_TYPE_F16) || + set_rows->src[1]->type != GGML_TYPE_I64 || !ggml_is_contiguous(set_rows->src[1]) || + set_rows->nb[0] != ggml_type_size(set_rows->type) || get_misalign_bytes(ctx, set_rows->src[1]) != 0) { + return false; + } + // As with the ROPE epilogue, each ne2 slice supplies one flattened row and ne3 is not encoded. + if (rms->ne[3] != 1 || !ggml_is_contiguous(rms->src[0]) || !ggml_is_contiguous(view)) { + return false; + } + if (view->ne[0] != rms->ne[0] * rms->ne[1] || view->ne[1] != rms->ne[2] || + view->ne[2] != 1 || view->ne[3] != 1 || + ggml_nelements(set_rows->src[1]) != rms->ne[2]) { return false; } @@ -17474,7 +17897,6 @@ static bool ggml_vk_tensors_overlap(const ggml_tensor * a, const ggml_tensor * b static bool ggml_vk_can_fuse_rms_norm_mul_rope(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx) { - GGML_UNUSED(ctx); const ggml_tensor *rms = cgraph->nodes[node_idx + 0]; const ggml_tensor *mul = cgraph->nodes[node_idx + 1]; const ggml_tensor *rope = cgraph->nodes[node_idx + 2]; @@ -17731,6 +18153,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; ctx->fused_topk_qsa = false; + ctx->fused_rms_norm_mode = RMS_NORM_COUNT; const char *fusion_string {}; if (!ctx->device->disable_fusion) { uint32_t num_adds = ggml_vk_fuse_multi_add(ctx, cgraph, i); @@ -17765,27 +18188,47 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg fusion_string = "MUL_MAT_ID_MUL"; op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = true; - } else if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, { i + 4 }) && + } else if (ggml_can_fuse_subgraph(cgraph, i, rms_norm_mul_rope_view_set_rows_pattern, { i + 4 }) && ggml_check_edges(cgraph, i, rms_norm_mul_rope_view_set_rows_edges) && ggml_vk_can_fuse_rms_norm_mul_rope(ctx, cgraph, i) && ggml_vk_can_fuse_rope_set_rows(ctx, cgraph, i + 2)) { ctx->num_additional_fused_ops = 4; + ctx->fused_rms_norm_mode = RMS_NORM_MUL_ROPE_VIEW_SET_ROWS; fusion_string = "RMS_NORM_MUL_ROPE_VIEW_SET_ROWS"; op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = false; op_srcs_fused_elementwise[2] = false; op_srcs_fused_elementwise[3] = false; op_srcs_fused_elementwise[4] = false; - } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE })&& + } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE }) && ggml_vk_can_fuse_rms_norm_mul_rope(ctx, cgraph, i)) { ctx->num_additional_fused_ops = 2; + ctx->fused_rms_norm_mode = RMS_NORM_MUL_ROPE; fusion_string = "RMS_NORM_MUL_ROPE"; // rope is approximately elementwise - whole rows are done by a single workgroup and it's row-wise op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = true; op_srcs_fused_elementwise[2] = true; + } else if (ggml_vk_can_fuse(ctx, cgraph, i, rms_norm_mul_add_mul_pattern)) { + ctx->num_additional_fused_ops = 3; + ctx->fused_rms_norm_mode = RMS_NORM_MUL_ADD_MUL; + fusion_string = "RMS_NORM_MUL_ADD_MUL"; + std::fill_n(op_srcs_fused_elementwise, 4, true); + } else if (ggml_vk_can_fuse(ctx, cgraph, i, rms_norm_mul_add_pattern)) { + ctx->num_additional_fused_ops = 2; + ctx->fused_rms_norm_mode = RMS_NORM_MUL_ADD; + fusion_string = "RMS_NORM_MUL_ADD"; + std::fill_n(op_srcs_fused_elementwise, 3, true); + } else if (ggml_can_fuse_subgraph(cgraph, i, rms_norm_view_set_rows_pattern, { i + 2 }) && + ggml_check_edges(cgraph, i, rms_norm_view_set_rows_edges) && + ggml_vk_can_fuse_rms_norm_set_rows(ctx, cgraph, i)) { + ctx->num_additional_fused_ops = 2; + ctx->fused_rms_norm_mode = RMS_NORM_VIEW_SET_ROWS; + fusion_string = "RMS_NORM_VIEW_SET_ROWS"; + std::fill_n(op_srcs_fused_elementwise, 3, false); } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) { ctx->num_additional_fused_ops = 1; + ctx->fused_rms_norm_mode = RMS_NORM_MUL; fusion_string = "RMS_NORM_MUL"; // rms_norm is not elementwise, but whole rows must be consumed and the scale factor computed before // they are overwritten, and one workgroup per row. So close enough. @@ -17804,7 +18247,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg fusion_string = "SSM_CONV_SILU"; op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = true; - } else if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, { i + 2 }) && + } else if (ggml_can_fuse_subgraph(cgraph, i, rope_view_set_rows_pattern, { i + 2 }) && ggml_check_edges(cgraph, i, rope_view_set_rows_edges) && ggml_vk_can_fuse_rope_set_rows(ctx, cgraph, i)) { ctx->num_additional_fused_ops = 2; @@ -17942,6 +18385,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; ctx->fused_topk_qsa = false; + ctx->fused_rms_norm_mode = RMS_NORM_COUNT; } } @@ -18142,6 +18586,22 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * continue; } + if (keep_pattern(rms_norm_mul_add_mul_pattern)) { + continue; + } + if (keep_pattern(rms_norm_mul_add_pattern)) { + continue; + } + if (keep_pattern(rms_norm_mul_rope_view_set_rows_pattern)) { + continue; + } + if (keep_pattern(rms_norm_view_set_rows_pattern)) { + continue; + } + if (keep_pattern(rope_view_set_rows_pattern)) { + continue; + } + // First, grab the next unused node. current_set.push_back(first_unused); @@ -18175,7 +18635,12 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * match_pattern(topk_moe_early_softmax, j) || match_pattern(topk_moe_late_softmax, j) || match_pattern(snake_pattern, j) || - in_qsa_pattern(j)) { + in_qsa_pattern(j) || + match_pattern(rms_norm_mul_add_mul_pattern, j) || + match_pattern(rms_norm_mul_add_pattern, j) || + match_pattern(rms_norm_mul_rope_view_set_rows_pattern, j) || + match_pattern(rms_norm_view_set_rows_pattern, j) || + match_pattern(rope_view_set_rows_pattern, j)) { continue; } bool ok = true; @@ -18215,30 +18680,41 @@ static void ggml_vk_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * } } } - // Look for ROPE + VIEW + SET_ROWS and make them consecutive - if (graph->nodes[rope_idx]->op == GGML_OP_ROPE) { + // Look for ROPE/RMS_NORM + VIEW + SET_ROWS and make them consecutive + if (graph->nodes[rope_idx]->op == GGML_OP_ROPE || graph->nodes[rope_idx]->op == GGML_OP_RMS_NORM) { int view_idx = -1; int set_rows_idx = -1; - for (int k = rope_idx+1; k < std::min(rope_idx + 10, graph->n_nodes); ++k) { - if (view_idx == -1 && - graph->nodes[k]->op == GGML_OP_VIEW && - graph->nodes[k]->src[0] == graph->nodes[rope_idx]) { + for (int k = rope_idx + 1; k < std::min(rope_idx + 15, graph->n_nodes); ++k) { + if (used[k]) { + continue; + } + if (view_idx == -1 && graph->nodes[k]->op == GGML_OP_VIEW && graph->nodes[k]->src[0] == graph->nodes[rope_idx]) { view_idx = k; continue; } - if (view_idx != -1 && - set_rows_idx == -1 && - graph->nodes[k]->op == GGML_OP_SET_ROWS && - graph->nodes[k]->src[0] == graph->nodes[view_idx]) { + if (view_idx != -1 && graph->nodes[k]->op == GGML_OP_SET_ROWS && graph->nodes[k]->src[0] == graph->nodes[view_idx]) { set_rows_idx = k; break; } } if (set_rows_idx != -1) { - current_set.push_back(view_idx); - current_set.push_back(set_rows_idx); - used[view_idx] = true; - used[set_rows_idx] = true; + const int node_idxs[] = { rope_idx, view_idx, set_rows_idx }; + const ggml_op ops[] = { graph->nodes[rope_idx]->op, GGML_OP_VIEW, GGML_OP_SET_ROWS }; + bool can_pull = ggml_can_fuse_subgraph_ext(graph, node_idxs, 3, ops, &set_rows_idx, 1); + + for (int c = rope_idx + 1; can_pull && c < set_rows_idx; ++c) { + if (!used[c] && c != view_idx && !is_empty(graph->nodes[c]) && + is_src_of(graph->nodes[set_rows_idx], graph->nodes[c])) { + can_pull = false; + } + } + + if (can_pull) { + current_set.push_back(view_idx); + current_set.push_back(set_rows_idx); + used[view_idx] = true; + used[set_rows_idx] = true; + } } } // Look for MUL_MAT_ID + ADD_ID + MUL @@ -18706,6 +19182,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: break; default: return false; @@ -18812,6 +19289,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: case GGML_TYPE_I32: return true; default: @@ -19034,6 +19512,31 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm } return false; } + case GGML_OP_DSV4_HC_COMB: + case GGML_OP_DSV4_HC_PRE: + case GGML_OP_DSV4_HC_POST: + { + if (op->type != GGML_TYPE_F32) { + return false; + } + for (uint32_t i = 0; i < GGML_MAX_SRC; ++i) { + if (op->src[i] && op->src[i]->type != GGML_TYPE_F32) { + return false; + } + } + // hc is hardcoded to 4 in the shaders. ggml only constrains it + // to 4 for COMB, so PRE/POST have to be checked here. + if (op->op == GGML_OP_DSV4_HC_PRE && op->src[0]->ne[1] != 4) { + return false; + } + if (op->op == GGML_OP_DSV4_HC_POST && op->src[1]->ne[1] != 4) { + return false; + } + if (op->op == GGML_OP_DSV4_HC_COMB) { + return device->pipeline_dsv4_hc_comb_f32 != nullptr; + } + return true; + } case GGML_OP_SOLVE_TRI: { if (op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32) { @@ -20022,6 +20525,13 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * tensor_clone = ggml_sum_rows(ggml_ctx, src_clone[0]); } else if (tensor->op == GGML_OP_CUMSUM) { tensor_clone = ggml_cumsum(ggml_ctx, src_clone[0]); + } else if (tensor->op == GGML_OP_DSV4_HC_COMB) { + tensor_clone = ggml_dsv4_hc_comb(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], + ggml_get_op_params_f32(tensor, 0), ggml_get_op_params_i32(tensor, 1)); + } else if (tensor->op == GGML_OP_DSV4_HC_PRE) { + tensor_clone = ggml_dsv4_hc_pre(ggml_ctx, src_clone[0], src_clone[1]); + } else if (tensor->op == GGML_OP_DSV4_HC_POST) { + tensor_clone = ggml_dsv4_hc_post(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]); } else if (tensor->op == GGML_OP_MEAN) { tensor_clone = ggml_mean(ggml_ctx, src_clone[0]); } else if (tensor->op == GGML_OP_ARGMAX) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl index 627932bd3..9df66cb44 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl @@ -608,6 +608,21 @@ vec2 get_dm(uint ib, uint a_offset) { } #endif +#if defined(DATA_A_TQ1_0) +float tq1_0_val(uint ib, uint e, uint a_offset) { + const uint bidx = tq1_0_byte_of(e); + const uint qbyte = uint(bidx < 48u ? data_a[a_offset + ib].qs[bidx] + : data_a[a_offset + ib].qh[bidx - 48u]); + return float(tq1_0_trit(qbyte, tq1_0_digit_of(e))) - 1.0; +} +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + return vec2(tq1_0_val(ib, iqs, a_offset), tq1_0_val(ib, iqs + 1u, a_offset)); +} +vec2 get_dm(uint ib, uint a_offset) { + return vec2(float(data_a[a_offset + ib].d), 0); +} +#endif + #if defined(DATA_A_TQ2_0) vec2 dequantize(uint ib, uint iqs, uint a_offset) { // elem e -> byte qs[(e/128)*32 + e%32], bits 2*((e%128)/32); w = q - 1 (d applied via get_dm) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl index 46cc69cb2..ef53264a7 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl @@ -247,6 +247,19 @@ f16vec4 dequantFuncQ8_0_v(const in decodeBufQ8_0 bl, const in uint blockCoords[2 return f16vec4(vec4(qi) * vec4(float(d))); } +layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufTQ1_0 { + block_tq1_0 block; +}; + +float16_t dequantFuncTQ1_0(const in decodeBufTQ1_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint e = coordInBlock[1]; + const uint bidx = tq1_0_byte_of(e); + const uint qbyte = uint(bidx < 48u ? bl.block.qs[bidx] : bl.block.qh[bidx - 48u]); + const uint xi = tq1_0_trit(qbyte, tq1_0_digit_of(e)); + return bl.block.d * (float16_t(int(xi)) - float16_t(1.0)); +} + layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufTQ2_0 { block_tq2_0 block; }; @@ -1406,6 +1419,8 @@ f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords #elif defined(DATA_A_Q8_0) #define dequantFuncA dequantFuncQ8_0 #define dequantFuncA_v dequantFuncQ8_0_v +#elif defined(DATA_A_TQ1_0) +#define dequantFuncA dequantFuncTQ1_0 #elif defined(DATA_A_TQ2_0) #define dequantFuncA dequantFuncTQ2_0 #define dequantFuncA_v dequantFuncTQ2_0_v diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_tq1_0.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_tq1_0.comp new file mode 100644 index 000000000..1632e7463 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_tq1_0.comp @@ -0,0 +1,28 @@ +#version 450 + +#include "dequant_head.glsl" + +layout (local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_tq1_0 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_GlobalInvocationID.x * 4; + + if (i >= p.nel) { + return; + } + + const uint ib = i / QUANT_K_TQ1_0; + const float d = float(data_a[ib].d); + + [[unroll]] for (uint j = 0; j < 4 && (i + j) < p.nel; ++j) { + const uint e = (i + j) % QUANT_K_TQ1_0; + const uint bidx = tq1_0_byte_of(e); + const uint qbyte = uint(bidx < 48u ? data_a[ib].qs[bidx] + : data_a[ib].qh[bidx - 48u]); + const uint xi = tq1_0_trit(qbyte, tq1_0_digit_of(e)); + data_b[i + j] = D_TYPE(d * (float(xi) - 1.0f)); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp new file mode 100644 index 000000000..f4ac0378a --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp @@ -0,0 +1,90 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require + +// 16 lanes per token, indexed idst + hc*isrc: idst in bits 0..1, isrc in bits 2..3, +// so subgroupShuffleXor by 1|2 reduces a row and by 4|8 a column. + +layout(constant_id = 0) const uint SUBGROUP_SIZE = 32; + +layout(local_size_x_id = 0, local_size_y = 4, local_size_z = 1) in; + +layout(push_constant) uniform parameter +{ + uint n_tokens; + + uint nbm0; uint nbm1; // mixes + uint nbs0; // scale + uint nbb0; // base + uint nbd0; uint nbd1; uint nbd2; // dst + + uint m_offset; + uint s_offset; + uint b_offset; + uint d_offset; + + float eps; + uint n_iter; +}; + +layout(binding = 0, std430) readonly buffer M { float data_m[]; }; +layout(binding = 1, std430) readonly buffer S { float data_s[]; }; +layout(binding = 2, std430) readonly buffer B { float data_b[]; }; +layout(binding = 3, std430) writeonly buffer D { float data_d[]; }; + +const uint hc = 4; +const uint comb_offset = 2 * hc; + +const uint TOKENS_PER_SUBGROUP = SUBGROUP_SIZE / 16; + +void main() { + const uint lane = gl_SubgroupInvocationID; + const uint blk = lane >> 4; // which 16-lane block, i.e. which token + const uint idx = lane & 15; // idst + hc*isrc + + const uint sg = gl_WorkGroupID.x * gl_WorkGroupSize.y + gl_SubgroupID; + const uint it = sg * TOKENS_PER_SUBGROUP + blk; + + // no early return, the shuffles need every lane; out-of-range blocks compute a discarded value + const bool in_range = it < n_tokens; + + const float scale_comb = data_s[s_offset + 2 * nbs0]; + + float v = 0.0f; + if (in_range) { + v = data_m[m_offset + (comb_offset + idx) * nbm0 + it * nbm1] * scale_comb + + data_b[b_offset + (comb_offset + idx) * nbb0]; + } + + // Softmax across destinations: the four lanes sharing an isrc. + float vmax = max(v, subgroupShuffleXor(v, 1)); + vmax = max(vmax, subgroupShuffleXor(vmax, 2)); + v = exp(v - vmax); + + float sum = v + subgroupShuffleXor(v, 1); + sum += subgroupShuffleXor(sum, 2); + v = v / sum + eps; + + // Normalize columns: equal destination indices are four lanes apart. + sum = v + subgroupShuffleXor(v, 4); + sum += subgroupShuffleXor(sum, 8); + v /= sum + eps; + + for (uint i = 1; i < n_iter; ++i) { + sum = v + subgroupShuffleXor(v, 1); + sum += subgroupShuffleXor(sum, 2); + v /= sum + eps; + + sum = v + subgroupShuffleXor(v, 4); + sum += subgroupShuffleXor(sum, 8); + v /= sum + eps; + } + + if (in_range) { + const uint idst = idx & 3; + const uint isrc = idx >> 2; + data_d[d_offset + idst * nbd0 + isrc * nbd1 + it * nbd2] = v; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp new file mode 100644 index 000000000..bab6f8767 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp @@ -0,0 +1,83 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +// Fan one stream back out to hc streams and add the combination-weighted +// residuals: +// +// dst[i0, idst, it] = x[i0, it]*post[idst, it] +// + sum_isrc residual[i0, isrc, it]*comb[idst, isrc, it] + +layout(constant_id = 0) const uint BLOCK_SIZE = 256; + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform parameter +{ + uint n_embd; + uint n_tokens; + + uint nbx0; uint nbx1; // x + uint nbr0; uint nbr1; uint nbr2; // residual + uint nbp0; uint nbp1; // post + uint nbc0; uint nbc1; uint nbc2; // comb + uint nbd0; uint nbd1; uint nbd2; // dst + + uint x_offset; + uint r_offset; + uint p_offset; + uint c_offset; + uint d_offset; +}; + +layout(binding = 0, std430) readonly buffer X { float data_x[]; }; +layout(binding = 1, std430) readonly buffer R { float data_r[]; }; +layout(binding = 2, std430) readonly buffer P { float data_p[]; }; +layout(binding = 3, std430) readonly buffer C { float data_c[]; }; +layout(binding = 4, std430) writeonly buffer D { float data_d[]; }; + +const uint hc = 4; + +shared float post_s[hc]; +shared float comb_s[hc * hc]; + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint it = gl_WorkGroupID.y; + + if (tid < hc) { + post_s[tid] = data_p[p_offset + tid * nbp0 + it * nbp1]; + } + if (tid < hc * hc) { + const uint idst = tid & 3; + const uint isrc = tid >> 2; + comb_s[tid] = data_c[c_offset + idst * nbc0 + isrc * nbc1 + it * nbc2]; + } + barrier(); + + // After the barrier, so every invocation reaches it. + const uint i0 = gl_WorkGroupID.x * BLOCK_SIZE + tid; + if (i0 >= n_embd) { + return; + } + + const float xv = data_x[x_offset + i0 * nbx0 + it * nbx1]; + + const uint rb = r_offset + i0 * nbr0 + it * nbr2; + + float r[hc]; + [[unroll]] + for (uint isrc = 0; isrc < hc; ++isrc) { + r[isrc] = data_r[rb + isrc * nbr1]; + } + + [[unroll]] + for (uint idst = 0; idst < hc; ++idst) { + float result = xv * post_s[idst]; + [[unroll]] + for (uint isrc = 0; isrc < hc; ++isrc) { + result = fma(r[isrc], comb_s[idst + hc * isrc], result); + } + data_d[d_offset + i0 * nbd0 + idst * nbd1 + it * nbd2] = result; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp new file mode 100644 index 000000000..51deabbac --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp @@ -0,0 +1,59 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +// Collapse the hc residual streams of a token into one, weighted per stream: +// +// dst[i0, it] = sum_ih x[i0, ih, it] * weights[ih, it] + +layout(constant_id = 0) const uint BLOCK_SIZE = 256; + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform parameter +{ + uint n_embd; + uint n_tokens; + + uint nbx0; uint nbx1; uint nbx2; // x + uint nbw0; uint nbw1; // weights + uint nbd0; uint nbd1; // dst + + uint x_offset; + uint w_offset; + uint d_offset; +}; + +layout(binding = 0, std430) readonly buffer X { float data_x[]; }; +layout(binding = 1, std430) readonly buffer W { float data_w[]; }; +layout(binding = 2, std430) writeonly buffer D { float data_d[]; }; + +const uint hc = 4; + +shared float w[hc]; + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint it = gl_WorkGroupID.y; + + if (tid < hc) { + w[tid] = data_w[w_offset + tid * nbw0 + it * nbw1]; + } + barrier(); + + // After the barrier, so every invocation reaches it. + const uint i0 = gl_WorkGroupID.x * BLOCK_SIZE + tid; + if (i0 >= n_embd) { + return; + } + + const uint xb = x_offset + i0 * nbx0 + it * nbx2; + + float result = 0.0f; + [[unroll]] + for (uint ih = 0; ih < hc; ++ih) { + result = fma(data_x[xb + ih * nbx1], w[ih], result); + } + + data_d[d_offset + i0 * nbd0 + it * nbd1] = result; +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp b/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp index 9dba437ed..19af30ac9 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_quant.comp @@ -27,10 +27,10 @@ void main() { const uint i11 = gid_z / p.ne12; const uint i12 = gid_z % p.ne12; - const uint i01 = data_b[i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; + const uint i01 = data_b[get_boffset() + i10*p.nb10 + i11*p.nb11 + i12*p.nb12]; - const uint a_offset = i01*p.nb01 + i11*p.nb02 + i12*p.nb03; - const uint d_offset = i10*p.nb21 + i11*p.nb22 + i12*p.nb23; + const uint a_offset = get_aoffset() + i01*p.nb01 + i11*p.nb02 + i12*p.nb03; + const uint d_offset = get_doffset() + i10*p.nb21 + i11*p.nb22 + i12*p.nb23; const uint ib = a_offset + i00/QUANT_K; // block index const uint iqs = (i00%QUANT_K)/QUANT_R; // quant index diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_tq1_0.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_tq1_0.comp new file mode 100644 index 000000000..2c99a268e --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_tq1_0.comp @@ -0,0 +1,85 @@ +#version 450 +#extension GL_EXT_shader_explicit_arithmetic_types : require + +#include "mul_mat_vec_base.glsl" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + +// Walks the packed bytes directly (byte m, digit t) rather than via +// tq1_0_byte_of()/tq1_0_digit_of(): one byte per thread, expanded in place. +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + const uint tid = gl_LocalInvocationID.x; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint i = 0; i < NUM_ROWS; ++i) { + temp[j][i] = FLOAT_TYPE(0); + } + } + + for (uint nrow = 0; nrow < num_rows; ++nrow) { + const uint ib0 = a_offset + (first_row + nrow) * num_blocks_per_row; + for (uint jcol = 0; jcol < NUM_COLS; ++jcol) { + const uint b_base = (jcol * p.batch_stride_b); + for (uint i = tid/8; i < num_blocks_per_row; i += gl_WorkGroupSize.x/8) { + const FLOAT_TYPE d = float(data_a[ib0 + i].d); + + // First qs chunk: 32 bytes (5*32 elements) + [[unroll]] for (uint m = tid%8; m < 32; m += 8) { + const uint q_byte = uint(data_a[ib0 + i].qs[m]); + [[unroll]] for (uint t = 0; t < 5; ++t) { + const uint xi = tq1_0_trit(q_byte, t); + const FLOAT_TYPE dequant_val = FLOAT_TYPE(d * (float(xi) - 1.0f)); + const uint elem = t * 32u + m; + const uint b_idx = i * QUANT_K + elem; + temp[jcol][nrow] += dequant_val * FLOAT_TYPE(data_b[b_base + b_offset + b_idx]); + } + } + + // Second qs chunk: 16 bytes (5*16 elements) + [[unroll]] for (uint m = tid%8; m < 16; m += 8) { + const uint q_byte = uint(data_a[ib0 + i].qs[32u + m]); + [[unroll]] for (uint t = 0; t < 5; ++t) { + const uint xi = tq1_0_trit(q_byte, t); + const FLOAT_TYPE dequant_val = FLOAT_TYPE(d * (float(xi) - 1.0f)); + const uint elem = 160u + t * 16u + m; + const uint b_idx = i * QUANT_K + elem; + temp[jcol][nrow] += dequant_val * FLOAT_TYPE(data_b[b_base + b_offset + b_idx]); + } + } + + // qh bytes: 4 bytes (4*4 elements) + [[unroll]] for (uint j = tid%8; j < 4; j += 8) { + const uint qh_byte = uint(data_a[ib0 + i].qh[j]); + [[unroll]] for (uint t = 0; t < 4; ++t) { + const uint xi = tq1_0_trit(qh_byte, t); + const FLOAT_TYPE dequant_val = FLOAT_TYPE(d * (float(xi) - 1.0f)); + const uint elem = 240u + t * 4u + j; + const uint b_idx = i * QUANT_K + elem; + temp[jcol][nrow] += dequant_val * FLOAT_TYPE(data_b[b_base + b_offset + b_idx]); + } + } + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, p.stride_d - first_row); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl index 7d852dced..bdc70af14 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl @@ -197,6 +197,24 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint k_pair = row * LOAD_VEC_A / 2; store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); +#elif defined(DATA_A_TQ1_0) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = (idx % 128) * 2; // element 0,2,4..254 + + const float d = float(data_a[ib].d); + vec2 v; + for (uint kk = 0u; kk < 2u; ++kk) { + const uint e = iqs + kk; + const uint bidx = tq1_0_byte_of(e); + const uint qbyte = uint(bidx < 48u ? data_a[ib].qs[bidx] + : data_a[ib].qh[bidx - 48u]); + v[kk] = d * (float(tq1_0_trit(qbyte, tq1_0_digit_of(e))) - 1.0); + } + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); #elif defined(DATA_A_TQ2_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp index 55b89f19a..ee813842c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp @@ -27,12 +27,24 @@ layout (binding = 6) readonly buffer R_I {uvec2 rope_data_i[];}; // indices for #define GGML_ROPE_TYPE_MROPE 8 #define GGML_ROPE_TYPE_VISION 24 +#elif RMS_NORM_ADD_FUSION + +layout (binding = 3) readonly buffer C {float data_c[];}; +layout (binding = 4) readonly buffer E {float data_e[];}; + +#elif RMS_NORM_SET_ROWS_FUSION + +layout (binding = 3) readonly buffer I {uvec2 data_i[];}; + #endif #extension GL_EXT_control_flow_attributes : enable #define BLOCK_SIZE 512 layout (constant_id = 1) const bool do_multiply = false; +#if RMS_NORM_ADD_FUSION +layout (constant_id = 2) const bool do_post_multiply = false; +#endif layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; @@ -57,6 +69,8 @@ void rms_norm(uint num_iters) { #if RMS_NORM_ROPE_FUSION // Per-row offset in shared memory uint32_t d_offset = 0; +#elif RMS_NORM_SET_ROWS_FUSION + uint32_t d_offset = data_i[channel].x*p.nb21 + row*ncols + get_doffset(); #else uint32_t d_offset = ((samp*nchannels + channel)*nrows + row)*ncols + get_doffset(); #endif @@ -91,14 +105,28 @@ void rms_norm(uint num_iters) { if (col >= ncols) { continue; } - data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)])); + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)]); +#if RMS_NORM_ADD_FUSION + value += FLOAT_TYPE(data_c[d_offset + col]); + if (do_post_multiply) { + value *= FLOAT_TYPE(data_e[0]); + } +#endif + data_d[d_offset + col] = D_TYPE(value); } } else { [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { if (col >= ncols) { continue; } - data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col])); + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col]); +#if RMS_NORM_ADD_FUSION + value += FLOAT_TYPE(data_c[d_offset + col]); + if (do_post_multiply) { + value *= FLOAT_TYPE(data_e[0]); + } +#endif + data_d[d_offset + col] = D_TYPE(value); } } } else { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_partials.comp b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_partials.comp index 4618b2c7e..cf7ab21f2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_partials.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_partials.comp @@ -10,11 +10,19 @@ #define BLOCK_SIZE 128 layout (constant_id = 1) const bool do_multiply = false; +#if RMS_NORM_ADD_FUSION +layout (constant_id = 2) const bool do_post_multiply = false; +#endif layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; layout (binding = 3, std430) readonly buffer PartialsBuf {float partial_sums[];}; +#if RMS_NORM_ADD_FUSION +layout (binding = 4) readonly buffer C {float data_c[];}; +layout (binding = 5) readonly buffer E {float data_e[];}; +#endif + shared FLOAT_TYPE sumsh[BLOCK_SIZE]; void main() { @@ -55,9 +63,23 @@ void main() { if (do_multiply) { if (ncols > p.ne10) { - data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)])); + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)]); +#if RMS_NORM_ADD_FUSION + value += FLOAT_TYPE(data_c[d_offset + col]); + if (do_post_multiply) { + value *= FLOAT_TYPE(data_e[0]); + } +#endif + data_d[d_offset + col] = D_TYPE(value); } else { - data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col])); + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col]); +#if RMS_NORM_ADD_FUSION + value += FLOAT_TYPE(data_c[d_offset + col]); + if (do_post_multiply) { + value *= FLOAT_TYPE(data_e[0]); + } +#endif + data_d[d_offset + col] = D_TYPE(value); } } else { data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col])); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl index adb1bb8b3..a19c7f2f4 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl @@ -303,6 +303,41 @@ struct block_q2_K_packed32 #define DATA_A_QUANT_K #endif +#define QUANT_K_TQ1_0 256 + +// TQ1_0: base-3 packed trits, 5 per byte in `qs` (48B) and 4 in `qh` (4B). +struct block_tq1_0 +{ + uint8_t qs[(QUANT_K_TQ1_0 - 4 * QUANT_K_TQ1_0 / 64) / 5]; + uint8_t qh[QUANT_K_TQ1_0 / 64]; + float16_t d; +}; + +// Element e in [0,255] -> its packed byte (0..47 qs, 48..51 qh) and digit. +uint tq1_0_byte_of(uint e) { + return e < 160u ? (e % 32u) + : e < 240u ? 32u + ((e - 160u) % 16u) + : 48u + ((e - 240u) % 4u); +} +uint tq1_0_digit_of(uint e) { + return e < 160u ? (e / 32u) + : e < 240u ? ((e - 160u) / 16u) + : ((e - 240u) / 4u); +} +// The 8-bit truncation below is part of the format, not an optimisation: +// the C reference does `uint8_t q = qs[..] * pow3[n]`. +uint tq1_0_trit(uint qbyte, uint t) { + const uint POW3_PACKED = (1u << 28) | (3u << 21) | (9u << 14) | (27u << 7) | 81u; + return ((((qbyte * ((POW3_PACKED >> (7u * (4u - t))) & 0x7Fu)) & 255u) * 3u) >> 8); +} + +#if defined(DATA_A_TQ1_0) +#define QUANT_K QUANT_K_TQ1_0 +#define QUANT_R 1 +#define A_TYPE block_tq1_0 +#define DATA_A_QUANT_K +#endif + #define QUANT_K_TQ2_0 256 // ternary (BitNet): 2-bit codes, w = (q - 1) * d; qs layout matches q2_K's diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 86bbdb0e3..fdcc13943 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -86,6 +86,7 @@ const std::vector type_names = { "iq4_nl", "mxfp4", "nvfp4", + "tq1_0", "tq2_0", "bf16", }; @@ -760,7 +761,7 @@ void process_shaders() { for (const auto& tname : type_names) { // mul mat vec std::string data_a_key = "DATA_A_" + to_uppercase(tname); - std::string shader = (string_ends_with(tname, "_k") || string_starts_with(tname, "iq1_") || string_starts_with(tname, "iq2_") || string_starts_with(tname, "iq3_") || tname == "tq2_0") ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp"; + std::string shader = (string_ends_with(tname, "_k") || string_starts_with(tname, "iq1_") || string_starts_with(tname, "iq2_") || string_starts_with(tname, "iq3_") || tname == "tq2_0" || tname == "tq1_0") ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp"; string_to_spv("mul_mat_vec_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}})); string_to_spv("mul_mat_vec_" + tname + "_f16_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}})); @@ -831,6 +832,10 @@ void process_shaders() { string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("rms_norm_mul_add_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_ADD_FUSION", "1"}})); + string_to_spv("rms_norm_mul_add_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_ADD_FUSION", "1"}})); + string_to_spv("rms_norm_set_rows_f32_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"RMS_NORM_SET_ROWS_FUSION", "1"}})); + string_to_spv("rms_norm_set_rows_f32_f16", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float16_t"}, {"RMS_NORM_SET_ROWS_FUSION", "1"}})); string_to_spv("rms_norm_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("rms_norm_mul_rope_f32_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float"}, {"RMS_NORM_ROPE_FUSION", "1"}})); string_to_spv("rms_norm_mul_rope_f32_f16", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float16_t"}, {"RMS_NORM_ROPE_FUSION", "1"}})); @@ -1063,6 +1068,9 @@ void process_shaders() { string_to_spv("fwht_f32", "fwht.comp", {}); string_to_spv("fwht_shmem_f32", "fwht.comp", {{"FWHT_SHMEM", "1"}}); string_to_spv("count_equal_i32", "count_equal.comp", merge_maps(base_dict, {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}})); + string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {}); + string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {}); + string_to_spv("dsv4_hc_post_f32", "dsv4_hc_post.comp", {}); string_to_spv("cumsum_f32", "cumsum.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("cumsum_multipass1_f32", "cumsum_multipass1.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("cumsum_multipass2_f32", "cumsum_multipass2.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 399d31f1d..d3a639f37 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -215,6 +215,7 @@ class Keys: KV_LORA_RANK_SWA = "{arch}.attention.kv_lora_rank_swa" SHARED_KV_LAYERS = "{arch}.attention.shared_kv_layers" SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern" + RECURRENT_LAYERS = "{arch}.attention.recurrent_layers" TEMPERATURE_SCALE = "{arch}.attention.temperature_scale" ROPE_PATTERN = "{arch}.attention.rope_pattern" @@ -619,6 +620,7 @@ class MODEL_ARCH(IntEnum): PADDLEOCR = auto() MIMO2 = auto() STEP35 = auto() + SPARK2_5 = auto() LLAMA_EMBED = auto() MAINCODER = auto() KIMI_LINEAR = auto() @@ -1373,6 +1375,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.PADDLEOCR: "paddleocr", MODEL_ARCH.MIMO2: "mimo2", MODEL_ARCH.STEP35: "step35", + MODEL_ARCH.SPARK2_5: "spark2_5", MODEL_ARCH.LLAMA_EMBED: "llama-embed", MODEL_ARCH.MAINCODER: "maincoder", MODEL_ARCH.KIMI_LINEAR: "kimi-linear", @@ -2294,6 +2297,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2314,6 +2318,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2337,6 +2342,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2357,6 +2363,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2402,6 +2409,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2504,6 +2512,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.TOKEN_TYPES, MODEL_TENSOR.ATTN_NORM_2, MODEL_TENSOR.ATTN_OUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -2532,6 +2541,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2561,6 +2571,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2573,6 +2584,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2600,6 +2612,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2631,6 +2644,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2646,6 +2660,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2661,6 +2676,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2675,6 +2691,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2689,6 +2706,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -2709,6 +2727,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -2725,6 +2744,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -2780,6 +2800,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -2796,6 +2817,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -2936,6 +2958,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3069,6 +3092,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3084,6 +3108,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3102,6 +3127,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.ROPE_FACTORS_LONG, MODEL_TENSOR.ROPE_FACTORS_SHORT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3139,6 +3165,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3151,6 +3178,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_ARCH.GEMMA2: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3167,6 +3195,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -3185,6 +3214,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -3221,6 +3251,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -3276,6 +3307,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.DENSE_2_OUT, MODEL_TENSOR.DENSE_3_OUT, MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -3296,6 +3328,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3459,6 +3492,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3488,6 +3522,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3502,6 +3537,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3516,6 +3552,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3567,6 +3604,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_ARCH.OLMO: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3579,6 +3617,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3594,6 +3633,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_ARCH.SEED_OSS: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3610,6 +3650,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3660,6 +3701,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3681,6 +3723,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3743,6 +3786,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_A, MODEL_TENSOR.ATTN_Q_B, @@ -3865,6 +3909,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -3941,6 +3986,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4086,6 +4132,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4100,6 +4147,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4121,6 +4169,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.SSM_D, MODEL_TENSOR.SSM_NORM, MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4140,6 +4189,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.SSM_D, MODEL_TENSOR.SSM_NORM, MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4170,6 +4220,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4185,6 +4236,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4210,6 +4262,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4241,6 +4294,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4255,6 +4309,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4280,6 +4335,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.SSM_D, MODEL_TENSOR.SSM_NORM, MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4343,6 +4399,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4382,6 +4439,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4473,6 +4531,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4536,6 +4595,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4551,6 +4611,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4602,6 +4663,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4616,6 +4678,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4633,6 +4696,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.ATTN_NORM, # Attention components + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, # Query projection MODEL_TENSOR.ATTN_K, # Key projection MODEL_TENSOR.ATTN_V, # Value projection @@ -4665,6 +4729,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4685,6 +4750,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4701,6 +4767,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4791,6 +4858,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4807,6 +4875,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4830,6 +4899,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.ATTN_NORM, # operator_norm MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4850,6 +4920,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.ATTN_NORM, # operator_norm MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4865,6 +4936,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4884,6 +4956,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4901,6 +4974,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -4918,6 +4992,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -4956,6 +5031,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -5019,6 +5095,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -5036,6 +5113,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -5051,6 +5129,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -5204,6 +5283,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -5231,12 +5311,26 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.SPARK2_5: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + ], MODEL_ARCH.LLAMA_EMBED: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, @@ -5256,6 +5350,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_Q_NORM, MODEL_TENSOR.ATTN_K, @@ -5272,6 +5367,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_QKV, MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 50e4d7c53..ed5a185b3 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -841,6 +841,9 @@ class GGUFWriter: else: self.add_array(key, value) + def add_recurrent_layers(self, value: Sequence[bool]) -> None: + self.add_array(Keys.Attention.RECURRENT_LAYERS.format(arch=self.arch), value) + def add_rope_pattern(self, value: Sequence[bool]) -> None: self.add_array(Keys.Attention.ROPE_PATTERN.format(arch=self.arch), value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index d644d502e..d2dfeece5 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -385,6 +385,7 @@ class TensorNameMap: MODEL_TENSOR.ATTN_SINKS: ( "model.layers.{bid}.self_attn.sinks", # openai-moe "model.layers.{bid}.self_attn.attention_sink_bias", # mimov2 + "model.layers.{bid}.self_attn.learnable_sink_param", # hy-v4 ), MODEL_TENSOR.ATTN_GATE: ( @@ -392,6 +393,7 @@ class TensorNameMap: "model.layers.{bid}.linear_attn.in_proj_z", # qwen3.5 "model.layers.{bid}.self_attn.g_proj", # step3.5 head-wise attention gate "model.layers.{bid}.self_attn.output_gate", # minimax-01 + "model.layers.{bid}.self_attn.linear_gate", # hy-v4 ), # Feed-forward norm @@ -1329,6 +1331,42 @@ class TensorNameMap: "model.layers.{bid}.self_attn.index_q_norm", # MSA ), + MODEL_TENSOR.HC_ATTN_FN: ( + "model.layers.{bid}.hc_attn_layer.hc_pre.hc_fn", # hy-v4 + ), + + MODEL_TENSOR.HC_ATTN_BASE: ( + "model.layers.{bid}.hc_attn_layer.hc_pre.hc_base", # hy-v4 + ), + + MODEL_TENSOR.HC_ATTN_SCALE: ( + "model.layers.{bid}.hc_attn_layer.hc_pre.hc_scale", # hy-v4 + ), + + MODEL_TENSOR.HC_FFN_FN: ( + "model.layers.{bid}.hc_mlp_layer.hc_pre.hc_fn", # hy-v4 + ), + + MODEL_TENSOR.HC_FFN_BASE: ( + "model.layers.{bid}.hc_mlp_layer.hc_pre.hc_base", # hy-v4 + ), + + MODEL_TENSOR.HC_FFN_SCALE: ( + "model.layers.{bid}.hc_mlp_layer.hc_pre.hc_scale", # hy-v4 + ), + + MODEL_TENSOR.HC_HEAD_FN: ( + "model.hc_head.hc_head_fn", # hy-v4 + ), + + MODEL_TENSOR.HC_HEAD_BASE: ( + "model.hc_head.hc_head_base", # hy-v4 + ), + + MODEL_TENSOR.HC_HEAD_SCALE: ( + "model.hc_head.hc_head_scale", # hy-v4 + ), + ############################################################################ # TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg MODEL_TENSOR.ENC_OUTPUT_NORM: ( diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index d06be641a..15f651919 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -146,6 +146,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_PADDLEOCR, "paddleocr" }, { LLM_ARCH_MIMO2, "mimo2" }, { LLM_ARCH_STEP35, "step35" }, + { LLM_ARCH_SPARK2_5, "spark2_5" }, { LLM_ARCH_LLAMA_EMBED, "llama-embed" }, { LLM_ARCH_MAINCODER, "maincoder" }, { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 62dfa5d81..f1d173a57 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -147,6 +147,7 @@ enum llm_arch { LLM_ARCH_PADDLEOCR, LLM_ARCH_MIMO2, LLM_ARCH_STEP35, + LLM_ARCH_SPARK2_5, LLM_ARCH_LLAMA_EMBED, LLM_ARCH_MAINCODER, LLM_ARCH_KIMI_LINEAR, diff --git a/src/llama-grammar.cpp b/src/llama-grammar.cpp index bb683e749..afac045d8 100644 --- a/src/llama-grammar.cpp +++ b/src/llama-grammar.cpp @@ -517,7 +517,7 @@ const char * llama_grammar_parser::parse_sequence( total_rules = min_times; } - if (n_prev_rules * total_rules >= MAX_REPETITION_THRESHOLD) { + if (n_prev_rules * total_rules > MAX_REPETITION_THRESHOLD) { throw std::runtime_error("number of rules that are going to be repeated multiplied by the new repetition exceeds sane defaults, please reduce the number of repetitions or rule complexity"); } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 87b791f3f..12033dcc0 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1624,8 +1624,26 @@ llm_graph_qkv llm_graph_context::build_qkv( int64_t n_head, int64_t n_head_kv, int il) const { - const int64_t n_embd_q = n_embd_head * n_head; - const int64_t n_embd_kv = n_embd_head * n_head_kv; + return build_qkv(layer, cur, + n_embd_head, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il); +} + +llm_graph_qkv llm_graph_context::build_qkv( + const llama_layer & layer, + ggml_tensor * cur, + int64_t n_embd_head_q, + int64_t n_head_q, + int64_t n_embd_head_k, + int64_t n_head_k, + int64_t n_embd_head_v, + int64_t n_head_v, + int il, + bool reshape) const { + const int64_t n_embd_q = n_embd_head_q * n_head_q; + const int64_t n_embd_k = n_embd_head_k * n_head_k; ggml_tensor * Qcur, * Kcur, * Vcur; @@ -1636,59 +1654,93 @@ llm_graph_qkv llm_graph_context::build_qkv( if (layer.wqkv_b) { qkv = ggml_add(ctx0, qkv, layer.wqkv_b); cb(qkv, "wqkv_b", il); + } else if (layer.wq_b && layer.wk_b && layer.wv_b) { + // Fused weights may coexist with separate Q/K/V biases in legacy or custom GGUFs. + ggml_tensor * qkv_b = ggml_concat(ctx0, ggml_concat(ctx0, layer.wq_b, layer.wk_b, 0), layer.wv_b, 0); + qkv = ggml_add(ctx0, qkv, qkv_b); + cb(qkv, "wqkv_b", il); } - if (hparams.f_clamp_kqv > 0.0f) { + if (reshape && hparams.f_clamp_kqv > 0.0f) { qkv = ggml_clamp(ctx0, qkv, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(qkv, "wqkv_clamped", il); } - Qcur = ggml_view_3d(ctx0, qkv, n_embd_head, n_head, n_tokens, - ggml_row_size(qkv->type, n_embd_head), qkv->nb[1], 0); - Kcur = ggml_view_3d(ctx0, qkv, n_embd_head, n_head_kv, n_tokens, - ggml_row_size(qkv->type, n_embd_head), qkv->nb[1], - ggml_row_size(qkv->type, n_embd_q)); - Vcur = ggml_view_3d(ctx0, qkv, n_embd_head, n_head_kv, n_tokens, - ggml_row_size(qkv->type, n_embd_head), qkv->nb[1], - ggml_row_size(qkv->type, n_embd_q + n_embd_kv)); + if (reshape) { + Qcur = ggml_view_3d(ctx0, qkv, n_embd_head_q, n_head_q, n_tokens, + ggml_row_size(qkv->type, n_embd_head_q), qkv->nb[1], 0); + Kcur = ggml_view_3d(ctx0, qkv, n_embd_head_k, n_head_k, n_tokens, + ggml_row_size(qkv->type, n_embd_head_k), qkv->nb[1], + ggml_row_size(qkv->type, n_embd_q)); + Vcur = ggml_view_3d(ctx0, qkv, n_embd_head_v, n_head_v, n_tokens, + ggml_row_size(qkv->type, n_embd_head_v), qkv->nb[1], + ggml_row_size(qkv->type, n_embd_q + n_embd_k)); + } else { + Qcur = ggml_view_2d(ctx0, qkv, n_embd_q, n_tokens, qkv->nb[1], 0); + Kcur = ggml_view_2d(ctx0, qkv, n_embd_k, n_tokens, qkv->nb[1], + ggml_row_size(qkv->type, n_embd_q)); + Vcur = ggml_view_2d(ctx0, qkv, n_embd_head_v * n_head_v, n_tokens, qkv->nb[1], + ggml_row_size(qkv->type, n_embd_q + n_embd_k)); + } + if (!reshape) { + Qcur = ggml_cont(ctx0, Qcur); + Kcur = ggml_cont(ctx0, Kcur); + Vcur = ggml_cont(ctx0, Vcur); + } } else { // separate Q/K/V path Qcur = build_lora_mm(layer.wq, cur, layer.wq_s); - cb(Qcur, "Qcur", il); - if (layer.wq_b) { - Qcur = ggml_add(ctx0, Qcur, layer.wq_b); + if (reshape) { cb(Qcur, "Qcur", il); } - if (hparams.f_clamp_kqv > 0.0f) { + if (layer.wq_b) { + Qcur = ggml_add(ctx0, Qcur, layer.wq_b); + if (reshape) { + cb(Qcur, "Qcur", il); + } + } + if (reshape && hparams.f_clamp_kqv > 0.0f) { Qcur = ggml_clamp(ctx0, Qcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(Qcur, "Qcur_clamped", il); } Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); - cb(Kcur, "Kcur", il); - if (layer.wk_b) { - Kcur = ggml_add(ctx0, Kcur, layer.wk_b); + if (reshape) { cb(Kcur, "Kcur", il); } - if (hparams.f_clamp_kqv > 0.0f) { + if (layer.wk_b) { + Kcur = ggml_add(ctx0, Kcur, layer.wk_b); + if (reshape) { + cb(Kcur, "Kcur", il); + } + } + if (reshape && hparams.f_clamp_kqv > 0.0f) { Kcur = ggml_clamp(ctx0, Kcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(Kcur, "Kcur_clamped", il); } Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); - cb(Vcur, "Vcur", il); - if (layer.wv_b) { - Vcur = ggml_add(ctx0, Vcur, layer.wv_b); + if (reshape) { cb(Vcur, "Vcur", il); } - if (hparams.f_clamp_kqv > 0.0f) { + if (layer.wv_b) { + Vcur = ggml_add(ctx0, Vcur, layer.wv_b); + if (reshape) { + cb(Vcur, "Vcur", il); + } + } + if (reshape && hparams.f_clamp_kqv > 0.0f) { Vcur = ggml_clamp(ctx0, Vcur, -hparams.f_clamp_kqv, hparams.f_clamp_kqv); cb(Vcur, "Vcur_clamped", il); } - Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); - Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); - Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + if (reshape) { + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_q, n_head_q, n_tokens); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head_k, n_head_k, n_tokens); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head_v, n_head_v, n_tokens); + } } - cb(Qcur, "Qcur", il); - cb(Kcur, "Kcur", il); - cb(Vcur, "Vcur", il); + if (reshape) { + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + } return { Qcur, Kcur, Vcur }; } diff --git a/src/llama-graph.h b/src/llama-graph.h index dddfdac7b..b486578c1 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1079,6 +1079,19 @@ struct llm_graph_context { int64_t n_head_kv, int il) const; + // Set reshape to false to return contiguous projections before clamp/reshape. + llm_graph_qkv build_qkv( + const llama_layer & layer, + ggml_tensor * cur, + int64_t n_embd_head_q, + int64_t n_head_q, + int64_t n_embd_head_k, + int64_t n_head_k, + int64_t n_embd_head_v, + int64_t n_head_v, + int il, + bool reshape = true) const; + ggml_tensor * build_ffn( ggml_tensor * cur, ggml_tensor * up, diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index df2a46d93..66f8bdec3 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -27,6 +27,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) { case LLM_ARCH_APERTUS: case LLM_ARCH_MIMO2: case LLM_ARCH_STEP35: + case LLM_ARCH_SPARK2_5: case LLM_ARCH_MUSE_GLIMMER: case LLM_ARCH_MELLUM: case LLM_ARCH_LAGUNA: diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 7961e6689..e56e7c309 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -182,6 +182,7 @@ #include "models/seed-oss.cpp" #include "models/smallthinker.cpp" #include "models/smollm3.cpp" +#include "models/spark2-5.cpp" #include "models/stablelm.cpp" #include "models/starcoder.cpp" #include "models/starcoder2.cpp" @@ -491,6 +492,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_kimi_k3(params); case LLM_ARCH_STEP35: return new llama_model_step35(params); + case LLM_ARCH_SPARK2_5: + return new llama_model_spark2_5(params); default: throw std::runtime_error(std::string("unsupported model architecture: '") + llm_arch_name(arch) + "'"); } @@ -3152,6 +3155,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_MIMO2: case LLM_ARCH_STEP35: + case LLM_ARCH_SPARK2_5: case LLM_ARCH_TALKIE: case LLM_ARCH_MELLUM: return LLAMA_ROPE_TYPE_NEOX; @@ -3386,6 +3390,12 @@ void llama_model_base::create_tensor_qkv(llama_layer & layer, int bid, layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", bid), {n_embd_, n_embd_qkv}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); if (layer.wqkv) { layer.wqkv_b = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias", bid), {n_embd_qkv}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); + // Fused weights may coexist with separate Q/K/V biases in legacy or custom GGUFs. + if (!layer.wqkv_b) { + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q, "bias", bid), {n_embd_q_}, TENSOR_NOT_REQUIRED); + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K, "bias", bid), {n_embd_k_}, TENSOR_NOT_REQUIRED); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V, "bias", bid), {n_embd_v_}, TENSOR_NOT_REQUIRED); + } } else { layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", bid), {n_embd_, n_embd_q_}, flags); layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", bid), {n_embd_, n_embd_k_}, flags); diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 79a92a9a4..2157cb7b9 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -550,6 +550,14 @@ struct llm_tokenizer_bpe : llm_tokenizer { "[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\r\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+[\r\n]*|\\s*[\r\n]+|\\s+(?!\\S)|\\s+", }; break; + case LLAMA_VOCAB_PRE_TYPE_SPARK2_5: + regex_exprs = { + "\\p{N}{1,3}", + "[一-龥぀-ゟ゠-ヿ]+", + "[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\r\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+|[\r\n]|\\s+(?!\\S)|\\s+", + "\\p{N}", + }; + break; case LLAMA_VOCAB_PRE_TYPE_YOUTU: regex_exprs = { "[가-힣ㄱ-ㆎ]+|[!…“”‘’—:;,、-〿︰-﹏]+|[ㄅ-ㄯ]+|[一-龥぀-ゟ゠-ヿ]+", @@ -2406,6 +2414,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "deepseek-v3") { pre_type = LLAMA_VOCAB_PRE_TYPE_DEEPSEEK3_LLM; clean_spaces = false; + } else if ( + tokenizer_pre == "spark2_5") { + pre_type = LLAMA_VOCAB_PRE_TYPE_SPARK2_5; + clean_spaces = false; } else if ( tokenizer_pre == "youtu") { pre_type = LLAMA_VOCAB_PRE_TYPE_YOUTU; diff --git a/src/llama-vocab.h b/src/llama-vocab.h index 9a10ecd01..c2c822594 100644 --- a/src/llama-vocab.h +++ b/src/llama-vocab.h @@ -67,6 +67,7 @@ enum llama_vocab_pre_type { LLAMA_VOCAB_PRE_TYPE_MELLUM2 = 55, LLAMA_VOCAB_PRE_TYPE_LAGUNA = 56, LLAMA_VOCAB_PRE_TYPE_HY_V4 = 57, + LLAMA_VOCAB_PRE_TYPE_SPARK2_5 = 58, }; struct LLM_KV; diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp index 1f2592cfa..e208c7d5a 100644 --- a/src/models/bailingmoe3.cpp +++ b/src/models/bailingmoe3.cpp @@ -280,8 +280,8 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); beta = ggml_sigmoid(ctx0, ggml_reshape_4d(ctx0, beta, 1, n_head, n_seq_tokens, n_seqs)); - q = ggml_l2_norm(ctx0, q, hparams.f_norm_rms_eps); - k = ggml_l2_norm(ctx0, k, hparams.f_norm_rms_eps); + q = build_gdn_l2_norm(ctx0, q, hparams.f_norm_rms_eps); + k = build_gdn_l2_norm(ctx0, k, hparams.f_norm_rms_eps); ggml_tensor * states_all = mctx_cur->get_s_l(il); ggml_tensor * state = build_rs(inp_rs, states_all, hparams.n_embd_s(), n_seqs); diff --git a/src/models/deepseek2.cpp b/src/models/deepseek2.cpp index 4628ff4da..deca86527 100644 --- a/src/models/deepseek2.cpp +++ b/src/models/deepseek2.cpp @@ -475,21 +475,12 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p const int ocr_rope_type = GGML_ROPE_TYPE_NEOX; GGML_ASSERT(n_embed_head == n_embd_head_k && n_embed_head == n_embd_head_v); - ggml_tensor * Qcur = NULL; - ggml_tensor * Kcur = NULL; - ggml_tensor * Vcur = NULL; - - Qcur = ggml_mul_mat(ctx0, model.layers[il].wq, cur); - Kcur = ggml_mul_mat(ctx0, model.layers[il].wk, cur); - Vcur = ggml_mul_mat(ctx0, model.layers[il].wv, cur); + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embed_head, n_head, n_head, il); cb(Qcur, "q", il); cb(Kcur, "k", il); cb(Vcur, "v", il); - Qcur = ggml_reshape_3d(ctx0, Qcur, n_embed_head, n_head, n_tokens); - Kcur = ggml_reshape_3d(ctx0, Kcur, n_embed_head, n_head, n_tokens); - Vcur = ggml_reshape_3d(ctx0, Vcur, n_embed_head, n_head, n_tokens); - GGML_ASSERT(fabs(freq_base - 10000.0) < 1e-4); Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_embed_head, ocr_rope_type, 0, freq_base, 1, 0, 1, 0, 0); Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_embed_head, ocr_rope_type, 0, freq_base, 1, 0, 1, 0, 0); diff --git a/src/models/deepseek2ocr.cpp b/src/models/deepseek2ocr.cpp index 1c5c452e9..3d630699e 100644 --- a/src/models/deepseek2ocr.cpp +++ b/src/models/deepseek2ocr.cpp @@ -40,9 +40,7 @@ void llama_model_deepseek2ocr::load_arch_tensors(llama_model_loader &) { for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; - layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd}, 0); - layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd}, 0); - layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd}, 0); + create_tensor_qkv(layer, i, n_embd, n_embd, n_embd, n_embd, 0); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0); // norm diff --git a/src/models/gemma3n.cpp b/src/models/gemma3n.cpp index 83eb8250a..ea616db3b 100644 --- a/src/models/gemma3n.cpp +++ b/src/models/gemma3n.cpp @@ -176,7 +176,14 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par hparams.f_attention_scale, il); } else { // reuse KV cache of earlier layers - ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur); + ggml_tensor * Qcur; + if (model.layers[il].wqkv) { + ggml_tensor * qkv = build_lora_mm(model.layers[il].wqkv, cur); + const int64_t q_dim = n_embd_head * n_head; + Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, q_dim, n_tokens, qkv->nb[1], 0)); + } else { + Qcur = build_lora_mm(model.layers[il].wq, cur); + } cb(Qcur, "Qcur", il); Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index d93990eb5..b8ae9623e 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -75,9 +75,13 @@ void llama_model_gemma4::load_arch_tensors(llama_model_loader &) { layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); // note: use_alternative_attention (v_proj is optional, if it's not present, use k_proj) - layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head * n_head}, 0); - layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k}, kv_flags); - layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v}, TENSOR_NOT_REQUIRED); + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), + {n_embd, n_embd_head * n_head + n_embd_k + n_embd_v}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); + if (!layer.wqkv) { + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head * n_head}, 0); + layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k}, kv_flags); + layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v}, TENSOR_NOT_REQUIRED); + } layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head * n_head, n_embd}, 0); layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head}, 0); @@ -202,9 +206,17 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para // Q projection (shared for both non-KV and KV layers) // this is to mirror Gemma4Attention in pytorch code + ggml_tensor * qkv_fused = nullptr; ggml_tensor * Qcur; - { + if (model.layers[il].wqkv) { + qkv_fused = build_lora_mm(model.layers[il].wqkv, cur, model.layers[il].wqkv_s); + cb(qkv_fused, "wqkv", il); + const int64_t q_dim = n_embd_head * n_head; + Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, q_dim, n_tokens, qkv_fused->nb[1], 0)); + } else { Qcur = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); + } + { cb(Qcur, "Qcur", il); Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens); @@ -219,12 +231,22 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para // self-attention if (hparams.has_kv(il)) { - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); + ggml_tensor * Kcur; + ggml_tensor * Vcur; + if (qkv_fused) { + const int64_t q_dim = n_embd_head * n_head; + const int64_t k_dim = n_embd_head * n_head_kv; + const int64_t v_dim = n_embd_head * n_head_kv; + const size_t esize = ggml_element_size(qkv_fused); + Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, k_dim, n_tokens, qkv_fused->nb[1], q_dim * esize)); + Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, v_dim, n_tokens, qkv_fused->nb[1], (q_dim + k_dim) * esize)); + } else { + Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); + Vcur = model.layers[il].wv + ? build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s) + : Kcur; // if v_proj is not present, use Kcur as Vcur + } cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = model.layers[il].wv - ? build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s) - : Kcur; // if v_proj is not present, use Kcur as Vcur cb(Vcur, "Vcur", il); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); diff --git a/src/models/jais2.cpp b/src/models/jais2.cpp index 8610fcc9f..64813b7b6 100644 --- a/src/models/jais2.cpp +++ b/src/models/jais2.cpp @@ -29,15 +29,9 @@ void llama_model_jais2::load_arch_tensors(llama_model_loader &) { layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0); - layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head}, 0); - layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, 0); - layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, 0); + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0); - // attention biases - all have shape n_embd (output dimension of projections) - layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q, "bias", i), {n_embd}, 0); - layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K, "bias", i), {n_embd}, 0); - layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V, "bias", i), {n_embd}, 0); layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, 0); layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); diff --git a/src/models/kimi-k3.cpp b/src/models/kimi-k3.cpp index b061093eb..b7604cbf2 100644 --- a/src/models/kimi-k3.cpp +++ b/src/models/kimi-k3.cpp @@ -441,9 +441,9 @@ ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer( ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs); state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head_kda, n_seqs); - const float eps = hparams.f_norm_rms_eps; - Qcur = ggml_l2_norm(ctx0, Qcur, eps); - Kcur = ggml_l2_norm(ctx0, Kcur, eps); + const float eps_norm = hparams.f_norm_rms_eps; + Qcur = build_gdn_l2_norm(ctx0, Qcur, eps_norm); + Kcur = build_gdn_l2_norm(ctx0, Kcur, eps_norm); auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il); diff --git a/src/models/kimi-linear.cpp b/src/models/kimi-linear.cpp index 601d1d9be..b9cf28d85 100644 --- a/src/models/kimi-linear.cpp +++ b/src/models/kimi-linear.cpp @@ -195,7 +195,7 @@ static ggml_tensor * causal_conv1d(ggml_cgraph * gf, ggml_context * ctx0, ggml_t // Causal Conv1d function for Q,K,V // When qkv is 0, it is Q, 1 is K, 2 is V // Step 1: Q, K, V projections -> [d_inner, n_tokens] - ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + ggml_tensor * x_proj = proj_w ? ggml_mul_mat(ctx0, proj_w, x) : x; // Reshape input: {d_inner, n_tokens} -> {d_inner, n_seq_tokens, n_seqs} ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); @@ -295,9 +295,20 @@ llama_model_kimi_linear::graph::graph(const llama_model & model, const llm_graph ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); cb(conv_states_all, "conv_states_all", il); ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); - ggml_tensor * Qcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); - ggml_tensor * Kcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); - ggml_tensor * Vcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * q_in = cur, * k_in = cur, * v_in = cur; + ggml_tensor * q_w = layer.wq, * k_w = layer.wk, * v_w = layer.wv; + if (layer.wqkv) { + ggml_tensor * qkv = ggml_mul_mat(ctx0, layer.wqkv, cur); + const int64_t d_inner = head_dim * n_head; + const size_t esize = ggml_element_size(qkv); + q_in = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, d_inner, n_tokens, qkv->nb[1], 0)); + k_in = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, d_inner, n_tokens, qkv->nb[1], d_inner * esize)); + v_in = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, d_inner, n_tokens, qkv->nb[1], 2 * d_inner * esize)); + q_w = nullptr; k_w = nullptr; v_w = nullptr; + } + ggml_tensor * Qcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, q_in, q_w, layer.ssm_q_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Kcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, k_in, k_w, layer.ssm_k_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Vcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, v_in, v_w, layer.ssm_v_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); // g1 = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias) ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); @@ -331,10 +342,11 @@ llama_model_kimi_linear::graph::graph(const llama_model & model, const llm_graph ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs); state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs); + const float eps_norm = hparams.f_norm_rms_eps; - Qcur = ggml_l2_norm(ctx0, Qcur, eps_norm); - Kcur = ggml_l2_norm(ctx0, Kcur, eps_norm); + Qcur = build_gdn_l2_norm(ctx0, Qcur, eps_norm); + Kcur = build_gdn_l2_norm(ctx0, Kcur, eps_norm); // Choose between build_delta_net_chunking and build_delta_net_recurrent based on n_tokens auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il); diff --git a/src/models/llada.cpp b/src/models/llada.cpp index 87d4259f9..ae3d6925c 100644 --- a/src/models/llada.cpp +++ b/src/models/llada.cpp @@ -36,12 +36,7 @@ void llama_model_llada::load_arch_tensors(llama_model_loader &) { layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0); - // Use separate Q, K, V projections without bias, matching LLaDALlamaBlock - layer.wq = - create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head_k * n_head }, 0); - layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), { n_embd, n_embd_k_gqa }, 0); - layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, 0); - // No bias for QKV projections as per config: include_bias=false, include_qkv_bias=false + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0); layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), { n_embd }, TENSOR_NOT_REQUIRED); diff --git a/src/models/minimax-m2.cpp b/src/models/minimax-m2.cpp index c2e69bfaa..7a22af036 100644 --- a/src/models/minimax-m2.cpp +++ b/src/models/minimax-m2.cpp @@ -71,14 +71,13 @@ llama_model_minimax_m2::graph::graph(const llama_model & model, const llm_graph_ cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "attn_norm", il); - // compute Q and K and RoPE them - ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur); + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur, "Qcur", il); - - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur); cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur); cb(Vcur, "Vcur", il); Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, diff --git a/src/models/models.h b/src/models/models.h index 93a6b3494..87195fddd 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -10,6 +10,13 @@ class llama_memory_hybrid_idx_context; +// ref: https://github.com/ggml-org/llama.cpp/pull/28068 +static inline ggml_tensor * build_gdn_l2_norm(ggml_context * ctx, ggml_tensor * x, float eps) { + const float n = x->ne[0]; + + return ggml_scale(ctx, ggml_rms_norm(ctx, x, eps/n), 1.0f/sqrtf(n)); +} + // // base classes // @@ -2606,3 +2613,16 @@ struct llama_model_step35 : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; + + +struct llama_model_spark2_5 : public llama_model_base { + llama_model_spark2_5(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; +}; diff --git a/src/models/olmo2.cpp b/src/models/olmo2.cpp index cb52cdef7..05b9394b8 100644 --- a/src/models/olmo2.cpp +++ b/src/models/olmo2.cpp @@ -93,14 +93,13 @@ llama_model_olmo2::graph::graph(const llama_model & model, const llm_graph // self_attention { - // compute Q and K and RoPE them - ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur); + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur, "Qcur", il); - - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur); cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur); cb(Vcur, "Vcur", il); Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, diff --git a/src/models/olmoe.cpp b/src/models/olmoe.cpp index 1e2baeb20..11c53f3f4 100644 --- a/src/models/olmoe.cpp +++ b/src/models/olmoe.cpp @@ -79,14 +79,13 @@ llama_model_olmoe::graph::graph(const llama_model & model, const llm_graph_param // self_attention { - // compute Q and K and RoPE them - ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur); + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur, "Qcur", il); - - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur); cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur); cb(Vcur, "Vcur", il); Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index 0b9210981..a1e263500 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -263,8 +263,14 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn( // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention // Qwen3Next uses a single Q projection that outputs query + gate - ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] + auto [Qcur_full, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "Qcur_full", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, ggml_element_size(Qcur_full) * n_embd_head * 2, @@ -275,12 +281,6 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn( Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); cb(Qcur, "Qcur_normed", il); - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); - cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s); - cb(Vcur, "Vcur", il); - // Apply K normalization Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); @@ -423,10 +423,11 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( cb(k_conv, "k_conv", il); cb(v_conv, "v_conv", il); + const float eps_norm = hparams.f_norm_rms_eps; - q_conv = ggml_l2_norm(ctx0, q_conv, eps_norm); - k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm); + q_conv = build_gdn_l2_norm(ctx0, q_conv, eps_norm); + k_conv = build_gdn_l2_norm(ctx0, k_conv, eps_norm); //q_conv = ggml_cont_4d(ctx0, q_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); //k_conv = ggml_cont_4d(ctx0, k_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); @@ -553,7 +554,11 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "mtp_attn_norm", il); - ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s); + auto [Qcur_full, Kcur, Vcur] = build_qkv(layer, cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "mtp_Qcur_full", il); ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, @@ -572,12 +577,10 @@ llama_model_qwen35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); cb(gate, "mtp_gate", il); - ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); cb(Kcur, "mtp_Kcur_normed", il); - ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); cb(Vcur, "mtp_Vcur", il); diff --git a/src/models/qwen35moe.cpp b/src/models/qwen35moe.cpp index ed4083f12..bdf772625 100644 --- a/src/models/qwen35moe.cpp +++ b/src/models/qwen35moe.cpp @@ -287,8 +287,14 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn( // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention // Qwen3Next uses a single Q projection that outputs query + gate - ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] + auto [Qcur_full, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "Qcur_full", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, ggml_element_size(Qcur_full) * n_embd_head * 2, @@ -299,12 +305,6 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn( Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); cb(Qcur, "Qcur_normed", il); - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); - cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s); - cb(Vcur, "Vcur", il); - // Apply K normalization Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); @@ -447,10 +447,11 @@ ggml_tensor * llama_model_qwen35moe::graph::build_layer_attn_linear( cb(k_conv, "k_conv", il); cb(v_conv, "v_conv", il); + const float eps_norm = hparams.f_norm_rms_eps; - q_conv = ggml_l2_norm(ctx0, q_conv, eps_norm); - k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm); + q_conv = build_gdn_l2_norm(ctx0, q_conv, eps_norm); + k_conv = build_gdn_l2_norm(ctx0, k_conv, eps_norm); //q_conv = ggml_cont_4d(ctx0, q_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); //k_conv = ggml_cont_4d(ctx0, k_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); @@ -617,7 +618,11 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "mtp_attn_norm", il); - ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s); + auto [Qcur_full, Kcur, Vcur] = build_qkv(layer, cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "mtp_Qcur_full", il); ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, @@ -636,12 +641,10 @@ llama_model_qwen35moe::graph_mtp::graph_mtp(const llama_model & model, const llm gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); cb(gate, "mtp_gate", il); - ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); cb(Kcur, "mtp_Kcur_normed", il); - ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); cb(Vcur, "mtp_Vcur", il); diff --git a/src/models/qwen3next.cpp b/src/models/qwen3next.cpp index eb823b8ea..b63fc9c6a 100644 --- a/src/models/qwen3next.cpp +++ b/src/models/qwen3next.cpp @@ -244,8 +244,14 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn( // Order: joint QG projection, QG split, Q norm, KV projection, K norm, RoPE, attention // Qwen3Next uses a single Q projection that outputs query + gate - ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); + auto [Qcur_full, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "Qcur_full", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); Qcur_full = ggml_reshape_4d(ctx0, Qcur_full, n_embd_head * 2, n_head, n_tokens, 1); @@ -260,12 +266,6 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn( Qcur_full->nb[1], Qcur_full->nb[2], Qcur_full->nb[3], n_embd_head * ggml_element_size(Qcur_full)); cb(gate, "gate", il); - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); - cb(Kcur, "Kcur", il); - - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s); - cb(Vcur, "Vcur", il); - Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); @@ -503,10 +503,11 @@ ggml_tensor * llama_model_qwen3next::graph::build_layer_attn_linear( cb(k_conv, "k_conv", il); cb(v_conv, "v_conv", il); + const float eps_norm = hparams.f_norm_rms_eps; - q_conv = ggml_l2_norm(ctx0, q_conv, eps_norm); - k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm); + q_conv = build_gdn_l2_norm(ctx0, q_conv, eps_norm); + k_conv = build_gdn_l2_norm(ctx0, k_conv, eps_norm); //q_conv = ggml_cont_4d(ctx0, q_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); //k_conv = ggml_cont_4d(ctx0, k_conv, head_k_dim, num_k_heads, n_seq_tokens, n_seqs); @@ -691,7 +692,11 @@ llama_model_qwen3next::graph_mtp::graph_mtp(const llama_model & model, const llm cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "mtp_attn_norm", il); - ggml_tensor * Qcur_full = build_lora_mm(layer.wq, cur, layer.wq_s); + auto [Qcur_full, Kcur, Vcur] = build_qkv(layer, cur, + n_embd_head * 2, n_head, + n_embd_head, n_head_kv, + n_embd_head, n_head_kv, + il, false); cb(Qcur_full, "mtp_Qcur_full", il); ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, @@ -702,12 +707,10 @@ llama_model_qwen3next::graph_mtp::graph_mtp(const llama_model & model, const llm Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il); cb(Qcur, "mtp_Qcur_normed", il); - ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); cb(Kcur, "mtp_Kcur_normed", il); - ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 1484c9b07..8ace95f73 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -936,10 +936,11 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( cb(k_conv, "k_conv", il); cb(v_conv, "v_conv", il); + const float eps_norm = hparams.f_norm_rms_eps; - q_conv = ggml_l2_norm(ctx0, q_conv, eps_norm); - k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm); + q_conv = build_gdn_l2_norm(ctx0, q_conv, eps_norm); + k_conv = build_gdn_l2_norm(ctx0, k_conv, eps_norm); // repeat to match shapes when head keys != value keys; unneeded with the fused GDN if (num_k_heads != num_v_heads && (!cparams.fused_gdn_ar || !cparams.fused_gdn_ch)) { diff --git a/src/models/spark2-5.cpp b/src/models/spark2-5.cpp new file mode 100644 index 000000000..107448777 --- /dev/null +++ b/src/models/spark2-5.cpp @@ -0,0 +1,146 @@ +#include "models.h" + +void llama_model_spark2_5::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl); + + hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; + hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); + + switch (hparams.n_layer()) { + case 28: type = LLM_TYPE_1_7B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_spark2_5::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); + if (output == nullptr) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + const int64_t n_head_i = hparams.n_head(i); + const int64_t n_head_kv_i = hparams.n_head_kv(i); + const int64_t n_embd_q = hparams.n_embd_head_k(i) * n_head_i; + const int64_t n_embd_k = hparams.n_embd_head_k(i) * n_head_kv_i; + const int64_t n_embd_v = hparams.n_embd_head_v(i) * n_head_kv_i; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + create_tensor_qkv(layer, i, n_embd, n_embd_q, n_embd_k, n_embd_v, 0); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_head_i}, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_q, n_embd}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + } +} + +std::unique_ptr llama_model_spark2_5::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_spark2_5::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_STANDARD); + + ggml_tensor * inpL = build_inp_embd(model.tok_embd); + ggml_tensor * inp_pos = build_inp_pos(); + auto * inp_attn = build_attn_inp_kv_iswa(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const float kq_scale = 1.0f / sqrtf(float(n_embd_head)); + + for (int il = 0; il < n_layer; ++il) { + ggml_tensor * inpSA = inpL; + ggml_tensor * cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + const int64_t n_head_i = hparams.n_head(il); + const int64_t n_head_kv_i = hparams.n_head_kv(il); + const int64_t n_rot_i = hparams.n_rot(il); + const float freq_base_i = model.get_rope_freq_base(cparams, il); + const float freq_scale_i = model.get_rope_freq_scale(cparams, il); + + ggml_tensor * attn_inp = cur; + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, n_embd_head, n_head_i, n_head_kv_i, il); + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, + n_rot_i, rope_type, n_ctx_orig, freq_base_i, freq_scale_i, + ext_factor, attn_factor, beta_fast, beta_slow); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, + n_rot_i, rope_type, n_ctx_orig, freq_base_i, freq_scale_i, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(Qcur, "Qcur_rope", il); + cb(Kcur, "Kcur_rope", il); + + cur = build_attn(inp_attn, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "attn_out", il); + + ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp); + gate = ggml_sigmoid(ctx0, gate); + cb(gate, "attn_gate", il); + + const int64_t n_tokens_i = cur->ne[1]; + cur = ggml_reshape_3d(ctx0, cur, n_embd_head, n_head_i, n_tokens_i); + gate = ggml_reshape_3d(ctx0, gate, 1, n_head_i, n_tokens_i); + cur = ggml_mul(ctx0, cur, gate); + cur = ggml_reshape_2d(ctx0, cur, n_embd_head * n_head_i, n_tokens_i); + cb(cur, "attn_gated", il); + + cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s); + cb(cur, "attn_out_proj", il); + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + model.layers[il].ffn_up, nullptr, nullptr, + model.layers[il].ffn_gate, nullptr, nullptr, + model.layers[il].ffn_down, nullptr, nullptr, + nullptr, + LLM_FFN_GELU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + inpL = cur; + } + + ggml_tensor * cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/step35.cpp b/src/models/step35.cpp index 53f3179c6..946a36960 100644 --- a/src/models/step35.cpp +++ b/src/models/step35.cpp @@ -216,9 +216,11 @@ llama_model_step35::graph::graph(const llama_model & model, const llm_graph_para { cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "attn_norm", il); - ggml_tensor * Qcur = build_lora_mm(model.layers[il].wq, cur); - ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur); - ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur); + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head_k, n_head_l, + n_embd_head_k, n_head_kv_l, + n_embd_head_v, n_head_kv_l, + il, false); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); @@ -425,9 +427,11 @@ llama_model_step35::graph_mtp::graph_mtp(const llama_model & model, const llm_gr cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "mtp_attn_norm", il); - ggml_tensor * Qcur = build_lora_mm(layer.wq, cur, layer.wq_s); - ggml_tensor * Kcur = build_lora_mm(layer.wk, cur, layer.wk_s); - ggml_tensor * Vcur = build_lora_mm(layer.wv, cur, layer.wv_s); + auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, + n_embd_head_k, n_head_l, + n_embd_head_k, n_head_kv_l, + n_embd_head_v, n_head_kv_l, + il, false); cb(Qcur, "mtp_Qcur", il); cb(Kcur, "mtp_Kcur", il); cb(Vcur, "mtp_Vcur", il); diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index db0fac995..4d2592b25 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -80,18 +80,19 @@ struct server_lru_sched { } // returns "" if no model can be given up - std::string pick_victim(std::unique_lock & lk, const std::string & exclude) { + std::string pick_victim(std::unique_lock & lk) { check_lock(lk); std::string victim; int64_t victim_last_used = 0; for (const auto & m : models.mapping) { - if (m.first == exclude) { - continue; - } // a busy model is mid-request, one still coming up has no request to finish if (m.second.req_count != 0 || !m.second.meta.is_ready_or_sleep()) { continue; } + // already on its way out, or a queued request wants it + if (models.stopping_models.count(m.first) || find(m.first)) { + continue; + } if (victim.empty() || m.second.meta.last_used < victim_last_used) { victim = m.first; victim_last_used = m.second.meta.last_used; @@ -109,7 +110,7 @@ struct server_lru_sched { SRV_INF("request for name=%s joined the queue, %d waiting\n", model_id.c_str(), e->n_waiters); return; } - queue.push_back({ model_id, 1, false, false }); + queue.push_back({ model_id, 1, false }); SRV_INF("models_max reached, request for name=%s queued at position %zu\n", model_id.c_str(), queue.size()); } @@ -144,85 +145,67 @@ struct server_lru_sched { return true; } - // ok means the model is up: drop the entry, the other waiters just watch its status now + // on failure the entry is back in line; on success it stays until its waiters leave, + // so the model coming up is never picked as a victim before they use it void claim_done(std::unique_lock & lk, const std::string & model_id, bool ok) { check_lock(lk); + if (ok) { + return; + } for (auto it = queue.begin(); it != queue.end(); ++it) { if (it->model_id == model_id) { - if (ok) { - queue.erase(it); - } else { - it->loading = false; - } + it->loading = false; return; } } } - // a model is on its way out for this entry, so other requests do not also give up one - void mark_slot_pending(std::unique_lock & lk, const std::string & model_id) { + // evict idle models while queued requests outnumber the slots that are free or being freed + // caller must hold models.mutex; never blocks, so it is safe from any thread + void tick(std::unique_lock & lk) { check_lock(lk); - if (entry_t * e = find(model_id)) { - e->slot_pending = true; + if (models.base_params.models_max <= 0 || queue.empty()) { + return; } - } - - // model_id went idle: give up its slot if a queued request needs one - // thread-safe, caller must NOT hold models.mutex - void on_model_idle(const std::string & model_id) { - if (models.base_params.models_max <= 0) { - return; // no limit, nothing is ever queued - } - { - std::unique_lock lk(models.mutex); - if (queue.empty()) { - return; - } - size_t promised = 0; - bool has_unserved = false; - for (const auto & e : queue) { - if (e.needs_slot()) { - has_unserved = true; - } else { - promised++; - } - } - if (!has_unserved) { - return; - } - if ((int) count_running() - (int) promised < models.base_params.models_max) { - return; // a slot is already on its way - } - // never give up a model that a queued request wants - for (const auto & e : queue) { - if (e.model_id == model_id) { - return; - } - } - auto it = models.mapping.find(model_id); - if (it == models.mapping.end() || it->second.req_count != 0 || !it->second.meta.is_ready_or_sleep()) { - return; - } - for (auto & e : queue) { - if (!e.slot_pending) { - e.slot_pending = true; - break; + int n_running = 0; + int n_stopping = 0; + for (const auto & m : models.mapping) { + if (m.second.meta.is_running()) { + n_running++; + if (models.stopping_models.count(m.first)) { + n_stopping++; } } } - SRV_INF("model name=%s went idle, giving up its slot to a queued request\n", model_id.c_str()); - models.unload(model_id); + int n_needed = 0; + int n_claimed = 0; // claimed the slot, but load() has not spawned yet + for (const auto & e : queue) { + if (!e.loading) { + n_needed++; + continue; + } + auto it = models.mapping.find(e.model_id); + if (it != models.mapping.end() && !it->second.meta.is_running()) { + n_claimed++; + } + } + int n_free = models.base_params.models_max - n_running + n_stopping - n_claimed; + while (n_free < n_needed) { + std::string victim = pick_victim(lk); + if (victim.empty()) { + return; // all remaining models are busy, wait for a request to end + } + SRV_INF("evicting idle LRU name=%s for a queued request\n", victim.c_str()); + models.request_stop(victim); + n_free++; + } } private: struct entry_t { std::string model_id; - int n_waiters; // requests waiting for this model - bool slot_pending; // a model is already being evicted for this entry - bool loading; // one of the waiters is doing the load right now - - // a slot is already coming, or already taken by the load in flight - bool needs_slot() const { return !slot_pending && !loading; } + int n_waiters; // requests waiting for this model + bool loading; // one of the waiters is doing the load right now }; entry_t * find(const std::string & model_id) { @@ -946,7 +929,7 @@ void server_models::unload_lru() { if (sched->has_capacity(lk)) { return; } - lru_model_name = sched->pick_victim(lk, ""); + lru_model_name = sched->pick_victim(lk); } if (!lru_model_name.empty()) { SRV_INF("models_max limit reached, removing LRU name=%s\n", lru_model_name.c_str()); @@ -1169,6 +1152,11 @@ void server_models::load(const std::string & name, const load_options & opts) { cv.notify_all(); } +void server_models::request_stop(const std::string & name) { + stopping_models.insert(name); + cv_stop.notify_all(); +} + void server_models::unload(const std::string & name) { std::unique_lock lk(mutex); auto it = mapping.find(name); @@ -1182,13 +1170,12 @@ void server_models::unload(const std::string & name) { }); } else if (it->second.meta.is_running()) { SRV_INF("stopping model instance name=%s\n", name.c_str()); - stopping_models.insert(name); if (it->second.meta.status == SERVER_MODEL_STATUS_LOADING) { // special case: if model is in loading state, unloading means force-killing it SRV_WRN("model name=%s is still loading, force-killing\n", name.c_str()); it->second.subproc->terminate(); } - cv_stop.notify_all(); + request_stop(name); // status change will be handled by the managing thread } else { SRV_WRN("model instance name=%s is not running\n", name.c_str()); @@ -1206,8 +1193,7 @@ void server_models::unload_all() { inst.subproc->stopped.store(true, std::memory_order_relaxed); } else if (inst.meta.is_running()) { SRV_INF("stopping model instance name=%s\n", name.c_str()); - stopping_models.insert(name); - cv_stop.notify_all(); + request_stop(name); // status change will be handled by the managing thread } // moving the thread to join list to avoid deadlock @@ -1234,6 +1220,8 @@ void server_models::update_status(const std::string & name, const update_status_ if (!args.progress.is_null()) { meta.progress = args.progress; } + // a model that comes up idle or goes down changes the slot count for queued requests + sched->tick(lk); } // broadcast status change to SSE { @@ -1380,13 +1368,11 @@ bool server_models::ensure_model_ready(const std::string & name, const std::func bool queued = false; bool did_load = false; - std::string victim; { std::unique_lock lk(mutex); auto it = mapping.find(name); if (it != mapping.end() && it->second.meta.status == SERVER_MODEL_STATUS_UNLOADED) { - bool has_capacity = sched->has_capacity(lk); - if (has_capacity && sched->queue_empty(lk)) { + if (sched->has_capacity(lk) && sched->queue_empty(lk)) { lk.unlock(); SRV_INF("model name=%s is not loaded, loading...\n", name.c_str()); load(name); @@ -1394,21 +1380,11 @@ bool server_models::ensure_model_ready(const std::string & name, const std::func } else { // also queue when a slot looks free but others wait already, else they starve sched->join(lk, name); + sched->tick(lk); queued = true; - if (!has_capacity) { - // an idle model may sit here right now, do not wait for a request to end - victim = sched->pick_victim(lk, name); - if (!victim.empty()) { - sched->mark_slot_pending(lk, name); - } - } } } } - if (!victim.empty()) { - SRV_INF("evicting idle LRU name=%s to make room for name=%s\n", victim.c_str(), name.c_str()); - unload(victim); - } // while queued, this is also where the load happens: the head of the queue does it SRV_INF("waiting until model name=%s is fully loaded...\n", name.c_str()); @@ -1470,9 +1446,7 @@ bool server_models::ensure_model_ready(const std::string & name, const std::func } lk.lock(); sched->claim_done(lk, name, ok); - if (ok) { - queued = false; // entry is gone, the other waiters watch the status now - } + sched->tick(lk); continue; } @@ -1480,6 +1454,7 @@ bool server_models::ensure_model_ready(const std::string & name, const std::func } } catch (...) { leave_queue(); + sched->tick(lk); // a slot freed for this waiter goes to the next one throw; } leave_queue(); @@ -1529,18 +1504,14 @@ server_http_res_ptr server_models::proxy_request(const server_http_req & req, co ); proxy->cleanup = [this, name]() { - bool went_idle = false; - { - std::unique_lock lk(mutex); - auto it = mapping.find(name); - if (it != mapping.end() && it->second.req_count > 0) { - it->second.req_count--; - went_idle = it->second.req_count == 0; + std::unique_lock lk(mutex); + auto it = mapping.find(name); + if (it != mapping.end() && it->second.req_count > 0) { + it->second.req_count--; + if (it->second.req_count == 0) { + sched->tick(lk); } } - if (went_idle) { - sched->on_model_idle(name); - } }; return proxy; diff --git a/tools/server/server-models.h b/tools/server/server-models.h index 5cbb6a801..7f6c26b35 100644 --- a/tools/server/server-models.h +++ b/tools/server/server-models.h @@ -216,6 +216,10 @@ private: // not thread-safe, caller must hold mutex void add_model(server_model_meta && meta); + // ask the monitoring thread to stop a running instance + // not thread-safe, caller must hold mutex + void request_stop(const std::string & name); + // notify SSE clients void notify_sse(const std::string & event, const std::string & model_id, const json & data = nullptr); diff --git a/tools/server/tests/unit/test_router.py b/tools/server/tests/unit/test_router.py index 96eb87978..e4b7f9fe4 100644 --- a/tools/server/tests/unit/test_router.py +++ b/tools/server/tests/unit/test_router.py @@ -297,6 +297,26 @@ def test_router_queue_is_fifo(): assert first.done_at < second.done_at, "queue was not served in arrival order" +def test_router_queue_two_waiters_share_one_eviction(): + """two requests that both find the same idle model must both be served in the end""" + global server + server.models_max = 1 + server.start() + + _load_model_and_wait(MODEL_A, timeout=120) + + # both arrive while MODEL_A is idle, so both want its slot; only one eviction can happen + first = _Bg(lambda: _tokenize(MODEL_B)).start() + second = _Bg(lambda: _tokenize(MODEL_C)).start() + + first.join(90) + second.join(90) + + first.assert_ok("first queued request") + second.assert_ok("second queued request") + assert _get_model_status(MODEL_A) == "unloaded" + + def test_router_no_models_autoload(): global server server.no_models_autoload = True diff --git a/tools/ui/CMakeLists.txt b/tools/ui/CMakeLists.txt index 208b46a5c..79ffe9fc1 100644 --- a/tools/ui/CMakeLists.txt +++ b/tools/ui/CMakeLists.txt @@ -36,60 +36,11 @@ endif() set(UI_CPP "${CMAKE_CURRENT_BINARY_DIR}/ui.cpp") set(UI_H "${CMAKE_CURRENT_BINARY_DIR}/ui.h") -if(CMAKE_CROSSCOMPILING) - find_program(HOST_CXX_COMPILER NAMES g++ clang++ NO_CMAKE_FIND_ROOT_PATH) - if(NOT HOST_CXX_COMPILER) - message(FATAL_ERROR "UI: no host C++ compiler (g++/clang++) found to build llama-ui-embed; set -DHOST_CXX_COMPILER=") - endif() - message(STATUS "UI: building llama-ui-embed with host compiler ${HOST_CXX_COMPILER}") - - if(CMAKE_HOST_WIN32) - set(LLAMA_UI_EMBED_EXE "${CMAKE_CURRENT_BINARY_DIR}/llama-ui-embed-host.exe") - else() - set(LLAMA_UI_EMBED_EXE "${CMAKE_CURRENT_BINARY_DIR}/llama-ui-embed-host") - endif() - - add_custom_command( - OUTPUT "${LLAMA_UI_EMBED_EXE}" - COMMAND "${HOST_CXX_COMPILER}" -O2 -std=c++17 - -o "${LLAMA_UI_EMBED_EXE}" "${CMAKE_CURRENT_SOURCE_DIR}/embed.cpp" - DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/embed.cpp" - COMMENT "Building llama-ui-embed (host)" - VERBATIM - ) - - # phony target to tie it into the dependency graph - add_custom_target(llama-ui-embed DEPENDS "${LLAMA_UI_EMBED_EXE}") -else() - # exclude llama-ui-embed from sanitizer flags, - # it's a build-time-only tool, no need to instrument it - # this is to fix TSan "memory layout is incompatible" error on CI - get_directory_property(_llama_ui_dir_co COMPILE_OPTIONS) - get_directory_property(_llama_ui_dir_ll LINK_LIBRARIES) - set(_llama_ui_embed_co ${_llama_ui_dir_co}) - set(_llama_ui_embed_ll ${_llama_ui_dir_ll}) - list(FILTER _llama_ui_embed_co EXCLUDE REGEX ".*-fsanitize=.*") - list(FILTER _llama_ui_embed_ll EXCLUDE REGEX ".*-fsanitize=.*") - set_directory_properties(PROPERTIES - COMPILE_OPTIONS "${_llama_ui_embed_co}" - LINK_LIBRARIES "${_llama_ui_embed_ll}") - - add_executable(llama-ui-embed embed.cpp) - target_compile_features(llama-ui-embed PRIVATE cxx_std_17) - set_target_properties(llama-ui-embed PROPERTIES - RUNTIME_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" - ) - set(LLAMA_UI_EMBED_EXE "$") - - # restore so the llama-ui library below keeps sanitizer instrumentation - set_directory_properties(PROPERTIES - COMPILE_OPTIONS "${_llama_ui_dir_co}" - LINK_LIBRARIES "${_llama_ui_dir_ll}") -endif() - -# Run the provisioning script every build so source changes in tools/ui/ are -# always picked up. The script uses copy_if_different for ui.cpp/ui.h, so the -# library only recompiles when contents actually change. +# Provision assets and generate ui.cpp/ui.h natively in CMake at build time. +# The generated sources are compiled by the regular target toolchain; no +# build-time host executable is needed (works in any cross-compile setup). +# The script uses copy_if_different semantics, so the library below only +# recompiles when the generated contents actually change. add_custom_target(llama-ui-assets ALL BYPRODUCTS ${UI_CPP} ${UI_H} COMMAND ${CMAKE_COMMAND} @@ -101,15 +52,12 @@ add_custom_target(llama-ui-assets ALL "-DHF_VERSION=${HF_UI_VERSION}" "-DHF_ENABLED=${LLAMA_USE_PREBUILT_UI}" "-DBUILD_UI=${LLAMA_BUILD_UI}" - "-DLLAMA_UI_EMBED=${LLAMA_UI_EMBED_EXE}" "-DLLAMA_UI_GZIP=${LLAMA_UI_GZIP}" -P "${PROJECT_SOURCE_DIR}/scripts/ui-assets.cmake" COMMENT "Provisioning UI assets" VERBATIM ) -add_dependencies(llama-ui-assets llama-ui-embed) - set_source_files_properties(${UI_CPP} ${UI_H} PROPERTIES GENERATED TRUE) add_library(${TARGET} STATIC ${UI_CPP} ${UI_H}) diff --git a/tools/ui/embed.cpp b/tools/ui/embed.cpp deleted file mode 100644 index b76c9047f..000000000 --- a/tools/ui/embed.cpp +++ /dev/null @@ -1,308 +0,0 @@ -// llama-ui-embed: generate ui.cpp / ui.h that embed UI assets as C arrays. -// -// Usage: -// llama-ui-embed [] -// -// Recursively embeds every regular file under . -// Asset names are relative paths from (e.g. "_app/immutable/bundle.HASH.js"). -// Without , emits an empty asset table. - -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include - - -static const char * mime_from_ext(const std::string & name) { - auto ext = name.rfind('.'); - if (ext == std::string::npos) return "application/octet-stream"; - std::string e = name.substr(ext + 1); - if (e == "html") return "text/html; charset=utf-8"; - if (e == "css") return "text/css"; - if (e == "js") return "application/javascript"; - if (e == "json") return "application/json"; - if (e == "webmanifest") return "application/manifest+json"; - if (e == "svg") return "image/svg+xml"; - if (e == "png") return "image/png"; - if (e == "jpg" || - e == "jpeg") return "image/jpeg"; - if (e == "ico") return "image/x-icon"; - if (e == "woff") return "font/woff"; - if (e == "woff2") return "font/woff2"; - return "application/octet-stream"; -} - -// Computes FNV-1a hash of the data -static uint64_t fnv_hash(const uint8_t * data, size_t len) { - const uint64_t fnv_prime = 0x100000001b3ULL; - uint64_t hash = 0xcbf29ce484222325ULL; - - for (size_t i = 0; i < len; ++i) { - hash ^= data[i]; - hash *= fnv_prime; - } - return hash; -} - -static bool read_file(const std::filesystem::path & path, std::vector & out) { - std::ifstream f(path, std::ios::binary | std::ios::ate); - if (!f) { - fprintf(stderr, "embed: cannot open %s\n", path.string().c_str()); - return false; - } - const auto sz = f.tellg(); - if (sz < 0) { - return false; - } - f.seekg(0); - out.resize(static_cast(sz)); - if (sz > 0 && !f.read(reinterpret_cast(out.data()), sz)) { - return false; - } - return true; -} - -static void append_bytes_hex(std::string & out, const std::vector & bytes) { - static const char hex[] = "0123456789abcdef"; - out.reserve(out.size() + bytes.size() * 5); - for (unsigned char b : bytes) { - out += '0'; - out += 'x'; - out += hex[b >> 4]; - out += hex[b & 0xf]; - out += ','; - } -} - -static bool write_if_different(const std::string & path, const std::string & content) { - std::ifstream f(path, std::ios::binary | std::ios::ate); - if (f) { - const auto sz = f.tellg(); - if (sz >= 0 && static_cast(sz) == content.size()) { - std::string existing(static_cast(sz), '\0'); - f.seekg(0); - if (sz == 0 || f.read(existing.data(), sz)) { - if (existing == content) { - return true; - } - } - } - } - - std::ofstream out(path, std::ios::binary | std::ios::trunc); - if (!out) { - fprintf(stderr, "embed: cannot write %s\n", path.c_str()); - return false; - } - if (!content.empty()) { - out.write(content.data(), static_cast(content.size())); - } - bool ok = out.good(); - if (ok) { - printf("embed: write output file %s\n", path.c_str()); - } - return ok; -} - -static std::string path_basename(const std::string & name) { - const size_t p = name.rfind('/'); - return p == std::string::npos ? name : name.substr(p + 1); -} -static bool str_starts_with(const std::string & s, const char * prefix) { - const size_t n = strlen(prefix); - return s.size() >= n && s.compare(0, n, prefix) == 0; -} -static bool str_ends_with(const std::string & s, const char * suffix) { - const size_t n = strlen(suffix); - return s.size() >= n && s.compare(s.size() - n, n, suffix) == 0; -} - -static std::string fmt(const char * pattern, ...) { - char tmp[512]; - va_list ap; - va_start(ap, pattern); - const int n = vsnprintf(tmp, sizeof(tmp), pattern, ap); - va_end(ap); - return (n > 0) ? std::string(tmp, static_cast(n)) : std::string(); -} - -struct asset_entry { - std::string name; - std::filesystem::path path; -}; - -int main(int argc, char ** argv) { - if (argc < 3 || argc > 4) { - fprintf(stderr, "usage: %s []\n", argv[0]); - return 1; - } - - const std::string out_cpp = argv[1]; - const std::string out_h = argv[2]; - const std::string asset_dir = (argc >= 4) ? argv[3] : std::string(); - - const bool use_gzip = !asset_dir.empty() && std::filesystem::exists(asset_dir + "/_gzip"); - const std::string in_dir = use_gzip ? (asset_dir + "/_gzip") : asset_dir; - - std::vector assets; - if (!in_dir.empty()) { - const std::filesystem::path dir = in_dir; - - std::error_code ec; - std::filesystem::recursive_directory_iterator it(dir, ec); - if (ec) { - fprintf(stderr, "embed: cannot iterate %s: %s\n", argv[3], ec.message().c_str()); - return 1; - } - for (const auto & entry : it) { - if (!entry.is_regular_file()) { - continue; - } - // name is the relative path from dir, with forward slashes - const std::string name = entry.path().lexically_relative(dir).generic_string(); - assets.push_back({ name, entry.path() }); - } - - // directory iteration order is unspecified; sort for reproducible output - std::sort(assets.begin(), assets.end(), - [](const asset_entry & a, const asset_entry & b) { return a.name < b.name; }); - } - - const int n_assets = static_cast(assets.size()); - - if (n_assets > 0) { - using match_fn = std::function; - auto exact = [](const char * name) -> match_fn { - return [name](const std::string & base) { return base == name; }; - }; - - struct required_check { const char * label; match_fn match; bool found; }; - required_check checks[] = { - { "index.html", exact("index.html"), false }, - { "manifest.webmanifest", exact("manifest.webmanifest"), false }, - { "sw.js", exact("sw.js"), false }, - { "build.json", exact("build.json"), false }, - { "version.json", exact("version.json"), false }, - { "bundle[hash].js", [](const std::string & b) { - return str_starts_with(b, "bundle") && str_ends_with(b, ".js"); - }, false }, - { "bundle[hash].css", [](const std::string & b) { - return str_starts_with(b, "bundle") && str_ends_with(b, ".css"); - }, false }, - { "workbox[hash].js", [](const std::string & b) { - return str_starts_with(b, "workbox") && str_ends_with(b, ".js"); - }, false }, - }; - - for (const auto & a : assets) { - const std::string base = path_basename(a.name); - for (auto & c : checks) { - if (!c.found) { c.found = c.match(base); } - } - } - - std::vector missing; - for (const auto & c : checks) { - if (!c.found) { missing.push_back(c.label); } - } - if (!missing.empty()) { - fprintf(stderr, "\ncurrent asset files:\n"); - for (const auto & a : assets) { - fprintf(stderr, " %s\n", a.name.c_str()); - } - fprintf(stderr, "missing required asset(s):\n"); - for (const char * m : missing) { - fprintf(stderr, " %s\n", m); - } - fprintf(stderr, "hint: try cleaning your build directory: %s\n", in_dir.c_str()); - return 1; - } - } - - std::string h; - h += "#pragma once\n\n#include \n#include \n\n"; - if (n_assets > 0) { - h += "#define LLAMA_UI_HAS_ASSETS 1\n\n"; - } - h += - "struct llama_ui_asset {\n" - " std::string name;\n" - " const unsigned char * data;\n" - " std::size_t size;\n" - " std::string etag;\n" - " std::string type;\n" - "};\n\n" - "const llama_ui_asset * llama_ui_find_asset(const std::string & name);\n" - "bool llama_ui_use_gzip();\n"; - h += fmt("const std::array & llama_ui_get_assets();\n", n_assets); - - std::string cpp; - cpp += "#include \"ui.h\"\n\n"; - - if (n_assets > 0) { - for (int i = 0; i < n_assets; i++) { - std::vector bytes; - if (!read_file(assets[i].path, bytes)) { - return 1; - } - if (bytes.empty()) { - fprintf(stderr, "embed: empty file: %s\n", assets[i].path.generic_string().c_str()); - return 1; - } - cpp += fmt("static const unsigned char asset_%d_data[] = {", i); - append_bytes_hex(cpp, bytes); - - // note: this is a simple hash for cache busting, not a cryptographic hash; fnv is enough here - const auto hash = fnv_hash(bytes.data(), bytes.size()); - - cpp += fmt("};\nstatic const std::size_t asset_%d_size = %zu;\n", - i, bytes.size()); - cpp += fmt("static const char asset_%d_etag[] = \"\\\"0x%016" PRIx64 "\\\"\";\n\n", - i, hash); - } - - cpp += fmt("static const std::array g_assets = {{\n", n_assets); - for (int i = 0; i < n_assets; i++) { - const std::string & name = assets[i].name; - cpp += fmt(" { \"%s\", asset_%d_data, asset_%d_size, asset_%d_etag, \"%s\" },\n", - name.c_str(), i, i, i, mime_from_ext(name)); - } - cpp += "}};\n\n"; - - cpp += - "const llama_ui_asset * llama_ui_find_asset(const std::string & name) {\n" - " for (const auto & a : g_assets) {\n" - " if (a.name == name) {\n" - " return &a;\n" - " }\n" - " }\n" - " return nullptr;\n" - "}\n"; - cpp += fmt("const std::array & llama_ui_get_assets() {\n", n_assets); - cpp += " return g_assets;\n" - "}\n"; - } else { - cpp += - "const llama_ui_asset * llama_ui_find_asset(const std::string &) {\n" - " return nullptr;\n" - "}\n" - "const std::array & llama_ui_get_assets() {\n" - " static const std::array empty{};\n" - " return empty;\n" - "}\n"; - } - cpp += fmt("bool llama_ui_use_gzip() { return %s; }\n", use_gzip ? "true" : "false"); - - bool ok = true; - ok = write_if_different(out_h, h) && ok; - ok = write_if_different(out_cpp, cpp) && ok; - return ok ? 0 : 1; -} diff --git a/tools/ui/src/app.d.ts b/tools/ui/src/app.d.ts index 5309dce8f..639a16df2 100644 --- a/tools/ui/src/app.d.ts +++ b/tools/ui/src/app.d.ts @@ -137,7 +137,6 @@ declare global { declare global { interface Window { - idxThemeStyle?: number; idxCodeBlock?: number; // File System Access API - not in the DOM lib and unavailable in some browsers diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte index fa2a50bc5..46d05338b 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte @@ -404,7 +404,7 @@ } -
+
{#if message.role === MessageRole.SYSTEM} {:else if mcpPromptExtra} @@ -425,25 +425,3 @@ /> {/if}
- - diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte index a2c742f0f..dac55caff 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageAssistant/ChatMessageAssistant.svelte @@ -82,8 +82,11 @@ let lastUserMessageHeight = $state(0); let assistantMarginTop = $state(0); + // The measured CSS vars feed the :last-child min-height rule only, so only + // the last assistant message needs them. Reading isLastAssistantMessage + // here also re-runs the effect when this message stops being the last. $effect(() => { - if (!assistantEl) return; + if (!assistantEl || !isLastAssistantMessage) return; assistantMarginTop = Math.round(parseFloat(getComputedStyle(assistantEl).marginTop)); diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte index a604a97e3..cc2b4a562 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlock.svelte @@ -13,7 +13,12 @@ import ChatMessageToolCallBlockWriteFile from './ChatMessageToolCallBlockWriteFile.svelte'; import { BuiltInTool } from '$lib/enums'; import type { AgenticSection, DatabaseMessageExtra } from '$lib/types'; - import { extractSearchQuery, extractSearchResults, isWebSearchToolName } from '$lib/utils'; + import { + extractSearchQuery, + extractSearchResults, + isWebSearchToolName, + looksLikeSearchResult + } from '$lib/utils'; interface Props { section: AgenticSection; @@ -26,11 +31,16 @@ let { attachments, isExecuting, isStreaming, onToggle, open, section }: Props = $props(); - const searchResults = $derived(extractSearchResults(section.toolResult)); - const searchQuery = $derived(extractSearchQuery(section.toolArgs)); - const isSearchCall = $derived( - searchResults.length > 0 || (searchQuery.length > 0 && isWebSearchToolName(section.toolName)) - ); + // Runs for every tool block on mount, before the body renders: the cheap + // content prefilter and the tool-name allow-list come first so blobs from + // exec/file tools are never line-split or JSON-parsed here + const isSearchCall = $derived.by(() => { + if (looksLikeSearchResult(section.toolResult)) { + return extractSearchResults(section.toolResult).length > 0; + } + + return isWebSearchToolName(section.toolName) && extractSearchQuery(section.toolArgs).length > 0; + }); {#if isSearchCall} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte index 2067e4268..22ffc256b 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockEditFile.svelte @@ -1,5 +1,5 @@ @@ -45,11 +49,11 @@ {meta.errorMessage}
- {:else if meta && meta.edits.length > 0} + {:else if meta && editFileBody && editFileBody.edits.length > 0} {#each editDiffs as diffLines, ei (ei)}
- Edit {ei + 1} of {meta.edits.length} + Edit {ei + 1} of {editFileBody.edits.length}
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockWriteFile.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockWriteFile.svelte index 178c479d9..cafa5280b 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockWriteFile.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ChatMessageToolCallBlockWriteFile.svelte @@ -1,5 +1,5 @@ @@ -45,7 +49,7 @@
{:else if meta} | null { } } +// Compiled per key on first use; the key set is tiny and fixed. +const toolArgStringRegexes = new Map(); + +/** + * Extract a string field from a JSON tool-args blob without parsing the + * whole document. write_file and edit_file args embed full file contents, + * yet the block title needs only the path; a targeted key match plus a + * JSON.parse of the captured string literal alone keeps title rendering + * O(path) instead of O(blob). Returns undefined when the key is missing + * or its value is not a string; callers fall back to the full parse. + */ +export function extractToolArgString( + toolArgs: string, + keys: readonly string[] +): string | undefined { + for (const key of keys) { + let pattern = toolArgStringRegexes.get(key); + + if (!pattern) { + pattern = new RegExp(TOOL_ARG_STRING_FIELD_PATTERN_TEMPLATE.replace('{key}', key)); + toolArgStringRegexes.set(key, pattern); + } + + const match = pattern.exec(toolArgs); + + if (!match) continue; + + try { + const value: unknown = JSON.parse(`"${match[1]}"`); + + if (typeof value === 'string') return value; + } catch { + // fall through to the next key; the full parse is the fallback + } + } + + return undefined; +} + /** * Parse a section's toolArgs against an expected tool name. Returns * `null` when: diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file.ts b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file.ts index 9ed6f92bc..d711466cb 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file.ts +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file.ts @@ -3,26 +3,12 @@ // rendering), plus the result blob for `result` / `edits_applied` / // `error` fields. -import { parseToolArgs } from './_shared'; -import { FILE_PATH_SEPARATOR_REGEX } from '$lib/constants'; +import { extractToolArgString, parseToolArgs } from './_shared'; +import { FILE_PATH_SEPARATOR_REGEX, TOOL_ARG_PATH_KEYS } from '$lib/constants'; import { BuiltInTool } from '$lib/enums'; -import type { AgenticSection } from '$lib/types'; +import type { AgenticSection, EditFileEdit, EditFileMeta, EditFileTitleMeta } from '$lib/types'; import { tryParseToolResultObject } from '$lib/utils'; -export type EditFileEdit = { - oldText: string; - newText: string; -}; - -export type EditFileMeta = { - fileName: string; - filePath: string; - edits: EditFileEdit[]; - resultMessage?: string; - editsApplied?: number; - errorMessage?: string; -}; - export function parseEditFileMeta(section: AgenticSection): EditFileMeta | null { const args = parseToolArgs(BuiltInTool.SERVER_EDIT_FILE, section, { partial: true }); @@ -79,3 +65,45 @@ export function parseEditFileMeta(section: AgenticSection): EditFileMeta | null resultMessage }; } + +/** + * Title-tier meta for edit_file blocks: everything the header and status + * pill render, obtained without parsing the embedded edit strings. The path + * comes from a targeted key extraction; the full parse runs only as a + * fallback for arg shapes the extraction can't see. + */ +export function parseEditFileTitleMeta(section: AgenticSection): EditFileTitleMeta | null { + if (section.toolName !== BuiltInTool.SERVER_EDIT_FILE || !section.toolArgs) return null; + + let rawPath: string | undefined = extractToolArgString(section.toolArgs, TOOL_ARG_PATH_KEYS); + + if (!rawPath) { + const args = parseToolArgs(BuiltInTool.SERVER_EDIT_FILE, section, { partial: true }); + const fallbackPath = args?.path ?? args?.file_path ?? args?.filePath; + + if (typeof fallbackPath === 'string' && fallbackPath) rawPath = fallbackPath; + } + + if (!rawPath) return null; + + const fileName = rawPath.split(FILE_PATH_SEPARATOR_REGEX).pop() || rawPath; + const resultObj = tryParseToolResultObject(section.toolResult); + + let resultMessage: string | undefined; + let editsApplied: number | undefined; + let errorMessage: string | undefined; + + if (typeof resultObj?.error === 'string') { + errorMessage = resultObj.error; + } else if (resultObj) { + if (typeof resultObj.result === 'string') { + resultMessage = resultObj.result; + } + + if (Number.isFinite(Number(resultObj.edits_applied))) { + editsApplied = Number(resultObj.edits_applied); + } + } + + return { editsApplied, errorMessage, fileName, filePath: rawPath, resultMessage }; +} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/run-javascript.ts b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/run-javascript.ts index 440a1f5d6..bd97cd2fe 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/run-javascript.ts +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/run-javascript.ts @@ -6,6 +6,7 @@ // are handled. import { parseToolArgs } from './_shared'; +import { JSON_ARRAY_OPEN, JSON_OBJECT_OPEN } from '$lib/constants'; import { BuiltInTool } from '$lib/enums'; import type { AgenticSection } from '$lib/types'; @@ -38,14 +39,21 @@ export function parseRunJavascriptMeta(section: AgenticSection): RunJavascriptMe // do we scan raw lines for the `Error:` prefix. let parsedObject: Record | null = null; - try { - const parsed: unknown = JSON.parse(toolResultString); + // Successful sandbox output is a JSON array, errors are objects; plain + // text (huge console logs) fails the parse below anyway, so only try + // when the blob starts with a JSON container + const trimmedResult = toolResultString.trimStart(); - if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) { - parsedObject = parsed as Record; + if (trimmedResult[0] === JSON_OBJECT_OPEN || trimmedResult[0] === JSON_ARRAY_OPEN) { + try { + const parsed: unknown = JSON.parse(trimmedResult); + + if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) { + parsedObject = parsed as Record; + } + } catch { + parsedObject = null; } - } catch { - parsedObject = null; } if (typeof parsedObject?.error === 'string') { diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/write-file.ts b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/write-file.ts index 5b9bf9f88..4a8e1a9c9 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/write-file.ts +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/write-file.ts @@ -3,22 +3,12 @@ // finishes) and surfaces `bytes`, `result`, and `error` from the // result blob. -import { parseToolArgs } from './_shared'; -import { CODE_BLOCK, FILE_PATH_SEPARATOR_REGEX } from '$lib/constants'; +import { extractToolArgString, parseToolArgs } from './_shared'; +import { CODE_BLOCK, FILE_PATH_SEPARATOR_REGEX, TOOL_ARG_PATH_KEYS } from '$lib/constants'; import { BuiltInTool } from '$lib/enums'; -import type { AgenticSection } from '$lib/types'; +import type { AgenticSection, WriteFileMeta, WriteFileTitleMeta } from '$lib/types'; import { getFileTypeByExtension, tryParseToolResultObject } from '$lib/utils'; -export type WriteFileMeta = { - fileName: string; - filePath: string; - language: string; - content: string; - bytesWritten?: number; - resultMessage?: string; - errorMessage?: string; -}; - export function parseWriteFileMeta(section: AgenticSection): WriteFileMeta | null { const args = parseToolArgs(BuiltInTool.SERVER_WRITE_FILE, section, { partial: true }); @@ -51,3 +41,43 @@ export function parseWriteFileMeta(section: AgenticSection): WriteFileMeta | nul resultMessage }; } + +/** + * Title-tier meta for write_file blocks: everything the header and status + * pill render, obtained without parsing the embedded file content. The path + * comes from a targeted key extraction; the full parse runs only as a + * fallback for arg shapes the extraction can't see. + */ +export function parseWriteFileTitleMeta(section: AgenticSection): WriteFileTitleMeta | null { + if (section.toolName !== BuiltInTool.SERVER_WRITE_FILE || !section.toolArgs) return null; + + let rawPath: string | undefined = extractToolArgString(section.toolArgs, TOOL_ARG_PATH_KEYS); + + if (!rawPath) { + const args = parseToolArgs(BuiltInTool.SERVER_WRITE_FILE, section, { partial: true }); + const fallbackPath = args?.path ?? args?.file_path ?? args?.filePath; + + if (typeof fallbackPath === 'string' && fallbackPath) rawPath = fallbackPath; + } + + if (!rawPath) return null; + + const fileName = rawPath.split(FILE_PATH_SEPARATOR_REGEX).pop() || rawPath; + const language = + getFileTypeByExtension(rawPath)?.replace(CODE_BLOCK.TEXT_LANGUAGE_PREFIX_REGEX, '') ?? + CODE_BLOCK.DEFAULT_LANGUAGE; + const resultObj = tryParseToolResultObject(section.toolResult); + const bytesWritten = + resultObj && Number.isFinite(Number(resultObj.bytes)) ? Number(resultObj.bytes) : undefined; + const resultMessage = typeof resultObj?.result === 'string' ? resultObj.result : undefined; + const errorMessage = typeof resultObj?.error === 'string' ? resultObj.error : undefined; + + return { + bytesWritten, + errorMessage, + fileName, + filePath: rawPath, + language, + resultMessage + }; +} diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte index 5137e261f..ea9428e07 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageAgenticContent.svelte @@ -46,49 +46,44 @@ isLastAssistantMessage ? !!agenticStore.getLastError(message.convId) : false ); - let permissionDismissed = $state(false); - const pendingPermission = $derived( isStreaming && isLastAssistantMessage ? agenticStore.getPendingPermissionRequest(message.convId) : null ); - let prevPendingRef: typeof pendingPermission = null; - $effect(() => { - if (pendingPermission !== prevPendingRef) { - prevPendingRef = pendingPermission; + // dismissal applies to the request object, so the next request ( new + // identity ) shows the card again without any reset bookkeeping + let dismissedPermission: typeof pendingPermission = $state(null); - if (pendingPermission) { - permissionDismissed = false; - } - } - }); + const visiblePermission = $derived( + pendingPermission && dismissedPermission !== pendingPermission ? pendingPermission : null + ); function handlePermission(decision: ToolPermissionDecision) { - permissionDismissed = true; + dismissedPermission = pendingPermission; agenticStore.resolvePermission(message.convId, decision); } - let continueDismissed = $state(false); - const pendingContinue = $derived( isStreaming && isLastAssistantMessage ? agenticStore.getPendingContinueRequest(message.convId) : false ); - let prevContinueRef = false; - $effect(() => { - if (pendingContinue !== prevContinueRef) { - prevContinueRef = pendingContinue; + let continueDismissed = $state(false); - if (pendingContinue) { - continueDismissed = false; - } + // the continue request is a plain boolean, so there is no identity to + // compare against; clear the dismissal whenever no request is pending so + // the next one starts from a clean state + $effect(() => { + if (!pendingContinue) { + continueDismissed = false; } }); + const showContinue = $derived(Boolean(pendingContinue) && !continueDismissed); + function handleContinue(shouldContinue: boolean) { continueDismissed = true; agenticStore.resolveContinue(message.convId, shouldContinue); @@ -238,15 +233,15 @@ {/each} {/if} - {#if pendingPermission && !permissionDismissed} + {#if visiblePermission} {/if} - {#if pendingContinue && !continueDismissed} + {#if showContinue} {/if}
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte index 4750a9f7c..0078225c0 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessages.svelte @@ -1,5 +1,6 @@ -
- {#each displayMessages as { isLastAssistantMessage, isLastUserMessage, message, nextAssistantMessage, siblingInfo, toolMessages } (message.id)} - - {/each} - - {#if conversationsStore.activeConversation && agenticStore.getPendingSteeringMessageContent(conversationsStore.activeConversation!.id)} - {@const convId = conversationsStore.activeConversation!.id} - {@const pendingContent = agenticStore.getPendingSteeringMessageContent(convId)} - - {#if pendingContent} - agenticStore.clearSteeringMessage(convId)} - onEdit={(newContent, extras) => - agenticStore.injectSteeringMessage(convId, newContent, extras)} - onSendImmediately={() => chatStore.abortCurrentFlow(convId)} + +{#key conversationsStore.activeConversation?.id ?? 'new'} +
+ {#each displayMessages as { isLastAssistantMessage, isLastUserMessage, message, nextAssistantMessage, siblingInfo, toolMessages } (message.id)} + - {/if} - {:else if conversationsStore.activeConversation && chatStore.getPendingMessageContent(conversationsStore.activeConversation!.id)} - {@const convId = conversationsStore.activeConversation!.id} - {@const pendingContent = chatStore.getPendingMessageContent(convId)} + {/each} - {#if pendingContent} - chatStore.clearPendingMessage(convId)} - onEdit={(newContent, extras) => chatStore.injectPendingMessage(convId, newContent, extras)} - onSendImmediately={() => chatStore.abortCurrentFlow(convId)} - /> + {#if conversationsStore.activeConversation && agenticStore.getPendingSteeringMessageContent(conversationsStore.activeConversation!.id)} + {@const convId = conversationsStore.activeConversation!.id} + {@const pendingContent = agenticStore.getPendingSteeringMessageContent(convId)} + + {#if pendingContent} + agenticStore.clearSteeringMessage(convId)} + onEdit={(newContent, extras) => + agenticStore.injectSteeringMessage(convId, newContent, extras)} + onSendImmediately={() => chatStore.abortCurrentFlow(convId)} + /> + {/if} + {:else if conversationsStore.activeConversation && chatStore.getPendingMessageContent(conversationsStore.activeConversation!.id)} + {@const convId = conversationsStore.activeConversation!.id} + {@const pendingContent = chatStore.getPendingMessageContent(convId)} + + {#if pendingContent} + chatStore.clearPendingMessage(convId)} + onEdit={(newContent, extras) => + chatStore.injectPendingMessage(convId, newContent, extras)} + onSendImmediately={() => chatStore.abortCurrentFlow(convId)} + /> + {/if} {/if} - {/if} -
+
+{/key} + + diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/LazyChatMessage.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/LazyChatMessage.svelte new file mode 100644 index 000000000..f9667bbbb --- /dev/null +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/LazyChatMessage.svelte @@ -0,0 +1,105 @@ + + +
+ {#if mounted} + + {/if} +
+ + diff --git a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte index 3ad3f2468..6cea95d0d 100644 --- a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte @@ -315,13 +315,18 @@
bottomed move with transform, not bottom: + // layout-property transitions need the main thread every frame and + // stutter while a long conversation loads; transform transitions + // run on the compositor and stay smooth + 'pointer-events-none md:sticky fixed mt-auto transition-transform duration-200', deviceStore.isStandalone ? 'bottom-6 right-4 left-4' : deviceStore.isIOSSafari ? 'bottom-1 left-2 right-2' : 'bottom-2 right-2 left-2', - isEmpty ? 'md:bottom-[calc(50dvh-7rem)] 2xl:bottom-[calc(50dvh-4rem)]' : 'md:bottom-4' + 'md:bottom-4', + isEmpty ? 'md:translate-y-[calc(-50dvh+8rem)] 2xl:translate-y-[calc(-50dvh+5rem)]' : '' ]} > diff --git a/tools/ui/src/lib/components/app/content/MarkdownContent/MarkdownContent.svelte b/tools/ui/src/lib/components/app/content/MarkdownContent/MarkdownContent.svelte index 87b41bd00..c217a769a 100644 --- a/tools/ui/src/lib/components/app/content/MarkdownContent/MarkdownContent.svelte +++ b/tools/ui/src/lib/components/app/content/MarkdownContent/MarkdownContent.svelte @@ -1,23 +1,12 @@ diff --git a/tools/ui/src/lib/components/app/content/MarkdownContent/markdown-processor.ts b/tools/ui/src/lib/components/app/content/MarkdownContent/markdown-processor.ts new file mode 100644 index 000000000..e973a6a4b --- /dev/null +++ b/tools/ui/src/lib/components/app/content/MarkdownContent/markdown-processor.ts @@ -0,0 +1,112 @@ +// Shared remark/rehype pipeline factory for MarkdownContent. +// +// The frozen plugin chain is expensive to build ( ~15 plugin instances ), +// and MarkdownContent used to rebuild it on every processMarkdown call: +// once per block at mount, and again on every coalesced chunk while +// streaming. Pipelines without attachments are shared process-wide per +// math flag; attachment-bearing pipelines are cached by the attachments +// array identity, which changes whenever extras are updated. + +import { rehypeEnhanceCodeBlocks } from './plugins/rehype/enhance-code-blocks'; +import { rehypeEnhanceLinks } from './plugins/rehype/enhance-links'; +import { rehypeEnhanceMermaidBlocks } from './plugins/rehype/enhance-mermaid-blocks'; +import { rehypeEnhanceSvgBlocks } from './plugins/rehype/enhance-svg-blocks'; +import { rehypeFileBadge } from './plugins/rehype/file-badge'; +import { rehypeMermaidPre } from './plugins/rehype/mermaid-pre'; +import { rehypeRtlSupport } from './plugins/rehype/rehype-rtl-support'; +import { rehypeResolveAttachmentImages } from './plugins/rehype/resolve-attachment-images'; +import { rehypeSvgPre } from './plugins/rehype/svg-pre'; +import { rehypeRestoreTableHtml } from './plugins/rehype/table-html-restorer'; +import { remarkLiteralHtml } from './plugins/remark/literal-html'; +import { FileTypeText } from '$lib/enums/files.enums'; +import type { DatabaseMessageExtra } from '$lib/types/database'; +import type { Root as HastRoot } from 'hast'; +import { all as lowlightAll } from 'lowlight'; +import type { Root as MdastRoot } from 'mdast'; +import rehypeHighlight from 'rehype-highlight'; +import rehypeKatex from 'rehype-katex'; +import rehypeStringify from 'rehype-stringify'; +import { remark } from 'remark'; +import remarkBreaks from 'remark-breaks'; +import remarkGfm from 'remark-gfm'; +import remarkMath from 'remark-math'; +import remarkRehype from 'remark-rehype'; + +export interface MarkdownProcessor { + parse(markdown: string): MdastRoot; + run(tree: MdastRoot): Promise; + stringify(tree: HastRoot): string; +} + +export interface MarkdownProcessorOptions { + attachments?: DatabaseMessageExtra[]; + disableMath?: boolean; +} + +const sharedPipelines = new Map(); +const attachmentPipelines = new WeakMap(); + +function buildPipeline({ + attachments, + disableMath = false +}: MarkdownProcessorOptions): MarkdownProcessor { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + let proc: any = remark().use(remarkGfm); // GitHub Flavored Markdown + + if (!disableMath) { + proc = proc.use(remarkMath); // Parse $inline$ and $$block$$ math + } + + proc = proc + .use(remarkBreaks) // Convert line breaks to
+ // Treat raw HTML as literal text with preserved indentation + .use(remarkLiteralHtml) + .use(remarkRehype); // Convert Markdown AST to rehype + + if (!disableMath) { + proc = proc.use(rehypeKatex); // Render math using KaTeX + } + + const pipeline = proc + .use(rehypeHighlight, { + aliases: { [FileTypeText.XML]: [FileTypeText.SVELTE, FileTypeText.VUE] }, + languages: lowlightAll + }) // Add syntax highlighting + .use(rehypeRestoreTableHtml) // Restore limited HTML (e.g.
,
    ) inside Markdown tables + .use(rehypeEnhanceLinks) // Add target="_blank" to links + .use(rehypeFileBadge) // Render file:// anchors as inline badge chips + .use(rehypeMermaidPre) // Convert mermaid blocks to
    +		.use(rehypeSvgPre) // Convert svg blocks to 
    +		.use(rehypeEnhanceCodeBlocks) // Wrap code blocks with header and actions
    +		.use(rehypeEnhanceMermaidBlocks) // Wrap mermaid blocks with header and actions
    +		.use(rehypeEnhanceSvgBlocks) // Wrap svg blocks with header and actions
    +		.use(rehypeResolveAttachmentImages, { attachments })
    +		.use(rehypeRtlSupport) // Add bidirectional text support
    +		.use(rehypeStringify, { allowDangerousHtml: true }); // Convert to HTML string
    +
    +	return pipeline as MarkdownProcessor;
    +}
    +
    +export function getMarkdownProcessor(options: MarkdownProcessorOptions): MarkdownProcessor {
    +	if (options.attachments && options.attachments.length > 0) {
    +		let cached = attachmentPipelines.get(options.attachments);
    +
    +		if (!cached) {
    +			cached = buildPipeline(options);
    +			attachmentPipelines.set(options.attachments, cached);
    +		}
    +
    +		return cached;
    +	}
    +
    +	const key = String(Boolean(options.disableMath));
    +
    +	let cached = sharedPipelines.get(key);
    +
    +	if (!cached) {
    +		cached = buildPipeline(options);
    +		sharedPipelines.set(key, cached);
    +	}
    +
    +	return cached;
    +}
    diff --git a/tools/ui/src/lib/constants/index.ts b/tools/ui/src/lib/constants/index.ts
    index e3241373e..d93ae6429 100644
    --- a/tools/ui/src/lib/constants/index.ts
    +++ b/tools/ui/src/lib/constants/index.ts
    @@ -16,6 +16,7 @@ export * from './context-gauge-popup.constants';
     export * from './conversation-import.constants';
     export * from './binary-detection.constants';
     export * from './content-detection.constants';
    +export * from './tool-call-args.constants';
     export * from './tool-ui.constants';
     export * from './cache.constants';
     export * from './chat-form.constants';
    diff --git a/tools/ui/src/lib/constants/tool-call-args.constants.ts b/tools/ui/src/lib/constants/tool-call-args.constants.ts
    new file mode 100644
    index 000000000..e74260be2
    --- /dev/null
    +++ b/tools/ui/src/lib/constants/tool-call-args.constants.ts
    @@ -0,0 +1,23 @@
    +// Tool-args and tool-result parsing helpers: the file tools' path field
    +// aliases, the JSON container gates for result blobs, and the targeted
    +// string-field pattern used for cheap title-tier extraction.
    +
    +/**
    + * Field aliases the file tools accept for the path argument. Tool contracts
    + * drifted over time: some models emit `file_path` / `filePath`.
    + */
    +export const TOOL_ARG_PATH_KEYS: readonly string[] = ['path', 'file_path', 'filePath'];
    +
    +/** Opening character of a JSON object; only an object root can carry fields. */
    +export const JSON_OBJECT_OPEN = '{';
    +
    +/** Opening character of a JSON array; successful sandbox output is one. */
    +export const JSON_ARRAY_OPEN = '[';
    +
    +/**
    + * Matches `"": ""` in a JSON args blob ( whitespace between
    + * tokens allowed ), capturing the raw string literal so only that literal
    + * gets decoded; escaped quotes stay inside the value group. `{key}` is
    + * replaced with the field name before use.
    + */
    +export const TOOL_ARG_STRING_FIELD_PATTERN_TEMPLATE = '"{key}"\\s*:\\s*"((?:[^"\\\\]|\\\\.)*)"';
    diff --git a/tools/ui/src/lib/stores/chat/index.svelte.ts b/tools/ui/src/lib/stores/chat/index.svelte.ts
    index 296c2cca5..4bdcc6845 100644
    --- a/tools/ui/src/lib/stores/chat/index.svelte.ts
    +++ b/tools/ui/src/lib/stores/chat/index.svelte.ts
    @@ -55,7 +55,6 @@ class ChatStore implements ChatStreamHost, ChatFlowsHost {
     		string,
     		{ response: string; messageId: string; model?: string | null }
     	>();
    -	currentResponse = $state('');
     	errorDialogState = $state(null);
     	// true while the active conversation has a local pipe (send, attach or resume-wait)
     	isLoading = $derived(this.activity.isLocal(conversationsStore.activeConversation?.id ?? ''));
    @@ -256,8 +255,6 @@ class ChatStore implements ChatStreamHost, ChatFlowsHost {
     		}
     
     		this.chatStreamingStates.delete(convId);
    -
    -		if (convId === conversationsStore.activeConversation?.id) this.currentResponse = '';
     	}
     	clearEditMode(): void {
     		this.isEditModeActive = false;
    @@ -272,11 +269,6 @@ class ChatStore implements ChatStreamHost, ChatFlowsHost {
     		this.pendingMessages.delete(convId);
     	}
     
    -	/** Reset per-view state when (re)mounting the empty chat screen. */
    -	clearUIState(): void {
    -		this.currentResponse = '';
    -	}
    -
     	consumePendingDraft(): { message: string; files: ChatUploadedFile[] } | null {
     		if (!this.pendingDraftMessage && this.pendingDraftFiles.length === 0) return null;
     
    @@ -766,8 +758,6 @@ class ChatStore implements ChatStreamHost, ChatFlowsHost {
     			model: model ?? this.chatStreamingStates.get(convId)?.model,
     			response
     		});
    -
    -		if (convId === conversationsStore.activeConversation?.id) this.currentResponse = response;
     	}
     
     	setEditModeActive(handler: (files: File[]) => void): void {
    @@ -1244,7 +1234,6 @@ class ChatStore implements ChatStreamHost, ChatFlowsHost {
     	syncLoadingStateForChat(convId: string): void {
     		const s = this.chatStreamingStates.get(convId);
     
    -		this.currentResponse = s?.response || '';
     		this.processing.setActiveConversation(convId);
     
     		// Sync streaming content to activeMessages so UI displays current content
    diff --git a/tools/ui/src/lib/stores/conversations/index.svelte.ts b/tools/ui/src/lib/stores/conversations/index.svelte.ts
    index df5b1ecef..c4fea2e4e 100644
    --- a/tools/ui/src/lib/stores/conversations/index.svelte.ts
    +++ b/tools/ui/src/lib/stores/conversations/index.svelte.ts
    @@ -52,6 +52,13 @@ class ConversationsStore implements ConversationsPreferencesHost {
     	/** In-flight init run; shared by concurrent callers, reset on failure to allow retry */
     	private initPromise: Promise | null = null;
     
    +	/**
    +	 * Messages loadConversation just read, handed off once so the chat
    +	 * screen can reuse them for sibling info instead of re-fetching the
    +	 * whole conversation a second time.
    +	 */
    +	private lastLoadedMessages: { convId: string; messages: DatabaseMessage[] } | null = null;
    +
     	/**
     	 * Memo of the last findMessageIndex() lookup. Streaming calls it once per
     	 * chunk for the same message, so a validated cache hit keeps that O(1)
    @@ -88,7 +95,13 @@ class ConversationsStore implements ConversationsPreferencesHost {
     		}
     
     		if (this.activeConversation?.id === id) {
    -			this.activeConversation = { ...this.activeConversation, ...updates };
    +			// field-wise, not object replacement: effects that track the active
    +			// conversation identity would otherwise refire on every rename or pin
    +			const target = this.activeConversation as unknown as Record;
    +
    +			for (const [key, value] of Object.entries(updates)) {
    +				if (target[key] !== value) target[key] = value;
    +			}
     		}
     	}
     
    @@ -202,11 +215,8 @@ class ConversationsStore implements ConversationsPreferencesHost {
     			const updates = await DatabaseService.bulkToggleConversationPins(convIds);
     			const activeId = this.activeConversation?.id;
     
    -			if (activeId && updates.has(activeId)) {
    -				this.activeConversation = {
    -					...this.activeConversation!,
    -					pinned: updates.get(activeId)!
    -				};
    +			if (this.activeConversation && activeId && updates.has(activeId)) {
    +				this.activeConversation.pinned = updates.get(activeId)!;
     			}
     
     			for (let i = 0; i < this.conversations.length; i++) {
    @@ -236,6 +246,17 @@ class ConversationsStore implements ConversationsPreferencesHost {
     		this.preferences.resetPending();
     	}
     
    +	/** One-shot handoff of the messages the last loadConversation read. */
    +	consumeLastLoadedMessages(convId: string): DatabaseMessage[] | null {
    +		if (this.lastLoadedMessages?.convId !== convId) return null;
    +
    +		const messages = this.lastLoadedMessages.messages;
    +
    +		this.lastLoadedMessages = null;
    +
    +		return messages;
    +	}
    +
     	/**
     	 * Creates a new conversation and navigates to it
     	 * @param name - Optional name for the conversation
    @@ -509,22 +530,15 @@ class ConversationsStore implements ConversationsPreferencesHost {
     			// it doesn't belong to this conversation.
     			this.preferences.pendingCwd = null;
     
    +			const allMessages = await DatabaseService.getConversationMessages(convId);
    +
    +			// set conversation and messages in one sync block so effects never see
    +			// the new conversation with the previous conversation's messages
    +			this.lastLoadedMessages = { convId, messages: allMessages };
     			this.activeConversation = conversation;
    -
    -			if (conversation.currNode) {
    -				const allMessages = await DatabaseService.getConversationMessages(convId);
    -				const filteredMessages = filterByLeafNodeId(
    -					allMessages,
    -					conversation.currNode,
    -					false
    -				) as DatabaseMessage[];
    -
    -				this.activeMessages = filteredMessages;
    -			} else {
    -				const messages = await DatabaseService.getConversationMessages(convId);
    -
    -				this.activeMessages = messages;
    -			}
    +			this.activeMessages = conversation.currNode
    +				? (filterByLeafNodeId(allMessages, conversation.currNode, false) as DatabaseMessage[])
    +				: allMessages;
     
     			return true;
     		} catch (error) {
    @@ -558,7 +572,7 @@ class ConversationsStore implements ConversationsPreferencesHost {
     		const currentLeafNodeId = findLeafNode(allMessages, siblingId);
     
     		await DatabaseService.updateCurrentNode(this.activeConversation.id, currentLeafNodeId);
    -		this.activeConversation = { ...this.activeConversation, currNode: currentLeafNodeId };
    +		this.activeConversation.currNode = currentLeafNodeId;
     		await this.refreshActiveMessages();
     
     		if (rootMessage && this.activeMessages.length > 0) {
    @@ -694,7 +708,7 @@ class ConversationsStore implements ConversationsPreferencesHost {
     		}
     
     		if (this.activeConversation?.id === targetId) {
    -			this.activeConversation = { ...this.activeConversation, lastModified: now };
    +			this.activeConversation.lastModified = now;
     		}
     
     		DatabaseService.updateConversation(targetId, { lastModified: now }).catch((error) =>
    @@ -710,7 +724,7 @@ class ConversationsStore implements ConversationsPreferencesHost {
     		if (!this.activeConversation) return;
     
     		await DatabaseService.updateCurrentNode(this.activeConversation.id, nodeId);
    -		this.activeConversation = { ...this.activeConversation, currNode: nodeId };
    +		this.activeConversation.currNode = nodeId;
     	}
     
     	/**
    diff --git a/tools/ui/src/lib/types/index.ts b/tools/ui/src/lib/types/index.ts
    index d91c2811a..333c1bd3c 100644
    --- a/tools/ui/src/lib/types/index.ts
    +++ b/tools/ui/src/lib/types/index.ts
    @@ -209,7 +209,16 @@ export type {
     export type { DesktopIconStripItem } from './navigation';
     
     // Tools types
    -export type { ToolEntry, ToolGroup, ToolUiEntry } from './tools';
    +export type {
    +	EditFileEdit,
    +	EditFileMeta,
    +	EditFileTitleMeta,
    +	ToolEntry,
    +	ToolGroup,
    +	ToolUiEntry,
    +	WriteFileMeta,
    +	WriteFileTitleMeta
    +} from './tools';
     
     // Reasoning
     export type { ReasoningEffortLevel } from './reasoning';
    diff --git a/tools/ui/src/lib/types/tools.d.ts b/tools/ui/src/lib/types/tools.d.ts
    index edcec65c7..fa8963bd1 100644
    --- a/tools/ui/src/lib/types/tools.d.ts
    +++ b/tools/ui/src/lib/types/tools.d.ts
    @@ -31,3 +31,50 @@ export interface ToolGroup {
     	serverId?: string;
     	tools: ToolEntry[];
     }
    +
    +export interface WriteFileMeta {
    +	fileName: string;
    +	filePath: string;
    +	language: string;
    +	content: string;
    +	bytesWritten?: number;
    +	resultMessage?: string;
    +	errorMessage?: string;
    +}
    +
    +/** Everything the write_file block title and status pill show; the full meta
    + *  ( with the embedded file content ) stays body-only so collapsed blocks
    + *  never parse the content blob. */
    +export interface WriteFileTitleMeta {
    +	fileName: string;
    +	filePath: string;
    +	language: string;
    +	bytesWritten?: number;
    +	resultMessage?: string;
    +	errorMessage?: string;
    +}
    +
    +export interface EditFileEdit {
    +	oldText: string;
    +	newText: string;
    +}
    +
    +export interface EditFileMeta {
    +	fileName: string;
    +	filePath: string;
    +	edits: EditFileEdit[];
    +	resultMessage?: string;
    +	editsApplied?: number;
    +	errorMessage?: string;
    +}
    +
    +/** Everything the edit_file block title and status pill show; the full meta
    + *  ( with the embedded edit strings ) stays body-only so collapsed blocks
    + *  never parse the args blob. */
    +export interface EditFileTitleMeta {
    +	fileName: string;
    +	filePath: string;
    +	resultMessage?: string;
    +	editsApplied?: number;
    +	errorMessage?: string;
    +}
    diff --git a/tools/ui/src/lib/utils/agentic.ts b/tools/ui/src/lib/utils/agentic.ts
    index cd150c5ef..28b3f43ee 100644
    --- a/tools/ui/src/lib/utils/agentic.ts
    +++ b/tools/ui/src/lib/utils/agentic.ts
    @@ -109,6 +109,89 @@ function deriveSingleTurnSections(
     	return sections;
     }
     
    +interface TurnSectionsCacheEntry {
    +	content: string | undefined;
    +	extra: DatabaseMessageExtra[] | undefined;
    +	reasoningContent: string | undefined;
    +	toolCalls: string | undefined;
    +	toolMessageContents: (string | undefined)[];
    +	toolMessageExtras: (DatabaseMessageExtra[] | undefined)[];
    +	toolMessages: DatabaseMessage[];
    +	sections: AgenticSection[];
    +}
    +
    +const turnSectionsCache = new WeakMap();
    +
    +function isTurnCacheValid(
    +	entry: TurnSectionsCacheEntry,
    +	message: DatabaseMessage,
    +	toolMessages: DatabaseMessage[]
    +): boolean {
    +	if (
    +		entry.content !== message.content ||
    +		entry.reasoningContent !== message.reasoningContent ||
    +		entry.toolCalls !== message.toolCalls ||
    +		entry.extra !== message.extra
    +	) {
    +		return false;
    +	}
    +
    +	if (entry.toolMessages.length !== toolMessages.length) return false;
    +
    +	for (let i = 0; i < toolMessages.length; i++) {
    +		if (entry.toolMessages[i] !== toolMessages[i]) return false;
    +
    +		if (entry.toolMessageContents[i] !== toolMessages[i].content) return false;
    +
    +		if (entry.toolMessageExtras[i] !== toolMessages[i].extra) return false;
    +	}
    +
    +	return true;
    +}
    +
    +/**
    + * deriveSingleTurnSections with structural reuse for completed turns.
    + *
    + * deriveAgenticSections runs in a $derived invalidated per streamed chunk, but
    + * only the last turn actually changes. Messages mutate in place and are never
    + * replaced, so a WeakMap keyed by the turn's assistant message plus reference
    + * checks on every field deriveSingleTurnSections reads detects any change. A
    + * cache hit also returns the same section objects, keeping downstream props
    + * stable so tool blocks skip their per-chunk re-derive. The streaming turn
    + * recomputes uncached on every chunk.
    + */
    +function deriveTurnSections(
    +	message: DatabaseMessage,
    +	toolMessages: DatabaseMessage[],
    +	streamingToolCalls: ApiChatCompletionToolCall[],
    +	isStreaming: boolean
    +): AgenticSection[] {
    +	if (isStreaming || streamingToolCalls.length > 0) {
    +		return deriveSingleTurnSections(message, toolMessages, streamingToolCalls, isStreaming);
    +	}
    +
    +	const cached = turnSectionsCache.get(message);
    +
    +	if (cached && isTurnCacheValid(cached, message, toolMessages)) {
    +		return cached.sections;
    +	}
    +
    +	const sections = deriveSingleTurnSections(message, toolMessages, [], false);
    +
    +	turnSectionsCache.set(message, {
    +		content: message.content,
    +		extra: message.extra,
    +		reasoningContent: message.reasoningContent,
    +		sections,
    +		toolCalls: message.toolCalls,
    +		toolMessageContents: toolMessages.map((tm) => tm.content),
    +		toolMessageExtras: toolMessages.map((tm) => tm.extra),
    +		toolMessages
    +	});
    +
    +	return sections;
    +}
    +
     /**
      * Derives display sections from structured message data.
      *
    @@ -132,13 +215,13 @@ export function deriveAgenticSections(
     	const hasAssistantContinuations = toolMessages.some((m) => m.role === MessageRole.ASSISTANT);
     
     	if (!hasAssistantContinuations) {
    -		return deriveSingleTurnSections(message, toolMessages, streamingToolCalls, isStreaming);
    +		return deriveTurnSections(message, toolMessages, streamingToolCalls, isStreaming);
     	}
     
     	const sections: AgenticSection[] = [];
     	const firstTurnToolMsgs = collectToolMessages(toolMessages, 0);
     
    -	sections.push(...deriveSingleTurnSections(message, firstTurnToolMsgs));
    +	sections.push(...deriveTurnSections(message, firstTurnToolMsgs, [], false));
     
     	let i = firstTurnToolMsgs.length;
     
    @@ -150,7 +233,7 @@ export function deriveAgenticSections(
     			const isLastTurn = i + 1 + turnToolMsgs.length >= toolMessages.length;
     
     			sections.push(
    -				...deriveSingleTurnSections(
    +				...deriveTurnSections(
     					msg,
     					turnToolMsgs,
     					isLastTurn ? streamingToolCalls : [],
    diff --git a/tools/ui/src/lib/utils/branching.ts b/tools/ui/src/lib/utils/branching.ts
    index 6c2c895cb..43d33d424 100644
    --- a/tools/ui/src/lib/utils/branching.ts
    +++ b/tools/ui/src/lib/utils/branching.ts
    @@ -105,18 +105,34 @@ export function filterByLeafNodeId(
      */
     function findLeafNodeInMap(
     	nodeMap: ReadonlyMap,
    -	messageId: string
    +	messageId: string,
    +	leafCache?: Map
     ): string {
    +	const path: string[] = [];
    +
     	let currentNode: DatabaseMessage | undefined = nodeMap.get(messageId);
     
     	while (currentNode && currentNode.children.length > 0) {
     		// Follow the last child (most recent branch)
    +		const cached = leafCache?.get(currentNode.id);
    +
    +		if (cached !== undefined) {
    +			for (const id of path) leafCache?.set(id, cached);
    +
    +			return cached;
    +		}
    +
    +		path.push(currentNode.id);
     		const lastChildId = currentNode.children[currentNode.children.length - 1];
     
     		currentNode = nodeMap.get(lastChildId);
     	}
     
    -	return currentNode?.id ?? messageId;
    +	const leafId = currentNode?.id ?? messageId;
    +
    +	for (const id of path) leafCache?.set(id, leafId);
    +
    +	return leafId;
     }
     
     /**
    @@ -176,7 +192,8 @@ export function findDescendantMessages(
      */
     export function getMessageSiblings(
     	nodeMap: ReadonlyMap,
    -	messageId: string
    +	messageId: string,
    +	leafCache?: Map
     ): ChatMessageSiblingInfo | null {
     	const message = nodeMap.get(messageId);
     
    @@ -212,7 +229,7 @@ export function getMessageSiblings(
     	// Convert sibling message IDs to their corresponding leaf node IDs
     	// This allows navigation between different conversation branches
     	const siblingLeafIds = siblingIds.map((siblingId: string) =>
    -		findLeafNodeInMap(nodeMap, siblingId)
    +		findLeafNodeInMap(nodeMap, siblingId, leafCache)
     	);
     	// Find current message's position among siblings
     	const currentIndex = siblingIds.indexOf(messageId);
    @@ -236,9 +253,12 @@ export function buildSiblingInfoMap(
     ): Map {
     	const nodeMap = new Map(messages.map((msg) => [msg.id, msg] as const));
     	const siblingMap = new Map();
    +	// Leaf walks repeat along the same child chains for every message; memoize
    +	// them per build so each edge is walked once instead of O(messages^2)
    +	const leafCache = new Map();
     
     	for (const msg of messages) {
    -		const info = getMessageSiblings(nodeMap, msg.id);
    +		const info = getMessageSiblings(nodeMap, msg.id, leafCache);
     
     		if (info) {
     			siblingMap.set(msg.id, info);
    diff --git a/tools/ui/src/lib/utils/index.ts b/tools/ui/src/lib/utils/index.ts
    index 079cdc871..721618c48 100644
    --- a/tools/ui/src/lib/utils/index.ts
    +++ b/tools/ui/src/lib/utils/index.ts
    @@ -285,7 +285,8 @@ export {
     	extractSearchResults,
     	extractSearchQuery,
     	faviconForUrl,
    -	isWebSearchToolName
    +	isWebSearchToolName,
    +	looksLikeSearchResult
     } from './search-results';
     
     // Cache utilities
    diff --git a/tools/ui/src/lib/utils/parse-exec-shell-error.ts b/tools/ui/src/lib/utils/parse-exec-shell-error.ts
    index 42d2ee254..a7b2eb5c8 100644
    --- a/tools/ui/src/lib/utils/parse-exec-shell-error.ts
    +++ b/tools/ui/src/lib/utils/parse-exec-shell-error.ts
    @@ -3,8 +3,14 @@ export function parseExecShellCommandError(
     ): string | undefined {
     	if (!toolResultString) return undefined;
     
    +	// Exec results are usually large plain-text stdout; only a JSON object
    +	// root can carry an error field, so skip the parse otherwise
    +	const trimmed = toolResultString.trimStart();
    +
    +	if (trimmed[0] !== '{') return undefined;
    +
     	try {
    -		const parsed: unknown = JSON.parse(toolResultString);
    +		const parsed: unknown = JSON.parse(trimmed);
     
     		if (
     			parsed &&
    diff --git a/tools/ui/src/lib/utils/parse-exec-shell-status.ts b/tools/ui/src/lib/utils/parse-exec-shell-status.ts
    index 1f7ec557e..71dd110bd 100644
    --- a/tools/ui/src/lib/utils/parse-exec-shell-status.ts
    +++ b/tools/ui/src/lib/utils/parse-exec-shell-status.ts
    @@ -15,15 +15,18 @@ export interface ExecShellExitStatus {
     }
     
     // Anchor to the absolute end so intermediate "[exit code: N]" string content
    -// (e.g. a shell echo) doesn't false-positive.
    +// (e.g. a shell echo) doesn't false-positive. The marker is at most ~50 chars
    +// with the timed-out suffix, so matching a tail slice keeps the cost constant
    +// for megabyte exec outputs instead of scanning the whole blob.
     const EXIT_CODE_TAIL_REGEX = /\[exit code: (-?\d+)\](?: \[exit due to timed out\])?\s*$/;
    +const EXIT_CODE_TAIL_SCAN = 128;
     
     export function parseExecShellCommandExitStatus(
     	toolResultString: string | undefined
     ): ExecShellExitStatus | undefined {
     	if (!toolResultString) return undefined;
     
    -	const match = toolResultString.match(EXIT_CODE_TAIL_REGEX);
    +	const match = toolResultString.slice(-EXIT_CODE_TAIL_SCAN).match(EXIT_CODE_TAIL_REGEX);
     
     	if (!match) return undefined;
     
    diff --git a/tools/ui/src/lib/utils/search-results.ts b/tools/ui/src/lib/utils/search-results.ts
    index facf7766d..0fe861d94 100644
    --- a/tools/ui/src/lib/utils/search-results.ts
    +++ b/tools/ui/src/lib/utils/search-results.ts
    @@ -156,6 +156,20 @@ function parseChunk(chunk: string): SearchResult | null {
     	return result;
     }
     
    +const EMPTY_SEARCH_RESULTS: SearchResult[] = [];
    +
    +/**
    + * Cheap prefilter for the wire format: a parseable result needs both a
    + * `Title:` and a `URL:` field line, so a blob missing either substring can
    + * never yield a result. Two substring scans cost far less than the
    + * line-split parse for the megabyte tool results exec and file tools emit.
    + */
    +export function looksLikeSearchResult(text: string | undefined | null): boolean {
    +	if (!text) return false;
    +
    +	return text.includes('Title:') && text.includes('URL:');
    +}
    +
     /** Bounded cache for extractSearchResults results. */
     const SEARCH_RESULTS_CACHE_MAX_SIZE = 32;
     const searchResultsCache = new Map();
    @@ -168,7 +182,7 @@ const searchResultsCache = new Map();
      * tool result strings.
      */
     export function extractSearchResults(text: string | undefined | null): SearchResult[] {
    -	if (!text) return [];
    +	if (!text || !looksLikeSearchResult(text)) return EMPTY_SEARCH_RESULTS;
     
     	const cached = searchResultsCache.get(text);
     
    diff --git a/tools/ui/src/lib/utils/tool-call-meta.ts b/tools/ui/src/lib/utils/tool-call-meta.ts
    index b64bca786..2c035446d 100644
    --- a/tools/ui/src/lib/utils/tool-call-meta.ts
    +++ b/tools/ui/src/lib/utils/tool-call-meta.ts
    @@ -4,6 +4,8 @@
     // Each tool needs to surface fields like `error`, `result`, `bytes`,
     // `edits_applied` without repeating the try/JSON.parse/object guard inline.
     
    +import { JSON_OBJECT_OPEN } from '$lib/constants';
    +
     /**
      * Parse a tool-result blob into a JSON object, or `null` if it isn't
      * one. Returns null for:
    @@ -16,8 +18,14 @@ export function tryParseToolResultObject(
     ): Record | null {
     	if (!toolResultString) return null;
     
    +	// Tool results are usually large plain text (file contents, stdout); only
    +	// a JSON object root can carry fields, so skip the parse otherwise
    +	const trimmed = toolResultString.trimStart();
    +
    +	if (trimmed[0] !== JSON_OBJECT_OPEN) return null;
    +
     	try {
    -		const parsed: unknown = JSON.parse(toolResultString);
    +		const parsed: unknown = JSON.parse(trimmed);
     
     		if (parsed && typeof parsed === 'object' && !Array.isArray(parsed)) {
     			return parsed as Record;
    diff --git a/tools/ui/src/routes/(chat)/+page.svelte b/tools/ui/src/routes/(chat)/+page.svelte
    index 08a6b11ad..53975d7b3 100644
    --- a/tools/ui/src/routes/(chat)/+page.svelte
    +++ b/tools/ui/src/routes/(chat)/+page.svelte
    @@ -3,7 +3,7 @@
     	import { page } from '$app/state';
     	import { DialogModelNotAvailable } from '$lib/components/app';
     	import { APP_NAME, URL_PARAMS } from '$lib/constants';
    -	import { chatStore, conversationsStore, modelsStore, serverStore } from '$lib/stores';
    +	import { conversationsStore, modelsStore, serverStore } from '$lib/stores';
     	import { onMount } from 'svelte';
     
     	let qParam = $derived(page.url.searchParams.get(URL_PARAMS.QUERY));
    @@ -77,7 +77,6 @@
     		}
     
     		conversationsStore.clearActiveConversation();
    -		chatStore.clearUIState();
     
     		await modelsStore.fetch();
     
    diff --git a/tools/ui/tests/unit/agentic-sections.test.ts b/tools/ui/tests/unit/agentic-sections.test.ts
    index 4096a1710..fdb3b2217 100644
    --- a/tools/ui/tests/unit/agentic-sections.test.ts
    +++ b/tools/ui/tests/unit/agentic-sections.test.ts
    @@ -290,3 +290,114 @@ describe('hasAgenticContent', () => {
     		expect(hasAgenticContent(msg)).toBe(false);
     	});
     });
    +
    +// The turn-section cache: completed turns are immutable, so repeated
    +// derivations return the same section objects - which is what keeps tool
    +// block props stable while another turn streams. Every field the cache
    +// compares must invalidate it; a miss here renders stale content.
    +
    +describe('completed turn section reuse', () => {
    +	const toolCallsJson = JSON.stringify([
    +		{ function: { arguments: '{"path":"/a"}', name: 'test' }, id: 'call_1', type: 'function' }
    +	]);
    +
    +	function makeSession() {
    +		return {
    +			anchor: makeAssistant({
    +				content: 'answer',
    +				reasoningContent: 'thinking',
    +				toolCalls: toolCallsJson
    +			}),
    +			tools: [makeToolMsg({ content: 'tool result', extra: [{ type: 'file' } as never] })]
    +		};
    +	}
    +
    +	it('returns the same section objects for unchanged inputs', () => {
    +		const { anchor, tools } = makeSession();
    +		const first = deriveAgenticSections(anchor, tools, [], false);
    +		const second = deriveAgenticSections(anchor, tools, [], false);
    +
    +		expect(second[0]).toBe(first[0]);
    +		expect(second[1]).toBe(first[1]);
    +	});
    +
    +	it('recomputes when the assistant content changes', () => {
    +		const { anchor, tools } = makeSession();
    +		const first = deriveAgenticSections(anchor, tools, [], false);
    +
    +		anchor.content = 'edited';
    +		const second = deriveAgenticSections(anchor, tools, [], false);
    +
    +		expect(second).not.toBe(first);
    +		expect(second.some((s) => s.type === AgenticSectionType.TEXT && s.content === 'edited')).toBe(
    +			true
    +		);
    +	});
    +
    +	it('recomputes when reasoning content changes', () => {
    +		const { anchor, tools } = makeSession();
    +		const first = deriveAgenticSections(anchor, tools, [], false);
    +
    +		anchor.reasoningContent = 'new thinking';
    +		const second = deriveAgenticSections(anchor, tools, [], false);
    +
    +		expect(second).not.toBe(first);
    +	});
    +
    +	it('recomputes when toolCalls change', () => {
    +		const { anchor, tools } = makeSession();
    +		const first = deriveAgenticSections(anchor, tools, [], false);
    +
    +		anchor.toolCalls = '[]';
    +		const second = deriveAgenticSections(anchor, tools, [], false);
    +
    +		expect(second).not.toBe(first);
    +	});
    +
    +	it('recomputes when a tool result or its extras change', () => {
    +		const { anchor, tools } = makeSession();
    +		const first = deriveAgenticSections(anchor, tools, [], false);
    +
    +		tools[0].content = 'new tool result';
    +		expect(deriveAgenticSections(anchor, tools, [], false)).not.toBe(first);
    +
    +		const firstAfterContent = deriveAgenticSections(anchor, tools, [], false);
    +
    +		tools[0].extra = [{ type: 'image' } as never];
    +		expect(deriveAgenticSections(anchor, tools, [], false)).not.toBe(firstAfterContent);
    +	});
    +
    +	it('never reuses the streaming turn', () => {
    +		const { anchor, tools } = makeSession();
    +		const first = deriveAgenticSections(anchor, tools, [], true);
    +		const second = deriveAgenticSections(anchor, tools, [], true);
    +
    +		expect(second).not.toBe(first);
    +	});
    +
    +	it('keeps completed turns stable while the last turn streams', () => {
    +		const anchor = makeAssistant({
    +			content: 'turn one',
    +			id: 'ast-1',
    +			toolCalls: JSON.stringify([
    +				{ function: { arguments: '{}', name: 'test' }, id: 'call_1', type: 'function' }
    +			])
    +		});
    +		const continuation = makeAssistant({ content: 'turn two', id: 'ast-2' });
    +		const tools = [
    +			makeToolMsg({ content: 'r1', id: 'tool-1', toolCallId: 'call_1' }),
    +			continuation,
    +			makeToolMsg({ content: 'r2', id: 'tool-2', toolCallId: 'call_2' })
    +		];
    +		const first = deriveAgenticSections(anchor, tools, [], true);
    +		const second = deriveAgenticSections(anchor, tools, [], true);
    +
    +		// turn one is complete: identical section objects across derivations
    +		expect(second.slice(0, 2)).toEqual(first.slice(0, 2));
    +		expect(second[0]).toBe(first[0]);
    +		expect(second[1]).toBe(first[1]);
    +
    +		// the streaming last turn recomputed: fresh section objects
    +		expect(second[second.length - 1]).not.toBe(first[first.length - 1]);
    +	});
    +});
    diff --git a/tools/ui/tests/unit/branching.test.ts b/tools/ui/tests/unit/branching.test.ts
    new file mode 100644
    index 000000000..8a752ae2f
    --- /dev/null
    +++ b/tools/ui/tests/unit/branching.test.ts
    @@ -0,0 +1,95 @@
    +// Sibling-info correctness for buildSiblingInfoMap, including the memoized
    +// leaf resolution. A wrong leaf id here breaks branch navigation, so the
    +// deep-chain and multi-branch cases below pin the resolution down.
    +
    +import { MessageRole, MessageType } from '$lib/enums';
    +import type { DatabaseMessage } from '$lib/types/database';
    +import { buildSiblingInfoMap, findLeafNode } from '$lib/utils/branching';
    +import { describe, expect, it } from 'vitest';
    +
    +function msg(id: string, parent: string | null, children: string[] = []): DatabaseMessage {
    +	return {
    +		children,
    +		content: '',
    +		convId: 'c1',
    +		id,
    +		parent,
    +		role: MessageRole.USER,
    +		timestamp: 0,
    +		type: MessageType.TEXT
    +	} as DatabaseMessage;
    +}
    +
    +/** root -> m1 -> ... -> m depth, each node with a single child. */
    +function linearChain(depth: number): DatabaseMessage[] {
    +	const messages = [msg('m0', null, ['m1'])];
    +
    +	for (let i = 1; i <= depth; i++) {
    +		messages.push(msg(`m${i}`, `m${i - 1}`, i < depth ? [`m${i + 1}`] : []));
    +	}
    +
    +	return messages;
    +}
    +
    +describe('buildSiblingInfoMap', () => {
    +	it('resolves the deepest leaf for every node of a long single chain', () => {
    +		const messages = linearChain(50);
    +		const map = buildSiblingInfoMap(messages);
    +		const leafId = messages[messages.length - 1].id;
    +
    +		// every non-root message of the chain is an only child, and its
    +		// navigation target is the chain's deepest leaf
    +		for (const m of messages.slice(1)) {
    +			const info = map.get(m.id);
    +
    +			expect(info?.totalSiblings).toBe(1);
    +			expect(info?.siblingIds).toEqual([leafId]);
    +		}
    +	});
    +
    +	it('reports sibling position and leaf targets on a branched tree', () => {
    +		// m0 -> m1, m4 ; m1 -> m2 ; m2 -> m3, m6 ; m4 -> m5
    +		const root = msg('m0', null, ['m1', 'm4']);
    +		const m1 = msg('m1', 'm0', ['m2']);
    +		const m2 = msg('m2', 'm1', ['m3', 'm6']);
    +		const m3 = msg('m3', 'm2');
    +		const m4 = msg('m4', 'm0', ['m5']);
    +		const m5 = msg('m5', 'm4');
    +		const m6 = msg('m6', 'm2');
    +		const map = buildSiblingInfoMap([root, m1, m2, m3, m4, m5, m6]);
    +
    +		// m1 and m4 share the root as parent; their nav targets are the
    +		// leaves of their subtrees ( m6 for the first branch, m5 for the second )
    +		expect(map.get(m1.id)).toMatchObject({
    +			currentIndex: 0,
    +			siblingIds: [m6.id, m5.id],
    +			totalSiblings: 2
    +		});
    +		expect(map.get(m4.id)).toMatchObject({
    +			currentIndex: 1,
    +			siblingIds: [m6.id, m5.id],
    +			totalSiblings: 2
    +		});
    +
    +		// m3 and m6 are siblings under m2; both are leaves
    +		expect(map.get(m3.id)?.siblingIds).toEqual([m3.id, m6.id]);
    +		expect(map.get(m6.id)?.currentIndex).toBe(1);
    +
    +		// the root has no parent and reports itself
    +		expect(map.get(root.id)).toMatchObject({
    +			currentIndex: 0,
    +			siblingIds: [root.id],
    +			totalSiblings: 1
    +		});
    +	});
    +
    +	it('agrees with findLeafNode for arbitrary nodes', () => {
    +		const messages = linearChain(20);
    +		const leafId = messages[messages.length - 1].id;
    +
    +		// every node of the chain resolves to the deepest leaf
    +		for (const m of messages) {
    +			expect(findLeafNode(messages, m.id), `leaf of ${m.id}`).toBe(leafId);
    +		}
    +	});
    +});
    diff --git a/tools/ui/tests/unit/conversations-store.test.ts b/tools/ui/tests/unit/conversations-store.test.ts
    new file mode 100644
    index 000000000..e06546597
    --- /dev/null
    +++ b/tools/ui/tests/unit/conversations-store.test.ts
    @@ -0,0 +1,90 @@
    +// Field updates to the active conversation must keep the object identity
    +// stable: effects that track the identity ( the chat screen's sibling-info
    +// refresh ) refire on every identity change, which used to trigger a full
    +// message refetch on every send and tool result.
    +
    +import { beforeEach, describe, expect, it, vi } from 'vitest';
    +
    +vi.mock('$lib/services/database.service', () => ({
    +	DatabaseService: {
    +		getConversation: vi.fn(),
    +		getConversationMessages: vi.fn(),
    +		updateConversation: vi.fn(),
    +		updateCurrentNode: vi.fn()
    +	}
    +}));
    +
    +import { DatabaseService } from '$lib/services/database.service';
    +import { conversationsStore } from '$lib/stores/conversations/index.svelte';
    +import type { DatabaseConversation, DatabaseMessage } from '$lib/types/database';
    +
    +const getConversationMock = vi.mocked(DatabaseService.getConversation);
    +const getMessagesMock = vi.mocked(DatabaseService.getConversationMessages);
    +const updateCurrentNodeMock = vi.mocked(DatabaseService.updateCurrentNode);
    +
    +function makeConversation(overrides: Partial = {}): DatabaseConversation {
    +	return {
    +		currNode: 'node-1',
    +		id: 'conv-1',
    +		lastModified: 1000,
    +		name: 'conversation',
    +		...overrides
    +	};
    +}
    +
    +async function loadActive(conversation: DatabaseConversation, messages: DatabaseMessage[]) {
    +	getConversationMock.mockResolvedValue(conversation);
    +	getMessagesMock.mockResolvedValue(messages);
    +
    +	expect(await conversationsStore.loadConversation(conversation.id)).toBe(true);
    +}
    +
    +beforeEach(() => {
    +	getConversationMock.mockReset();
    +	getMessagesMock.mockReset();
    +	updateCurrentNodeMock.mockReset();
    +	updateCurrentNodeMock.mockResolvedValue(undefined);
    +	vi.mocked(DatabaseService.updateConversation).mockReset();
    +	vi.mocked(DatabaseService.updateConversation).mockResolvedValue(undefined);
    +});
    +
    +describe('active conversation identity', () => {
    +	it('hands the load read off exactly once', async () => {
    +		await loadActive(makeConversation(), []);
    +
    +		expect(conversationsStore.consumeLastLoadedMessages('conv-1')).toEqual([]);
    +		// a second consume is a miss: branch actions must fall back to a refetch
    +		expect(conversationsStore.consumeLastLoadedMessages('conv-1')).toBeNull();
    +	});
    +
    +	it('writes currNode in place on updateCurrentNode', async () => {
    +		await loadActive(makeConversation(), []);
    +		const before = conversationsStore.activeConversation;
    +
    +		await conversationsStore.updateCurrentNode('node-2');
    +
    +		expect(conversationsStore.activeConversation).toBe(before);
    +		expect(conversationsStore.activeConversation?.currNode).toBe('node-2');
    +	});
    +
    +	it('writes renamed and pinned fields in place on applyConversationUpdate', async () => {
    +		await loadActive(makeConversation(), []);
    +		const before = conversationsStore.activeConversation;
    +
    +		conversationsStore.applyConversationUpdate('conv-1', { name: 'renamed', pinned: true });
    +
    +		expect(conversationsStore.activeConversation).toBe(before);
    +		expect(conversationsStore.activeConversation?.name).toBe('renamed');
    +		expect(conversationsStore.activeConversation?.pinned).toBe(true);
    +	});
    +
    +	it('writes lastModified in place on updateConversationTimestamp', async () => {
    +		await loadActive(makeConversation(), []);
    +		const before = conversationsStore.activeConversation;
    +
    +		conversationsStore.updateConversationTimestamp('conv-1');
    +
    +		expect(conversationsStore.activeConversation).toBe(before);
    +		expect(conversationsStore.activeConversation?.lastModified).toBeGreaterThan(1000);
    +	});
    +});
    diff --git a/tools/ui/tests/unit/parse-exec-shell-status.test.ts b/tools/ui/tests/unit/parse-exec-shell-status.test.ts
    index ed499d078..7e22bf9ee 100644
    --- a/tools/ui/tests/unit/parse-exec-shell-status.test.ts
    +++ b/tools/ui/tests/unit/parse-exec-shell-status.test.ts
    @@ -71,3 +71,21 @@ describe('isExitCodeSummaryLine', () => {
     		expect(isExitCodeSummaryLine('[exit code: 7]', undefined)).toBe(false);
     	});
     });
    +
    +describe('parseExecShellCommandExitStatus tail scan', () => {
    +	it('finds the marker at the end of a blob larger than the tail window', () => {
    +		// the parser matches only the last ~128 chars; a marker past that
    +		// window must still parse, and an earlier fake must not match
    +		const blob = `${'the shell prints [exit code: 1] mid-stream\n'.repeat(2000)}[exit code: 0]`;
    +		const status = parseExecShellCommandExitStatus(blob);
    +
    +		expect(status?.code).toBe(0);
    +		expect(status?.timedOut).toBe(false);
    +	});
    +
    +	it('keeps rejecting markers that are not at the absolute end', () => {
    +		const blob = `${'stdout\n'.repeat(2000)}[exit code: 0]\nsome trailing log line`;
    +
    +		expect(parseExecShellCommandExitStatus(blob)).toBeUndefined();
    +	});
    +});
    diff --git a/tools/ui/tests/unit/search-results.test.ts b/tools/ui/tests/unit/search-results.test.ts
    index c168dec25..561ab935a 100644
    --- a/tools/ui/tests/unit/search-results.test.ts
    +++ b/tools/ui/tests/unit/search-results.test.ts
    @@ -2,7 +2,8 @@ import {
     	extractSearchQuery,
     	extractSearchResults,
     	faviconForUrl,
    -	isWebSearchToolName
    +	isWebSearchToolName,
    +	looksLikeSearchResult
     } from '$lib/utils/search-results';
     import { describe, expect, it } from 'vitest';
     
    @@ -119,3 +120,27 @@ describe('isWebSearchToolName', () => {
     		expect(isWebSearchToolName('exec_shell_command')).toBe(false);
     	});
     });
    +
    +describe('extractSearchResults prefilter', () => {
    +	it('returns the shared empty array for blobs without the wire format', () => {
    +		// exec/file tool results never carry Title:/URL: field lines; the
    +		// cheap prefilter must skip the line-split parse for them
    +		const stdout = `${'make[1]: entering directory\n'.repeat(5000)}`;
    +
    +		expect(extractSearchResults(stdout)).toEqual([]);
    +	});
    +
    +	it('returns an empty result when only one required field is present', () => {
    +		expect(extractSearchResults('URL: https://example.com')).toEqual([]);
    +		expect(extractSearchResults('Title: only a title')).toEqual([]);
    +	});
    +});
    +
    +describe('looksLikeSearchResult', () => {
    +	it('requires both Title and URL field markers', () => {
    +		expect(looksLikeSearchResult('Title: a\nURL: https://b')).toBe(true);
    +		expect(looksLikeSearchResult('URL: https://b')).toBe(false);
    +		expect(looksLikeSearchResult('plain stdout')).toBe(false);
    +		expect(looksLikeSearchResult(undefined)).toBe(false);
    +	});
    +});
    diff --git a/tools/ui/tests/unit/tool-call-meta.test.ts b/tools/ui/tests/unit/tool-call-meta.test.ts
    index bb28e3830..f94d2279f 100644
    --- a/tools/ui/tests/unit/tool-call-meta.test.ts
    +++ b/tools/ui/tests/unit/tool-call-meta.test.ts
    @@ -28,3 +28,15 @@ describe('tryParseToolResultObject', () => {
     		expect(tryParseToolResultObject('{bad')).toBeNull();
     	});
     });
    +
    +describe('tryParseToolResultObject gating', () => {
    +	it('parses JSON objects that start after leading whitespace', () => {
    +		expect(tryParseToolResultObject('\n  {"result":"ok"}')).toEqual({ result: 'ok' });
    +	});
    +
    +	it('skips the parse for large plain-text results', () => {
    +		// most tool results are file contents or stdout; the gate avoids a
    +		// doomed JSON.parse over the whole blob
    +		expect(tryParseToolResultObject(`${'stdout line\n'.repeat(2000)}`)).toBeNull();
    +	});
    +});
    diff --git a/tools/ui/tests/unit/tool-calls.test.ts b/tools/ui/tests/unit/tool-calls.test.ts
    index f84a2405e..a2274f9d2 100644
    --- a/tools/ui/tests/unit/tool-calls.test.ts
    +++ b/tools/ui/tests/unit/tool-calls.test.ts
    @@ -1,5 +1,8 @@
     import { parseToolArgs } from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/_shared';
    -import { parseEditFileMeta } from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file';
    +import {
    +	parseEditFileMeta,
    +	parseEditFileTitleMeta
    +} from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file';
     import { parseExecShellCommandMeta } from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/exec-shell-command';
     import { parseFileGlobSearchMeta } from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/file-glob-search';
     import { parseGrepSearchMeta } from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/grep-search';
    @@ -7,10 +10,10 @@ import { parseReadFileMeta } from '$lib/components/app/chat/ChatMessages/ChatMes
     import { parseRunJavascriptMeta } from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/run-javascript';
     import {
     	parseWriteFileMeta,
    -	type WriteFileMeta
    +	parseWriteFileTitleMeta
     } from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/write-file';
     import { AgenticSectionType, BuiltInTool } from '$lib/enums';
    -import type { AgenticSection } from '$lib/types';
    +import type { AgenticSection, WriteFileMeta } from '$lib/types';
     import { abbreviateHome, formatCwdMessage, lastPathSegment, parseCwdMessage } from '$lib/utils';
     import { describe, expect, it } from 'vitest';
     
    @@ -223,6 +226,113 @@ describe('parseWriteFileMeta', () => {
     	});
     });
     
    +describe('parseWriteFileTitleMeta', () => {
    +	it('matches the full meta for path, language and result fields', () => {
    +		const args = JSON.stringify({ content: 'x'.repeat(50_000), path: '/foo.ts' });
    +		const toolResult = '{"result":"wrote","bytes":42}';
    +		const section = makeSection(
    +			{ toolArgs: args, toolName: BuiltInTool.SERVER_WRITE_FILE, toolResult },
    +			BuiltInTool.SERVER_WRITE_FILE
    +		);
    +		const full = parseWriteFileMeta(section);
    +		const title = parseWriteFileTitleMeta(section);
    +
    +		expect(title?.filePath).toBe(full?.filePath);
    +		expect(title?.fileName).toBe(full?.fileName);
    +		expect(title?.language).toBe(full?.language);
    +		expect(title?.bytesWritten).toBe(full?.bytesWritten);
    +		expect(title?.resultMessage).toBe(full?.resultMessage);
    +		expect(title?.errorMessage).toBe(full?.errorMessage);
    +	});
    +
    +	it('extracts a path with escaped characters without parsing the content blob', () => {
    +		const section = makeSection(
    +			{
    +				toolArgs: '{"path":"/a\\nb\\"c/d.ts","content":"x"}',
    +				toolName: BuiltInTool.SERVER_WRITE_FILE
    +			},
    +			BuiltInTool.SERVER_WRITE_FILE
    +		);
    +
    +		expect(parseWriteFileTitleMeta(section)?.filePath).toBe('/a\nb"c/d.ts');
    +	});
    +
    +	it('falls back to the full parse for args the extractor can not see', () => {
    +		const section = makeSection(
    +			{
    +				// key written with an escaped unicode escape sequence in the name
    +				toolArgs: '{"\\u0070ath":"/foo.ts","content":"x"}',
    +				toolName: BuiltInTool.SERVER_WRITE_FILE
    +			},
    +			BuiltInTool.SERVER_WRITE_FILE
    +		);
    +
    +		expect(parseWriteFileTitleMeta(section)?.filePath).toBe('/foo.ts');
    +	});
    +
    +	it('accepts partial args like the full parser', () => {
    +		const section = makeSection(
    +			{ toolArgs: '{"path":"/foo.t', toolName: BuiltInTool.SERVER_WRITE_FILE },
    +			BuiltInTool.SERVER_WRITE_FILE
    +		);
    +
    +		expect(parseWriteFileTitleMeta(section)?.filePath).toBe('/foo.t');
    +	});
    +
    +	it('returns null for sections with a different tool name', () => {
    +		expect(
    +			parseWriteFileTitleMeta(
    +				makeSection({
    +					toolArgs: '{"path":"/x","content":"y"}',
    +					toolName: BuiltInTool.SERVER_READ_FILE
    +				})
    +			)
    +		).toBeNull();
    +	});
    +});
    +
    +describe('parseEditFileTitleMeta', () => {
    +	it('matches the full meta for path and result fields', () => {
    +		const section = makeSection(
    +			{
    +				toolArgs: '{"path":"/foo.ts","edits":[{"old_text":"a","new_text":"b"}]}' + ' '.repeat(0),
    +				toolName: BuiltInTool.SERVER_EDIT_FILE,
    +				toolResult: '{"result":"ok","edits_applied":1}'
    +			},
    +			BuiltInTool.SERVER_EDIT_FILE
    +		);
    +		const full = parseEditFileMeta(section);
    +		const title = parseEditFileTitleMeta(section);
    +
    +		expect(title?.filePath).toBe(full?.filePath);
    +		expect(title?.fileName).toBe(full?.fileName);
    +		expect(title?.editsApplied).toBe(full?.editsApplied);
    +		expect(title?.resultMessage).toBe(full?.resultMessage);
    +		expect(title?.errorMessage).toBe(full?.errorMessage);
    +	});
    +
    +	it('surfaces errorMessage from the result blob without parsing args', () => {
    +		const section = makeSection(
    +			{
    +				toolArgs: '{"path":"/foo.ts","edits":[]}',
    +				toolName: BuiltInTool.SERVER_EDIT_FILE,
    +				toolResult: '{"error":"permission denied"}'
    +			},
    +			BuiltInTool.SERVER_EDIT_FILE
    +		);
    +
    +		expect(parseEditFileTitleMeta(section)?.errorMessage).toBe('permission denied');
    +	});
    +
    +	it('returns null when args have no path-like field', () => {
    +		expect(
    +			parseEditFileTitleMeta(
    +				makeSection({ toolArgs: '{"edits":[]}', toolName: BuiltInTool.SERVER_EDIT_FILE })
    +			)
    +		).toBeNull();
    +	});
    +});
    +
     describe('parseEditFileMeta', () => {
     	it('parses edits array and applies editsApplied from the result', () => {
     		const section = makeSection(
    diff --git a/tools/ui/ui.cpp.in b/tools/ui/ui.cpp.in
    new file mode 100644
    index 000000000..7f91ef2a2
    --- /dev/null
    +++ b/tools/ui/ui.cpp.in
    @@ -0,0 +1,36 @@
    +// Generated by scripts/ui-assets.cmake - do not edit.
    +
    +#include "ui.h"
    +
    +@ASSET_ARRAYS@
    +#if defined(LLAMA_UI_HAS_ASSETS)
    +static const std::array g_assets = {{
    +@ASSET_TABLE@
    +}};
    +#endif
    +
    +const llama_ui_asset * llama_ui_find_asset(const std::string & name) {
    +#if defined(LLAMA_UI_HAS_ASSETS)
    +    for (const auto & a : g_assets) {
    +        if (a.name == name) {
    +            return &a;
    +        }
    +    }
    +#else
    +    (void) name;
    +#endif
    +    return nullptr;
    +}
    +
    +const std::array & llama_ui_get_assets() {
    +#if defined(LLAMA_UI_HAS_ASSETS)
    +    return g_assets;
    +#else
    +    static const std::array empty{};
    +    return empty;
    +#endif
    +}
    +
    +bool llama_ui_use_gzip() {
    +    return @USE_GZIP@;
    +}
    diff --git a/tools/ui/ui.h.in b/tools/ui/ui.h.in
    new file mode 100644
    index 000000000..4555b0dd5
    --- /dev/null
    +++ b/tools/ui/ui.h.in
    @@ -0,0 +1,21 @@
    +// Generated by scripts/ui-assets.cmake - do not edit.
    +
    +#pragma once
    +
    +#include 
    +#include 
    +
    +// Defined as 1 only when assets were embedded (tools/server checks defined()).
    +#cmakedefine LLAMA_UI_HAS_ASSETS 1
    +
    +struct llama_ui_asset {
    +    std::string           name;
    +    const unsigned char * data;
    +    std::size_t           size;
    +    std::string           etag;
    +    std::string           type;
    +};
    +
    +const llama_ui_asset * llama_ui_find_asset(const std::string & name);
    +bool llama_ui_use_gzip();
    +const std::array & llama_ui_get_assets();