From 3e390af18de1131dda27913c00511725b4274c8a Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Thu, 15 Oct 2026 05:26:10 +0200 Subject: [PATCH 001/423] vulkan: Use L2-friendly workgroup scheduling for cm2 mulmat (#30215) --- .../vulkan-shaders/mul_mm_cm2.comp | 29 +++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp index 89022f047848..10dda7768295 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp @@ -45,6 +45,7 @@ const uint BNover2 = BN / 2; const uint BNover4 = enable_smaller_matrices ? (BN / 4) : (BN / 2); layout (constant_id = 5) const uint ALIGNED = 0; layout (constant_id = 6) const uint subgroup_size = 32; +#define GROUP_A_BUDGET (8u * 1024u * 1024u) #ifdef MULMAT_QUANT @@ -258,8 +259,6 @@ void main() { const uint row_begin = tile_idx * BN; _ne1 = data_row_map[2 + 2 * tile_idx]; const uint ic = 0; -#else - const uint ic = gl_WorkGroupID.y; #endif #ifdef MUL_MAT_ID @@ -289,9 +288,33 @@ void main() { #endif const uint blocks_m = (p.M + BM - 1) / BM; - const uint ir = gl_WorkGroupID.x % blocks_m; const uint ik = gl_WorkGroupID.x / blocks_m; +#ifdef MUL_MAT_ID + const uint ir = gl_WorkGroupID.x % blocks_m; +#else + // L2-friendly workgroup scheduling + const uint blocks_n = (p.N + BN - 1) / BN; +#if defined(MULMAT_QUANT) + const uint a_panel_bytes = (BM * p.K / mm_quant_k()) * MmABlockBytes; +#elif defined(DATA_A_F32) + const uint a_panel_bytes = BM * p.K * 4; +#elif QUANT_K == 1 + const uint a_panel_bytes = BM * p.K * 2; +#else + const uint a_panel_bytes = (BM * p.K) / 2 + (BM * p.K) / 16; +#endif + const uint group_m = clamp(GROUP_A_BUDGET / max(a_panel_bytes, 1u), 1u, min(blocks_m, 32u)); + const uint tiles_per_group = group_m * blocks_n; + const uint lin = gl_WorkGroupID.y * blocks_m + (gl_WorkGroupID.x % blocks_m); + const uint group_id = lin / tiles_per_group; + const uint first_m = group_id * group_m; + const uint gsize = min(blocks_m - first_m, group_m); + const uint in_group = lin - group_id * tiles_per_group; + const uint ir = first_m + in_group % gsize; + const uint ic = in_group / gsize; +#endif + #ifdef MUL_MAT_ID load_row_ids(row_begin); #endif From 843e7dcc15098dd6b3aaf33b4847052293d8ef1a Mon Sep 17 00:00:00 2001 From: Kartik Gulia Date: Thu, 15 Oct 2026 09:19:25 +0530 Subject: [PATCH 002/423] spec : add lilicorr support (#30506) * spec : add lilicorr support * convert : use convention for hparam assignment * add todo for hparam use in llama-batch.cpp Co-authored-by: Georgi Gerganov * typo fix --------- Co-authored-by: Georgi Gerganov --- common/speculative.cpp | 59 +++++-- conversion/__init__.py | 1 + conversion/nemotron.py | 101 +++++++++++- gguf-py/gguf/constants.py | 49 ++++++ gguf-py/gguf/gguf_writer.py | 24 +++ gguf-py/gguf/tensor_mapping.py | 47 ++++++ include/llama.h | 3 +- src/llama-arch.cpp | 36 +++++ src/llama-arch.h | 21 +++ src/llama-batch.cpp | 4 + src/llama-hparams.h | 6 + src/llama-model.h | 44 ++++++ src/models/dflash.cpp | 276 +++++++++++++++++++++++++++++++-- 13 files changed, 649 insertions(+), 22 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 50647dbb20aa..701093308292 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -999,6 +999,12 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { int32_t block_size = 0; llama_token mask_token_id = 0; + bool is_lilicorr = false; + int32_t lilicorr_top_k = 0; + size_t lilicorr_rows = 0; + std::vector lilicorr_embd; + std::vector lilicorr_pos; + bool is_dflash2 = false; bool is_mrope = false; int32_t selector_top_k = 0; @@ -1053,6 +1059,16 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { selector_top_k = llama_model_dflash_selector_top_k(model_dft); is_dflash2 = selector_top_k > 0; + char lilicorr_meta[32] = {}; + lilicorr_top_k = llama_model_meta_val_str(model_dft, "lilicorr.candidate_topk", lilicorr_meta, sizeof(lilicorr_meta)) >= 0 + ? std::atoi(lilicorr_meta) : 0; + is_lilicorr = lilicorr_top_k > 0; + if (is_lilicorr) { + lilicorr_rows = (size_t) block_size + 1; + const size_t lilicorr_slots = (size_t) n_seq * lilicorr_rows; + lilicorr_embd.resize(lilicorr_slots * (size_t) n_embd_dec); + lilicorr_pos.assign(lilicorr_slots, -1); + } mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); if (is_dspark && this->params.p_min > 0.0f) { @@ -1098,7 +1114,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // offload draft sampling to the backend backend_chains.assign(n_seq, nullptr); - if (this->params.backend_sampling && !is_dflash2) { + if (this->params.backend_sampling && !is_dflash2 && !is_lilicorr) { for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params()); llama_sampler_chain_add(chain, llama_sampler_init_top_k(10)); @@ -1118,7 +1134,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits. - llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2); + llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2 && !is_lilicorr); llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise } @@ -1233,6 +1249,24 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { __func__, rc, (int) n_chunk, (int) offset); return false; } + if (is_lilicorr) { + // Ordinary DFlash can write to the KV cache, but LiliCorr needs to store embeddings across sequences. + const float * projected = llama_get_embeddings_nextn(ctx_dft); + if (projected) { + const size_t seq_offset = (size_t) seq_id * lilicorr_rows; + const llama_pos pos_min = std::max(0, batch_inject.tokens.back().pos[0] - block_size); + for (int32_t i = 0; i < n_chunk; ++i) { + const llama_pos pos = batch_inject.tokens[i].pos[0]; + if (pos < pos_min) { + continue; + } + const size_t slot = seq_offset + (size_t) pos % lilicorr_rows; + std::memcpy(lilicorr_embd.data() + slot * (size_t) n_embd_dec, + projected + (size_t) i * n_embd_dec, (size_t) n_embd_dec * sizeof(float)); + lilicorr_pos[slot] = pos; + } + } + } } } @@ -1260,12 +1294,17 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { const int32_t n = (int32_t) dp.pos0; const int32_t n_draft = params.n_max; - - const int32_t n_block_tokens = n_draft + (is_dspark && sample_from_anchor ? 0 : 1); + const int32_t n_block_tokens = is_lilicorr ? block_size : n_draft + (is_dspark && sample_from_anchor ? 0 : 1); i_block_beg[seq_id] = batch.size(); n_block [seq_id] = n_block_tokens; for (int32_t i = 0; i < n_block_tokens; ++i) { - batch.add(i == 0 ? dp.id_last : mask_token_id, n + i, seq_id, !is_dflash2); + const int32_t idx = batch.add(i == 0 ? dp.id_last : mask_token_id, n + i, seq_id, !is_dflash2 && !is_lilicorr); + if (is_lilicorr) { + // For LiliCorr drafting batch, set state embeddings for use in the associated transformer head. + const llama_pos pos = dp.pos0 - 1; + const size_t slot = (size_t) seq_id * lilicorr_rows + (size_t) pos % lilicorr_rows; + batch.set_embd_state(idx, { lilicorr_embd.data() + slot * (size_t) n_embd_dec, 1, (size_t) n_embd_dec }); + } } } @@ -1293,21 +1332,23 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { auto & result = *dp.result; - if (is_dflash2) { + if (is_dflash2 || is_lilicorr) { + // Both DFlash and LiliCorr greedily walk the candidate lattice. const float * lattice = llama_get_embeddings_nextn(ctx_dft); GGML_ASSERT(lattice && "DFlash2 selector produced no lattice"); int32_t predecessor = 0; for (int32_t i = 1; i < n_block_tokens; ++i) { const float * row = lattice + (size_t) (beg + i) * n_embd_dec; - const float * scores = row + selector_top_k + (size_t) predecessor * selector_top_k; + const int32_t top_k = is_dflash2 ? selector_top_k : lilicorr_top_k; + const float * scores = row + top_k + (size_t) predecessor * top_k; predecessor = (int32_t) std::distance(scores, - std::max_element(scores, scores + selector_top_k)); + std::max_element(scores, scores + top_k)); if (params.p_min > 0.0f) { // softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max)) float sum = 0.0f; - for (int32_t k = 0; k < selector_top_k; ++k) { + for (int32_t k = 0; k < top_k; ++k) { sum += std::exp(scores[k] - scores[predecessor]); } if (1.0f / sum < params.p_min) { diff --git a/conversion/__init__.py b/conversion/__init__.py index 7f32a87f062a..d5b480a9bf41 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -59,6 +59,7 @@ "DeepseekV32ForCausalLM": "deepseek", "DFlashDraftModel": "qwen", "DFlash2DraftModel": "qwen", + "LiLiCorrDraftModel": "nemotron", "Qwen3DSparkModel": "qwen", "DSparkDraftModel": "qwen", "DSparkSpeculator": "qwen", diff --git a/conversion/nemotron.py b/conversion/nemotron.py index f7eefcf8d0ea..581f41a0230b 100644 --- a/conversion/nemotron.py +++ b/conversion/nemotron.py @@ -2,15 +2,114 @@ from typing import Any, Callable, Iterable, TYPE_CHECKING +import numpy as np import torch if TYPE_CHECKING: from pathlib import Path from torch import Tensor -from .base import MmprojModel, ModelBase, TextModel, gguf, logger +from .base import MmprojModel, ModelBase, ModelType, TextModel, get_model_architecture, gguf, logger from .granite import GraniteHybridModel +from .qwen import DFlashModel + + +@ModelBase.register("LiLiCorrDraftModel") +class LiLiCorrDraftModel(DFlashModel): + model_arch = gguf.MODEL_ARCH.DFLASH + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + prefix = "model.lilicorr." + if not any(name.startswith(prefix) for name in self.model_tensors): + raise ValueError("no lilicorr.* tensors found") + dflash_config = self.hparams["dflash_config"] + self.lilicorr_geometry = { + "hidden_size": int(dflash_config["lilicorr_hidden_size"]), + "candidate_topk": int(dflash_config["lilicorr_candidate_topk"]), + "num_heads": int(dflash_config["lilicorr_num_heads"]), + "num_layers": int(dflash_config["lilicorr_num_layers"]), + "factor_dim": int(dflash_config["lilicorr_factor_dim"]), + "mlp_hidden_size": self.model_tensors[prefix + "layers.0.mlp.up_proj.weight"]().shape[0], + } + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count + self.lilicorr_geometry["num_layers"] + 1) + block = int(dflash_config.get("block_size", self.hparams.get("block_size", 16))) + topk = self.lilicorr_geometry["candidate_topk"] + heads = self.lilicorr_geometry["num_heads"] + nodes = (block - 1) * topk + relative = gguf.LazyBase.to_eager(self.model_tensors.pop(prefix + "relative_slot_bias")()).float().cpu().numpy() + same = gguf.LazyBase.to_eager(self.model_tensors.pop(prefix + "same_slot_bias")()).float().cpu().numpy() + bias = np.empty((heads, nodes, nodes), dtype=np.float32) + for q in range(nodes): + q_slot = q // topk + for k in range(nodes): + k_slot = k // topk + bias[:, q, k] = relative[:, q_slot - k_slot + block - 1] + if q_slot == k_slot: + bias[:, q, k] += same + self.model_tensors[prefix + "attention_bias.bias"] = lambda bias=bias: torch.from_numpy(bias) + for name in ("slot_embedding", "rank_embedding"): + gen = self.model_tensors.pop(prefix + name) + self.model_tensors[prefix + name + ".weight"] = lambda gen=gen: gen().reshape(-1, self.lilicorr_geometry["hidden_size"]) + for i in range(self.lilicorr_geometry["num_layers"]): + layer_prefix = prefix + f"layers.{i}.attn.in_proj" + for suffix in ("weight", "bias"): + self.model_tensors[layer_prefix + "." + suffix] = self.model_tensors.pop(layer_prefix + "_" + suffix) + dflash_config["selector_rank"] = 0 + dflash_config["selector_top_k"] = 0 + + def set_vocab(self): + if self.target_model_dir is None: + raise ValueError("LiLiCorr draft requires --target-model-dir") + from . import get_model_class + target_text_config = ModelBase.load_hparams(self.target_model_dir, False) + target_text_config = {**target_text_config, **target_text_config.get("text_config", {})} + target_arch = get_model_architecture(target_text_config, ModelType.TEXT) + target_cls = get_model_class(target_arch) + original_dir = self.dir_model + self.dir_model = self.target_model_dir + try: + if target_arch == "NemotronHForCausalLM": + self.is_moe = "num_experts_per_tok" in target_text_config + target_cls.set_vocab(self) # ty: ignore[unresolved-attribute] + finally: + self.dir_model = original_dir + self.gguf_writer.add_mask_token_id(self.hparams["dflash_config"]["mask_token_id"]) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + dflash_config = self.hparams["dflash_config"] + self.gguf_writer.add_lilicorr_hidden_size(int(self.lilicorr_geometry["hidden_size"])) + self.gguf_writer.add_lilicorr_candidate_topk(int(self.lilicorr_geometry["candidate_topk"])) + self.gguf_writer.add_lilicorr_num_heads(int(self.lilicorr_geometry["num_heads"])) + self.gguf_writer.add_lilicorr_num_layers(int(self.lilicorr_geometry["num_layers"])) + self.gguf_writer.add_lilicorr_factor_dim(int(self.lilicorr_geometry["factor_dim"])) + self.gguf_writer.add_lilicorr_mlp_hidden_size(int(self.lilicorr_geometry["mlp_hidden_size"])) + self.gguf_writer.add_lilicorr_vector_eps(float(dflash_config["lilicorr_vector_eps"])) + self.gguf_writer.add_lilicorr_logit_scale(float(dflash_config["lilicorr_logit_scale"])) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.startswith("model.lilicorr."): + if name.startswith("model.lilicorr.layers."): + parts = name.split(".", 4) + if not parts[4].startswith("attn.in_proj."): + parts[3] = str(self.block_count + int(parts[3])) + name = ".".join(parts) + elif name.startswith("model.lilicorr.feature_norm."): + feature_bid = self.block_count + self.lilicorr_geometry["num_layers"] + name = name.replace("model.lilicorr.feature_norm.", f"model.lilicorr.feature_mlp.{feature_bid}.norm.") + elif name.startswith("model.lilicorr.feature_mlp."): + feature_bid = self.block_count + self.lilicorr_geometry["num_layers"] + name = name.replace("model.lilicorr.feature_mlp.", f"model.lilicorr.feature_mlp.{feature_bid}.") + yield self.map_tensor_name(name), data_torch + return + yield from super().modify_tensors(data_torch, name, bid) + + def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool: + if name.startswith("model.lilicorr."): + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) @ModelBase.register( diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index c31ab4a42e89..f9baf687ee16 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -309,6 +309,16 @@ class KDA: class WKV: HEAD_SIZE = "{arch}.wkv.head_size" + class LiliCorr: + HIDDEN_SIZE = "lilicorr.hidden_size" + CANDIDATE_TOPK = "lilicorr.candidate_topk" + NUM_HEADS = "lilicorr.num_heads" + NUM_LAYERS = "lilicorr.num_layers" + FACTOR_DIM = "lilicorr.factor_dim" + MLP_HIDDEN_SIZE = "lilicorr.mlp_hidden_size" + VECTOR_EPS = "lilicorr.vector_eps" + LOGIT_SCALE = "lilicorr.logit_scale" + class PosNet: EMBEDDING_LENGTH = "{arch}.posnet.embedding_length" BLOCK_COUNT = "{arch}.posnet.block_count" @@ -1279,6 +1289,19 @@ class MODEL_TENSOR(IntEnum): DFLASH_SELECTOR_PREV = auto() DFLASH_SELECTOR_NEXT = auto() DFLASH_SELECTOR_HIDDEN = auto() + LILICORR_TOKEN_PROJ = auto() + LILICORR_PASS_HIDDEN_PROJ = auto() + LILICORR_SLOT_EMBD = auto() + LILICORR_RANK_EMBD = auto() + LILICORR_ATTN_BIAS = auto() + LILICORR_CONTEXT_PROJ = auto() + LILICORR_ATTN_IN_PROJ = auto() + LILICORR_OUTPUT_NORM = auto() + LILICORR_ANCHOR_NORM = auto() + LILICORR_FACTOR_INPUT_PROJ = auto() + LILICORR_OUT_HEAD = auto() + LILICORR_IN_HEAD = auto() + LILICORR_ANCHOR_OUT_HEAD = auto() # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -2093,6 +2116,19 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor", MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor", MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden", + MODEL_TENSOR.LILICORR_TOKEN_PROJ: "lilicorr.token_proj", + MODEL_TENSOR.LILICORR_PASS_HIDDEN_PROJ: "lilicorr.pass_hidden_proj", + MODEL_TENSOR.LILICORR_SLOT_EMBD: "lilicorr.slot_embedding", + MODEL_TENSOR.LILICORR_RANK_EMBD: "lilicorr.rank_embedding", + MODEL_TENSOR.LILICORR_ATTN_BIAS: "lilicorr.attention_bias", + MODEL_TENSOR.LILICORR_CONTEXT_PROJ: "lilicorr.context_proj", + MODEL_TENSOR.LILICORR_ATTN_IN_PROJ: "lilicorr.layers.{bid}.attn.in_proj", + MODEL_TENSOR.LILICORR_OUTPUT_NORM: "lilicorr.output_norm", + MODEL_TENSOR.LILICORR_ANCHOR_NORM: "lilicorr.anchor_norm", + MODEL_TENSOR.LILICORR_FACTOR_INPUT_PROJ: "lilicorr.factor_input_proj", + MODEL_TENSOR.LILICORR_OUT_HEAD: "lilicorr.out_head", + MODEL_TENSOR.LILICORR_IN_HEAD: "lilicorr.in_head", + MODEL_TENSOR.LILICORR_ANCHOR_OUT_HEAD: "lilicorr.anchor_out_head", MODEL_TENSOR.D2T: "d2t", MODEL_TENSOR.ATTN_V_GATE: "blk.{bid}.attn_v_gate", MODEL_TENSOR.ATTN_V_EXP: "blk.{bid}.attn_v_exps", @@ -5548,6 +5584,19 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DFLASH_SELECTOR_PREV, MODEL_TENSOR.DFLASH_SELECTOR_NEXT, MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN, + MODEL_TENSOR.LILICORR_TOKEN_PROJ, + MODEL_TENSOR.LILICORR_PASS_HIDDEN_PROJ, + MODEL_TENSOR.LILICORR_SLOT_EMBD, + MODEL_TENSOR.LILICORR_RANK_EMBD, + MODEL_TENSOR.LILICORR_ATTN_BIAS, + MODEL_TENSOR.LILICORR_CONTEXT_PROJ, + MODEL_TENSOR.LILICORR_ATTN_IN_PROJ, + MODEL_TENSOR.LILICORR_OUTPUT_NORM, + MODEL_TENSOR.LILICORR_ANCHOR_NORM, + MODEL_TENSOR.LILICORR_FACTOR_INPUT_PROJ, + MODEL_TENSOR.LILICORR_OUT_HEAD, + MODEL_TENSOR.LILICORR_IN_HEAD, + MODEL_TENSOR.LILICORR_ANCHOR_OUT_HEAD, ], MODEL_ARCH.MISTRAL4: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 02ac3e823a9d..cc2bd8e7b3b8 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1236,6 +1236,30 @@ def add_kda_safe_gate(self, value: bool) -> None: def add_kda_gate_lower_bound(self, value: float) -> None: self.add_float32(Keys.KDA.GATE_LOWER_BOUND.format(arch=self.arch), value) + def add_lilicorr_hidden_size(self, value: int) -> None: + self.add_uint32(Keys.LiliCorr.HIDDEN_SIZE, value) + + def add_lilicorr_candidate_topk(self, value: int) -> None: + self.add_uint32(Keys.LiliCorr.CANDIDATE_TOPK, value) + + def add_lilicorr_num_heads(self, value: int) -> None: + self.add_uint32(Keys.LiliCorr.NUM_HEADS, value) + + def add_lilicorr_num_layers(self, value: int) -> None: + self.add_uint32(Keys.LiliCorr.NUM_LAYERS, value) + + def add_lilicorr_factor_dim(self, value: int) -> None: + self.add_uint32(Keys.LiliCorr.FACTOR_DIM, value) + + def add_lilicorr_mlp_hidden_size(self, value: int) -> None: + self.add_uint32(Keys.LiliCorr.MLP_HIDDEN_SIZE, value) + + def add_lilicorr_vector_eps(self, value: float) -> None: + self.add_float32(Keys.LiliCorr.VECTOR_EPS, value) + + def add_lilicorr_logit_scale(self, value: float) -> None: + self.add_float32(Keys.LiliCorr.LOGIT_SCALE, value) + def add_tokenizer_model(self, model: str) -> None: self.add_string(Keys.Tokenizer.MODEL, model) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 6445784a359f..5eab5b1fdcb2 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -182,11 +182,48 @@ class TensorNameMap: MODEL_TENSOR.A_QF_PROJ_LINEAR: ( "projector.linear", ), + MODEL_TENSOR.LILICORR_TOKEN_PROJ: ( + "model.lilicorr.token_proj", + ), + MODEL_TENSOR.LILICORR_PASS_HIDDEN_PROJ: ( + "model.lilicorr.pass_hidden_proj", + ), + MODEL_TENSOR.LILICORR_SLOT_EMBD: ( + "model.lilicorr.slot_embedding", + ), + MODEL_TENSOR.LILICORR_RANK_EMBD: ( + "model.lilicorr.rank_embedding", + ), + MODEL_TENSOR.LILICORR_ATTN_BIAS: ( + "model.lilicorr.attention_bias", + ), + MODEL_TENSOR.LILICORR_CONTEXT_PROJ: ( + "model.lilicorr.context_proj", + ), + MODEL_TENSOR.LILICORR_OUTPUT_NORM: ( + "model.lilicorr.output_norm", + ), + MODEL_TENSOR.LILICORR_ANCHOR_NORM: ( + "model.lilicorr.anchor_norm", + ), + MODEL_TENSOR.LILICORR_FACTOR_INPUT_PROJ: ( + "model.lilicorr.factor_input_proj", + ), + MODEL_TENSOR.LILICORR_OUT_HEAD: ( + "model.lilicorr.out_head", + ), + MODEL_TENSOR.LILICORR_IN_HEAD: ( + "model.lilicorr.in_head", + ), + MODEL_TENSOR.LILICORR_ANCHOR_OUT_HEAD: ( + "model.lilicorr.anchor_out_head", + ), } block_mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = { # Attention norm MODEL_TENSOR.ATTN_NORM: ( + "model.lilicorr.layers.{bid}.attn_norm", # lilicorr "gpt_neox.layers.{bid}.input_layernorm", # gptneox "transformer.h.{bid}.ln_1", # gpt2 gpt-j refact qwen jais exaone "transformer.blocks.{bid}.norm_1", # mpt @@ -324,6 +361,7 @@ class TensorNameMap: # Attention output MODEL_TENSOR.ATTN_OUT: ( + "model.lilicorr.layers.{bid}.attn.out_proj", # lilicorr "gpt_neox.layers.{bid}.attention.dense", # gptneox "transformer.h.{bid}.attn.c_proj", # gpt2 refact qwen jais "transformer.blocks.{bid}.attn.out_proj", # mpt @@ -407,6 +445,8 @@ class TensorNameMap: # Feed-forward norm MODEL_TENSOR.FFN_NORM: ( + "model.lilicorr.layers.{bid}.mlp_norm", # lilicorr + "model.lilicorr.feature_mlp.{bid}.norm", # lilicorr "gpt_neox.layers.{bid}.post_attention_layernorm", # gptneox "transformer.h.{bid}.ln_2", # gpt2 refact qwen jais exaone "h.{bid}.post_attention_layernorm", # bloom @@ -504,6 +544,8 @@ class TensorNameMap: # Feed-forward up MODEL_TENSOR.FFN_UP: ( + "model.lilicorr.layers.{bid}.mlp.up_proj", # lilicorr + "model.lilicorr.feature_mlp.{bid}.up_proj", # lilicorr "gpt_neox.layers.{bid}.mlp.dense_h_to_4h", # gptneox "transformer.h.{bid}.mlp.c_fc", # gpt2 jais "transformer.blocks.{bid}.ffn.up_proj", # mpt @@ -640,6 +682,8 @@ class TensorNameMap: # Feed-forward down MODEL_TENSOR.FFN_DOWN: ( + "model.lilicorr.layers.{bid}.mlp.down_proj", # lilicorr + "model.lilicorr.feature_mlp.{bid}.down_proj", # lilicorr "gpt_neox.layers.{bid}.mlp.dense_4h_to_h", # gptneox "transformer.h.{bid}.mlp.c_proj", # gpt2 refact qwen jais "transformer.blocks.{bid}.ffn.down_proj", # mpt @@ -1554,6 +1598,9 @@ class TensorNameMap: MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: ( "model.candidate_selector.hidden_projection", ), + MODEL_TENSOR.LILICORR_ATTN_IN_PROJ: ( + "model.lilicorr.layers.{bid}.attn.in_proj", + ), MODEL_TENSOR.CLS: ( "classifier", # jina diff --git a/include/llama.h b/include/llama.h index f15ca9f75d7e..3af6e22233ef 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1056,7 +1056,8 @@ extern "C" { // "state" here means extra hidden state carried over from a previous stage, e.g.: // - MTP: state from N layers of the target model // - Qwen3 VL (deepstack): state from N layers of the vision encoder - // Returns false if the context does not take a state embedding (currently only MTP contexts do) + // - LiliCorr: context embeddings for the transformers + // Returns false if the context does not take a state embedding (MTP and LiliCorr contexts do) LLAMA_API bool llama_batch_ext_set_embd_state( struct llama_batch_ext * batch, int32_t idx, diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 1acda79d864e..5744c5fb844d 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -389,6 +389,15 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_DFLASH_SELECTOR_RANK, "%s.selector_rank" }, { LLM_KV_DFLASH_SELECTOR_TOP_K, "%s.selector_top_k" }, + { LLM_KV_LILICORR_HIDDEN_SIZE, "lilicorr.hidden_size" }, + { LLM_KV_LILICORR_CANDIDATE_TOPK, "lilicorr.candidate_topk" }, + { LLM_KV_LILICORR_NUM_HEADS, "lilicorr.num_heads" }, + { LLM_KV_LILICORR_NUM_LAYERS, "lilicorr.num_layers" }, + { LLM_KV_LILICORR_VECTOR_EPS, "lilicorr.vector_eps" }, + { LLM_KV_LILICORR_LOGIT_SCALE, "lilicorr.logit_scale" }, + { LLM_KV_LILICORR_FACTOR_DIM, "lilicorr.factor_dim" }, + { LLM_KV_LILICORR_MLP_HIDDEN_SIZE, "lilicorr.mlp_hidden_size" }, + { LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" }, // sentence-transformers dense modules feature dims { LLM_KV_DENSE_2_FEAT_IN, "%s.dense_2_feat_in" }, @@ -759,6 +768,19 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_DFLASH_SELECTOR_PREV, "selector_predecessor" }, { LLM_TENSOR_DFLASH_SELECTOR_NEXT, "selector_successor" }, { LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "selector_hidden" }, + { LLM_TENSOR_LILICORR_TOKEN_PROJ, "lilicorr.token_proj" }, + { LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "lilicorr.pass_hidden_proj" }, + { LLM_TENSOR_LILICORR_SLOT_EMBD, "lilicorr.slot_embedding" }, + { LLM_TENSOR_LILICORR_RANK_EMBD, "lilicorr.rank_embedding" }, + { LLM_TENSOR_LILICORR_ATTN_BIAS, "lilicorr.attention_bias" }, + { LLM_TENSOR_LILICORR_CONTEXT_PROJ, "lilicorr.context_proj" }, + { LLM_TENSOR_LILICORR_ATTN_IN_PROJ, "lilicorr.layers.%d.attn.in_proj" }, + { LLM_TENSOR_LILICORR_OUTPUT_NORM, "lilicorr.output_norm" }, + { LLM_TENSOR_LILICORR_ANCHOR_NORM, "lilicorr.anchor_norm" }, + { LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, "lilicorr.factor_input_proj" }, + { LLM_TENSOR_LILICORR_OUT_HEAD, "lilicorr.out_head" }, + { LLM_TENSOR_LILICORR_IN_HEAD, "lilicorr.in_head" }, + { LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, "lilicorr.anchor_out_head" }, { LLM_TENSOR_ATTN_V_GATE, "blk.%d.attn_v_gate" }, { LLM_TENSOR_ATTN_V_EXPS, "blk.%d.attn_v_exps" }, }; @@ -1076,6 +1098,20 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_DFLASH_SELECTOR_PREV, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, {LLM_TENSOR_DFLASH_SELECTOR_NEXT, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, {LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_LILICORR_TOKEN_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_LILICORR_SLOT_EMBD, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_ADD}}, + {LLM_TENSOR_LILICORR_RANK_EMBD, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_ADD}}, + {LLM_TENSOR_LILICORR_ATTN_BIAS, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_ADD}}, + {LLM_TENSOR_LILICORR_CONTEXT_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_LILICORR_ATTN_IN_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_LILICORR_OUTPUT_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_LILICORR_ANCHOR_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_LILICORR_OUT_HEAD, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_LILICORR_IN_HEAD, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_V_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_V_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, }; diff --git a/src/llama-arch.h b/src/llama-arch.h index e6e9bd4b94c4..94771c2a1c32 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -432,6 +432,14 @@ enum llm_kv { LLM_KV_DFLASH_CONV_GROUP_SIZE, LLM_KV_DFLASH_SELECTOR_RANK, LLM_KV_DFLASH_SELECTOR_TOP_K, + LLM_KV_LILICORR_HIDDEN_SIZE, + LLM_KV_LILICORR_CANDIDATE_TOPK, + LLM_KV_LILICORR_NUM_HEADS, + LLM_KV_LILICORR_NUM_LAYERS, + LLM_KV_LILICORR_VECTOR_EPS, + LLM_KV_LILICORR_LOGIT_SCALE, + LLM_KV_LILICORR_FACTOR_DIM, + LLM_KV_LILICORR_MLP_HIDDEN_SIZE, LLM_KV_NORM_BEFORE_RESIDUAL, LLM_KV_NORM_BEFORE_FC, @@ -766,6 +774,19 @@ enum llm_tensor { LLM_TENSOR_DFLASH_SELECTOR_PREV, LLM_TENSOR_DFLASH_SELECTOR_NEXT, LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, + LLM_TENSOR_LILICORR_TOKEN_PROJ, + LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, + LLM_TENSOR_LILICORR_SLOT_EMBD, + LLM_TENSOR_LILICORR_RANK_EMBD, + LLM_TENSOR_LILICORR_ATTN_BIAS, + LLM_TENSOR_LILICORR_CONTEXT_PROJ, + LLM_TENSOR_LILICORR_ATTN_IN_PROJ, + LLM_TENSOR_LILICORR_OUTPUT_NORM, + LLM_TENSOR_LILICORR_ANCHOR_NORM, + LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, + LLM_TENSOR_LILICORR_OUT_HEAD, + LLM_TENSOR_LILICORR_IN_HEAD, + LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, LLM_TENSOR_ATTN_V_GATE, LLM_TENSOR_ATTN_V_EXPS, }; diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 3c4b52008646..321aa9d23647 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1147,6 +1147,10 @@ size_t llama_batch_ext_select_n_embd_state(llama_context_type ctx_type, const ll if (ctx_type == LLAMA_CONTEXT_TYPE_MTP) { return hparams.n_embd_out(); } + // TODO: this probably should become a function of `llm_arch` + if (hparams.lilicorr_candidate_topk > 0) { + return hparams.n_embd_inp(); + } return 0; } diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 064de68f36c2..922164a544b8 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -264,6 +264,12 @@ struct llama_hparams { uint32_t dflash_conv_group_size = 0; uint32_t dflash_selector_rank = 0; uint32_t dflash_selector_top_k = 0; + uint32_t lilicorr_hidden_size = 0; + uint32_t lilicorr_candidate_topk = 0; + uint32_t lilicorr_num_heads = 0; + uint32_t lilicorr_layers = 0; + float lilicorr_vector_eps = 0.0f; + float lilicorr_logit_scale = 0.0f; // llama4 smallthinker uint32_t n_moe_layer_step = 0; diff --git a/src/llama-model.h b/src/llama-model.h index bc640e0ddb2e..bcfbf1e2cb39 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -386,6 +386,20 @@ struct llama_layer { struct ggml_tensor * dflash_ffn_conv_base = nullptr; struct ggml_tensor * dflash_ffn_conv_proj = nullptr; + struct ggml_tensor * lilicorr_attn_norm = nullptr; + struct ggml_tensor * lilicorr_mlp_norm = nullptr; + struct ggml_tensor * lilicorr_attn_in_proj = nullptr; + struct ggml_tensor * lilicorr_attn_in_proj_b = nullptr; + struct ggml_tensor * lilicorr_attn_out_proj = nullptr; + struct ggml_tensor * lilicorr_attn_out_proj_b = nullptr; + struct ggml_tensor * lilicorr_attn_out_proj_s = nullptr; + struct ggml_tensor * lilicorr_mlp_0 = nullptr; + struct ggml_tensor * lilicorr_mlp_0_b = nullptr; + struct ggml_tensor * lilicorr_mlp_0_s = nullptr; + struct ggml_tensor * lilicorr_mlp_2 = nullptr; + struct ggml_tensor * lilicorr_mlp_2_b = nullptr; + struct ggml_tensor * lilicorr_mlp_2_s = nullptr; + // mamba proj struct ggml_tensor * ssm_in = nullptr; struct ggml_tensor * ssm_x = nullptr; @@ -764,6 +778,36 @@ struct llama_model { struct ggml_tensor * dflash_selector_next = nullptr; struct ggml_tensor * dflash_selector_hidden = nullptr; + struct ggml_tensor * lilicorr_token_proj = nullptr; + struct ggml_tensor * lilicorr_token_proj_b = nullptr; + struct ggml_tensor * lilicorr_token_proj_s = nullptr; + struct ggml_tensor * lilicorr_pass_hidden_proj = nullptr; + struct ggml_tensor * lilicorr_pass_hidden_proj_b = nullptr; + struct ggml_tensor * lilicorr_pass_hidden_proj_s = nullptr; + struct ggml_tensor * lilicorr_feature_norm = nullptr; + struct ggml_tensor * lilicorr_feature_norm_b = nullptr; + struct ggml_tensor * lilicorr_feature_mlp_1 = nullptr; + struct ggml_tensor * lilicorr_feature_mlp_1_b = nullptr; + struct ggml_tensor * lilicorr_feature_mlp_3 = nullptr; + struct ggml_tensor * lilicorr_feature_mlp_3_b = nullptr; + struct ggml_tensor * lilicorr_slot_embedding = nullptr; + struct ggml_tensor * lilicorr_rank_embedding = nullptr; + struct ggml_tensor * lilicorr_attention_bias = nullptr; + struct ggml_tensor * lilicorr_context_proj = nullptr; + struct ggml_tensor * lilicorr_context_proj_b = nullptr; + struct ggml_tensor * lilicorr_context_proj_s = nullptr; + struct ggml_tensor * lilicorr_output_norm = nullptr; + struct ggml_tensor * lilicorr_anchor_norm = nullptr; + struct ggml_tensor * lilicorr_factor_input_proj = nullptr; + struct ggml_tensor * lilicorr_factor_input_proj_b = nullptr; + struct ggml_tensor * lilicorr_out_head = nullptr; + struct ggml_tensor * lilicorr_out_head_b = nullptr; + struct ggml_tensor * lilicorr_in_head = nullptr; + struct ggml_tensor * lilicorr_in_head_b = nullptr; + struct ggml_tensor * lilicorr_anchor_out_head = nullptr; + struct ggml_tensor * lilicorr_anchor_out_head_b = nullptr; + + // unified vector to store target-model extracted layer ids in eagle3, dflash, etc. std::vector target_layer_ids; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index aeeadbbd15ea..982351dc860a 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -33,6 +33,19 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_DFLASH_CONV_GROUP_SIZE, hparams.dflash_conv_group_size, false); ml.get_key(LLM_KV_DFLASH_SELECTOR_RANK, hparams.dflash_selector_rank, false); ml.get_key(LLM_KV_DFLASH_SELECTOR_TOP_K, hparams.dflash_selector_top_k, false); + const bool has_lilicorr = ml.get_key(LLM_KV_LILICORR_HIDDEN_SIZE, hparams.lilicorr_hidden_size, false); + ml.get_key(LLM_KV_LILICORR_CANDIDATE_TOPK, hparams.lilicorr_candidate_topk, has_lilicorr); + ml.get_key(LLM_KV_LILICORR_NUM_HEADS, hparams.lilicorr_num_heads, has_lilicorr); + ml.get_key(LLM_KV_LILICORR_NUM_LAYERS, hparams.lilicorr_layers, has_lilicorr); + ml.get_key(LLM_KV_LILICORR_VECTOR_EPS, hparams.lilicorr_vector_eps, has_lilicorr); + ml.get_key(LLM_KV_LILICORR_LOGIT_SCALE, hparams.lilicorr_logit_scale, has_lilicorr); + if (has_lilicorr && (hparams.dflash_block_size < 2 || hparams.lilicorr_hidden_size == 0 || + hparams.lilicorr_num_heads == 0 || hparams.lilicorr_hidden_size % hparams.lilicorr_num_heads || + hparams.lilicorr_candidate_topk == 0 || hparams.lilicorr_layers == 0 || + hparams.dflash_selector_top_k > 0 || + (uint64_t) hparams.lilicorr_candidate_topk * (hparams.lilicorr_candidate_topk + 1) > hparams.n_embd)) { + throw std::runtime_error("invalid LiLiCorr DFlash metadata"); + } if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) { throw std::runtime_error("DFlash model requires 'target_layers' in GGUF metadata"); @@ -137,6 +150,10 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { } const struct ggml_tensor * selector_meta = ml->get_tensor_meta("selector_hidden.weight"); + const bool conv_meta = hparams.dflash_conv_kernel_size > 0 && hparams.dflash_conv_group_size > 0; + if (conv_meta && n_embd % hparams.dflash_conv_group_size != 0) { + throw std::runtime_error("DFlash hidden size must be divisible by conv_group_size"); + } if (selector_meta) { const int64_t rank = hparams.dflash_selector_rank; if (rank <= 0 || hparams.dflash_block_size <= 0 || hparams.dflash_selector_top_k <= 0 || @@ -159,6 +176,66 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { hparams.dflash_selector_rank, hparams.dflash_selector_top_k); } + if (hparams.lilicorr_layers > 0) { + const int64_t hidden = hparams.lilicorr_hidden_size; + const int64_t top_k = hparams.lilicorr_candidate_topk; + const int64_t heads = hparams.lilicorr_num_heads; + uint32_t factor = 0, mlp = 0; + ml->get_key(LLM_KV_LILICORR_FACTOR_DIM, factor); + ml->get_key(LLM_KV_LILICORR_MLP_HIDDEN_SIZE, mlp); + const int64_t slots = hparams.dflash_block_size - 1; + const int feature_bid = n_layer_all + hparams.lilicorr_layers; + if (factor == 0 || mlp == 0 || hparams.lilicorr_layers > layers.size()) { + throw std::runtime_error("invalid LiLiCorr head dimensions"); + } + if (hidden != n_embd) { + lilicorr_token_proj = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "weight"), { n_embd, hidden }, 0); + lilicorr_token_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "bias"), { hidden }, 0); + lilicorr_token_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_TOKEN_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED); + } + lilicorr_pass_hidden_proj = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "weight"), { hparams.n_embd, hidden }, 0); + lilicorr_pass_hidden_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "bias"), { hidden }, 0); + lilicorr_pass_hidden_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_PASS_HIDDEN_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED); + lilicorr_feature_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", feature_bid), { 5 }, 0); + lilicorr_feature_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", feature_bid), { 5 }, 0); + lilicorr_feature_mlp_1 = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", feature_bid), { 5, hidden }, 0); + lilicorr_feature_mlp_1_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", feature_bid), { hidden }, 0); + lilicorr_feature_mlp_3 = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", feature_bid), { hidden, hidden }, 0); + lilicorr_feature_mlp_3_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", feature_bid), { hidden }, 0); + lilicorr_slot_embedding = create_tensor(tn(LLM_TENSOR_LILICORR_SLOT_EMBD, "weight"), { hidden, slots }, 0); + lilicorr_rank_embedding = create_tensor(tn(LLM_TENSOR_LILICORR_RANK_EMBD, "weight"), { hidden, top_k }, 0); + lilicorr_attention_bias = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_BIAS, "bias"), { slots * top_k, slots * top_k, heads }, 0); + lilicorr_context_proj = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "weight"), { n_embd, hidden }, 0); + lilicorr_context_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "bias"), { hidden }, 0); + lilicorr_context_proj_s = create_tensor(tn(LLM_TENSOR_LILICORR_CONTEXT_PROJ, "scale"), { 1 }, TENSOR_NOT_REQUIRED); + for (uint32_t i = 0; i < hparams.lilicorr_layers; ++i) { + const int bid = n_layer_all + i; + layers[i].lilicorr_attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", bid), { hidden }, 0); + layers[i].lilicorr_mlp_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", bid), { hidden }, 0); + layers[i].lilicorr_attn_in_proj = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_IN_PROJ, "weight", i), { hidden, 3 * hidden }, 0); + layers[i].lilicorr_attn_in_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_ATTN_IN_PROJ, "bias", i), { 3 * hidden }, 0); + layers[i].lilicorr_attn_out_proj = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", bid), { hidden, hidden }, 0); + layers[i].lilicorr_attn_out_proj_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", bid), { hidden }, 0); + layers[i].lilicorr_attn_out_proj_s = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED); + layers[i].lilicorr_mlp_0 = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", bid), { hidden, mlp }, 0); + layers[i].lilicorr_mlp_0_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", bid), { mlp }, 0); + layers[i].lilicorr_mlp_0_s = create_tensor(tn(LLM_TENSOR_FFN_UP, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED); + layers[i].lilicorr_mlp_2 = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", bid), { mlp, hidden }, 0); + layers[i].lilicorr_mlp_2_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", bid), { hidden }, 0); + layers[i].lilicorr_mlp_2_s = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "scale", bid), { 1 }, TENSOR_NOT_REQUIRED); + } + lilicorr_output_norm = create_tensor(tn(LLM_TENSOR_LILICORR_OUTPUT_NORM, "weight"), { hidden }, 0); + lilicorr_anchor_norm = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_NORM, "weight"), { hidden }, 0); + lilicorr_factor_input_proj = create_tensor(tn(LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, "weight"), { 3 * hidden, hidden }, 0); + lilicorr_factor_input_proj_b = create_tensor(tn(LLM_TENSOR_LILICORR_FACTOR_INPUT_PROJ, "bias"), { hidden }, 0); + lilicorr_out_head = create_tensor(tn(LLM_TENSOR_LILICORR_OUT_HEAD, "weight"), { hidden, factor }, 0); + lilicorr_out_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_OUT_HEAD, "bias"), { factor }, 0); + lilicorr_in_head = create_tensor(tn(LLM_TENSOR_LILICORR_IN_HEAD, "weight"), { hidden, factor }, 0); + lilicorr_in_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_IN_HEAD, "bias"), { factor }, 0); + lilicorr_anchor_out_head = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, "weight"), { hidden, factor }, 0); + lilicorr_anchor_out_head_b = create_tensor(tn(LLM_TENSOR_LILICORR_ANCHOR_OUT_HEAD, "bias"), { factor }, 0); + } + fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED); output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) @@ -244,7 +321,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0); layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0); - if (selector_meta) { + if (conv_meta) { const int64_t kernel = hparams.dflash_conv_kernel_size; const int64_t groups = n_embd / hparams.dflash_conv_group_size; const int64_t projected = 2 * kernel * groups; @@ -561,6 +638,168 @@ static void build_dflash2_selector(llm_graph_context & g, const llama_model & mo ggml_build_forward_expand(g.gf, packed); } +static void build_lilicorr_scores(llm_graph_context & g, const llama_model & model, + ggml_tensor * tok_embd, ggml_tensor * anchors) { + ggml_context * ctx0 = g.ctx0; + auto & res = g.res; + + const auto & hparams = g.hparams; + const int64_t n_tokens = g.n_tokens; + const int64_t n_embd = g.n_embd; + + const int64_t hidden = hparams.lilicorr_hidden_size; + const int64_t top_k = hparams.lilicorr_candidate_topk; + const int64_t heads = hparams.lilicorr_num_heads; + const int64_t n_blocks = g.ubatch.n_seqs_unq; + GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0); + GGML_ASSERT(res->t_logits->ne[1] == n_tokens); + GGML_ASSERT(anchors && tok_embd); + + const int64_t tokens_per_block = n_tokens / n_blocks; + const int64_t block_size = hparams.dflash_block_size; + const int64_t row_used = top_k + top_k * top_k; + const int64_t slots = block_size - 1; + const int64_t nodes = slots * top_k; + const int64_t factor_dim = model.lilicorr_out_head->ne[1]; + GGML_ASSERT(tokens_per_block == block_size); + + ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k); + ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens); + ggml_tensor * candidate_logits = ggml_reshape_2d(ctx0, + ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens); + // Rank features and embeddings require candidates in descending logit order. + ggml_tensor * order = ggml_argsort(ctx0, candidate_logits, GGML_SORT_ORDER_DESC); + ggml_tensor * candidate_rows = ggml_reshape_3d(ctx0, candidates, 1, top_k, n_tokens); + candidates = ggml_reshape_2d(ctx0, ggml_get_rows(ctx0, candidate_rows, order), top_k, n_tokens); + ggml_tensor * prob_rows = ggml_reshape_3d(ctx0, ggml_soft_max(ctx0, res->t_logits), 1, res->t_logits->ne[0], n_tokens); + ggml_tensor * probs = ggml_reshape_2d(ctx0, + ggml_get_rows(ctx0, prob_rows, candidates), top_k, n_tokens); + + ggml_tensor * cand_run = ggml_reshape_3d(ctx0, candidates, top_k, n_tokens / n_blocks, n_blocks); + ggml_tensor * probs_run = ggml_reshape_3d(ctx0, probs, top_k, n_tokens / n_blocks, n_blocks); + cand_run = ggml_cont(ctx0, ggml_view_3d(ctx0, cand_run, top_k, slots, n_blocks, + cand_run->nb[1], cand_run->nb[2], cand_run->nb[1])); + probs_run = ggml_cont(ctx0, ggml_view_3d(ctx0, probs_run, top_k, slots, n_blocks, + probs_run->nb[1], probs_run->nb[2], probs_run->nb[1])); + + ggml_tensor * lp = ggml_log(ctx0, probs_run); + ggml_tensor * lp_blk = ggml_reshape_3d(ctx0, lp, top_k, slots, n_blocks); + ggml_tensor * top_lp = ggml_view_3d(ctx0, lp_blk, 1, slots, n_blocks, + lp_blk->nb[1], lp_blk->nb[2], 0); + ggml_tensor * relative = ggml_sub(ctx0, lp_blk, top_lp); + ggml_tensor * feature = ggml_reshape_4d(ctx0, lp_blk, 1, top_k, slots, n_blocks); + + ggml_tensor * rank = ggml_arange(ctx0, 0, top_k, 1); + rank = ggml_repeat(ctx0, ggml_reshape_3d(ctx0, rank, top_k, 1, 1), lp_blk); + ggml_tensor * first = ggml_scale_bias(ctx0, ggml_step(ctx0, rank), -1.0f, 1.0f); + + for (ggml_tensor * item : { probs_run, relative, + ggml_scale(ctx0, rank, top_k == 1 ? 0.0f : 1.0f / (top_k - 1)), first }) { + feature = ggml_concat(ctx0, feature, ggml_reshape_4d(ctx0, item, 1, top_k, slots, n_blocks), 0); + } + feature = ggml_reshape_2d(ctx0, feature, 5, nodes * n_blocks); + + ggml_tensor * embeds = ggml_get_rows(ctx0, tok_embd, ggml_reshape_1d(ctx0, cand_run, nodes * n_blocks)); + ggml_tensor * token_states = n_embd == hidden ? embeds : ggml_add(ctx0, g.build_lora_mm(model.lilicorr_token_proj, embeds, model.lilicorr_token_proj_s), model.lilicorr_token_proj_b); + + ggml_tensor * pass_blk = ggml_reshape_3d(ctx0, res->t_embd, n_embd, block_size, n_blocks); + ggml_tensor * passes = ggml_reshape_2d(ctx0, ggml_cont(ctx0, ggml_view_3d(ctx0, pass_blk, n_embd, slots, n_blocks, + pass_blk->nb[1], pass_blk->nb[2], pass_blk->nb[1])), + n_embd, slots * n_blocks); + + ggml_tensor * pass_states = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_pass_hidden_proj, passes, model.lilicorr_pass_hidden_proj_s), model.lilicorr_pass_hidden_proj_b); + + pass_states = ggml_reshape_4d(ctx0, pass_states, hidden, 1, slots, n_blocks); + pass_states = ggml_reshape_2d(ctx0, ggml_repeat_4d(ctx0, pass_states, hidden, top_k, slots, n_blocks), hidden, nodes * n_blocks); + + ggml_tensor * feature_states = ggml_norm(ctx0, feature, 1.0e-5f); + feature_states = ggml_add(ctx0, ggml_mul(ctx0, feature_states, model.lilicorr_feature_norm), model.lilicorr_feature_norm_b); + feature_states = ggml_silu(ctx0, ggml_add(ctx0, g.build_lora_mm(model.lilicorr_feature_mlp_1, feature_states), model.lilicorr_feature_mlp_1_b)); + feature_states = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_feature_mlp_3, feature_states), model.lilicorr_feature_mlp_3_b); + + ggml_tensor * slot = ggml_reshape_4d(ctx0, model.lilicorr_slot_embedding, hidden, 1, slots, 1); + ggml_tensor * rank_embd = ggml_reshape_4d(ctx0, model.lilicorr_rank_embedding, hidden, top_k, 1, 1); + ggml_tensor * pos = ggml_add(ctx0, ggml_repeat_4d(ctx0, slot, hidden, top_k, slots, n_blocks), ggml_repeat_4d(ctx0, rank_embd, hidden, top_k, slots, n_blocks)); + pos = ggml_reshape_2d(ctx0, pos, hidden, nodes * n_blocks); + + ggml_tensor * x = ggml_add(ctx0, ggml_add(ctx0, token_states, pass_states), feature_states); + x = ggml_add(ctx0, x, pos); + ggml_tensor * anchor_blk = ggml_reshape_3d(ctx0, anchors, n_embd, block_size, n_blocks); + ggml_tensor * anchor = ggml_reshape_2d(ctx0, + ggml_cont(ctx0, ggml_view_3d(ctx0, anchor_blk, n_embd, 1, n_blocks, + anchor_blk->nb[1], anchor_blk->nb[2], 0)), n_embd, n_blocks); + ggml_tensor * anchor_state = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_context_proj, anchor, model.lilicorr_context_proj_s), model.lilicorr_context_proj_b); + + const int64_t head_dim = hidden / heads; + for (uint32_t layer = 0; layer < hparams.lilicorr_layers; ++layer) { + const auto & weights = model.layers[layer]; + ggml_tensor * attn_input = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), weights.lilicorr_attn_norm); + ggml_tensor * qkv = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_attn_in_proj, attn_input), weights.lilicorr_attn_in_proj_b); + qkv = ggml_reshape_3d(ctx0, qkv, 3 * hidden, nodes, n_blocks); + ggml_tensor * q = ggml_cont(ctx0, ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks, + qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], 0)); + ggml_tensor * k = ggml_cont(ctx0, ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks, + qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], hidden * qkv->nb[0])); + ggml_tensor * v = ggml_view_4d(ctx0, qkv, head_dim, nodes, heads, n_blocks, + qkv->nb[1], head_dim * qkv->nb[0], qkv->nb[2], 2 * hidden * qkv->nb[0]); + ggml_tensor * attn = ggml_soft_max_ext(ctx0, ggml_mul_mat(ctx0, k, q), model.lilicorr_attention_bias, + 1.0f / std::sqrt(static_cast(head_dim)), 0.0f); + ggml_tensor * joined = ggml_out_prod(ctx0, v, ggml_transpose(ctx0, attn)); + joined = ggml_permute(ctx0, joined, 0, 2, 1, 3); + ggml_tensor * attn_output = ggml_cont_2d(ctx0, joined, hidden, nodes * n_blocks); + attn_output = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_attn_out_proj, attn_output, weights.lilicorr_attn_out_proj_s), weights.lilicorr_attn_out_proj_b); + x = ggml_add(ctx0, x, attn_output); + ggml_tensor * mlp = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), weights.lilicorr_mlp_norm); + mlp = ggml_silu(ctx0, ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_mlp_0, mlp, weights.lilicorr_mlp_0_s), weights.lilicorr_mlp_0_b)); + mlp = ggml_add(ctx0, g.build_lora_mm(weights.lilicorr_mlp_2, mlp, weights.lilicorr_mlp_2_s), weights.lilicorr_mlp_2_b); + x = ggml_add(ctx0, x, mlp); + } + + x = ggml_mul(ctx0, ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps), model.lilicorr_output_norm); + anchor_state = ggml_mul(ctx0, ggml_rms_norm(ctx0, anchor_state, hparams.f_norm_rms_eps), model.lilicorr_anchor_norm); + ggml_tensor * anchor_nodes = ggml_reshape_3d(ctx0, anchor_state, hidden, 1, n_blocks); + anchor_nodes = ggml_reshape_2d(ctx0, ggml_repeat_4d(ctx0, anchor_nodes, hidden, nodes, n_blocks, 1), hidden, nodes * n_blocks); + ggml_tensor * factor_in = ggml_concat(ctx0, x, anchor_nodes, 0); + factor_in = ggml_concat(ctx0, factor_in, ggml_mul(ctx0, x, anchor_nodes), 0); + ggml_tensor * factors = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_factor_input_proj, factor_in), model.lilicorr_factor_input_proj_b); + factors = ggml_silu(ctx0, factors); + + ggml_tensor * out = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_out_head, factors), model.lilicorr_out_head_b); + out = ggml_reshape_4d(ctx0, ggml_l2_norm(ctx0, out, hparams.lilicorr_vector_eps), factor_dim, top_k, slots, n_blocks); + + ggml_tensor * in = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_in_head, factors), model.lilicorr_in_head_b); + in = ggml_reshape_4d(ctx0, ggml_l2_norm(ctx0, in, hparams.lilicorr_vector_eps), factor_dim, top_k, slots, n_blocks); + + ggml_tensor * anchor_out = ggml_add(ctx0, g.build_lora_mm(model.lilicorr_anchor_out_head, anchor_state), model.lilicorr_anchor_out_head_b); + anchor_out = ggml_l2_norm(ctx0, anchor_out, hparams.lilicorr_vector_eps); + anchor_out = ggml_reshape_3d(ctx0, anchor_out, factor_dim, 1, n_blocks); + + ggml_tensor * first_in = ggml_cont(ctx0, ggml_view_4d(ctx0, in, factor_dim, top_k, 1, n_blocks, + in->nb[1], in->nb[2], in->nb[3], 0)); + ggml_tensor * first_score = ggml_scale(ctx0, + ggml_mul_mat(ctx0, first_in, ggml_reshape_4d(ctx0, anchor_out, factor_dim, 1, 1, n_blocks)), + hparams.lilicorr_logit_scale); + ggml_tensor * scores = ggml_pad(ctx0, + ggml_reshape_3d(ctx0, first_score, top_k, 1, n_blocks), top_k * top_k - top_k, 0, 0, 0); + if (slots > 1) { + ggml_tensor * next = ggml_cont(ctx0, ggml_view_4d(ctx0, in, factor_dim, top_k, slots - 1, n_blocks, + in->nb[1], in->nb[2], in->nb[3], in->nb[2])); + ggml_tensor * previous = ggml_cont(ctx0, ggml_view_4d(ctx0, out, factor_dim, top_k, slots - 1, n_blocks, + out->nb[1], out->nb[2], out->nb[3], 0)); + ggml_tensor * score = ggml_scale(ctx0, ggml_mul_mat(ctx0, next, previous), hparams.lilicorr_logit_scale); + scores = ggml_concat(ctx0, scores, + ggml_reshape_3d(ctx0, score, top_k * top_k, slots - 1, n_blocks), 1); + } + + ggml_tensor * row = ggml_concat(ctx0, ggml_cast(ctx0, ggml_reshape_3d(ctx0, cand_run, top_k, slots, n_blocks), + GGML_TYPE_F32), scores, 0); + ggml_tensor * packed = ggml_pad_ext(ctx0, row, 0, n_embd - row_used, 1, 0, 0, 0, 0, 0); + packed = ggml_reshape_2d(ctx0, packed, n_embd, n_tokens); + g.cb(packed, "lilicorr_lattice", -1); + res->t_h_nextn = packed; + ggml_build_forward_expand(g.gf, packed); +} + // DFlash decoder, dual-mode by batch type: // * embd batch -> fused target features: project + inject K/V into the cache. // * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens @@ -667,6 +906,10 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra } res->t_embd = inp_g; + if (model.lilicorr_pass_hidden_proj) { + // Expose context features for retrieval by LiliCorr + res->t_h_nextn = inp_g; + } ggml_build_forward_expand(gf, inp_g); return; @@ -682,22 +925,31 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra tok_embd = model_other->tok_embd; } - auto inp = std::make_unique(n_embd); - - inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); - ggml_set_input(inp->tokens); - res->t_inp_tokens = inp->tokens; - - ggml_tensor * inp_tokens = inp->tokens; + ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp_tokens); + res->t_inp_tokens = inp_tokens; + + ggml_tensor * inp_anchor = nullptr; + if (model.lilicorr_pass_hidden_proj && ubatch.embd_state) { + // LiliCorr sets inp_anchor to pass the context features to the lattice scorer + auto inp = std::make_unique(n_embd, n_embd); + inp->tokens = inp_tokens; + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(inp->h); + inp_anchor = inp->h; + res->add_input(std::move(inp)); + } else { + auto inp = std::make_unique(n_embd); + inp->tokens = inp_tokens; + res->add_input(std::move(inp)); + } - ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens); + ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp_tokens); if (hparams.f_embedding_scale != 0.0f) { inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale); } cb(inpL, "inp_noise_embd", -1); - res->add_input(std::move(inp)); - for (int il = 0; il < n_layer; ++il) { const auto & layer = model.layers[il]; @@ -848,6 +1100,8 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra if (model.dflash_selector_hidden) { build_dflash2_selector(*this, model, inp_tokens); + } else if (model.lilicorr_pass_hidden_proj && ubatch.embd_state) { + build_lilicorr_scores(*this, model, tok_embd, inp_anchor); } } From 2639a3b4e029d1be80ca4a472d62dfea23609075 Mon Sep 17 00:00:00 2001 From: Pascal Date: Thu, 15 Oct 2026 06:30:17 +0200 Subject: [PATCH 003/423] server: save only restorable slot states to the prompt cache (#30498) * server: save only restorable slot states to the prompt cache The idle-slot save loop stores every idle slot whenever a task starts, while only completion tasks restore from the prompt cache. Rerank, embedding, decision and infill states fill the cache up to its limit without ever being read back, and slots already cached are offered again at every task start. Each slot now tracks with cache_pending whether its state is restorable and missing from the cache: a completion release or an explicit slot restore sets it, prompt_save consumes it, prompt_clear resets it. Both save paths share this rule, and the LRU path restores only for completion tasks. * server: test that decision slots skip the prompt cache Decision tasks never mark their slot as cacheable, so starting new tasks must not save their idle slots to the RAM prompt cache, with or without a unified KV cache. * server: restore any reusing task from the prompt cache Embedding and rerank tasks never reuse their state, so their slots are cleared on release like child slots and the cache only holds states that a later task can restore. Infill and decision tasks now restore from the cache too, which replaces the cache_pending flag. Decision models disable the prompt cache unless --cache-ram is given, since their states seldom repeat. --- common/arg.cpp | 3 ++- common/common.h | 1 + tools/cli/README.md | 2 +- tools/server/README.md | 2 +- tools/server/server-context.cpp | 18 ++++++++++------ tools/server/server-task.h | 5 +++++ tools/server/tests/unit/test_rerank.py | 25 ++++++++++++++++++++++ tools/server/tests/unit/test_systemone.py | 26 +++++++++++++++++++++++ 8 files changed, 73 insertions(+), 9 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 29f6df32fe5f..8b302d5d77f0 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1717,10 +1717,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_env("LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT").set_examples({LLAMA_EXAMPLE_SERVER})); add_opt(common_arg( {"-cram", "--cache-ram"}, "N", - string_format("set the maximum cache size in MiB (default: %d, -1 - no limit, 0 - disable)" + string_format("set the maximum cache size in MiB (default: %d, 0 for decision models, -1 - no limit, 0 - disable)" "[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)", params.cache_ram_mib), [](common_params & params, int value) { params.cache_ram_mib = value; + params.cache_ram_explicit = true; } ).set_env("LLAMA_ARG_CACHE_RAM").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); add_opt(common_arg( diff --git a/common/common.h b/common/common.h index 2048d6a79585..2cd6e2ec9858 100644 --- a/common/common.h +++ b/common/common.h @@ -640,6 +640,7 @@ struct common_params { int32_t kv_unified_per_slot = 0; // max context per parallel slot; 0 = unset int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. + bool cache_ram_explicit = false; // set by --cache-ram, decision models disable the cache otherwise std::string public_path = ""; // NOLINT std::string api_prefix = ""; // NOLINT diff --git a/tools/cli/README.md b/tools/cli/README.md index 1637741193df..be8fb09d238b 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -148,7 +148,7 @@ | `--display-prompt, --no-display-prompt` | whether to print prompt at generation (default: true) | | `-co, --color [on\|off\|auto]` | Colorize output to distinguish prompt and user input from generations ('on', 'off', or 'auto', default: 'auto')
'auto' enables colors when output is to a terminal | | `-ctxcp, --ctx-checkpoints, --swa-checkpoints N` | max number of context checkpoints to create per slot (default: 32)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)
(env: LLAMA_ARG_CTX_CHECKPOINTS) | -| `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) | +| `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, 0 for decision models, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) | | `--context-shift, --no-context-shift` | whether to use context shift on infinite text generation (default: disabled)
(env: LLAMA_ARG_CONTEXT_SHIFT) | | `-sys, --system-prompt PROMPT` | system prompt to use with model (if applicable, depending on chat template) | | `--show-timings, --no-show-timings` | whether to show timing information after each response (default: true)
(env: LLAMA_ARG_SHOW_TIMINGS) | diff --git a/tools/server/README.md b/tools/server/README.md index 37c41f1bb2ae..4e4472ec24bd 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -165,7 +165,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--kv-unified-per-slot N` | context limit per parallel slot (default: unset, behavior unchanged).
when set without -c/--ctx-size, the shared KV pool is sized to n_parallel*N
(env: LLAMA_ARG_KV_UNIFIED_PER_SLOT) | | `-ctxcp, --ctx-checkpoints, --swa-checkpoints N` | max number of context checkpoints to create per slot (default: 32)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)
(env: LLAMA_ARG_CTX_CHECKPOINTS) | | `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)
(env: LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT) | -| `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) | +| `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, 0 for decision models, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) | | `-kvu, --kv-unified, -no-kvu, --no-kv-unified` | use single unified KV buffer shared across all sequences (default: enabled if number of slots is auto)
(env: LLAMA_ARG_KV_UNIFIED) | | `--cache-idle-slots, --no-cache-idle-slots` | save idle slots to the prompt cache on new task, and clear them when using unified KV (default: enabled, requires cache-ram)
(env: LLAMA_ARG_CACHE_IDLE_SLOTS) | | `--context-shift, --no-context-shift` | whether to use context shift on infinite text generation (default: disabled)
(env: LLAMA_ARG_CONTEXT_SHIFT) | diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 60c563345399..d5b4a7029d4f 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -594,7 +594,8 @@ struct server_slot { state = SLOT_STATE_IDLE; // do not keep context of the child slots - the parent's context is enough - if (task->is_child()) { + // nor of the stateless slots - no task ever reuses it + if (task->is_child() || task->is_stateless()) { prompt_clear(); } @@ -1489,6 +1490,13 @@ struct server_context_impl { batch.init(ctx_tgt, std::max(n_batch, params_base.n_parallel), n_embd); } + // decision requests rarely repeat a state, the prompt cache is opt-in for decision models + if (decision.type != COMMON_DECISION_TYPE_NONE && !params_base.cache_ram_explicit) { + SRV_INF("%s", "decision model: prompt cache disabled by default, use `--cache-ram N` to enable it\n"); + params_base.cache_ram_mib = 0; + params_base.cache_idle_slots = false; + } + if (params_base.cache_ram_mib != 0) { if (params_base.cache_ram_mib < 0) { SRV_TRC("prompt cache is enabled, size limit: %s\n", "no limit"); @@ -1782,8 +1790,8 @@ struct server_context_impl { if (ret) { update_cache = update_cache && prompt_cache; - // cache prompts only for completion tasks - update_cache = update_cache && task.type == SERVER_TASK_TYPE_COMPLETION; + // cache prompts only for the tasks that reuse them + update_cache = update_cache && !task.is_stateless(); if (update_cache) { SRV_TRC("%s", "updating prompt cache\n"); @@ -3597,9 +3605,7 @@ struct server_context_impl { return; } - const bool is_stateless_task = slot.task->type == SERVER_TASK_TYPE_EMBEDDING || slot.task->type == SERVER_TASK_TYPE_RERANK; - - if (slot.task->params.cache_prompt && !is_stateless_task) { + if (slot.task->params.cache_prompt && !slot.task->is_stateless()) { // reuse any previously computed tokens that are common with the new prompt n_past = slot.prompt.tokens.get_common_prefix(input_tokens); diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 50926ec3ca8a..e7bb043988d7 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -300,6 +300,11 @@ struct server_task { bool is_child() const { return id_parent != -1; } + + // embedding and rerank tasks always process their whole prompt, their state is never reused + bool is_stateless() const { + return type == SERVER_TASK_TYPE_EMBEDDING || type == SERVER_TASK_TYPE_RERANK; + } }; struct result_prompt_progress { diff --git a/tools/server/tests/unit/test_rerank.py b/tools/server/tests/unit/test_rerank.py index ded826710968..53f15af098f6 100644 --- a/tools/server/tests/unit/test_rerank.py +++ b/tools/server/tests/unit/test_rerank.py @@ -41,6 +41,31 @@ def test_rerank(): assert least_relevant["index"] == 3 +@pytest.mark.parametrize("kv_unified", [False, True]) +def test_rerank_skips_prompt_cache(kv_unified, tmp_path): + global server + server.n_slots = 2 + server.cache_ram = 100 + server.kv_unified = kv_unified + server.debug = True + server.log_path = str(tmp_path / "server.log") + server.start() + + for i in range(4): + res = server.make_request("POST", "/rerank", data={ + "query": f"query {i}", + "documents": TEST_DOCUMENTS[i:], + }) + assert res.status_code == 200 + + server.stop() + + with open(server.log_path) as f: + log = f.read() + assert "__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__" in log + assert "__TEST_TAG_CACHE_IDLE_SLOT__" not in log + + def test_rerank_tei_format(): global server server.start() diff --git a/tools/server/tests/unit/test_systemone.py b/tools/server/tests/unit/test_systemone.py index d9d98841055f..61c2ab09897f 100644 --- a/tools/server/tests/unit/test_systemone.py +++ b/tools/server/tests/unit/test_systemone.py @@ -158,6 +158,32 @@ def test_systemone_shared_prompt(): assert abs(res.body["answers"]["angry"]["noul"] - res_single.body["answers"]["angry"]["noul"]) < 0.01 +@pytest.mark.parametrize("cache_ram", [None, 100]) +def test_systemone_prompt_cache(cache_ram, tmp_path): + global server + server.n_slots = 1 + server.cache_ram = cache_ram + server.debug = True + server.log_path = str(tmp_path / "server.log") + server.start() + + # the third request comes back to the state of the first one + for state in ["Ticket A: " + TEST_STATE, "Ticket B: " + TEST_STATE * 2, "Ticket A: " + TEST_STATE]: + res = server.make_request("POST", "/v1/systemone", data={ + "state": state, + "questions": TEST_QUESTIONS, + }) + assert res.status_code == 200 + + server.stop() + + with open(server.log_path) as f: + log = f.read() + # the prompt cache is opt-in for decision models, and restores their states once enabled + assert ("__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__" in log) == (cache_ram is not None) + assert ("found better prompt" in log) == (cache_ram is not None) + + def test_systemone_images(): global server server = ServerPreset.tinyopenjev() From 1ddc262558fe736f4e046dc6d203eefbf0bbe5d8 Mon Sep 17 00:00:00 2001 From: Pascal Date: Thu, 15 Oct 2026 06:38:12 +0200 Subject: [PATCH 004/423] graph: keep the graph non empty without embeddings (#30522) An embedding-only graph has no logits, and build_pooling() returns early when embeddings are off, so since the models no longer expand t_embd the graph ends up empty and its inputs are never allocated. Expand t_embd in that case. --- src/llama-graph.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index f69fe63f8f33..6cadda9716ff 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3978,6 +3978,10 @@ void llm_graph_context::build_pooling( ggml_tensor * cls_out_b, ggml_tensor * cls_norm) const { if (!cparams.embeddings) { + // a graph without logits still computes t_embd, so that it is never empty + if (res->t_logits == nullptr && res->t_embd != nullptr) { + ggml_build_forward_expand(gf, res->t_embd); + } return; } From e467e5b5e5ba43071d6151b3d881deb9300d000e Mon Sep 17 00:00:00 2001 From: Pascal Date: Thu, 15 Oct 2026 06:39:18 +0200 Subject: [PATCH 005/423] kv-cache: honor PARTIAL_ONLY in the plain KV cache (#30514) * llama: honor PARTIAL_ONLY in the plain KV cache A non-SWA llama_kv_cache rolls back with seq_rm, so it holds no partial state, yet state_write ignored the flag and serialized every cell. A draft context made of a plain KV cache (MTP head, dense -md draft) therefore copied its whole KV into each server checkpoint, growing it by the draft KV size per token. The cache now writes and reads nothing for PARTIAL_ONLY unless it is SWA, and the iSWA wrapper no longer gates its base cache since it filters itself. * Update src/llama-kv-cache.cpp Co-authored-by: Georgi Gerganov * llama: delegate PARTIAL_ONLY to the caches of every memory wrapper The hybrid, hybrid-idx and dsa-iswa wrappers skipped their attention caches on PARTIAL_ONLY, which also dropped an SWA window that is partial state. Each cache now applies the flag itself, so the wrappers only delegate. The indexer is marked stale only when a restore rewrites cells. test-llama-context gains a case checking that the partial state is not larger than the full state and, without SWA, does not grow with the tokens. --------- Co-authored-by: Georgi Gerganov --- src/llama-kv-cache-dsa-iswa.cpp | 10 ++------ src/llama-kv-cache-iswa.cpp | 14 +++------- src/llama-kv-cache.cpp | 8 ++++-- src/llama-memory-hybrid-idx.cpp | 21 ++++++--------- src/llama-memory-hybrid.cpp | 16 +++--------- tests/llama-context/cases-kv.cpp | 44 ++++++++++++++++++++++++++++++++ 6 files changed, 68 insertions(+), 45 deletions(-) diff --git a/src/llama-kv-cache-dsa-iswa.cpp b/src/llama-kv-cache-dsa-iswa.cpp index dc10342a19c0..784c78aa7839 100644 --- a/src/llama-kv-cache-dsa-iswa.cpp +++ b/src/llama-kv-cache-dsa-iswa.cpp @@ -234,18 +234,12 @@ bool llama_kv_cache_dsa_iswa::get_can_shift() const { } void llama_kv_cache_dsa_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { - if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { - kv_dsa->state_write(io, seq_id, flags); - } - + kv_dsa->state_write(io, seq_id, flags); kv_swa->state_write(io, seq_id, flags); } void llama_kv_cache_dsa_iswa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { - if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { - kv_dsa->state_read(io, seq_id, flags); - } - + kv_dsa->state_read(io, seq_id, flags); kv_swa->state_read(io, seq_id, flags); } diff --git a/src/llama-kv-cache-iswa.cpp b/src/llama-kv-cache-iswa.cpp index a6ce547e64e1..da440f6a2dca 100644 --- a/src/llama-kv-cache-iswa.cpp +++ b/src/llama-kv-cache-iswa.cpp @@ -255,19 +255,13 @@ bool llama_kv_cache_iswa::get_can_shift() const { } void llama_kv_cache_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { - if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { - kv_base->state_write(io, seq_id, flags); - } - - kv_swa->state_write(io, seq_id, flags); + kv_base->state_write(io, seq_id, flags); + kv_swa ->state_write(io, seq_id, flags); } void llama_kv_cache_iswa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { - if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { - kv_base->state_read(io, seq_id, flags); - } - - kv_swa->state_read(io, seq_id, flags); + kv_base->state_read(io, seq_id, flags); + kv_swa ->state_read(io, seq_id, flags); } llama_kv_cache * llama_kv_cache_iswa::get_base() const { diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index d5208d034811..60a7e855497e 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2076,7 +2076,9 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla return; } - GGML_UNUSED(flags); + if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) && swa_type == LLAMA_SWA_TYPE_NONE) { + return; + } io.write(&n_stream, sizeof(n_stream)); @@ -2155,7 +2157,9 @@ const slot_info_vec_t * sinfos_in) { return; } - GGML_UNUSED(flags); + if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) && swa_type == LLAMA_SWA_TYPE_NONE) { + return; + } // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG] GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())); diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 21ec0b2fe1f3..2c5a57680dc8 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -257,13 +257,9 @@ void llama_memory_hybrid_idx::state_write(llama_io_write_i & io, llama_seq_id se llama_memory_hybrid::state_write(io, seq_id, flags); // [TAG_HYBRID_IDX_STATE] the indexer section goes last, so it is a pure suffix: an old reader stops early instead of misparsing it - // The indexer mirrors the attention cache, so it uses the same PARTIAL_ONLY gate. - if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { - if (mem_idx) { - mem_idx->state_write(io, seq_id, flags); - } + if (mem_idx) { + mem_idx->state_write(io, seq_id, flags); } - } void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { @@ -276,17 +272,16 @@ void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_ llama_kv_cache::slot_info_vec_t sinfos_attn; try { - if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { - get_mem_attn()->state_read_sinfo(io, seq_id, flags, mem_idx ? &sinfos_attn : nullptr, nullptr); - } + // the indexer cache has the SWA type of the attention cache, so both read the same partial state + get_mem_attn()->state_read_sinfo(io, seq_id, flags, mem_idx ? &sinfos_attn : nullptr, nullptr); get_mem_recr()->state_read(io, seq_id, flags); // [TAG_HYBRID_IDX_STATE] must mirror the write order in state_write - if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { - if (mem_idx) { - mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn); - // the restore rewrites the cells behind the pool layout's back + if (mem_idx) { + mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn); + // a restore that rewrites cells does it behind the pool layout's back + if (!sinfos_attn.empty()) { mem_idx_stale_set(seq_id, 0); } } diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 62f12de954a0..bbd4bee0b5a2 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -188,26 +188,18 @@ std::map llama_memory_hybrid::memory_breakdo } void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { - if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { - mem_attn->state_write(io, seq_id, flags); - } + mem_attn->state_write(io, seq_id, flags); mem_recr->state_write(io, seq_id, flags); } void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { - const bool read_attn = (flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0; - - if (read_attn) { - mem_attn->state_read(io, seq_id, flags); - } + mem_attn->state_read(io, seq_id, flags); try { mem_recr->state_read(io, seq_id, flags); } catch (...) { - // the attention part is already restored - undo it - if (read_attn) { - mem_attn->state_clear(seq_id); - } + // a failed restore leaves the sequence empty in both caches + mem_attn->state_clear(seq_id); throw; } diff --git a/tests/llama-context/cases-kv.cpp b/tests/llama-context/cases-kv.cpp index 1529774688e1..27e6493e45d3 100644 --- a/tests/llama-context/cases-kv.cpp +++ b/tests/llama-context/cases-kv.cpp @@ -341,6 +341,49 @@ struct case_state_roundtrip : ctx_test { } }; +// the partial state holds what seq_rm cannot roll back: recurrent states and SWA windows +struct case_state_partial_size : ctx_test { + const char * name() const override { return "ps"; } + + const char * description() const override { + return "the partial state is not larger than the full state and, without SWA, does not grow with the tokens"; + } + + case_result run(model_run & mr) const override { + auto ctx = mr.ctx(); + if (!ctx) { + return case_result::fail(); + } + + const llama_pos n_tokens = mr.tokens.size(); + const llama_pos n_half = n_tokens / 2; + + if (!decode_seq(ctx.get(), mr.tokens, 0, n_half, 0, true)) { + return case_result::fail(); + } + + const size_t partial_a = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + const size_t full_a = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE); + + if (!decode_seq(ctx.get(), mr.tokens, n_half, n_tokens, 0, true)) { + return case_result::fail(); + } + + const size_t partial_b = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + const size_t full_b = llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_NONE); + + TLOG_INF("%s: partial %zu -> %zu, full %zu -> %zu\n", __func__, partial_a, partial_b, full_a, full_b); + + CTX_REQUIRE(partial_a <= full_a && partial_b <= full_b, "the partial state is larger than the full state"); + + if (llama_model_n_swa(mr.model) == 0) { + CTX_REQUIRE(partial_a == partial_b, "the partial state grows with the tokens"); + } + + return case_result::pass(); + } +}; + // a failed restore must leave the sequence empty and must not change the logits of other sequences struct case_state_restore_failure : ctx_test { const char * name() const override { return "rf"; } @@ -686,6 +729,7 @@ std::vector make_kv_cases() { cases.push_back(std::make_unique(false)); cases.push_back(std::make_unique(true)); cases.push_back(std::make_unique()); + cases.push_back(std::make_unique()); cases.push_back(std::make_unique()); cases.push_back(std::make_unique()); cases.push_back(std::make_unique()); From 6e66cb5d4e9a2b4f02820a9f3053422b92ec4d51 Mon Sep 17 00:00:00 2001 From: Ethan Guo Date: Thu, 15 Oct 2026 14:39:47 +0800 Subject: [PATCH 006/423] metal : add BF16 variant of the tensor API flash attention kernel (#30491) * metal : add BF16 variant of the tensor API flash attention kernel * cont : move tensor FA kernels to kernels/fa_tensor.metal --- ggml/src/ggml-metal/CMakeLists.txt | 9 +- ggml/src/ggml-metal/ggml-metal-device.cpp | 4 +- ggml/src/ggml-metal/ggml-metal-device.m | 1 + ggml/src/ggml-metal/ggml-metal-ops.cpp | 4 +- ggml/src/ggml-metal/kernels/fa_f16.metal | 288 ------------------- ggml/src/ggml-metal/kernels/fa_tensor.metal | 298 ++++++++++++++++++++ tests/test-backend-ops.cpp | 7 + 7 files changed, 316 insertions(+), 295 deletions(-) create mode 100644 ggml/src/ggml-metal/kernels/fa_tensor.metal diff --git a/ggml/src/ggml-metal/CMakeLists.txt b/ggml/src/ggml-metal/CMakeLists.txt index 05cdea156316..91691b223669 100644 --- a/ggml/src/ggml-metal/CMakeLists.txt +++ b/ggml/src/ggml-metal/CMakeLists.txt @@ -45,6 +45,7 @@ set(METALLIB_KERNEL_SOURCES kernels/fa_q5_0.metal kernels/fa_q5_1.metal kernels/fa_q8_0.metal + kernels/fa_tensor.metal kernels/fa_vec_f16.metal kernels/fa_vec_f32.metal kernels/fa_vec_q4_0.metal @@ -233,12 +234,12 @@ else() VERBATIM ) - set(AIR_FA_TENSOR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/fa_f16_tensor.air") + set(AIR_FA_TENSOR "${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/fa_tensor_tensor.air") add_custom_command( OUTPUT ${AIR_FA_TENSOR} - COMMAND xcrun -sdk ${METAL_SDK} metal ${XC_FLAGS_TENSOR} -DGGML_METAL_HAS_TENSOR -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/fa_f16.metal -o ${AIR_FA_TENSOR} - DEPENDS kernels/fa_f16.metal ${METALLIB_KERNELS_FA_SHARED} kernels/common.h kernels/dequantize.h ${METALLIB_COMMON} ggml-metal-impl.h - COMMENT "Compiling kernels/fa_f16.metal (tensor API)" + COMMAND xcrun -sdk ${METAL_SDK} metal ${XC_FLAGS_TENSOR} -DGGML_METAL_HAS_TENSOR -I ${CMAKE_RUNTIME_OUTPUT_DIRECTORY} -c ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/kernels/fa_tensor.metal -o ${AIR_FA_TENSOR} + DEPENDS kernels/fa_tensor.metal kernels/common.h ${METALLIB_COMMON} ggml-metal-impl.h + COMMENT "Compiling kernels/fa_tensor.metal (tensor API)" VERBATIM ) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index b70129de8f19..2c7972b0a5d6 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1866,7 +1866,9 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_t const int32_t dk = (int32_t) op->src[1]->ne[0]; const int32_t dv = (int32_t) op->src[2]->ne[0]; - snprintf(base, 256, "kernel_flash_attn_ext_tensor_f16_dk%d_dv%d", dk, dv); + const char * type = op->src[1]->type == GGML_TYPE_BF16 ? "bf16" : "f16"; + + snprintf(base, 256, "kernel_flash_attn_ext_tensor_%s_dk%d_dv%d", type, dk, dv); snprintf(name, 256, "%s_mask=%d_sinks=%d_bias=%d_scap=%d", base, has_mask, has_sinks, has_bias, has_scap); ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 13aecbdd320a..d1bce8bad85a 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -118,6 +118,7 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi X(FA_Q5_0, fa_q5_0) \ X(FA_Q5_1, fa_q5_1) \ X(FA_Q8_0, fa_q8_0) \ + X(FA_TENSOR, fa_tensor) \ X(FA_VEC_F16, fa_vec_f16) \ X(FA_VEC_F32, fa_vec_f32) \ X(FA_VEC_Q4_0, fa_vec_q4_0) \ diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 115d73695a0b..53c981618511 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -3161,7 +3161,7 @@ static bool ggml_metal_op_flash_attn_ext_use_tensor(const ggml_tensor * op, bool return false; } - if (op->src[1]->type != GGML_TYPE_F16 && !ggml_metal_op_flash_attn_ext_use_kv_f16(op)) { + if (op->src[1]->type != GGML_TYPE_F16 && op->src[1]->type != GGML_TYPE_BF16 && !ggml_metal_op_flash_attn_ext_use_kv_f16(op)) { return false; } @@ -3678,7 +3678,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { /*.logit_softcap =*/ logit_softcap, }; - // shared memory layout: queries (half), scores (float), probabilities (half), row scale (float), rescale flag (int) + // shared memory layout: queries (16-bit), scores (float), probabilities (16-bit), row scale (float), rescale flag (int) const size_t smem = GGML_PAD(nqptg*ne00*sizeof(ggml_fp16_t) + nqptg*ncpsg*(sizeof(float) + sizeof(ggml_fp16_t)) + nqptg*sizeof(float) + sizeof(int32_t), 16); auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_tensor(lib, op, has_mask, has_sinks, has_bias, has_scap); diff --git a/ggml/src/ggml-metal/kernels/fa_f16.metal b/ggml/src/ggml-metal/kernels/fa_f16.metal index 0f44b419bdea..6e30b09e3a5a 100644 --- a/ggml/src/ggml-metal/kernels/fa_f16.metal +++ b/ggml/src/ggml-metal/kernels/fa_f16.metal @@ -75,291 +75,3 @@ template [[host_name("kernel_flash_attn_ext_bf16_dk576_dv512")]] kernel flash_at #undef FA_TYPES #undef FA_TYPES_BF #undef FA_TYPES_F32 - -#ifdef GGML_METAL_HAS_TENSOR - -constant bool FC_flash_attn_ext_tensor_has_mask [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 0)]]; -constant bool FC_flash_attn_ext_tensor_has_sinks [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 1)]]; -constant bool FC_flash_attn_ext_tensor_has_bias [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 2)]]; -constant bool FC_flash_attn_ext_tensor_has_scap [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 3)]]; - -// ref: https://arxiv.org/pdf/2307.08691.pdf -template< - short DK, // K head size - short DV, // V head size - short Q = OP_FLASH_ATTN_EXT_TENSOR_NQPSG, // queries per threadgroup - short C = OP_FLASH_ATTN_EXT_TENSOR_NCPSG, // cache items per threadgroup - short NSG = OP_FLASH_ATTN_EXT_TENSOR_NSG> // number of simd groups -kernel void kernel_flash_attn_ext_tensor( - constant ggml_metal_kargs_flash_attn_ext & args, - device const char * q, - device const char * k, - device const char * v, - device const char * mask, - device const char * sinks, - device const char * blk, - device char * dst, - threadgroup char * shmem [[threadgroup(0)]], - uint3 tgpig [[threadgroup_position_in_grid]], - ushort tiisg [[thread_index_in_simdgroup]], - ushort sgitg [[simdgroup_index_in_threadgroup]]) { - constexpr short NW = N_SIMDWIDTH; - constexpr short NT = NW*NSG; - constexpr short NQ = Q/NSG; - constexpr short NC = C/NW; // columns per thread - - static_assert(DK % 4 == 0, "DK must be divisible by 4"); - static_assert(Q % NSG == 0, "Q must be divisible by NSG"); - static_assert(C % NW == 0, "C must be divisible by NW"); - - const int iq3 = tgpig[2]; - const int iq2 = tgpig[1]; - const int iq1 = tgpig[0]*Q; - - const short tiitg = sgitg*NW + tiisg; - - threadgroup half * sq = (threadgroup half *) shmem; // [Q, DK] queries - threadgroup float * ss = (threadgroup float *) (sq + Q*DK); // [Q, C] scores - threadgroup half * sp = (threadgroup half *) (ss + Q*C); // [Q, C] probabilities - threadgroup float * sr = (threadgroup float *) (sp + Q*C); // [Q] per-row scale of O - threadgroup int * sf = (threadgroup int *) (sr + Q); // [1] last iteration (ic0 + 1) that rescaled O - - q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03; - - { - const int ikv2 = iq2/(args.ne02/args.ne_12_2); - const int ikv3 = iq3/(args.ne03/args.ne_12_3); - - k += ikv2*args.nb12 + ikv3*args.nb13; - v += ikv2*args.nb22 + ikv3*args.nb23; - } - - // with softcap the scale is small (scale/softcap), so it is applied to the scores to keep the precision of Q - const float qscale = FC_flash_attn_ext_tensor_has_scap ? 1.0f : args.scale; - - // load the queries, with the scale folded in - for (int i = tiitg; i < Q*DK/4; i += NT) { - const int j = i/(DK/4); - - float4 q4 = 0.0f; - if (iq1 + j < args.ne01) { - q4 = ((device const float4 *) (q + j*args.nb01))[i%(DK/4)]; - } - - ((threadgroup half4 *) sq)[i] = (half4) (q4*qscale); - } - - device const half * pm[NQ]; - - FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { - const short j = jj*NSG + sgitg; - - pm[jj] = (device const half *) (mask + (iq1 + j)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); - } - - { - const int nblk1 = (args.ne01 + Q - 1)/Q; - const int nblk0 = (args.ne11 + C - 1)/C; - - blk += (((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*nblk1 + iq1/Q)*nblk0; - } - - float M[NQ]; - float S[NQ]; - - FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { - M[jj] = -FLT_MAX/2; - S[jj] = 0.0f; - } - - float slope = 1.0f; - - // ALiBi - if (FC_flash_attn_ext_tensor_has_bias) { - const short h = iq2; - - const float base = h < args.n_head_log2 ? args.m0 : args.m1; - const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1; - - slope = pow(base, exph); - } - - const int sk = args.ns10; - const int sv = args.ns20; - - auto tq = tensor(sq, dextents(DK, Q)); - auto ts = tensor(ss, dextents(C, Q)); - auto tp = tensor(sp, dextents(C, Q)); - - mpp::tensor_ops::matmul2d< - mpp::tensor_ops::matmul2d_descriptor(Q, C, DK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply), - execution_simdgroups> mm_qk; - - mpp::tensor_ops::matmul2d< - mpp::tensor_ops::matmul2d_descriptor(Q, DV, C, false, false, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate), - execution_simdgroups> mm_pv; - - auto tv0 = tensor((device half *) v, dextents(DV, C), array({1, sv})); - - // the O matrix from the paper - auto co = mm_pv.template get_destination_cooperative_tensor(); - - FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) { - if (co.is_valid_element(i)) { - co[i] = 0.0f; - } - } - - if (tiitg == 0) { - sf[0] = 0; - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - - // the host guarantees ne11 % C == 0 - for (int ic0 = 0, ic = 0; ic < args.ne11; ++ic0, ic += C) { - char blk_cur = 1; - - if (FC_flash_attn_ext_tensor_has_mask) { - blk_cur = blk[ic0]; - - if (blk_cur == 0) { - continue; - } - } - - // Q*K^T - { - auto tk = tensor((device half *) (k + (uint64_t) ic*args.nb11), dextents(DK, C), array({1, sk})); - - mm_qk.run(tq, tk, ts); - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - - // online softmax - FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { - const short j = jj*NSG + sgitg; - - float s[NC]; - - FOR_UNROLL (short ii = 0; ii < NC; ++ii) { - s[ii] = ss[j*C + ii*NW + tiisg]; - } - - if (FC_flash_attn_ext_tensor_has_scap) { - FOR_UNROLL (short ii = 0; ii < NC; ++ii) { - s[ii] = args.logit_softcap*precise::tanh(s[ii]*args.scale); - } - } - - if (FC_flash_attn_ext_tensor_has_mask && blk_cur != 2 && iq1 + j < args.ne31) { - FOR_UNROLL (short ii = 0; ii < NC; ++ii) { - s[ii] += slope*(float) pm[jj][ic + ii*NW + tiisg]; - } - } - - float m = M[jj]; - - FOR_UNROLL (short ii = 0; ii < NC; ++ii) { - m = max(m, s[ii]); - } - - m = simd_max(m); - - // lazy rescaling: move the running max only when it grows by more than 8 (e^8 fits in half) - float ms = 1.0f; - - if (m > M[jj] + 8.0f) { - ms = exp(M[jj] - m); - M[jj] = m; - - if (tiisg == 0) { - sf[0] = ic0 + 1; - } - } - - float sum = 0.0f; - - FOR_UNROLL (short ii = 0; ii < NC; ++ii) { - // the sum uses the same rounded values as P*V - const half p = (half) exp(s[ii] - M[jj]); - - sp[j*C + ii*NW + tiisg] = p; - - sum += (float) p; - } - - S[jj] = S[jj]*ms + simd_sum(sum); - - if (tiisg == 0) { - sr[j] = ms; - } - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - - // O = diag(ms)*O + P*V - if (sf[0] == ic0 + 1) { - FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) { - if (co.is_valid_element(i)) { - co[i] *= sr[co.get_multidimensional_index(i)[1]]; - } - } - } - - { - auto tv = tensor((device half *) (v + (uint64_t) ic*args.nb21), dextents(DV, C), array({1, sv})); - - mm_pv.run(tp, tv, co); - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - } - - FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { - const short j = jj*NSG + sgitg; - - // the sink only adds to the denominator - its rescale of O is folded into the final scale - float ms = 1.0f; - - if (FC_flash_attn_ext_tensor_has_sinks) { - const float s = ((device const float *) sinks)[iq2]; - const float m = max(M[jj], s); - - ms = exp(M[jj] - m); - - S[jj] = S[jj]*ms + exp(s - m); - } - - if (tiisg == 0) { - sr[j] = S[jj] == 0.0f ? 0.0f : ms/S[jj]; - } - } - - threadgroup_barrier(mem_flags::mem_threadgroup); - - FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) { - if (co.is_valid_element(i)) { - co[i] *= sr[co.get_multidimensional_index(i)[1]]; - } - } - - // store to global memory - rows past ne01 are clipped by the tensor extents - device float * pdst = (device float *) dst + ((uint64_t) iq3*args.ne2*args.ne1 + iq2 + (uint64_t) iq1*args.ne1)*DV; - - auto td = tensor(pdst, dextents(DV, args.ne01 - iq1), array({1, args.ne1*DV})); - - co.store(td); -} - -typedef decltype(kernel_flash_attn_ext_tensor<64, 64>) flash_attn_ext_tensor_t; - -template [[host_name("kernel_flash_attn_ext_tensor_f16_dk64_dv64" )]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<64, 64>; -template [[host_name("kernel_flash_attn_ext_tensor_f16_dk128_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<128, 128>; -template [[host_name("kernel_flash_attn_ext_tensor_f16_dk192_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<192, 128>; -template [[host_name("kernel_flash_attn_ext_tensor_f16_dk256_dv256")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<256, 256>; -template [[host_name("kernel_flash_attn_ext_tensor_f16_dk512_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<512, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>; -template [[host_name("kernel_flash_attn_ext_tensor_f16_dk576_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor<576, 512, OP_FLASH_ATTN_EXT_TENSOR_NQPSG_LARGE>; - -#endif // GGML_METAL_HAS_TENSOR diff --git a/ggml/src/ggml-metal/kernels/fa_tensor.metal b/ggml/src/ggml-metal/kernels/fa_tensor.metal new file mode 100644 index 000000000000..4a9ed30e3b25 --- /dev/null +++ b/ggml/src/ggml-metal/kernels/fa_tensor.metal @@ -0,0 +1,298 @@ +#include "common.h" + +#ifdef GGML_METAL_HAS_TENSOR + +constant bool FC_flash_attn_ext_tensor_has_mask [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 0)]]; +constant bool FC_flash_attn_ext_tensor_has_sinks [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 1)]]; +constant bool FC_flash_attn_ext_tensor_has_bias [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 2)]]; +constant bool FC_flash_attn_ext_tensor_has_scap [[function_constant(FC_FLASH_ATTN_EXT_TENSOR + 3)]]; + +// ref: https://arxiv.org/pdf/2307.08691.pdf +template< + typename T, // Q, K, V and P types + typename T4, + short DK, // K head size + short DV, // V head size + short Q = OP_FLASH_ATTN_EXT_TENSOR_NQPSG, // queries per threadgroup + short C = OP_FLASH_ATTN_EXT_TENSOR_NCPSG, // cache items per threadgroup + short NSG = OP_FLASH_ATTN_EXT_TENSOR_NSG> // number of simd groups +kernel void kernel_flash_attn_ext_tensor( + constant ggml_metal_kargs_flash_attn_ext & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * sinks, + device const char * blk, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig [[threadgroup_position_in_grid]], + ushort tiisg [[thread_index_in_simdgroup]], + ushort sgitg [[simdgroup_index_in_threadgroup]]) { + constexpr short NW = N_SIMDWIDTH; + constexpr short NT = NW*NSG; + constexpr short NQ = Q/NSG; + constexpr short NC = C/NW; // columns per thread + + static_assert(DK % 4 == 0, "DK must be divisible by 4"); + static_assert(Q % NSG == 0, "Q must be divisible by NSG"); + static_assert(C % NW == 0, "C must be divisible by NW"); + + const int iq3 = tgpig[2]; + const int iq2 = tgpig[1]; + const int iq1 = tgpig[0]*Q; + + const short tiitg = sgitg*NW + tiisg; + + threadgroup T * sq = (threadgroup T *) shmem; // [Q, DK] queries + threadgroup float * ss = (threadgroup float *) (sq + Q*DK); // [Q, C] scores + threadgroup T * sp = (threadgroup T *) (ss + Q*C); // [Q, C] probabilities + threadgroup float * sr = (threadgroup float *) (sp + Q*C); // [Q] per-row scale of O + threadgroup int * sf = (threadgroup int *) (sr + Q); // [1] last iteration (ic0 + 1) that rescaled O + + q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03; + + { + const int ikv2 = iq2/(args.ne02/args.ne_12_2); + const int ikv3 = iq3/(args.ne03/args.ne_12_3); + + k += ikv2*args.nb12 + ikv3*args.nb13; + v += ikv2*args.nb22 + ikv3*args.nb23; + } + + // with softcap the scale is small (scale/softcap), so it is applied to the scores to keep the precision of Q + const float qscale = FC_flash_attn_ext_tensor_has_scap ? 1.0f : args.scale; + + // load the queries, with the scale folded in + for (int i = tiitg; i < Q*DK/4; i += NT) { + const int j = i/(DK/4); + + float4 q4 = 0.0f; + if (iq1 + j < args.ne01) { + q4 = ((device const float4 *) (q + j*args.nb01))[i%(DK/4)]; + } + + ((threadgroup T4 *) sq)[i] = (T4) (q4*qscale); + } + + device const half * pm[NQ]; + + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + + pm[jj] = (device const half *) (mask + (iq1 + j)*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33); + } + + { + const int nblk1 = (args.ne01 + Q - 1)/Q; + const int nblk0 = (args.ne11 + C - 1)/C; + + blk += (((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*nblk1 + iq1/Q)*nblk0; + } + + float M[NQ]; + float S[NQ]; + + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + M[jj] = -FLT_MAX/2; + S[jj] = 0.0f; + } + + float slope = 1.0f; + + // ALiBi + if (FC_flash_attn_ext_tensor_has_bias) { + const short h = iq2; + + const float base = h < args.n_head_log2 ? args.m0 : args.m1; + const short exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1; + + slope = pow(base, exph); + } + + const int sk = args.ns10; + const int sv = args.ns20; + + auto tq = tensor(sq, dextents(DK, Q)); + auto ts = tensor(ss, dextents(C, Q)); + auto tp = tensor(sp, dextents(C, Q)); + + mpp::tensor_ops::matmul2d< + mpp::tensor_ops::matmul2d_descriptor(Q, C, DK, false, true, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply), + execution_simdgroups> mm_qk; + + mpp::tensor_ops::matmul2d< + mpp::tensor_ops::matmul2d_descriptor(Q, DV, C, false, false, false, mpp::tensor_ops::matmul2d_descriptor::mode::multiply_accumulate), + execution_simdgroups> mm_pv; + + auto tv0 = tensor((device T *) v, dextents(DV, C), array({1, sv})); + + // the O matrix from the paper + auto co = mm_pv.template get_destination_cooperative_tensor(); + + FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) { + if (co.is_valid_element(i)) { + co[i] = 0.0f; + } + } + + if (tiitg == 0) { + sf[0] = 0; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // the host guarantees ne11 % C == 0 + for (int ic0 = 0, ic = 0; ic < args.ne11; ++ic0, ic += C) { + char blk_cur = 1; + + if (FC_flash_attn_ext_tensor_has_mask) { + blk_cur = blk[ic0]; + + if (blk_cur == 0) { + continue; + } + } + + // Q*K^T + { + auto tk = tensor((device T *) (k + (uint64_t) ic*args.nb11), dextents(DK, C), array({1, sk})); + + mm_qk.run(tq, tk, ts); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // online softmax + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + + float s[NC]; + + FOR_UNROLL (short ii = 0; ii < NC; ++ii) { + s[ii] = ss[j*C + ii*NW + tiisg]; + } + + if (FC_flash_attn_ext_tensor_has_scap) { + FOR_UNROLL (short ii = 0; ii < NC; ++ii) { + s[ii] = args.logit_softcap*precise::tanh(s[ii]*args.scale); + } + } + + if (FC_flash_attn_ext_tensor_has_mask && blk_cur != 2 && iq1 + j < args.ne31) { + FOR_UNROLL (short ii = 0; ii < NC; ++ii) { + s[ii] += slope*(float) pm[jj][ic + ii*NW + tiisg]; + } + } + + float m = M[jj]; + + FOR_UNROLL (short ii = 0; ii < NC; ++ii) { + m = max(m, s[ii]); + } + + m = simd_max(m); + + // lazy rescaling: move the running max only when it grows by more than 8 (e^8 fits in half) + float ms = 1.0f; + + if (m > M[jj] + 8.0f) { + ms = exp(M[jj] - m); + M[jj] = m; + + if (tiisg == 0) { + sf[0] = ic0 + 1; + } + } + + float sum = 0.0f; + + FOR_UNROLL (short ii = 0; ii < NC; ++ii) { + // the sum uses the same rounded values as P*V + const T p = (T) exp(s[ii] - M[jj]); + + sp[j*C + ii*NW + tiisg] = p; + + sum += (float) p; + } + + S[jj] = S[jj]*ms + simd_sum(sum); + + if (tiisg == 0) { + sr[j] = ms; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // O = diag(ms)*O + P*V + if (sf[0] == ic0 + 1) { + FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) { + if (co.is_valid_element(i)) { + co[i] *= sr[co.get_multidimensional_index(i)[1]]; + } + } + } + + { + auto tv = tensor((device T *) (v + (uint64_t) ic*args.nb21), dextents(DV, C), array({1, sv})); + + mm_pv.run(tp, tv, co); + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + } + + FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { + const short j = jj*NSG + sgitg; + + // the sink only adds to the denominator - its rescale of O is folded into the final scale + float ms = 1.0f; + + if (FC_flash_attn_ext_tensor_has_sinks) { + const float s = ((device const float *) sinks)[iq2]; + const float m = max(M[jj], s); + + ms = exp(M[jj] - m); + + S[jj] = S[jj]*ms + exp(s - m); + } + + if (tiisg == 0) { + sr[j] = S[jj] == 0.0f ? 0.0f : ms/S[jj]; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + FOR_UNROLL (short i = 0; i < co.get_capacity(); ++i) { + if (co.is_valid_element(i)) { + co[i] *= sr[co.get_multidimensional_index(i)[1]]; + } + } + + // store to global memory - rows past ne01 are clipped by the tensor extents + device float * pdst = (device float *) dst + ((uint64_t) iq3*args.ne2*args.ne1 + iq2 + (uint64_t) iq1*args.ne1)*DV; + + auto td = tensor(pdst, dextents(DV, args.ne01 - iq1), array({1, args.ne1*DV})); + + co.store(td); +} + +typedef decltype(kernel_flash_attn_ext_tensor) flash_attn_ext_tensor_t; + +template [[host_name("kernel_flash_attn_ext_tensor_f16_dk64_dv64" )]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; +template [[host_name("kernel_flash_attn_ext_tensor_f16_dk128_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; +template [[host_name("kernel_flash_attn_ext_tensor_f16_dk192_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; +template [[host_name("kernel_flash_attn_ext_tensor_f16_dk256_dv256")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; +template [[host_name("kernel_flash_attn_ext_tensor_f16_dk512_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; +template [[host_name("kernel_flash_attn_ext_tensor_f16_dk576_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; + +template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk64_dv64" )]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; +template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk128_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; +template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk192_dv128")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; +template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk256_dv256")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; +template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk512_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; +template [[host_name("kernel_flash_attn_ext_tensor_bf16_dk576_dv512")]] kernel flash_attn_ext_tensor_t kernel_flash_attn_ext_tensor; + +#endif // GGML_METAL_HAS_TENSOR diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 0d7a7ed819d1..e61823084a4d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -11813,11 +11813,18 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {2, 1}, 1024, 32, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(512, 512, 4, {2, 1}, 1024, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // BF16 K/V for the other head sizes of the Metal tensor API kernel + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {12, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16)); + test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, { 8, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16)); + test_cases.emplace_back(new test_flash_attn_ext(512, 512, 4, { 4, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 4, { 4, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16, {0, 1, 2, 3}, true, true)); + // FLASH_ATTN_EXT: large logits test_cases.emplace_back(new test_flash_attn_ext_large_logits( 64, 64, 16, {4, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext_large_logits(128, 128, 8, {4, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext_large_logits(256, 256, 4, {4, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext_large_logits(256, 256, 4, {4, 1}, 1024, 75, true, false, 0, 10.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext_large_logits(256, 256, 4, {4, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16)); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); From 5a638aedb7c4352376d1e565a9a420b85f604628 Mon Sep 17 00:00:00 2001 From: cqderek Date: Thu, 15 Oct 2026 15:07:45 +0800 Subject: [PATCH 007/423] cpu: add F16 SWIGLU_OAI reference implementation and validation (#29200) * cpu: add F16 SWIGLU_OAI reference implementation and F16 activation op test coverage Adds ggml_compute_forward_swiglu_oai_f16() to the CPU backend, mirroring the existing F32 path, and extends test-backend-ops with F16 test cases for SWIGLU_OAI so downstream backends(e.g. Hexagon) can be CPU-diffed for correctness. * ggml-openvino: fix f16 SWIGLU_OAI type mismatch in alpha/one constants translate_glu_swiglu_oai creates alpha_const and one as f32 constants unconditionally, so f16 inputs hit a type mismatch in Swish/Add on GPU ("Swish inputs must have the same type but they are: f16 and f32"). translate_glu_swiglu_clamp already upcasts to f32 and converts back at the end; apply the same pattern here. --------- Co-authored-by: ravi9 --- ggml/src/ggml-cpu/ops.cpp | 71 +++++++++++++++++++ ggml/src/ggml-cuda/ggml-cuda.cu | 3 + .../ggml-openvino/openvino/op/glu_swiglu.cpp | 14 +++- tests/test-backend-ops.cpp | 22 ++++-- 4 files changed, 105 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index c7e7c47a796e..3b39683d5c33 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -3548,6 +3548,73 @@ static void ggml_compute_forward_swiglu_oai_f32( } } +static void ggml_compute_forward_swiglu_oai_f16( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; + + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); + + if (src1) { + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + } + + const int ith = params->ith; + const int nth = params->nth; + + const int nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + const int nr = ggml_nrows(src0); + + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); + + const int32_t swapped = ggml_get_op_params_i32(dst, 1); + const float alpha = ggml_get_op_params_f32(dst, 2); + const float limit = ggml_get_op_params_f32(dst, 3); + + // rows per thread + const int dr = (nr + nth - 1)/nth; + + // row range for this thread + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + for (int i1 = ir0; i1 < ir1; i1++) { + ggml_fp16_t * src0_p = (ggml_fp16_t *) (src0_d + i1*src0_o); + ggml_fp16_t * src1_p = (ggml_fp16_t *) (src1_d + i1*src1_o); + ggml_fp16_t * dst_p = (ggml_fp16_t *) ((char *) dst->data + i1*(dst->nb[1])); + + if (!src1) { + src0_p += swapped ? nc : 0; + src1_p += swapped ? 0 : nc; + } + + for (int k = 0; k < nc; k++) { + const float x = std::min(GGML_CPU_FP16_TO_FP32(src0_p[k]), limit); + const float y = std::clamp(GGML_CPU_FP16_TO_FP32(src1_p[k]), -limit, limit); + const float out_glu = x / (1.f + expf(alpha * (-x))); + dst_p[k] = GGML_CPU_FP32_TO_FP16(out_glu * (y + 1.f)); + } + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = GGML_CPU_FP16_TO_FP32(dst_p[k]); + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); + } +#endif // NDEBUG + } +} + static void ggml_compute_forward_swiglu_oai( const ggml_compute_params * params, ggml_tensor * dst) { @@ -3559,6 +3626,10 @@ static void ggml_compute_forward_swiglu_oai( { ggml_compute_forward_swiglu_oai_f32(params, dst); } break; + case GGML_TYPE_F16: + { + ggml_compute_forward_swiglu_oai_f16(params, dst); + } break; default: { GGML_ABORT("fatal error"); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 9fe618f855d1..7c45d0d1fb23 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5474,6 +5474,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g (ggml_get_glu_op(op) == GGML_GLU_OP_SWIGLU_OAI || ggml_get_glu_op(op) == GGML_GLU_OP_SWIGLU_CLAMP)) { return false; } + if (op->src[0]->type == GGML_TYPE_F16 && ggml_get_glu_op(op) == GGML_GLU_OP_SWIGLU_OAI) { + return false; + } return ggml_is_contiguous_1(op->src[0]); default: return false; diff --git a/ggml/src/ggml-openvino/openvino/op/glu_swiglu.cpp b/ggml/src/ggml-openvino/openvino/op/glu_swiglu.cpp index 7eea81d96e59..2c206a597db6 100644 --- a/ggml/src/ggml-openvino/openvino/op/glu_swiglu.cpp +++ b/ggml/src/ggml-openvino/openvino/op/glu_swiglu.cpp @@ -75,6 +75,15 @@ OutputVector translate_glu_swiglu_oai(const NodeContext & context) { const float alpha = reinterpret_cast(params)[2]; const float limit = reinterpret_cast(params)[3]; + // Compute in f32: alpha_const/one are f32, and Swish/Add need matching input types. + auto output_type = context.get_output_type(); + if (src0.get_element_type() != ov::element::f32) { + src0 = std::make_shared(src0, ov::element::f32); + } + if (src1.get_element_type() != ov::element::f32) { + src1 = std::make_shared(src1, ov::element::f32); + } + auto gate = std::make_shared(src0, -std::numeric_limits::infinity(), limit); auto alpha_const = ov::op::v0::Constant::create(ov::element::f32, {}, {alpha}); auto out_glu = std::make_shared(gate, alpha_const); @@ -82,7 +91,10 @@ OutputVector translate_glu_swiglu_oai(const NodeContext & context) { auto up = std::make_shared(src1, -limit, limit); auto one = ov::op::v0::Constant::create(ov::element::f32, {}, {1.0f}); auto up_plus_one = std::make_shared(up, one); - auto res = std::make_shared(out_glu, up_plus_one); + ov::Output res = std::make_shared(out_glu, up_plus_one); + if (res.get_element_type() != output_type) { + res = std::make_shared(res, output_type); + } return rename_outputs_with_suffix({res}, context.get_name()); } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e61823084a4d..93b30d8d25e8 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2308,6 +2308,12 @@ struct test_swiglu_oai : public test_case { init_tensor_uniform(t, -150.f, 150.f); } } + + double max_nmse_err() override { + // F16 epsilon ~ 9.7e-4: relax tolerance vs the F32 default to absorb the + // natural roundoff drift from the exp()/clamp() chain computed in F16. + return type == GGML_TYPE_F16 ? 5e-6 : 1e-7; + } }; struct test_swiglu_clamp : public test_case { @@ -2365,6 +2371,12 @@ struct test_swiglu_clamp : public test_case { init_tensor_uniform(t, -150.f, 150.f); } } + + double max_nmse_err() override { + // F16 epsilon ~ 9.7e-4: relax tolerance vs the F32 default to absorb the + // natural roundoff drift from the sigmoid()/clamp() chain computed in F16. + return type == GGML_TYPE_F16 ? 5e-6 : 1e-7; + } }; // GGML_OP_GET_ROWS @@ -9747,10 +9759,12 @@ static std::vector> make_test_cases_eval() { } } - for (int v : {0, 1}) { - for (float alpha : {.5f, 1.702f}) { - for (float limit : {2.0f, 7.0f}) { - test_cases.emplace_back(new test_swiglu_oai(GGML_TYPE_F32, { 128, 2, 2, 2 }, v, alpha, limit)); + for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16}) { + for (int v : {0, 1}) { + for (float alpha : {.5f, 1.702f}) { + for (float limit : {2.0f, 7.0f}) { + test_cases.emplace_back(new test_swiglu_oai(type, { 128, 2, 2, 2 }, v, alpha, limit)); + } } } } From 733f57a9994f7a237176aa9f4cdc31bc466d056d Mon Sep 17 00:00:00 2001 From: Nathan Wilson <67372905+Nathanw1014@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:39:02 +1030 Subject: [PATCH 008/423] vulkan : copy strided f16 KV once on AMD so reads spread across memory channels (#27703) * vulkan : copy strided f16 KV once so reads spread across memory channels Assisted-by: Claude (Opus 4.8 / Opus 5) * tests : cover the f16 KV copy to scratch Assisted-by: Claude (Opus 4.8 / Opus 5) * vulkan : gate the f16 KV scratch copy to AMD Assisted-by: Claude (Opus 4.8 / Opus 5) * vulkan : reuse the generic copy shader for the f16 KV scratch fill The f16 clause needs only a strided copy, not a dequant, so drop the dedicated shader and dispatch the existing copy pipeline instead. The copy is described in source memory order so the strided access lands on the writes, which stream, instead of the reads, which would concentrate on a few memory channels. Assisted-by: Claude (Opus 4.8 / Opus 5 / Fable 5) * vulkan : extend the strided KV scratch copy to bf16 Assisted-by: Claude Opus 5 * vulkan : clean up comments Assisted-by: Claude Opus 5 * consolidate copy shaders * disable on RDNA4 --------- Co-authored-by: Ruben Ortlam --- ggml/src/ggml-vulkan/ggml-vulkan-types.h | 4 +- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 119 +++++++++++------- .../vulkan-shaders/vulkan-shaders-gen.cpp | 8 +- tests/test-backend-ops.cpp | 7 ++ 4 files changed, 85 insertions(+), 53 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h index 7794c631062b..a41ce596625a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h @@ -835,8 +835,8 @@ struct vk_device_struct { vk_pipeline pipeline_roll_f32; vk_pipeline pipeline_repeat_i32, pipeline_repeat_back_f32; vk_pipeline pipeline_repeat_i16; - vk_pipeline pipeline_cpy_f32_f32, pipeline_cpy_f32_f16, pipeline_cpy_f16_f16, pipeline_cpy_f16_f32, pipeline_cpy_f32_bf16, pipeline_cpy_bf16_f32, pipeline_cpy_bf16_f16, pipeline_cpy_f32_i32, pipeline_cpy_i32_f32; - vk_pipeline pipeline_contig_cpy_f32_f32, pipeline_contig_cpy_f32_f16, pipeline_contig_cpy_f16_f16, pipeline_contig_cpy_f16_f32, pipeline_contig_cpy_f32_bf16, pipeline_contig_cpy_bf16_f32, pipeline_contig_cpy_bf16_f16, pipeline_contig_cpy_f32_i32, pipeline_contig_cpy_i32_f32; + vk_pipeline pipeline_cpy_32, pipeline_cpy_f32_f16, pipeline_cpy_16, pipeline_cpy_f16_f32, pipeline_cpy_f32_bf16, pipeline_cpy_bf16_f32, pipeline_cpy_bf16_f16, pipeline_cpy_f32_i32, pipeline_cpy_i32_f32; + vk_pipeline pipeline_contig_cpy_32, pipeline_contig_cpy_f32_f16, pipeline_contig_cpy_16, pipeline_contig_cpy_f16_f32, pipeline_contig_cpy_f32_bf16, pipeline_contig_cpy_bf16_f32, pipeline_contig_cpy_bf16_f16, pipeline_contig_cpy_f32_i32, pipeline_contig_cpy_i32_f32; vk_pipeline pipeline_cpy_f32_quant[GGML_TYPE_COUNT]; vk_pipeline pipeline_cpy_quant_f32[GGML_TYPE_COUNT]; vk_pipeline pipeline_cpy_transpose_16, pipeline_cpy_transpose_32; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8357f01fe28d..42057114071b 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3273,9 +3273,9 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_rms_norm_back_f32, "rms_norm_back_f32", rms_norm_back_f32_len, rms_norm_back_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_l2_norm_f32, "l2_norm_f32", l2_norm_f32_len, l2_norm_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f32, "cpy_f32_f32", cpy_f32_f32_len, cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_32, "cpy_32", cpy_32_len, cpy_32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_f16, "cpy_f32_f16", cpy_f32_f16_len, cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_cpy_f16_f16, "cpy_f16_f16", cpy_f16_f16_len, cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_16, "cpy_16", cpy_16_len, cpy_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f16_f32, "cpy_f16_f32", cpy_f16_f32_len, cpy_f16_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_bf16,"cpy_f32_bf16",cpy_f32_bf16_len,cpy_f32_bf16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_bf16_f32,"cpy_bf16_f32",cpy_bf16_f32_len,cpy_bf16_f32_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); @@ -3283,9 +3283,9 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_cpy_i32_f32, "cpy_i32_f32", cpy_i32_f32_len, cpy_i32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_i32, "cpy_f32_i32", cpy_f32_i32_len, cpy_f32_i32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f32, "contig_cpy_f32_f32", contig_cpy_f32_f32_len, contig_cpy_f32_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_32, "contig_cpy_32", contig_cpy_32_len, contig_cpy_32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_f16, "contig_cpy_f32_f16", contig_cpy_f32_f16_len, contig_cpy_f32_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f16_f16, "contig_cpy_f16_f16", contig_cpy_f16_f16_len, contig_cpy_f16_f16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_16, "contig_cpy_16", contig_cpy_16_len, contig_cpy_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f16_f32, "contig_cpy_f16_f32", contig_cpy_f16_f32_len, contig_cpy_f16_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_f32_bf16,"contig_cpy_f32_bf16",contig_cpy_f32_bf16_len,contig_cpy_f32_bf16_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_contig_cpy_bf16_f32,"contig_cpy_bf16_f32",contig_cpy_bf16_f32_len,contig_cpy_bf16_f32_data,"main", 2, sizeof(vk_op_unary_push_constants), {512, 1, 1}, {}, 1); @@ -6100,13 +6100,6 @@ vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_t } } - if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F32) { - if (contig) { - return ctx->device->pipeline_contig_cpy_f32_f32; - } else { - return ctx->device->pipeline_cpy_f32_f32; - } - } if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F16) { if (contig) { return ctx->device->pipeline_contig_cpy_f32_f16; @@ -6114,13 +6107,6 @@ vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_t return ctx->device->pipeline_cpy_f32_f16; } } - if (src->type == GGML_TYPE_F16 && to == GGML_TYPE_F16) { - if (contig) { - return ctx->device->pipeline_contig_cpy_f16_f16; - } else { - return ctx->device->pipeline_cpy_f16_f16; - } - } if (src->type == GGML_TYPE_F16 && to == GGML_TYPE_F32) { if (contig) { return ctx->device->pipeline_contig_cpy_f16_f32; @@ -6203,15 +6189,15 @@ vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_t GGML_ASSERT(ggml_is_quantized(to) || ggml_type_size(src->type) == 2 || ggml_type_size(src->type) == 4); if ((ggml_type_size(src->type) % 4) == 0) { if (contig) { - return ctx->device->pipeline_contig_cpy_f32_f32; + return ctx->device->pipeline_contig_cpy_32; } else { - return ctx->device->pipeline_cpy_f32_f32; + return ctx->device->pipeline_cpy_32; } } else { if (contig) { - return ctx->device->pipeline_contig_cpy_f16_f16; + return ctx->device->pipeline_contig_cpy_16; } else { - return ctx->device->pipeline_cpy_f16_f16; + return ctx->device->pipeline_cpy_16; } } } @@ -8134,19 +8120,30 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const }; const bool k_quant = k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_BF16 && k->type != GGML_TYPE_F32; const bool v_quant = v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_BF16 && v->type != GGML_TYPE_F32; - const bool use_dequant_kv = k_quant && v_quant && neq1 >= 64 && + const bool kv_raw = (k->type == GGML_TYPE_F16 || k->type == GGML_TYPE_BF16) && k->type == v->type; + // gated to AMD, except RDNA4, which regressed in benchmarks + const bool kv_raw_strided = ctx->device->vendor_id == VK_VENDOR_ID_AMD && + ctx->device->architecture != vk_device_architecture::AMD_RDNA4 && + neq1 >= 256 && + kv_raw && + (k->nb[1] != (uint64_t)HSK * ggml_type_size(k->type) || + v->nb[1] != (uint64_t)HSV * ggml_type_size(v->type)) && + (HSK % 8) == 0 && (HSV % 8) == 0; + const size_t kv_scratch_ts = kv_raw_strided ? ggml_type_size(k->type) : sizeof(ggml_fp16_t); + const bool use_dequant_kv = ((k_quant && v_quant) || kv_raw_strided) && neq1 >= 64 && is_dense_kv_cache(k) && is_dense_kv_cache(v) && - (uint64_t)ggml_nelements(k) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && - (uint64_t)ggml_nelements(v) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && - ctx->device->pipeline_dequant_transpose[k->type] != nullptr && - ctx->device->pipeline_dequant_transpose[v->type] != nullptr && + (uint64_t)ggml_nelements(k) * kv_scratch_ts <= ctx->device->properties.limits.maxStorageBufferRange && + (uint64_t)ggml_nelements(v) * kv_scratch_ts <= ctx->device->properties.limits.maxStorageBufferRange && + (kv_raw_strided || + (ctx->device->pipeline_dequant_transpose[k->type] != nullptr && + ctx->device->pipeline_dequant_transpose[v->type] != nullptr)) && // coopmat2 path does not benefit from the f16 scratch !ctx->device->coopmat2 && // Intel Xe1 regresses, see PR 25494 (ctx->device->vendor_id != VK_VENDOR_ID_INTEL || (ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1)); - const ggml_type k_type_eff = use_dequant_kv ? GGML_TYPE_F16 : k->type; - const ggml_type v_type_eff = use_dequant_kv ? GGML_TYPE_F16 : v->type; + const ggml_type k_type_eff = (use_dequant_kv && !kv_raw_strided) ? GGML_TYPE_F16 : k->type; + const ggml_type v_type_eff = (use_dequant_kv && !kv_raw_strided) ? GGML_TYPE_F16 : v->type; // For scalar/coopmat1 FA, we can use the "large" size to accommodate qga. // For coopmat2 FA, we always use the small size (which is still pretty large for gqa). @@ -8422,28 +8419,56 @@ void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const } if (use_dequant_kv) { - const uint64_t fp = sizeof(ggml_fp16_t); - const uint64_t k_f16_sz = (uint64_t)ggml_nelements(k) * fp; - const uint64_t v_f16_sz = (uint64_t)ggml_nelements(v) * fp; - if (ctx->prealloc_size_x < k_f16_sz + v_f16_sz) { - ctx->prealloc_size_x = k_f16_sz + v_f16_sz; + const uint64_t k_sz = (uint64_t)ggml_nelements(k) * kv_scratch_ts; + const uint64_t v_sz = (uint64_t)ggml_nelements(v) * kv_scratch_ts; + if (ctx->prealloc_size_x < k_sz + v_sz) { + ctx->prealloc_size_x = k_sz + v_sz; ggml_vk_preallocate_buffers(ctx, subctx); } - vk_pipeline tr_k = ctx->device->pipeline_dequant_transpose[k->type]; - vk_pipeline tr_v = ctx->device->pipeline_dequant_transpose[v->type]; - ggml_pipeline_request_descriptor_sets(ctx, tr_k, 1); - ggml_pipeline_request_descriptor_sets(ctx, tr_v, 1); - if (ctx->prealloc_x_need_sync) { - ggml_vk_sync_buffers(ctx, subctx); - } - vk_subbuffer k_dst = vk_subbuffer{ ctx->prealloc_x, 0, k_f16_sz }; - vk_subbuffer v_dst = vk_subbuffer{ ctx->prealloc_x, k_f16_sz, v_f16_sz }; + vk_subbuffer k_dst = vk_subbuffer{ ctx->prealloc_x, 0, k_sz }; + vk_subbuffer v_dst = vk_subbuffer{ ctx->prealloc_x, k_sz, v_sz }; const uint32_t k_nel = (uint32_t)ggml_nelements(k); const uint32_t v_nel = (uint32_t)ggml_nelements(v); - { const std::vector pc = { (uint32_t)HSK, (uint32_t)nek2, (uint32_t)KV, 0, k_nel }; - ggml_vk_dispatch_pipeline(ctx, subctx, tr_k, { k_buf, k_dst }, pc, { k_nel, 1, 1 }); } - { const std::vector pc = { (uint32_t)HSV, (uint32_t)nev2, (uint32_t)KV, 0, v_nel }; - ggml_vk_dispatch_pipeline(ctx, subctx, tr_v, { v_buf, v_dst }, pc, { v_nel, 1, 1 }); } + if (kv_raw_strided) { + vk_pipeline cp_k = ggml_vk_get_cpy_pipeline(ctx, k, nullptr, k->type); + vk_pipeline cp_v = ggml_vk_get_cpy_pipeline(ctx, v, nullptr, v->type); + ggml_pipeline_request_descriptor_sets(ctx, cp_k, 1); + ggml_pipeline_request_descriptor_sets(ctx, cp_v, 1); + if (ctx->prealloc_x_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + auto make_pc = [](uint32_t hs, uint32_t nh, uint32_t kv, uint32_t nel) { + vk_op_unary_push_constants pc{}; + pc.ne = nel; + pc.ne00 = hs; pc.ne01 = nh; pc.ne02 = kv; pc.ne03 = nel / (hs * nh * kv); + pc.nb00 = 1; pc.nb01 = hs; pc.nb02 = hs * nh; pc.nb03 = hs * nh * kv; + pc.ne10 = hs; pc.ne11 = nh; pc.ne12 = kv; pc.ne13 = pc.ne03; + pc.nb10 = 1; pc.nb11 = hs * kv; pc.nb12 = hs; pc.nb13 = hs * kv * nh; + init_pushconst_fastdiv(pc); + return pc; + }; + auto cpy_elems = [](uint32_t ne) -> std::array { + if (ne > 262144) { return { 512, 512, CEIL_DIV(ne, 262144) }; } + if (ne > 512) { return { 512, CEIL_DIV(ne, 512), 1 }; } + return { ne, 1, 1 }; + }; + ggml_vk_dispatch_pipeline(ctx, subctx, cp_k, { k_buf, k_dst }, + make_pc((uint32_t)HSK, (uint32_t)nek2, (uint32_t)KV, k_nel), cpy_elems(k_nel)); + ggml_vk_dispatch_pipeline(ctx, subctx, cp_v, { v_buf, v_dst }, + make_pc((uint32_t)HSV, (uint32_t)nev2, (uint32_t)KV, v_nel), cpy_elems(v_nel)); + } else { + vk_pipeline tr_k = ctx->device->pipeline_dequant_transpose[k->type]; + vk_pipeline tr_v = ctx->device->pipeline_dequant_transpose[v->type]; + ggml_pipeline_request_descriptor_sets(ctx, tr_k, 1); + ggml_pipeline_request_descriptor_sets(ctx, tr_v, 1); + if (ctx->prealloc_x_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + { const std::vector pc = { (uint32_t)HSK, (uint32_t)nek2, (uint32_t)KV, 0, k_nel }; + ggml_vk_dispatch_pipeline(ctx, subctx, tr_k, { k_buf, k_dst }, pc, { k_nel, 1, 1 }); } + { const std::vector pc = { (uint32_t)HSV, (uint32_t)nev2, (uint32_t)KV, 0, v_nel }; + ggml_vk_dispatch_pipeline(ctx, subctx, tr_v, { v_buf, v_dst }, pc, { v_nel, 1, 1 }); } + } ggml_vk_sync_buffers(ctx, subctx); k_buf = k_dst; v_buf = v_dst; 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 42ddf4858178..f545015464f1 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -864,18 +864,18 @@ void process_shaders() { string_to_spv("rms_norm_back_f32", "rms_norm_back.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("l2_norm_f32", "l2_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); - string_to_spv("cpy_f32_f32", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("cpy_32", "copy.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}}); string_to_spv("cpy_f32_f16", "copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); - string_to_spv("cpy_f16_f16", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); + string_to_spv("cpy_16", "copy.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}}); string_to_spv("cpy_f16_f32", "copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); string_to_spv("cpy_f32_bf16","copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "uint16_t"}, {"DATA_D_BF16", "1"}}); string_to_spv("cpy_bf16_f32","copy.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "float"}, {"DATA_A_BF16", "1"}}); string_to_spv("cpy_bf16_f16","copy.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "float16_t"}, {"DATA_A_BF16", "1"}}); - string_to_spv("contig_cpy_f32_f32", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float"}}); + string_to_spv("contig_cpy_32", "contig_copy.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}}); string_to_spv("contig_cpy_f32_i32", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "int"}}); string_to_spv("contig_cpy_i32_f32", "contig_copy.comp", {{"A_TYPE", "int"}, {"D_TYPE", "float"}}); string_to_spv("contig_cpy_f32_f16", "contig_copy.comp", {{"A_TYPE", "float"}, {"D_TYPE", "float16_t"}}); - string_to_spv("contig_cpy_f16_f16", "contig_copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); + string_to_spv("contig_cpy_16", "contig_copy.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}}); string_to_spv("contig_cpy_f16_f32", "contig_copy.comp", {{"A_TYPE", "float16_t"}, {"D_TYPE", "float"}, {"OPTIMIZATION_ERROR_WORKAROUND", "1"}}); string_to_spv("contig_cpy_f32_bf16","contig_copy.comp",{{"A_TYPE", "float"}, {"D_TYPE", "uint16_t"}, {"DATA_D_BF16", "1"}}); string_to_spv("contig_cpy_bf16_f32","contig_copy.comp",{{"A_TYPE", "uint16_t"}, {"D_TYPE", "float"}, {"DATA_A_BF16", "1"}}); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 93b30d8d25e8..ba7f7eb4e56c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -11817,6 +11817,13 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 1024, 320, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(96, 96, 8, {4, 1}, 512, 256, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {8, 1}, 512, 259, true, false, 8.0f, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 512, 288, true, false, 0, 30.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, {8, 1}, 1024, 320, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16, {0, 2, 1, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {8, 1}, 512, 259, true, false, 8.0f, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16, {0, 2, 1, 3}, false)); + // FLASH_ATTN_EXT MMA: non-pow2 head size and MLA K/V view. test_cases.emplace_back(new test_flash_attn_ext(192, 128, 8, {8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); From 77c1f39929d0cae9488d846fff0925cd0f1738c8 Mon Sep 17 00:00:00 2001 From: lhez Date: Thu, 15 Oct 2026 00:11:51 -0700 Subject: [PATCH 009/423] opencl: apply flash attn softcap before mask (#30546) --- .../src/ggml-opencl/kernels/flash_attn_f16.cl | 25 +++--- .../src/ggml-opencl/kernels/flash_attn_f32.cl | 25 +++--- .../ggml-opencl/kernels/flash_attn_f32_f16.cl | 77 +++++++++---------- .../kernels/flash_attn_f32_q4_0.cl | 50 ++++++------ .../kernels/flash_attn_f32_q8_0.cl | 44 +++++------ 5 files changed, 108 insertions(+), 113 deletions(-) diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl index f9797d346005..cb6fd89308f4 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f16.cl @@ -175,6 +175,12 @@ __kernel void flash_attn_f16( ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale; ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale; + if (logit_softcap > 0.0f) { + s0 = logit_softcap * tanh(s0 / logit_softcap); + s1 = logit_softcap * tanh(s1 / logit_softcap); + s2 = logit_softcap * tanh(s2 / logit_softcap); + s3 = logit_softcap * tanh(s3 / logit_softcap); + } if (is_causal) { const int causal_limit = n_kv - n_q + my_query_row; @@ -196,13 +202,6 @@ __kernel void flash_attn_f16( if (k_row3 < n_kv) s3 += slope * (ACC_TYPE)mask_ptr[k_row3]; } - if (logit_softcap > 0.0f) { - s0 = logit_softcap * tanh(s0 / logit_softcap); - s1 = logit_softcap * tanh(s1 / logit_softcap); - s2 = logit_softcap * tanh(s2 / logit_softcap); - s3 = logit_softcap * tanh(s3 / logit_softcap); - } - const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3))); const ACC_TYPE scale_prev = native_exp(m_i - m_new); const ACC_TYPE p0 = native_exp(s0 - m_new); @@ -331,13 +330,13 @@ __kernel void flash_attn_f16_q1( dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc); } ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global DATA_TYPE* mask_ptr = (const global DATA_TYPE*)(mask_base); score += slope * (ACC_TYPE)mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } m_i = max(m_i, score); } @@ -367,13 +366,13 @@ __kernel void flash_attn_f16_q1( dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc); } ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global DATA_TYPE* mask_ptr = (const global DATA_TYPE*)(mask_base); score += slope * (ACC_TYPE)mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } const ACC_TYPE p = exp(score - m_final); l_i += p; FA_UNROLL diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl index 5911524e156c..15b9cf30a740 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32.cl @@ -176,6 +176,12 @@ __kernel void flash_attn_f32( ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale; ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale; + if (logit_softcap > 0.0f) { + s0 = logit_softcap * tanh(s0 / logit_softcap); + s1 = logit_softcap * tanh(s1 / logit_softcap); + s2 = logit_softcap * tanh(s2 / logit_softcap); + s3 = logit_softcap * tanh(s3 / logit_softcap); + } if (is_causal) { const int causal_limit = n_kv - n_q + my_query_row; @@ -197,13 +203,6 @@ __kernel void flash_attn_f32( if (k_row3 < n_kv) s3 += slope * (ACC_TYPE)mask_ptr[k_row3]; } - if (logit_softcap > 0.0f) { - s0 = logit_softcap * tanh(s0 / logit_softcap); - s1 = logit_softcap * tanh(s1 / logit_softcap); - s2 = logit_softcap * tanh(s2 / logit_softcap); - s3 = logit_softcap * tanh(s3 / logit_softcap); - } - const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3))); const ACC_TYPE scale_prev = native_exp(m_i - m_new); const ACC_TYPE p0 = native_exp(s0 - m_new); @@ -332,13 +331,13 @@ __kernel void flash_attn_f32_q1( dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc); } ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base); score += slope * (ACC_TYPE)mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } m_i = max(m_i, score); } @@ -368,13 +367,13 @@ __kernel void flash_attn_f32_q1( dot_acc = mad(q_priv[k], CONVERT_ACC4(k_ptr[k]), dot_acc); } ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base); score += slope * (ACC_TYPE)mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } const ACC_TYPE p = exp(score - m_final); l_i += p; FA_UNROLL diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl index 27553c729af8..d1cc1480105d 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_f16.cl @@ -336,6 +336,10 @@ __kernel void FA_TILE_NAME( ACC_TYPE score0 = partial0 * scale; ACC_TYPE score1 = partial1 * scale; + if (logit_softcap > 0.0f) { + score0 = logit_softcap * tanh(score0 / logit_softcap); + score1 = logit_softcap * tanh(score1 / logit_softcap); + } if (!query_valid) { score0 = FA_M_INIT; score1 = FA_M_INIT; } if (is_causal) { @@ -359,11 +363,6 @@ __kernel void FA_TILE_NAME( } } - if (logit_softcap > 0.0f) { - score0 = logit_softcap * tanh(score0 / logit_softcap); - score1 = logit_softcap * tanh(score1 / logit_softcap); - } - const ACC_TYPE m_new = max(m_i, max(score0, score1)); // Whole tile masked (m_new == FA_M_INIT): force the exp() args // far negative so the tile contributes 0, not exp(0)=1. @@ -408,6 +407,9 @@ __kernel void FA_TILE_NAME( score += local_partial[j][q_lane * N_SPLIT + s]; } score *= scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (is_causal && k_row > (n_kv - n_q + my_query_row)) score = FA_M_INIT; if (k_row >= n_kv) score = FA_M_INIT; @@ -424,10 +426,6 @@ __kernel void FA_TILE_NAME( } } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } - m_new = max(m_new, score); local_p[q_lane][j] = score; } @@ -502,6 +500,12 @@ __kernel void FA_TILE_NAME( ACC_TYPE s1 = (dot_acc1.s0 + dot_acc1.s1 + dot_acc1.s2 + dot_acc1.s3) * scale; ACC_TYPE s2 = (dot_acc2.s0 + dot_acc2.s1 + dot_acc2.s2 + dot_acc2.s3) * scale; ACC_TYPE s3 = (dot_acc3.s0 + dot_acc3.s1 + dot_acc3.s2 + dot_acc3.s3) * scale; + if (logit_softcap > 0.0f) { + s0 = logit_softcap * tanh(s0 / logit_softcap); + s1 = logit_softcap * tanh(s1 / logit_softcap); + s2 = logit_softcap * tanh(s2 / logit_softcap); + s3 = logit_softcap * tanh(s3 / logit_softcap); + } if (is_causal) { const int causal_limit = n_kv - n_q + my_query_row; @@ -531,13 +535,6 @@ __kernel void FA_TILE_NAME( } } - if (logit_softcap > 0.0f) { - s0 = logit_softcap * tanh(s0 / logit_softcap); - s1 = logit_softcap * tanh(s1 / logit_softcap); - s2 = logit_softcap * tanh(s2 / logit_softcap); - s3 = logit_softcap * tanh(s3 / logit_softcap); - } - const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3))); // Whole tile masked (m_new == FA_M_INIT): force the exp() args // far negative so the tile contributes 0, not exp(0)=1. @@ -748,13 +745,13 @@ __kernel void flash_attn_f32_f16_q1( dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc); } ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base); score += slope * (ACC_TYPE)mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } m_i = max(m_i, score); } @@ -776,13 +773,13 @@ __kernel void flash_attn_f32_f16_q1( dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc); } ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base); score += slope * (ACC_TYPE)mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } const ACC_TYPE p = exp(score - m_final); l_i += p; FA_UNROLL @@ -1122,13 +1119,13 @@ __kernel void flash_attn_f32_f16_q1_local_tile( if (tid == 0) { float s = red[0] * scale; + if (logit_softcap > 0.0f) { + s = logit_softcap * tanh(s / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base; s += slope * (float) mask_ptr[kv_idx]; } - if (logit_softcap > 0.0f) { - s = logit_softcap * tanh(s / logit_softcap); - } score_shared = s; } barrier(CLK_LOCAL_MEM_FENCE); @@ -1323,13 +1320,13 @@ __kernel void flash_attn_f32_f16_q1_local_mq_split( contrib += q_shared[h * DK + d] * (float) k_tile[j * DK + d]; } float s = sub_group_reduce_add(contrib) * scale; + if (logit_softcap > 0.0f) { + s = logit_softcap * tanh(s / logit_softcap); + } if (mask_base[h] != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h]; s += slope[h] * (float) mask_ptr[kv_idx]; } - if (logit_softcap > 0.0f) { - s = logit_softcap * tanh(s / logit_softcap); - } score[h] = s; } @@ -1510,13 +1507,13 @@ __kernel void flash_attn_f32_f16_q1_vec_mq( for (int h = 0; h < MQ_GQA; ++h) { const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3; ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale; + if (logit_softcap > 0.0f) { + s = logit_softcap * tanh(s / logit_softcap); + } if (mask_base[h] != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h]; s += slope[h] * (ACC_TYPE) mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - s = logit_softcap * tanh(s / logit_softcap); - } score[h] = s; } @@ -2448,13 +2445,13 @@ __kernel void flash_attn_f32_f16_q1_vec_mq_split_k_img( for (int h = 0; h < MQ_GQA; ++h) { const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3; ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale; + if (logit_softcap > 0.0f) { + s = logit_softcap * tanh(s / logit_softcap); + } if (mask_base[h] != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h]; s += slope[h] * (ACC_TYPE) mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - s = logit_softcap * tanh(s / logit_softcap); - } score[h] = s; } @@ -2633,13 +2630,13 @@ __kernel void flash_attn_f32_f16_q1_split( dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc); } ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base); score += slope * (ACC_TYPE) mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } m_i = max(m_i, score); } @@ -2662,13 +2659,13 @@ __kernel void flash_attn_f32_f16_q1_split( dot_acc = mad(q_shared[k], CONVERT_KV_ACC4(k_ptr[k]), dot_acc); } ACC_TYPE score = (dot_acc.s0 + dot_acc.s1 + dot_acc.s2 + dot_acc.s3) * scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base); score += slope * (ACC_TYPE) mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } const ACC_TYPE p = exp(score - m_c); l_i += p; #pragma unroll diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl index 48adba4f725b..3e75dc942bfb 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q4_0.cl @@ -287,14 +287,14 @@ __kernel void flash_attn_f32_q4_0_q1( #endif } score *= scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base); score += slope * (ACC_TYPE)mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } // Online softmax step. const ACC_TYPE m_new = max(m_i, score); @@ -554,14 +554,14 @@ __kernel void flash_attn_f32_q4_0_q1_vec( ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3; ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale; #endif + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base; score += slope * (ACC_TYPE) mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } const ACC_TYPE m_new = max(m_i, score); const ACC_TYPE scale_prev = native_exp(m_i - m_new); @@ -745,13 +745,13 @@ __kernel void flash_attn_f32_q4_0_q1_split( #endif } score *= scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base); score += slope * (ACC_TYPE) mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } // Online softmax step. const ACC_TYPE m_new = max(m_i, score); @@ -1037,13 +1037,13 @@ __kernel void flash_attn_f32_q4_0_q1_vec_mq_split( #pragma unroll for (int h = 0; h < MQ_GQA; ++h) { ACC_TYPE s = sub_group_reduce_add(lane_contrib[h]) * scale; + if (logit_softcap > 0.0f) { + s = logit_softcap * tanh(s / logit_softcap); + } if (mask_base[h] != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h]; s += slope[h] * (ACC_TYPE) mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - s = logit_softcap * tanh(s / logit_softcap); - } score[h] = s; } #else @@ -1067,13 +1067,13 @@ __kernel void flash_attn_f32_q4_0_q1_vec_mq_split( for (int h = 0; h < MQ_GQA; ++h) { const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3; ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale; + if (logit_softcap > 0.0f) { + s = logit_softcap * tanh(s / logit_softcap); + } if (mask_base[h] != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h]; s += slope[h] * (ACC_TYPE) mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - s = logit_softcap * tanh(s / logit_softcap); - } score[h] = s; } #endif @@ -1385,13 +1385,13 @@ __kernel void flash_attn_f32_q4_0_q1_vec_mq_split_c8( s += sub_group_shuffle_xor(s, step); } s *= scale; + if (logit_softcap > 0.0f) { + s = logit_softcap * tanh(s / logit_softcap); + } if (mask_base[h] != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h]; s += slope[h] * (ACC_TYPE) mask_ptr[k_safe]; } - if (logit_softcap > 0.0f) { - s = logit_softcap * tanh(s / logit_softcap); - } score[h] = valid ? s : FA_M_INIT; } @@ -1852,9 +1852,15 @@ __kernel void flash_attn_f32_q4_0( s2 = sub_group_shuffle(s2, tri_base + 0) + sub_group_shuffle(s2, tri_base + 1) + sub_group_shuffle(s2, tri_base + 2); s3 = sub_group_shuffle(s3, tri_base + 0) + sub_group_shuffle(s3, tri_base + 1) + sub_group_shuffle(s3, tri_base + 2); #endif - if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; } #endif s0 *= scale; s1 *= scale; s2 *= scale; s3 *= scale; + if (logit_softcap > 0.0f) { + s0 = logit_softcap * tanh(s0 / logit_softcap); + s1 = logit_softcap * tanh(s1 / logit_softcap); + s2 = logit_softcap * tanh(s2 / logit_softcap); + s3 = logit_softcap * tanh(s3 / logit_softcap); + } + if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; } if (is_causal) { const int causal_limit = n_kv - n_q + my_query_row; @@ -1876,12 +1882,6 @@ __kernel void flash_attn_f32_q4_0( if (k_row2 < n_kv) s2 += slope * (ACC_TYPE) mask_ptr[k_row2]; if (k_row3 < n_kv) s3 += slope * (ACC_TYPE) mask_ptr[k_row3]; } - if (logit_softcap > 0.0f) { - s0 = logit_softcap * tanh(s0 / logit_softcap); - s1 = logit_softcap * tanh(s1 / logit_softcap); - s2 = logit_softcap * tanh(s2 / logit_softcap); - s3 = logit_softcap * tanh(s3 / logit_softcap); - } const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3))); // Whole tile masked (m_new == FA_M_INIT): force the exp() args diff --git a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl index f50912d21101..75d0d9b0a4df 100644 --- a/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl +++ b/ggml/src/ggml-opencl/kernels/flash_attn_f32_q8_0.cl @@ -232,14 +232,14 @@ __kernel void flash_attn_f32_q8_0_q1( #endif } score *= scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE* mask_ptr = (const global MASK_DATA_TYPE*)(mask_base); score += slope * (ACC_TYPE)mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } // Online softmax step. const ACC_TYPE m_new = max(m_i, score); @@ -437,14 +437,14 @@ __kernel void flash_attn_f32_q8_0_q1_vec( } ACC_TYPE dot_partial = dot4.s0 + dot4.s1 + dot4.s2 + dot4.s3; ACC_TYPE score = sub_group_reduce_add(dot_partial) * scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base; score += slope * (ACC_TYPE) mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } const ACC_TYPE m_new = max(m_i, score); const ACC_TYPE scale_prev = native_exp(m_i - m_new); @@ -627,13 +627,13 @@ __kernel void flash_attn_f32_q8_0_q1_split( #endif } score *= scale; + if (logit_softcap > 0.0f) { + score = logit_softcap * tanh(score / logit_softcap); + } if (mask_base != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) (mask_base); score += slope * (ACC_TYPE) mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - score = logit_softcap * tanh(score / logit_softcap); - } // Online softmax step. const ACC_TYPE m_new = max(m_i, score); @@ -882,13 +882,13 @@ __kernel void flash_attn_f32_q8_0_q1_vec_mq_split( for (int h = 0; h < MQ_GQA; ++h) { const ACC_TYPE dot_partial = dot4[h].s0 + dot4[h].s1 + dot4[h].s2 + dot4[h].s3; ACC_TYPE s = sub_group_reduce_add(dot_partial) * scale; + if (logit_softcap > 0.0f) { + s = logit_softcap * tanh(s / logit_softcap); + } if (mask_base[h] != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h]; s += slope[h] * (ACC_TYPE) mask_ptr[k_idx]; } - if (logit_softcap > 0.0f) { - s = logit_softcap * tanh(s / logit_softcap); - } score[h] = s; } @@ -1151,13 +1151,13 @@ __kernel void flash_attn_f32_q8_0_q1_vec_mq_split_c8( s += sub_group_shuffle_xor(s, step); } s *= scale; + if (logit_softcap > 0.0f) { + s = logit_softcap * tanh(s / logit_softcap); + } if (mask_base[h] != NULL) { const global MASK_DATA_TYPE * mask_ptr = (const global MASK_DATA_TYPE *) mask_base[h]; s += slope[h] * (ACC_TYPE) mask_ptr[k_safe]; } - if (logit_softcap > 0.0f) { - s = logit_softcap * tanh(s / logit_softcap); - } score[h] = valid ? s : FA_M_INIT; } @@ -1652,9 +1652,15 @@ __kernel void flash_attn_f32_q8_0( s2 = sub_group_shuffle(s2, tri_base + 0) + sub_group_shuffle(s2, tri_base + 1) + sub_group_shuffle(s2, tri_base + 2); s3 = sub_group_shuffle(s3, tri_base + 0) + sub_group_shuffle(s3, tri_base + 1) + sub_group_shuffle(s3, tri_base + 2); #endif - if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; } #endif s0 *= scale; s1 *= scale; s2 *= scale; s3 *= scale; + if (logit_softcap > 0.0f) { + s0 = logit_softcap * tanh(s0 / logit_softcap); + s1 = logit_softcap * tanh(s1 / logit_softcap); + s2 = logit_softcap * tanh(s2 / logit_softcap); + s3 = logit_softcap * tanh(s3 / logit_softcap); + } + if (!query_valid) { s0 = FA_M_INIT; s1 = FA_M_INIT; s2 = FA_M_INIT; s3 = FA_M_INIT; } if (is_causal) { const int causal_limit = n_kv - n_q + my_query_row; @@ -1676,12 +1682,6 @@ __kernel void flash_attn_f32_q8_0( if (k_row2 < n_kv) s2 += slope * (ACC_TYPE) mask_ptr[k_row2]; if (k_row3 < n_kv) s3 += slope * (ACC_TYPE) mask_ptr[k_row3]; } - if (logit_softcap > 0.0f) { - s0 = logit_softcap * tanh(s0 / logit_softcap); - s1 = logit_softcap * tanh(s1 / logit_softcap); - s2 = logit_softcap * tanh(s2 / logit_softcap); - s3 = logit_softcap * tanh(s3 / logit_softcap); - } const ACC_TYPE m_new = max(m_i, max(max(s0, s1), max(s2, s3))); // Whole tile masked (m_new == FA_M_INIT): force the exp() args From 68ac7657268db96d1beab80a065c424b3ae7f029 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrien=20Gallou=C3=ABt?= Date: Thu, 15 Oct 2026 09:12:45 +0200 Subject: [PATCH 010/423] llama : split llama_prefetch() into merge and prefetch steps (#30532) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * llama : split llama_prefetch() into merge and prefetch steps Merge the ranges first, then call the system functions. This fixes -Wunreachable-code-return on Windows. Signed-off-by: Adrien Gallouët * Use const_cast Signed-off-by: Adrien Gallouët --------- Signed-off-by: Adrien Gallouët --- src/llama-mmap.cpp | 39 ++++++++++++++++++--------------------- 1 file changed, 18 insertions(+), 21 deletions(-) diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index 80a719670b7b..4d1cd31e9863 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -865,23 +865,7 @@ void llama_prefetch(llama_memory_ranges mr) { }); uintptr_t begin = 0, end = 0; -#if defined(_WIN32) - // collect the mr and prefetch them in one call, so the reads can be issued concurrently - std::vector entries; - auto prefetch = [&]() { - entries.push_back({ (PVOID) begin, (SIZE_T) (end - begin) }); - return true; - }; -#else - auto prefetch = [&]() { - // TODO: in a follow-up, use posix_madvise(.., POSIX_MADV_WILLNEED) to match the rest of this file - if (madvise((void *) begin, end - begin, MADV_WILLNEED) != 0) { - LLAMA_LOG_WARN("llama_prefetch: madvise(MADV_WILLNEED) failed: %s\n", strerror(errno)); - return false; - } - return true; - }; -#endif + size_t n_ranges = 0; for (const auto & range : mr) { if (!range.addr || range.size == 0) { continue; @@ -890,9 +874,7 @@ void llama_prefetch(llama_memory_ranges mr) { const uintptr_t first = pointer / page * page; const uintptr_t last = (pointer + range.size + page - 1) / page * page; if (end && first > end) { - if (!prefetch()) { - return; - } + mr[n_ranges++] = { (const void *) begin, end - begin }; end = 0; } if (!end) { @@ -901,13 +883,28 @@ void llama_prefetch(llama_memory_ranges mr) { end = std::max(end, last); } if (end) { - prefetch(); + mr[n_ranges++] = { (const void *) begin, end - begin }; } + mr.resize(n_ranges); #if defined(_WIN32) + // prefetch the merged ranges in one call, so the reads can be issued concurrently + std::vector entries; + entries.reserve(mr.size()); + for (const auto & range : mr) { + entries.push_back({ const_cast(range.addr), (SIZE_T) range.size }); + } if (!entries.empty() && !pPrefetchVirtualMemory(GetCurrentProcess(), (ULONG_PTR) entries.size(), entries.data(), 0)) { LLAMA_LOG_WARN("llama_prefetch: PrefetchVirtualMemory failed: %s\n", llama_format_win_err(GetLastError()).c_str()); } +#else + for (const auto & range : mr) { + // TODO: in a follow-up, use posix_madvise(.., POSIX_MADV_WILLNEED) to match the rest of this file + if (madvise(const_cast(range.addr), range.size, MADV_WILLNEED) != 0) { + LLAMA_LOG_WARN("llama_prefetch: madvise(MADV_WILLNEED) failed: %s\n", strerror(errno)); + return; + } + } #endif #else GGML_UNUSED(mr); From 2ce20d22fdc10cd43bc7c3076559f028a4f45401 Mon Sep 17 00:00:00 2001 From: Simone Iacomino <92865212+simoneiacomino@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:12:56 +0200 Subject: [PATCH 011/423] unicode : reject overlong UTF-8 and out-of-range codepoints (#30485) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * unicode : reject overlong UTF-8 and out-of-range codepoints * Update tests/test-unicode.cpp Co-authored-by: Sigbjørn Skjæret * tests: document byte-wise UTF-8 error recovery * break out of first test on failure --------- Co-authored-by: Sigbjørn Skjæret --- src/unicode.cpp | 9 +++++++++ tests/test-unicode.cpp | 46 +++++++++++++++++++++++++++++++++++++++--- 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/src/unicode.cpp b/src/unicode.cpp index 8d8843e0e468..6c834cbced3c 100644 --- a/src/unicode.cpp +++ b/src/unicode.cpp @@ -42,6 +42,9 @@ uint32_t unicode_cpt_from_utf8(const std::string & utf8, size_t & offset) { throw std::invalid_argument("invalid character"); } auto result = ((utf8[offset + 0] & 0x1f) << 6) | (utf8[offset + 1] & 0x3f); + if (result < 0x80) { + throw std::invalid_argument("invalid character"); + } offset += 2; return result; } @@ -50,6 +53,9 @@ uint32_t unicode_cpt_from_utf8(const std::string & utf8, size_t & offset) { throw std::invalid_argument("invalid character"); } auto result = ((utf8[offset + 0] & 0x0f) << 12) | ((utf8[offset + 1] & 0x3f) << 6) | (utf8[offset + 2] & 0x3f); + if (result < 0x800) { + throw std::invalid_argument("invalid character"); + } offset += 3; return result; } @@ -58,6 +64,9 @@ uint32_t unicode_cpt_from_utf8(const std::string & utf8, size_t & offset) { throw std::invalid_argument("invalid character"); } auto result = ((utf8[offset + 0] & 0x07) << 18) | ((utf8[offset + 1] & 0x3f) << 12) | ((utf8[offset + 2] & 0x3f) << 6) | (utf8[offset + 3] & 0x3f); + if (result < 0x10000 || result > 0x10ffff) { + throw std::invalid_argument("invalid character"); + } offset += 4; return result; } diff --git a/tests/test-unicode.cpp b/tests/test-unicode.cpp index 561ff3e5452f..5ced7a28a7df 100644 --- a/tests/test-unicode.cpp +++ b/tests/test-unicode.cpp @@ -6,6 +6,7 @@ #include int main() { + bool success = true; const std::vector, std::vector>> normalization = { {{}, {}}, {{0x65, 0x301}, {0xE9}}, @@ -23,7 +24,46 @@ int main() { for (const auto & test : normalization) { if (unicode_cpts_normalize_nfc(test.first) != test.second || unicode_cpts_normalize_nfc(test.second) != test.second) { fprintf(stderr, "unexpected NFC normalization\n"); - return 1; + success = false; + break; + } + } + + const struct { + std::string text; + std::vector expected; + } utf8_tests[] = { + { "\x7f", { 0x7f } }, + { "\xc2\x80", { 0x80 } }, + { "\xdf\xbf", { 0x7ff } }, + { "\xe0\xa0\x80", { 0x800 } }, + { "\xef\xbf\xbf", { 0xffff } }, + { "\xf0\x90\x80\x80", { 0x10000 } }, + { "\xf4\x8f\xbf\xbf", { 0x10ffff } }, + // Invalid-input cases preserve the existing recovery behavior: one U+FFFD per byte. + { "\xc0\xa0", { 0xfffd, 0xfffd } }, + { "\xc1\xbf", { 0xfffd, 0xfffd } }, + { "\xe0\x80\xa0", { 0xfffd, 0xfffd, 0xfffd } }, + { "\xe0\x9f\xbf", { 0xfffd, 0xfffd, 0xfffd } }, + { "\xf0\x80\x80\xa0", { 0xfffd, 0xfffd, 0xfffd, 0xfffd } }, + { "\xf0\x8f\xbf\xbf", { 0xfffd, 0xfffd, 0xfffd, 0xfffd } }, + { "\xf4\x90\x80\x80", { 0xfffd, 0xfffd, 0xfffd, 0xfffd } }, + { "\xf5\x80\x80\x80", { 0xfffd, 0xfffd, 0xfffd, 0xfffd } }, + { "\xf7\xbf\xbf\xbf", { 0xfffd, 0xfffd, 0xfffd, 0xfffd } }, + { "\xe2\x80", { 0xfffd, 0xfffd } }, + }; + + for (const auto & test : utf8_tests) { + auto expected = test.expected; + expected.insert(expected.begin(), 'a'); + expected.push_back('b'); + if (unicode_cpts_from_utf8("a" + test.text + "b") != expected) { + fprintf(stderr, "unexpected UTF-8 decoding for bytes:"); + for (const unsigned char byte : test.text) { + fprintf(stderr, " %02x", byte); + } + fprintf(stderr, "\n"); + success = false; } } @@ -39,8 +79,8 @@ int main() { fprintf(stderr, " [%s]", piece.c_str()); } fprintf(stderr, "\n"); - return 1; + success = false; } - return 0; + return success ? 0 : 1; } From 23f83e056d88ddae0d56b01ccc65fa42ab6625f7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrien=20Gallou=C3=ABt?= Date: Thu, 15 Oct 2026 09:13:18 +0200 Subject: [PATCH 012/423] ggml-cpu : optimize pad_f32 (#30526) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Adrien Gallouët --- ggml/src/ggml-cpu/ops.cpp | 69 +++++++++++++++++++++------------------ 1 file changed, 38 insertions(+), 31 deletions(-) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 3b39683d5c33..7a73e138ef1f 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -8515,40 +8515,47 @@ static void ggml_compute_forward_pad_f32( const int32_t lp3 = ggml_get_op_params_i32(dst, 6); const int32_t rp3 = ggml_get_op_params_i32(dst, 7); - // TODO: optimize + const int64_t nrows = ne1 * ne2 * ne3; + const bool src_row_contig = src0->nb[0] == sizeof(float); + + for (int64_t row = ith; row < nrows; row += nth) { + const int64_t i1 = row % ne1; + const int64_t i2 = (row / ne1) % ne2; + const int64_t i3 = row / (ne1 * ne2); + + float * dst_row = dst_ptr + row * ne0; + + if constexpr (circular_t) { + // circular padding wraps around on a torus + const int64_t src_i1 = ggml_wrap_around(i1 - lp1, ne01); + const int64_t src_i2 = ggml_wrap_around(i2 - lp2, ne02); + const int64_t src_i3 = ggml_wrap_around(i3 - lp3, ne03); - for (int64_t i2 = 0; i2 < ne2; ++i2) { - for (int64_t i1 = ith; i1 < ne1; i1 += nth) { for (int64_t i0 = 0; i0 < ne0; ++i0) { - for (int64_t i3 = 0; i3 < ne3; ++i3) { - // circular means wrap around on a torus, so x and y loop around - if constexpr (circular_t) { - const int64_t dst_idx = i3*(ne0*ne1*ne2) + i2*(ne0*ne1) + i1*ne0 + i0; - const int64_t src_i0 = ggml_wrap_around(i0 - lp0, ne00); - const int64_t src_i1 = ggml_wrap_around(i1 - lp1, ne01); - const int64_t src_i2 = ggml_wrap_around(i2 - lp2, ne02); - const int64_t src_i3 = ggml_wrap_around(i3 - lp3, ne03); - - const int64_t src_idx = - src_i3*nb03 + - src_i2*nb02 + - src_i1*nb01 + - src_i0*nb00; - - const float * src_ptr = (const float *)((char *) src0->data + src_idx); - dst_ptr[dst_idx] = *src_ptr; + const int64_t src_i0 = ggml_wrap_around(i0 - lp0, ne00); + dst_row[i0] = *(const float *)((const char *) src0->data + src_i3*nb03 + src_i2*nb02 + src_i1*nb01 + src_i0*nb00); + } + } else { + const bool interior = (i1 >= lp1 && i1 < ne1 - rp1) + && (i2 >= lp2 && i2 < ne2 - rp2) + && (i3 >= lp3 && i3 < ne3 - rp3); + if (!interior) { + memset(dst_row, 0, ne0 * sizeof(float)); + continue; + } + + const char * src_row = (const char *) src0->data + (i3 - lp3)*nb03 + (i2 - lp2)*nb02 + (i1 - lp1)*nb01; + + if (src_row_contig && lp0 >= 0 && rp0 >= 0) { + memset(dst_row, 0, lp0 * sizeof(float)); + memcpy(dst_row + lp0, src_row, ne00 * sizeof(float)); + memset(dst_row + lp0 + ne00, 0, rp0 * sizeof(float)); + } else { + for (int64_t i0 = 0; i0 < ne0; ++i0) { + if (i0 >= lp0 && i0 < ne0 - rp0) { + dst_row[i0] = *(const float *)(src_row + (i0 - lp0)*nb00); } else { - const int64_t dst_idx = i3*(ne0*ne1*ne2) + i2*(ne0*ne1) + i1*ne0 + i0; - if ((i0 >= lp0 && i0 < ne0 - rp0) \ - && (i1 >= lp1 && i1 < ne1 - rp1) \ - && (i2 >= lp2 && i2 < ne2 - rp2) \ - && (i3 >= lp3 && i3 < ne3 - rp3)) { - const int64_t src_idx = (i3 - lp3)*nb03 + (i2 - lp2)*nb02 + (i1 - lp1)*nb01 + (i0 - lp0)*nb00; - const float * src_ptr = (const float *)((char *) src0->data + src_idx); - dst_ptr[dst_idx] = *src_ptr; - } else { - dst_ptr[dst_idx] = 0; - } + dst_row[i0] = 0.0f; } } } From 2130ca84764e75c7ae9106cec671caa73ddbe416 Mon Sep 17 00:00:00 2001 From: William Tan <1284324+Ninja3047@users.noreply.github.com> Date: Thu, 15 Oct 2026 03:14:01 -0400 Subject: [PATCH 013/423] jinja : fix loop body variable scoping (#30533) Assisted-by: Codex --- common/jinja/runtime.cpp | 7 ++++--- tests/test-jinja.cpp | 32 ++++++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/common/jinja/runtime.cpp b/common/jinja/runtime.cpp index b8a69995b0ab..a53acf84c4de 100644 --- a/common/jinja/runtime.cpp +++ b/common/jinja/runtime.cpp @@ -600,6 +600,7 @@ value for_statement::execute_impl(context & ctx) const { bool noIteration = true; for (size_t i = 0; i < filtered_items.size(); i++) { JJ_DEBUG("For loop iteration %zu/%zu", i + 1, filtered_items.size()); + context loop_scope(scope); value_object loop_obj = mk_val(); loop_obj->has_builtins = false; // loop object has no builtins loop_obj->insert("index", mk_val(i + 1)); @@ -611,11 +612,11 @@ value for_statement::execute_impl(context & ctx) const { loop_obj->insert("length", mk_val(filtered_items.size())); loop_obj->insert("previtem", i > 0 ? filtered_items[i - 1] : mk_val("previtem")); loop_obj->insert("nextitem", i < filtered_items.size() - 1 ? filtered_items[i + 1] : mk_val("nextitem")); - scope.set_val("loop", loop_obj); - scope_update_fns[i](scope); + loop_scope.set_val("loop", loop_obj); + scope_update_fns[i](loop_scope); try { for (auto & stmt : body) { - value val = stmt->execute(scope); + value val = stmt->execute(loop_scope); result->push_back(val); } } catch (const continue_statement::signal &) { diff --git a/tests/test-jinja.cpp b/tests/test-jinja.cpp index 572b30816e5b..9b53e8cec3d4 100644 --- a/tests/test-jinja.cpp +++ b/tests/test-jinja.cpp @@ -360,6 +360,38 @@ static void test_loops(testing & t) { "1x1y2x2y" ); + test_template(t, "loop assignments do not persist between iterations", + "{% for i in [1, 2] %}{% if i == 1 %}{% set x = 7 %}{% endif %}{{ x is defined }};{% endfor %}", + json::object(), + "True;False;" + ); + + test_template(t, "loop assignments shadow outer variables", + "{% set x = 9 %}{% for i in [1, 2] %}{% if i == 1 %}{% set x = 7 %}{% endif %}{{ x }};{% endfor %}{{ x }}", + json::object(), + "7;9;9" + ); + + test_template(t, "nested loop assignments do not persist between iterations", + "{% for i in [1, 2] %}{% for j in [1, 2] %}{% if j == 1 %}{% set x = i %}{% endif %}{{ x is defined }};{% endfor %}{% endfor %}", + json::object(), + "True;False;True;False;" + ); + + test_template(t, "reasoning is local to each message", + "{% for m in messages %}{% if m.reasoning_content is string %}{% set reasoning_content = m.reasoning_content %}{% elif '' in m.content %}{% set reasoning_content = m.content.split('')[0].split('')[-1] %}{% endif %}[{{ reasoning_content is defined }}:{{ reasoning_content if reasoning_content is defined else '' }}]{% endfor %}", + {{"messages", json::array({ + {{"reasoning_content", "first"}, {"content", "answer"}}, + {{"content", "answer"}}, + {{"reasoning_content", nullptr}, {"content", "answer"}}, + {{"reasoning_content", 1}, {"content", "answer"}}, + {{"reasoning_content", ""}, {"content", "answer"}}, + {{"content", "embeddedanswer"}}, + {{"content", "answer"}} + })}}, + "[True:first][False:][False:][False:][True:][True:embedded][False:]" + ); + test_template(t, "for with range", "{% for i in range(3) %}{{ i }}{% endfor %}", json::object(), From 9166ceaf16599406366cd2d01dc6a2126bbada4d Mon Sep 17 00:00:00 2001 From: lhez Date: Thu, 15 Oct 2026 00:21:18 -0700 Subject: [PATCH 014/423] opencl: fix OOB tail access for adreno noshuffle gemv kernels (#30421) Co-authored-by: Hongqiang Wang --- .../kernels/gemv_noshuffle_iq4_nl_f32.cl | 20 ++-- .../kernels/gemv_noshuffle_q1_0_f32.cl | 12 ++- .../kernels/gemv_noshuffle_q4_0_f32.cl | 50 +++++----- .../kernels/gemv_noshuffle_q4_0_f32_spec.cl | 20 ++-- .../kernels/gemv_noshuffle_q4_1_f32.cl | 54 ++++++----- .../kernels/gemv_noshuffle_q4_k_f32.cl | 91 +++++++++++-------- .../kernels/gemv_noshuffle_q4_k_f32_o4.cl | 5 +- .../kernels/gemv_noshuffle_q5_0_f32.cl | 28 +++--- .../kernels/gemv_noshuffle_q5_1_f32.cl | 30 +++--- .../kernels/gemv_noshuffle_q5_k_f32.cl | 80 ++++++++-------- .../kernels/gemv_noshuffle_q6_k_f32.cl | 86 ++++++++++-------- .../kernels/gemv_noshuffle_q6_k_f32_o4.cl | 5 +- .../kernels/gemv_noshuffle_q8_0_f32.cl | 40 ++++---- 13 files changed, 287 insertions(+), 234 deletions(-) diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_iq4_nl_f32.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_iq4_nl_f32.cl index 1f832cb253b7..e17ca09ba815 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_iq4_nl_f32.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_iq4_nl_f32.cl @@ -230,6 +230,8 @@ kernel void kernel_gemv_noshuffle_iq4_nl_f32( uint M = ne01; uint LINE_STRIDE_A = M / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = NSUBGROUPS * M; private uint4 regA; @@ -240,7 +242,7 @@ kernel void kernel_gemv_noshuffle_iq4_nl_f32( // loop along K in block granularity, skip 4 blocks every iter for (uint k = groupId; k < (K / QK4_NL); k += NSUBGROUPS) { - regS = src0_d[gid + k * LINE_STRIDE_A]; // each fiber loads scale of two rows + regS = src0_d[gid_s + k * LINE_STRIDE_A]; // each fiber loads scale of two rows // first 4 fibers in each wave load 8 B values to its private scope if (slid < 4) { regB.s0123 = read_imagef(src1, (slid * 2 + k * 8)); @@ -248,20 +250,20 @@ kernel void kernel_gemv_noshuffle_iq4_nl_f32( } // load half weights for two blocks in consecutive rows - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), regS, regB); #else dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), regS, regB); #endif // VECTOR_SUB_GROUP_BROADCAST - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), regS, regB); #else diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q1_0_f32.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q1_0_f32.cl index 9efede29411b..d935a489c12d 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q1_0_f32.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q1_0_f32.cl @@ -73,6 +73,8 @@ __kernel void kernel_gemv_noshuffle_q1_0_f32( uint M = ne01; uint LINE_STRIDE_A = M; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = 4 * M; uint4 regA; @@ -83,7 +85,7 @@ __kernel void kernel_gemv_noshuffle_q1_0_f32( #pragma unroll 1 for (uint kb = groupId; kb < (K / QK1_0); kb += N_SIMDGROUP) { - regS = src0_d[gid + kb * LINE_STRIDE_A]; // each fiber loads its row's scale + regS = src0_d[gid_s + kb * LINE_STRIDE_A]; // each fiber loads its row's scale // first 16 fibers load 8 B values each -> 128 activations for this block if (slid < 16) { @@ -92,10 +94,10 @@ __kernel void kernel_gemv_noshuffle_q1_0_f32( } // load this row's 4 uint32 (128 sign bits) - regA.s0 = read_imageui(src0_q, (gid + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA.s1 = read_imageui(src0_q, (gid + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA.s2 = read_imageui(src0_q, (gid + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA.s3 = read_imageui(src0_q, (gid + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA.s0 = read_imageui(src0_q, (gid_s + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA.s1 = read_imageui(src0_q, (gid_s + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA.s2 = read_imageui(src0_q, (gid_s + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA.s3 = read_imageui(src0_q, (gid_s + kb * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; float scale = (float)regS; dequantizeBlockAccum_q1(totalSum, regA.s0, scale, regB, 0); diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_0_f32.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_0_f32.cl index 0dcc929a6bea..6066ae3a9b40 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_0_f32.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_0_f32.cl @@ -224,6 +224,8 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32( uint M = ne01; uint LINE_STRIDE_A = M / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = N_SIMDGROUP * M; __private uint4 regA; @@ -234,7 +236,7 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32( // loop along K in block granularity, skip 4 blocks every iter for (uint k = groupId; k < (K / QK4_0); k += N_SIMDGROUP) { - regS = src0_d[gid + k * LINE_STRIDE_A]; // each fiber loads scale of two rows + regS = src0_d[gid_s + k * LINE_STRIDE_A]; // each fiber loads scale of two rows // first 4 fibers in each wave load 8 B values to its private scope if (slid < 4) { regB.s0123 = read_imagef(src1, (slid * 2 + k * 8)); @@ -242,20 +244,20 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32( } // load half weights for two blocks in consecutive rows - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), regS, regB); #else dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), regS, regB); #endif // VECTOR_SUB_GROUP_BROADCAST - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), regS, regB); #else @@ -329,6 +331,8 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32_mc3( uint M = ne01; uint LINE_STRIDE_A = M / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); // BLOCK_STRIDE_A is the LAYOUT stride between consecutive K-blocks = 4 uints // per q4_0 block * M (set by the trans4_ns convert). The "4" is uints/block, NOT // the subgroup count — keep it fixed so the K-split count (nsg) can vary. @@ -347,17 +351,17 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32_mc3( __private float2 ts3 = (float2)(0.0f); for (uint k = groupId; k < (K / QK4_0); k += nsg) { - regS = src0_d[gid + k * LINE_STRIDE_A]; + regS = src0_d[gid_s + k * LINE_STRIDE_A]; // weights loaded ONCE, reused across the columns - regA_hi.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA_hi.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA_hi.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA_hi.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; - regA_lo.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA_lo.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA_lo.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA_lo.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA_hi.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA_hi.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA_hi.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA_hi.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA_lo.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA_lo.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA_lo.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA_lo.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; MC_COL_Q40(ts0, 0); MC_COL_Q40(ts1, 1); @@ -381,10 +385,12 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32_mc3( } dst = (global float*)((global char*)dst + offsetd); // dst is column-major [M rows x n_cols cols]: (row, col) at col*M + row - vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2])); - vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2])); - if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2])); - if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2])); + if (gid < LINE_STRIDE_A) { + vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2])); + vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2])); + if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2])); + if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2])); + } } } #undef MC_COL_Q40 diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_0_f32_spec.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_0_f32_spec.cl index 0dca20f71f92..63990c414044 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_0_f32_spec.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_0_f32_spec.cl @@ -210,6 +210,8 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32( { uint groupId = get_local_id(1); uint gid = get_global_id(0); + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, (uint)LINE_STRIDE_A - 1u); ushort slid = get_sub_group_local_id(); __private uint4 regA; @@ -220,7 +222,7 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32( // loop along K in block granularity, skip 4 blocks every iter for (uint k = groupId; k < (K / QK4_0); k += N_SIMDGROUP) { - regS = src0_d[gid + k * LINE_STRIDE_A]; // each fiber loads scale of two rows + regS = src0_d[gid_s + k * LINE_STRIDE_A]; // each fiber loads scale of two rows // first 4 fibers in each wave load 8 B values to its private scope if (slid < 4) { regB.s0123 = read_imagef(src1, (slid * 2 + k * 8)); @@ -228,20 +230,20 @@ __kernel void kernel_gemv_noshuffle_q4_0_f32( } // load half weights for two blocks in consecutive rows - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), regS, regB); #else dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), regS, regB); #endif // VECTOR_SUB_GROUP_BROADCAST - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), regS, regB); #else diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_1_f32.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_1_f32.cl index 2ccf4214c0bb..b0b4793d7f18 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_1_f32.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_1_f32.cl @@ -209,6 +209,8 @@ kernel void kernel_gemv_noshuffle_q4_1_f32( uint M = ne01; uint LINE_STRIDE_A = M / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = NSUBGROUPS * M; private uint4 regA; @@ -220,8 +222,8 @@ kernel void kernel_gemv_noshuffle_q4_1_f32( // loop along K in block granularity, skip 4 blocks every iter for (uint k = groupId; k < (K / QK4_0); k += NSUBGROUPS) { - regS = src0_d[gid + k * LINE_STRIDE_A]; // each fiber loads scale of two rows - regM = src0_m[gid + k * LINE_STRIDE_A]; // each fiber loads min of two rows + regS = src0_d[gid_s + k * LINE_STRIDE_A]; // each fiber loads scale of two rows + regM = src0_m[gid_s + k * LINE_STRIDE_A]; // each fiber loads min of two rows // first 4 fibers in each wave load 8 B values to its private scope if (slid < 4) { regB.s0123 = read_imagef(src1, (slid * 2 + k * 8)); @@ -229,20 +231,20 @@ kernel void kernel_gemv_noshuffle_q4_1_f32( } // load half weights for two blocks in consecutive rows - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; #ifdef VECTOR_SUB_GROUP_BROADCAT dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), regS, regM, regB); #else dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), regS, regM, regB); #endif // VECTOR_SUB_GROUP_BROADCAT - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; #ifdef VECTOR_SUB_GROUP_BROADCAT dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), regS, regM, regB); #else @@ -326,6 +328,8 @@ kernel void kernel_gemv_noshuffle_q4_1_f32_mc3( uint M = ne01; uint LINE_STRIDE_A = M / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = NSUBGROUPS * M; uint COL_STRIDE = K / 4; // float4 pixels per activation column @@ -339,18 +343,18 @@ kernel void kernel_gemv_noshuffle_q4_1_f32_mc3( private float2 ts3 = (float2)(0.0f); for (uint k = groupId; k < (K / QK4_0); k += NSUBGROUPS) { - regS = src0_d[gid + k * LINE_STRIDE_A]; - regM = src0_m[gid + k * LINE_STRIDE_A]; + regS = src0_d[gid_s + k * LINE_STRIDE_A]; + regM = src0_m[gid_s + k * LINE_STRIDE_A]; // weights loaded ONCE, reused across the columns - regA_hi.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA_hi.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA_hi.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA_hi.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; - regA_lo.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA_lo.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA_lo.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA_lo.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA_hi.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA_hi.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA_hi.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA_hi.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA_lo.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA_lo.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA_lo.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA_lo.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; MC_COL_Q41(ts0, 0); MC_COL_Q41(ts1, 1); @@ -373,10 +377,12 @@ kernel void kernel_gemv_noshuffle_q4_1_f32_mc3( acc += reduceLM[SUBGROUP_SIZE * 2 + slid]; dst = (global float*)((global char*)dst + offsetd); // dst is column-major [M rows x n_cols cols]: (row, col) at col*M + row - vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2])); - vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2])); - if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2])); - if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2])); + if (gid < LINE_STRIDE_A) { + vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2])); + vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2])); + if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2])); + if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2])); + } } } #undef MC_COL_Q41 diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_k_f32.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_k_f32.cl index 916e00118639..3dc30db5889f 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_k_f32.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_k_f32.cl @@ -415,6 +415,8 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_glu( uint M = ne01; uint LINE_STRIDE_A = M / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = 4 * M; private uint4 regA; @@ -434,9 +436,9 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_glu( for (uint k = groupId; k < (K / 32); k += nsg) { \ uint sb = k / 8; \ uint j = k % 8; \ - half2 d = DD[gid + sb * LINE_STRIDE_A]; \ - half2 dm = MM[gid + sb * LINE_STRIDE_A]; \ - global const uchar * sc0 = SS + sb * 12 * M + 2 * gid; \ + half2 d = DD[gid_s + sb * LINE_STRIDE_A]; \ + half2 dm = MM[gid_s + sb * LINE_STRIDE_A]; \ + global const uchar * sc0 = SS + sb * 12 * M + 2 * gid_s; \ global const uchar * sc1 = sc0 + 1; \ uchar sv0, mn0, sv1, mn1; \ get_scale_min_k4(j, sc0, M, &sv0, &mn0, mask_d6, mask_d4, mask_hi2); \ @@ -447,15 +449,15 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_glu( regB.s0123 = read_imagef(src1, (slid * 2 + k * 8)); \ regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8)); \ } \ - regA.s0 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; \ - regA.s1 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; \ - regA.s2 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; \ - regA.s3 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; \ + regA.s0 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; \ + regA.s1 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; \ + regA.s2 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; \ + regA.s3 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; \ DEQ_HI(SUM, as_ushort8(regA), regS, regM, regB); \ - regA.s0 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; \ - regA.s1 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; \ - regA.s2 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; \ - regA.s3 = read_imageui(Q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; \ + regA.s0 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; \ + regA.s1 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; \ + regA.s2 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; \ + regA.s3 = read_imageui(Q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; \ DEQ_LO(SUM, as_ushort8(regA), regS, regM, regB); \ } @@ -489,8 +491,10 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_glu( upSum += p.zw; } dst = (global float*)((global char*)dst + offsetd); - dst[gid * 2 + 0] = glu_apply(glu_op, gateSum.s0, upSum.s0); - dst[gid * 2 + 1] = glu_apply(glu_op, gateSum.s1, upSum.s1); + if (gid < LINE_STRIDE_A) { + dst[gid * 2 + 0] = glu_apply(glu_op, gateSum.s0, upSum.s0); + dst[gid * 2 + 1] = glu_apply(glu_op, gateSum.s1, upSum.s1); + } } } @@ -533,6 +537,9 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_splitk( uint LINE_STRIDE_A = M / 2; uint BLOCK_STRIDE_A = 4 * M; // physical, independent of the K-split + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); + private uint4 regA; private half2 regS, regM; private float8 regB; @@ -542,9 +549,9 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_splitk( for (uint k = kslice * nsg + groupId; k < (K / 32); k += ksplit * nsg) { uint sb = k / 8; uint j = k % 8; - half2 d = src0_d[gid + sb * LINE_STRIDE_A]; - half2 dm = src0_m[gid + sb * LINE_STRIDE_A]; - global const uchar * sc0 = src0_s + sb * 12 * M + 2 * gid; + half2 d = src0_d[gid_s + sb * LINE_STRIDE_A]; + half2 dm = src0_m[gid_s + sb * LINE_STRIDE_A]; + global const uchar * sc0 = src0_s + sb * 12 * M + 2 * gid_s; global const uchar * sc1 = sc0 + 1; uchar sv0, mn0, sv1, mn1; get_scale_min_k4(j, sc0, M, &sv0, &mn0, mask_d6, mask_d4, mask_hi2); @@ -555,19 +562,19 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_splitk( regB.s0123 = read_imagef(src1, (slid * 2 + k * 8)); regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8)); } - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), regS, regM, regB); #else dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), regS, regM, regB); #endif - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), regS, regM, regB); #else @@ -584,7 +591,9 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_splitk( for (uint i = 0; i < nsg - 1; ++i) { totalSum += reduceLM[SUBGROUP_SIZE * i + slid]; } - vstore2(totalSum, 0, &(partial[kslice * M + gid * 2])); + // Keep tail stores inside this K-slice. + if (gid * 2 + 0 < M) partial[kslice * M + gid * 2 + 0] = totalSum.s0; + if (gid * 2 + 1 < M) partial[kslice * M + gid * 2 + 1] = totalSum.s1; } } @@ -689,6 +698,8 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_mc3( uint M = ne01; uint LINE_STRIDE_A = M / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = NSUBGROUPS * M; uint COL_STRIDE = K / 4; // float4 pixels per activation column @@ -713,10 +724,10 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_mc3( uint sb = k / 8; uint j = k % 8; - half2 d = src0_d[gid + sb * LINE_STRIDE_A]; - half2 dm = src0_m[gid + sb * LINE_STRIDE_A]; + half2 d = src0_d[gid_s + sb * LINE_STRIDE_A]; + half2 dm = src0_m[gid_s + sb * LINE_STRIDE_A]; - global const uchar * sc0 = src0_s + sb * 12 * M + 2 * gid; + global const uchar * sc0 = src0_s + sb * 12 * M + 2 * gid_s; global const uchar * sc1 = sc0 + 1; uchar sv0, mn0, sv1, mn1; @@ -727,14 +738,14 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_mc3( regM = convert_half2(convert_float2(dm) * convert_float2((uchar2)(mn0, mn1))); // weights loaded ONCE, reused across the 3 columns - regA_hi.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA_hi.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA_hi.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA_hi.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; - regA_lo.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA_lo.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA_lo.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA_lo.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA_hi.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA_hi.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA_hi.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA_hi.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA_lo.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA_lo.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA_lo.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA_lo.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; #ifdef Q4K_MC3_DEQUANT_ONCE // Dequant the 32 weights/row (16 hi + 16 lo) ONCE into half2[] (byte- @@ -827,8 +838,10 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_mc3( acc += reduceLM[SUBGROUP_SIZE * 2 + slid]; dst = (global float*)((global char*)dst + offsetd); // dst is column-major [M rows x 3 cols]: (row, col) at col*M + row - vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2])); - vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2])); - vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2])); + if (gid < LINE_STRIDE_A) { + vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2])); + vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2])); + vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2])); + } } } diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_k_f32_o4.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_k_f32_o4.cl index 02916bb91ffa..2c1c659a31e8 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_k_f32_o4.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q4_k_f32_o4.cl @@ -231,8 +231,9 @@ kernel void kernel_gemv_noshuffle_q4_k_f32_o4( // Two consecutive pair-indices (each the same access pattern the 2-output // kernel uses); together they cover 4 consecutive output rows. - uint gid_a = gid * 2; - uint gid_b = gid * 2 + 1; + // Keep tail lanes active for subgroup ops, but read valid weight pairs. + uint gid_a = min(gid * 2, (uint)(ne01 / 2) - 1u); + uint gid_b = min(gid * 2 + 1, (uint)(ne01 / 2) - 1u); uint K = ne00; uint M = ne01; diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_0_f32.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_0_f32.cl index 7dbf5a3bbbfb..d8537e158ee3 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_0_f32.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_0_f32.cl @@ -209,6 +209,8 @@ __kernel void kernel_gemv_noshuffle_q5_0_f32( uint M = ne01; uint LINE_STRIDE_A = M / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = NSUBGROUPS * M; private uint4 regA; @@ -218,13 +220,13 @@ __kernel void kernel_gemv_noshuffle_q5_0_f32( private float2 totalSum = (float2)(0.0f); for (uint k = groupId; k < (K / QK5_0); k += NSUBGROUPS) { - regS = src0_d[gid + k * LINE_STRIDE_A]; + regS = src0_d[gid_s + k * LINE_STRIDE_A]; ushort4 qh_raw; - qh_raw.s0 = src0_qh[gid + (4*k + 0) * LINE_STRIDE_A]; - qh_raw.s1 = src0_qh[gid + (4*k + 1) * LINE_STRIDE_A]; - qh_raw.s2 = src0_qh[gid + (4*k + 2) * LINE_STRIDE_A]; - qh_raw.s3 = src0_qh[gid + (4*k + 3) * LINE_STRIDE_A]; + qh_raw.s0 = src0_qh[gid_s + (4*k + 0) * LINE_STRIDE_A]; + qh_raw.s1 = src0_qh[gid_s + (4*k + 1) * LINE_STRIDE_A]; + qh_raw.s2 = src0_qh[gid_s + (4*k + 2) * LINE_STRIDE_A]; + qh_raw.s3 = src0_qh[gid_s + (4*k + 3) * LINE_STRIDE_A]; uchar8 raw = as_uchar8(qh_raw); uchar8 qh_bytes = (uchar8)(raw.s0, raw.s2, raw.s4, raw.s6, @@ -236,10 +238,10 @@ __kernel void kernel_gemv_noshuffle_q5_0_f32( regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8)); } - regA.s0 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA.s1 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA.s2 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA.s3 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA.s0 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA.s1 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA.s2 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA.s3 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_q5_0_sgbroadcast_8_hi(totalSum, as_ushort8(regA), qh_bytes, regS, regB); @@ -247,10 +249,10 @@ __kernel void kernel_gemv_noshuffle_q5_0_f32( dequantizeBlockAccum_ns_q5_0_sgbroadcast_1_hi(totalSum, as_ushort8(regA), qh_bytes, regS, regB); #endif // VECTOR_SUB_GROUP_BROADCAST - regA.s0 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA.s1 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA.s2 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA.s3 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA.s0 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA.s1 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA.s2 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA.s3 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_q5_0_sgbroadcast_8_lo(totalSum, as_ushort8(regA), qh_bytes, regS, regB); #else diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_1_f32.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_1_f32.cl index ba0e2a711565..82b4c741e804 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_1_f32.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_1_f32.cl @@ -210,6 +210,8 @@ __kernel void kernel_gemv_noshuffle_q5_1_f32( uint M = ne01; uint LINE_STRIDE_A = M / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = NSUBGROUPS * M; __private uint4 regA; @@ -220,14 +222,14 @@ __kernel void kernel_gemv_noshuffle_q5_1_f32( __private float2 totalSum = (float2)(0.0f); for (uint k = groupId; k < (K / QK5_1); k += NSUBGROUPS) { - regS = src0_d[gid + k * LINE_STRIDE_A]; - regM = src0_m[gid + k * LINE_STRIDE_A]; + regS = src0_d[gid_s + k * LINE_STRIDE_A]; + regM = src0_m[gid_s + k * LINE_STRIDE_A]; ushort4 qh_raw; - qh_raw.s0 = src0_qh[gid + (4*k + 0) * LINE_STRIDE_A]; - qh_raw.s1 = src0_qh[gid + (4*k + 1) * LINE_STRIDE_A]; - qh_raw.s2 = src0_qh[gid + (4*k + 2) * LINE_STRIDE_A]; - qh_raw.s3 = src0_qh[gid + (4*k + 3) * LINE_STRIDE_A]; + qh_raw.s0 = src0_qh[gid_s + (4*k + 0) * LINE_STRIDE_A]; + qh_raw.s1 = src0_qh[gid_s + (4*k + 1) * LINE_STRIDE_A]; + qh_raw.s2 = src0_qh[gid_s + (4*k + 2) * LINE_STRIDE_A]; + qh_raw.s3 = src0_qh[gid_s + (4*k + 3) * LINE_STRIDE_A]; uchar8 raw = as_uchar8(qh_raw); uchar8 qh_bytes = (uchar8)(raw.s0, raw.s2, raw.s4, raw.s6, @@ -239,10 +241,10 @@ __kernel void kernel_gemv_noshuffle_q5_1_f32( regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8)); } - regA.s0 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA.s1 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA.s2 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA.s3 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA.s0 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA.s1 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA.s2 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA.s3 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_q5_1_sgbroadcast_8_hi(totalSum, as_ushort8(regA), qh_bytes, regS, regM, regB); @@ -250,10 +252,10 @@ __kernel void kernel_gemv_noshuffle_q5_1_f32( dequantizeBlockAccum_ns_q5_1_sgbroadcast_1_hi(totalSum, as_ushort8(regA), qh_bytes, regS, regM, regB); #endif // VECTOR_SUB_GROUP_BROADCAST - regA.s0 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA.s1 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA.s2 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA.s3 = read_imageui(src0_qs, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA.s0 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA.s1 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA.s2 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA.s3 = read_imageui(src0_qs, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_q5_1_sgbroadcast_8_lo(totalSum, as_ushort8(regA), qh_bytes, regS, regM, regB); #else diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_k_f32.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_k_f32.cl index ae864b19ba9a..6570167c9a68 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_k_f32.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q5_k_f32.cl @@ -232,6 +232,8 @@ kernel void kernel_gemv_noshuffle_q5_k_f32( uint M = ne01; uint LINE_STRIDE_A = M / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = NSUBGROUPS * M; uint LINE_STRIDE_A_QH = M / 2; @@ -250,11 +252,11 @@ kernel void kernel_gemv_noshuffle_q5_k_f32( uint sb = k / 8; uint j = k % 8; - half2 d = src0_d[gid + sb * LINE_STRIDE_A]; - half2 dm = src0_m[gid + sb * LINE_STRIDE_A]; + half2 d = src0_d[gid_s + sb * LINE_STRIDE_A]; + half2 dm = src0_m[gid_s + sb * LINE_STRIDE_A]; - global const uchar * sc0 = src0_s + 2 * gid * scales_per_row + sb * 12; - global const uchar * sc1 = src0_s + (2 * gid + 1) * scales_per_row + sb * 12; + global const uchar * sc0 = src0_s + 2 * gid_s * scales_per_row + sb * 12; + global const uchar * sc1 = src0_s + (2 * gid_s + 1) * scales_per_row + sb * 12; uchar sv0, mn0, sv1, mn1; get_scale_min_k4(j, sc0, &sv0, &mn0, mask_d6, mask_d4, mask_hi2); @@ -268,25 +270,25 @@ kernel void kernel_gemv_noshuffle_q5_k_f32( regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8)); } - regH.s0 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 0)).x); - regH.s1 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 1)).x); - regH.s2 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 2)).x); - regH.s3 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 3)).x); + regH.s0 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 0)).x); + regH.s1 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 1)).x); + regH.s2 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 2)).x); + regH.s3 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 3)).x); - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_sgbroadcast_8_hi(totalSum, as_ushort8(regA), as_uchar8(regH), regS, regM, regB); #else dequantizeBlockAccum_ns_sgbroadcast_1_hi(totalSum, as_ushort8(regA), as_uchar8(regH), regS, regM, regB); #endif // VECTOR_SUB_GROUP_BROADCAST - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; #ifdef VECTOR_SUB_GROUP_BROADCAST dequantizeBlockAccum_ns_sgbroadcast_8_lo(totalSum, as_ushort8(regA), as_uchar8(regH), regS, regM, regB); #else @@ -373,6 +375,8 @@ kernel void kernel_gemv_noshuffle_q5_k_f32_mc3( uint M = ne01; uint LINE_STRIDE_A = M / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = NSUBGROUPS * M; uint LINE_STRIDE_A_QH = M / 2; uint BLOCK_STRIDE_A_QH = NSUBGROUPS * M / 2; @@ -393,11 +397,11 @@ kernel void kernel_gemv_noshuffle_q5_k_f32_mc3( uint sb = k / 8; uint j = k % 8; - half2 d = src0_d[gid + sb * LINE_STRIDE_A]; - half2 dm = src0_m[gid + sb * LINE_STRIDE_A]; + half2 d = src0_d[gid_s + sb * LINE_STRIDE_A]; + half2 dm = src0_m[gid_s + sb * LINE_STRIDE_A]; - global const uchar * sc0 = src0_s + 2 * gid * scales_per_row + sb * 12; - global const uchar * sc1 = src0_s + (2 * gid + 1) * scales_per_row + sb * 12; + global const uchar * sc0 = src0_s + 2 * gid_s * scales_per_row + sb * 12; + global const uchar * sc1 = src0_s + (2 * gid_s + 1) * scales_per_row + sb * 12; uchar sv0, mn0, sv1, mn1; get_scale_min_k4(j, sc0, &sv0, &mn0, mask_d6, mask_d4, mask_hi2); @@ -407,19 +411,19 @@ kernel void kernel_gemv_noshuffle_q5_k_f32_mc3( regM = convert_half2(convert_float2(dm) * convert_float2((uchar2)(mn0, mn1))); // high-bit plane + weights loaded ONCE, reused across the columns - regH.s0 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 0)).x); - regH.s1 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 1)).x); - regH.s2 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 2)).x); - regH.s3 = as_ushort(read_imageh(src0_qh, (gid + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 3)).x); - - regA_hi.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA_hi.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA_hi.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA_hi.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; - regA_lo.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA_lo.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA_lo.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA_lo.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regH.s0 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 0)).x); + regH.s1 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 1)).x); + regH.s2 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 2)).x); + regH.s3 = as_ushort(read_imageh(src0_qh, (gid_s + k * BLOCK_STRIDE_A_QH + LINE_STRIDE_A_QH * 3)).x); + + regA_hi.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA_hi.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA_hi.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA_hi.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA_lo.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA_lo.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA_lo.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA_lo.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; MC_COL_Q5K(ts0, 0); MC_COL_Q5K(ts1, 1); @@ -442,10 +446,12 @@ kernel void kernel_gemv_noshuffle_q5_k_f32_mc3( acc += reduceLM[SUBGROUP_SIZE * 2 + slid]; dst = (global float*)((global char*)dst + offsetd); // dst is column-major [M rows x n_cols cols]: (row, col) at col*M + row - vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2])); - vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2])); - if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2])); - if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2])); + if (gid < LINE_STRIDE_A) { + vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0 * M + gid * 2])); + vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1 * M + gid * 2])); + if (n_cols > 2) vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2 * M + gid * 2])); + if (n_cols > 3) vstore2((float2)(acc.s6, acc.s7), 0, &(dst[3 * M + gid * 2])); + } } } #undef MC_COL_Q5K diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q6_k_f32.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q6_k_f32.cl index 32624ac868fe..a6908a6a7089 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q6_k_f32.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q6_k_f32.cl @@ -219,26 +219,28 @@ kernel void kernel_gemv_noshuffle_q6_K_f32( float2 total_sum = 0.0f; int line_stride_a = ne01 / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + int gid_s = min(gid, line_stride_a - 1); int block_stride_a = NSUBGROUPS * ne01; for (int k = grp; k < nb; k += NSUBGROUPS) { - reg_d = src0_d[gid + k/8 * line_stride_a]; - reg_s = as_char4(src0_s[gid + k * line_stride_a]); + reg_d = src0_d[gid_s + k/8 * line_stride_a]; + reg_s = as_char4(src0_s[gid_s + k * line_stride_a]); if (slid < 4) { reg_b.s0123 = read_imagef(src1, 0 + slid*2 + k*8); reg_b.s4567 = read_imagef(src1, 1 + slid*2 + k*8); } - reg_a_l.s0 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*0).x; - reg_a_l.s1 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*1).x; - reg_a_l.s2 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*2).x; - reg_a_l.s3 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*3).x; + reg_a_l.s0 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*0).x; + reg_a_l.s1 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*1).x; + reg_a_l.s2 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*2).x; + reg_a_l.s3 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*3).x; - reg_a_h.s0 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*0).x); - reg_a_h.s1 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*1).x); - reg_a_h.s2 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*2).x); - reg_a_h.s3 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*3).x); + reg_a_h.s0 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*0).x); + reg_a_h.s1 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*1).x); + reg_a_h.s2 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*2).x); + reg_a_h.s3 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*3).x); #ifdef VECTOR_SUB_GROUP_BROADCAT dequantize_block_acc_bcast_8_hi(total_sum, as_ushort8(reg_a_l), as_uchar8(reg_a_h), reg_d, reg_s, reg_b); @@ -246,15 +248,15 @@ kernel void kernel_gemv_noshuffle_q6_K_f32( dequantize_block_acc_bcast_1_hi(total_sum, as_ushort8(reg_a_l), as_uchar8(reg_a_h), reg_d, reg_s, reg_b); #endif // VECTOR_SUB_GROUP_BROADCAT - reg_a_l.s0 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*4).x; - reg_a_l.s1 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*5).x; - reg_a_l.s2 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*6).x; - reg_a_l.s3 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*7).x; + reg_a_l.s0 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*4).x; + reg_a_l.s1 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*5).x; + reg_a_l.s2 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*6).x; + reg_a_l.s3 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*7).x; - reg_a_h.s0 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*4).x); - reg_a_h.s1 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*5).x); - reg_a_h.s2 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*6).x); - reg_a_h.s3 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*7).x); + reg_a_h.s0 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*4).x); + reg_a_h.s1 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*5).x); + reg_a_h.s2 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*6).x); + reg_a_h.s3 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*7).x); #ifdef VECTOR_SUB_GROUP_BROADCAT dequantize_block_acc_bcast_8_lo(total_sum, as_ushort8(reg_a_l), as_uchar8(reg_a_h), reg_d, reg_s, reg_b); @@ -323,6 +325,8 @@ kernel void kernel_gemv_noshuffle_q6_K_f32_mc3( int nb = ne00 / 32; int line_stride_a = ne01 / 2; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + int gid_s = min(gid, line_stride_a - 1); int block_stride_a = NSUBGROUPS * ne01; int COL_STRIDE = ne00 / 4; // float4 pixels per activation column @@ -335,27 +339,27 @@ kernel void kernel_gemv_noshuffle_q6_K_f32_mc3( float2 ts0 = 0.0f, ts1 = 0.0f, ts2 = 0.0f; for (int k = grp; k < nb; k += NSUBGROUPS) { - reg_d = src0_d[gid + k/8 * line_stride_a]; - reg_s = as_char4(src0_s[gid + k * line_stride_a]); + reg_d = src0_d[gid_s + k/8 * line_stride_a]; + reg_s = as_char4(src0_s[gid_s + k * line_stride_a]); // weights loaded ONCE (hi: blocks 0-3, lo: blocks 4-7), reused x3 cols - ql_hi.s0 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*0).x; - ql_hi.s1 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*1).x; - ql_hi.s2 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*2).x; - ql_hi.s3 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*3).x; - qh_hi.s0 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*0).x); - qh_hi.s1 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*1).x); - qh_hi.s2 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*2).x); - qh_hi.s3 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*3).x); - - ql_lo.s0 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*4).x; - ql_lo.s1 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*5).x; - ql_lo.s2 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*6).x; - ql_lo.s3 = read_imageui(src0_ql, gid + k*block_stride_a + line_stride_a*7).x; - qh_lo.s0 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*4).x); - qh_lo.s1 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*5).x); - qh_lo.s2 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*6).x); - qh_lo.s3 = as_ushort(read_imageh(src0_qh, gid + k*block_stride_a + line_stride_a*7).x); + ql_hi.s0 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*0).x; + ql_hi.s1 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*1).x; + ql_hi.s2 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*2).x; + ql_hi.s3 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*3).x; + qh_hi.s0 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*0).x); + qh_hi.s1 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*1).x); + qh_hi.s2 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*2).x); + qh_hi.s3 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*3).x); + + ql_lo.s0 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*4).x; + ql_lo.s1 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*5).x; + ql_lo.s2 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*6).x; + ql_lo.s3 = read_imageui(src0_ql, gid_s + k*block_stride_a + line_stride_a*7).x; + qh_lo.s0 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*4).x); + qh_lo.s1 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*5).x); + qh_lo.s2 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*6).x); + qh_lo.s3 = as_ushort(read_imageh(src0_qh, gid_s + k*block_stride_a + line_stride_a*7).x); // Per-column: load only this column's activation (single reg_b live) -> // 1/3 the activation register pressure, cutting the private-mem spill. @@ -402,8 +406,10 @@ kernel void kernel_gemv_noshuffle_q6_K_f32_mc3( acc += reduce_lm[SUBGROUP_SIZE*2 + slid]; dst = (global float*)((global char*)dst + offsetd); // dst column-major [ne01 rows x 3 cols]: (row, col) at col*ne01 + row - vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0*ne01 + gid*2])); - vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1*ne01 + gid*2])); - vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2*ne01 + gid*2])); + if (gid < line_stride_a) { + vstore2((float2)(acc.s0, acc.s1), 0, &(dst[0*ne01 + gid*2])); + vstore2((float2)(acc.s2, acc.s3), 0, &(dst[1*ne01 + gid*2])); + vstore2((float2)(acc.s4, acc.s5), 0, &(dst[2*ne01 + gid*2])); + } } } diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q6_k_f32_o4.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q6_k_f32_o4.cl index 84447e61bb6f..22efb6c3caa6 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q6_k_f32_o4.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q6_k_f32_o4.cl @@ -249,8 +249,9 @@ kernel void Q6K_O4_NAME( // ne01/4 apart) is slower because two distant cache-line streams have worse // locality than the adjacent pair whose reads interleave into the same lines // each iteration. - int gid_a = gid * 2; - int gid_b = gid * 2 + 1; + // Keep tail lanes active for subgroup ops, but read valid weight pairs. + int gid_a = min(gid * 2, ne01 / 2 - 1); + int gid_b = min(gid * 2 + 1, ne01 / 2 - 1); int nb = ne00 / 32; diff --git a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q8_0_f32.cl b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q8_0_f32.cl index 6f6d7425c656..6109c967aac1 100644 --- a/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q8_0_f32.cl +++ b/ggml/src/ggml-opencl/kernels/gemv_noshuffle_q8_0_f32.cl @@ -155,6 +155,8 @@ __kernel void kernel_gemv_noshuffle_q8_0_f32_splitk( uint M = ne01; uint LINE_STRIDE_A = M; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = 8 * M; // physical, independent of the K-split __private uint8 regA; @@ -164,19 +166,19 @@ __kernel void kernel_gemv_noshuffle_q8_0_f32_splitk( #pragma unroll 1 for (uint k = kslice * nsg + groupId; k < (K / QK8_0); k += ksplit * nsg) { - regS = src0_d[gid + k * LINE_STRIDE_A]; + regS = src0_d[gid_s + k * LINE_STRIDE_A]; if (slid < 4) { regB.s0123 = read_imagef(src1, (slid * 2 + k * 8)); regB.s4567 = read_imagef(src1, (1 + slid * 2 + k * 8)); } - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; - regA.s4 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA.s5 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA.s6 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA.s7 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA.s4 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA.s5 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA.s6 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA.s7 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; dequantizeBlockAccum_ns_sgbroadcast_1(totalSum, regA, convert_float(regS), regB); } @@ -227,6 +229,8 @@ __kernel void kernel_gemv_noshuffle_q8_0_f32( uint M = ne01; uint LINE_STRIDE_A = M; + // Keep tail lanes active for subgroup ops, but read a valid weight row. + uint gid_s = min(gid, LINE_STRIDE_A - 1); uint BLOCK_STRIDE_A = 8 * M; // 32 / 4 = 8 __private uint8 regA; @@ -238,7 +242,7 @@ __kernel void kernel_gemv_noshuffle_q8_0_f32( // loop along K in block granularity, skip 4 blocks every iter #pragma unroll 1 /* tell compiler not to unroll */ for (uint k = groupId; k < (K / QK8_0); k += N_SIMDGROUP) { - regS = src0_d[gid + k * LINE_STRIDE_A]; // each fiber loads scale of one rows + regS = src0_d[gid_s + k * LINE_STRIDE_A]; // each fiber loads the scale of one row // first 4 fibers in each wave load 8 B values to its private scope if (slid < 4) { regB.s0123 = read_imagef(src1, (slid * 2 + k * 8)); @@ -246,14 +250,14 @@ __kernel void kernel_gemv_noshuffle_q8_0_f32( } // load weights for one block in consecutive rows - regA.s0 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; - regA.s1 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; - regA.s2 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; - regA.s3 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; - regA.s4 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; - regA.s5 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; - regA.s6 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; - regA.s7 = read_imageui(src0_q, (gid + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; + regA.s0 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 0)).x; + regA.s1 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 1)).x; + regA.s2 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 2)).x; + regA.s3 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 3)).x; + regA.s4 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 4)).x; + regA.s5 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 5)).x; + regA.s6 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 6)).x; + regA.s7 = read_imageui(src0_q, (gid_s + k * BLOCK_STRIDE_A + LINE_STRIDE_A * 7)).x; dequantizeBlockAccum_ns_sgbroadcast_1(totalSum, regA, convert_float(regS), regB); } From e36df604703dec1b6ce82f214698725475a57035 Mon Sep 17 00:00:00 2001 From: Neo Zhang Date: Thu, 15 Oct 2026 15:26:01 +0800 Subject: [PATCH 015/423] sycl : replace oneAPI package by DLE+oneDNN package (#30131) * replace oneAPI package by DLE+oneDNN package * support to install oneDNN in windows, in release and docker * for internal test * remove test code * Update intel.Dockerfile --- .devops/intel.Dockerfile | 7 ++++++- .github/workflows/build-sycl.yml | 20 +++++++++++++++----- .github/workflows/release.yml | 21 ++++++++++++++++----- 3 files changed, 37 insertions(+), 11 deletions(-) diff --git a/.devops/intel.Dockerfile b/.devops/intel.Dockerfile index e8b7e6077045..ced0d641a889 100644 --- a/.devops/intel.Dockerfile +++ b/.devops/intel.Dockerfile @@ -19,7 +19,12 @@ RUN npm ci COPY tools/ui/ ./ RUN LLAMA_BUILD_NUMBER="$APP_VERSION" npm run build -FROM docker.io/intel/oneapi-toolkit:$ONEAPI_VERSION AS build +FROM docker.io/intel/deep-learning-essentials:$ONEAPI_VERSION AS build + +RUN echo "Download oneDNN installation package" && \ + wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/94c3dbac-0852-45be-a57d-21c204cada3e/intel-onednn-2026.0.2.46_offline.sh -O intel-onednn_offline.sh && \ + echo "Install oneDNN components" && \ + bash intel-onednn_offline.sh -s -a --silent --eula accept ARG GGML_SYCL_F16=ON ARG LEVEL_ZERO_VERSION=1.28.2 diff --git a/.github/workflows/build-sycl.yml b/.github/workflows/build-sycl.yml index 741c9ce77c1c..54380f40f50b 100644 --- a/.github/workflows/build-sycl.yml +++ b/.github/workflows/build-sycl.yml @@ -67,8 +67,15 @@ jobs: shell: bash run: | cd /tmp - wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/5996e26b-f48a-42b1-8db0-b002ad0bd8d7/intel-oneapi-toolkit-2026.1.1.33_offline.sh -O intel-oneapi-toolkit_offline.sh - sudo bash intel-oneapi-toolkit_offline.sh -s -a --silent --eula accept + echo "Download oneAPI Deep Learning Essentials installation package" + wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/b5c100a7-8052-4e88-90e2-63d9916a998b/intel-deep-learning-essentials-2026.1.4.17_offline.sh -O intel-deep-learning-essentials_offline.sh + echo "Install oneAPI Deep Learning Essentials components" + sudo bash intel-deep-learning-essentials_offline.sh -s -a --silent --eula accept + + echo "Download oneDNN installation package" + wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/94c3dbac-0852-45be-a57d-21c204cada3e/intel-onednn-2026.0.2.46_offline.sh -O intel-onednn_offline.sh + echo "Install oneDNN components" + sudo bash intel-onednn_offline.sh -s -a --silent --eula accept - name: Install Level Zero SDK shell: bash @@ -138,8 +145,10 @@ jobs: shell: bash env: - WINDOWS_BASEKIT_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/0cb67a0d-67f6-410b-868b-f4a0a17ff0cf/intel-oneapi-toolkit-2026.1.1.32_offline.exe - WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.dnnl:intel.oneapi.win.tbb.devel + WINDOWS_DLE_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/a51104eb-6f4a-4a0e-9dae-1b81b0c81ae1/intel-deep-learning-essentials-2026.1.4.16_offline.exe + WINDOWS_ONEDNN_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/34e28083-d04b-46a3-826c-8c1d00bf1e1c/intel-onednn-2026.0.2.47_offline.exe + WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.tbb.devel + WINDOWS_DNNL: intel.oneapi.win.dnnl LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.33.1/level-zero-win-sdk-1.33.1.zip ONEAPI_ROOT: "C:/Program Files (x86)/Intel/oneAPI" ONEAPI_INSTALLER_VERSION: "2026.1" @@ -151,7 +160,8 @@ jobs: - name: Download & Install oneAPI shell: bash run: | - scripts/install-oneapi.bat $WINDOWS_BASEKIT_URL $WINDOWS_DPCPP_MKL + scripts/install-oneapi.bat $WINDOWS_DLE_URL $WINDOWS_DPCPP_MKL + scripts/install-oneapi.bat $WINDOWS_ONEDNN_URL $WINDOWS_DNNL - name: Install Level Zero SDK shell: pwsh diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 9745230f500c..035ea958e26d 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1300,8 +1300,11 @@ jobs: shell: bash env: - WINDOWS_BASEKIT_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/0cb67a0d-67f6-410b-868b-f4a0a17ff0cf/intel-oneapi-toolkit-2026.1.1.32_offline.exe - WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.dnnl:intel.oneapi.win.tbb.devel + WINDOWS_DLE_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/a51104eb-6f4a-4a0e-9dae-1b81b0c81ae1/intel-deep-learning-essentials-2026.1.4.16_offline.exe + WINDOWS_ONEDNN_URL: https://registrationcenter-download.intel.com/akdlm/IRC_NAS/34e28083-d04b-46a3-826c-8c1d00bf1e1c/intel-onednn-2026.0.2.47_offline.exe + WINDOWS_DPCPP_MKL: intel.oneapi.win.cpp-dpcpp-common:intel.oneapi.win.mkl.devel:intel.oneapi.win.tbb.devel + WINDOWS_DNNL: intel.oneapi.win.dnnl + LEVEL_ZERO_SDK_URL: https://github.com/oneapi-src/level-zero/releases/download/v1.33.1/level-zero-win-sdk-1.33.1.zip ONEAPI_ROOT: "C:/Program Files (x86)/Intel/oneAPI" ONEAPI_INSTALLER_VERSION: "2026.1" @@ -1314,7 +1317,8 @@ jobs: - name: Download & Install oneAPI shell: bash run: | - scripts/install-oneapi.bat $WINDOWS_BASEKIT_URL $WINDOWS_DPCPP_MKL + scripts/install-oneapi.bat $WINDOWS_DLE_URL $WINDOWS_DPCPP_MKL + scripts/install-oneapi.bat $WINDOWS_ONEDNN_URL $WINDOWS_DNNL - name: Install Level Zero SDK shell: pwsh @@ -1425,8 +1429,15 @@ jobs: shell: bash run: | cd /tmp - wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/5996e26b-f48a-42b1-8db0-b002ad0bd8d7/intel-oneapi-toolkit-2026.1.1.33_offline.sh -O intel-oneapi-toolkit_offline.sh - sudo bash intel-oneapi-toolkit_offline.sh -s -a --silent --eula accept + echo "Download oneAPI Deep Learning Essentials installation package" + wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/b5c100a7-8052-4e88-90e2-63d9916a998b/intel-deep-learning-essentials-2026.1.4.17_offline.sh -O intel-deep-learning-essentials_offline.sh + echo "Install oneAPI Deep Learning Essentials components" + sudo bash intel-deep-learning-essentials_offline.sh -s -a --silent --eula accept + + echo "Download oneDNN installation package" + wget https://registrationcenter-download.intel.com/akdlm/IRC_NAS/94c3dbac-0852-45be-a57d-21c204cada3e/intel-onednn-2026.0.2.46_offline.sh -O intel-onednn_offline.sh + echo "Install oneDNN components" + sudo bash intel-onednn_offline.sh -s -a --silent --eula accept - name: Install Level Zero SDK shell: bash From 47f5a64ae12833b9092c6c47b54aa9f310c4f6f1 Mon Sep 17 00:00:00 2001 From: losewayy <1134285612@qq.com> Date: Thu, 15 Oct 2026 16:37:35 +0800 Subject: [PATCH 016/423] vulkan : fix concurrent first device init (#29766) * vulkan : fix concurrent first device init * tests : load models in parallel in test-thread-safety Loading them serially never exercised concurrent first backend init, so the crash reported in #29764 was invisible to this test. * vulkan : create fence before retaining device in buffer context --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 16 +++++++-- tests/test-thread-safety.cpp | 51 +++++++++++++++++----------- 2 files changed, 45 insertions(+), 22 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 42057114071b..436bfe3bb054 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3990,10 +3990,12 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { vk_device ggml_vk_get_device(size_t idx) { VK_LOG_DEBUG("ggml_vk_get_device(" << idx << ")"); + static std::mutex mutex; + std::lock_guard lock(mutex); + if (vk_instance.devices[idx] == nullptr) { VK_LOG_DEBUG("Initializing new vk_device"); vk_device device = std::make_shared(); - vk_instance.devices[idx] = device; device->memory_logger = std::unique_ptr(new vk_memory_logger()); @@ -4891,14 +4893,14 @@ vk_device ggml_vk_get_device(size_t idx) { device->async_use_transfer_queue = false; } + device->fence = device->device.createFence({}); + device->buffer_type = { /* .iface = */ ggml_backend_vk_buffer_type_interface, /* .device = */ ggml_backend_reg_dev_get(ggml_backend_vk_reg(), idx), /* .context = */ new ggml_backend_vk_buffer_type_context{ device->name, device }, }; - device->fence = device->device.createFence({}); - device->idx = idx; device->serialize_submissions = getenv("GGML_VK_SERIALIZE_SUBMISSIONS") != nullptr; @@ -4920,6 +4922,9 @@ vk_device ggml_vk_get_device(size_t idx) { device->mmvq_mode = 1; } + // publish the device only when it is fully built + vk_instance.devices[idx] = device; + return device; } @@ -16463,6 +16468,11 @@ void vk_queue_handle_unsynchronized::submit(vk::ArrayProxy } vk_device_struct::~vk_device_struct() { + // the build failed before the device was created, nothing to destroy + if (device == VK_NULL_HANDLE) { + return; + } + VK_LOG_DEBUG("destroy device " << name); device.destroyFence(fence); diff --git a/tests/test-thread-safety.cpp b/tests/test-thread-safety.cpp index 4fbb1a2700bc..0e735d684d4e 100644 --- a/tests/test-thread-safety.cpp +++ b/tests/test-thread-safety.cpp @@ -1,5 +1,5 @@ // thread safety test -// - Loads a copy of the same model on each GPU, plus a copy on the CPU +// - Loads a copy of the same model on each GPU, plus a copy on the CPU, all in parallel // - Creates n_parallel (--parallel) contexts per model // - Runs inference in parallel on each context @@ -51,30 +51,43 @@ int main(int argc, char ** argv) { //const int num_models = std::max(1, gpu_dev_count); const int num_contexts = std::max(1, params.n_parallel); - std::vector models; + std::vector models(num_models); + std::vector load_threads; std::vector threads; std::atomic failed = false; for (int m = 0; m < num_models; ++m) { - auto mparams = common_model_params_to_llama(params); - - if (m < gpu_dev_count) { - mparams.split_mode = LLAMA_SPLIT_MODE_NONE; - mparams.devices = gpus[m].data(); - } else if (m == gpu_dev_count) { - mparams.split_mode = LLAMA_SPLIT_MODE_NONE; - mparams.main_gpu = -1; // CPU model - } else { - mparams.split_mode = LLAMA_SPLIT_MODE_LAYER; - } + load_threads.emplace_back([&, m]() { + auto mparams = common_model_params_to_llama(params); + + if (m < gpu_dev_count) { + mparams.split_mode = LLAMA_SPLIT_MODE_NONE; + mparams.devices = gpus[m].data(); + } else if (m == gpu_dev_count) { + mparams.split_mode = LLAMA_SPLIT_MODE_NONE; + mparams.main_gpu = -1; // CPU model + } else { + mparams.split_mode = LLAMA_SPLIT_MODE_LAYER; + } + + llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams); + if (model == NULL) { + LOG_ERR("%s: failed to load model '%s'\n", __func__, params.model.path.c_str()); + failed.store(true); + return; + } + + models[m].reset(model); + }); + } - llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams); - if (model == NULL) { - LOG_ERR("%s: failed to load model '%s'\n", __func__, params.model.path.c_str()); - return 1; - } + for (auto & thread : load_threads) { + thread.join(); + } - models.emplace_back(model); + if (failed) { + LOG_ERR("One or more models failed to load.\n"); + return 1; } for (int m = 0; m < num_models; ++m) { From 82e1e2fd84b93c0dce5b312b333117f85dddc5b4 Mon Sep 17 00:00:00 2001 From: Yuri Khrustalev Date: Thu, 15 Oct 2026 04:49:40 -0400 Subject: [PATCH 017/423] jinja : write nan and inf in tojson as json.dumps does (#30483) tojson wrote nan, inf and -inf for these floats, as values and as dict keys; json.dumps writes NaN, Infinity and -Infinity. The string and safe filters of a float keep nan and inf, as Python's str() does, so they no longer use the tojson function. float_to_json is renamed float_repr, as it now gives str(); float_json gives the json.dumps text. Also say in a comment why a last digit 9 is not bumped at a power of 2. --- common/jinja/value.cpp | 34 +++++++++++++++++++++++++++------- tests/test-jinja.cpp | 7 +++++++ 2 files changed, 34 insertions(+), 7 deletions(-) diff --git a/common/jinja/value.cpp b/common/jinja/value.cpp index aed834ce460b..0ecbbf8a8f7f 100644 --- a/common/jinja/value.cpp +++ b/common/jinja/value.cpp @@ -244,6 +244,8 @@ static std::string json_ensure_ascii_preserving_format(const std::string & json_ return output; } +static std::string float_repr(double val); + static value tojson(const func_args & args) { args.ensure_count(1, 5); value val_ascii = args.get_kwarg_or_pos("ensure_ascii", 1); @@ -634,6 +636,11 @@ const func_builtins & value_int_t::get_builtins() const { const func_builtins & value_float_t::get_builtins() const { + // Python's str() of a float, nan and inf stay as they are (tojson writes NaN and Infinity) + static const func_handler tostring = [](const func_args & args) -> value { + args.ensure_vals(); + return mk_val(float_repr(args.get_pos(0)->as_float())); + }; static const func_builtins builtins = { {"default", default_value}, {"abs", [](const func_args & args) -> value { @@ -650,8 +657,8 @@ const func_builtins & value_float_t::get_builtins() const { args.ensure_vals(); return mk_val(args.get_pos(0)->as_float()); }}, - {"safe", tojson}, - {"string", tojson}, + {"safe", tostring}, + {"string", tostring}, {"tojson", tojson}, }; return builtins; @@ -1522,10 +1529,10 @@ void global_from_json(context & ctx, const common_json & json_obj, bool mark_inp } } -// a float as Python writes it (repr, json.dumps): the shortest digits that read back as the same value -static std::string float_to_json(double val) { +// a float as Python's repr() and str() write it: the shortest digits that read back as the same value +static std::string float_repr(double val) { if (!std::isfinite(val)) { - return std::isnan(val) ? "nan" : val < 0 ? "-inf" : "inf"; // as before (json.dumps writes NaN and Infinity, str() writes nan and inf) + return std::isnan(val) ? "nan" : val < 0 ? "-inf" : "inf"; } const double mag = std::fabs(val); @@ -1551,6 +1558,7 @@ static std::string float_to_json(double val) { if (prec == 16 || reads_back(d)) { digits = d; } else if (is_pow2 && d.back() != '9') { + // a last digit 9 is not bumped (it would carry): no power of 2, from 2^-1074 to 2^1023, needs it d.back()++; if (reads_back(d)) { digits = d; @@ -1581,6 +1589,17 @@ static std::string float_to_json(double val) { return out; } +// a float as json.dumps writes it: nan and inf as NaN and Infinity, which are not valid JSON +static std::string float_json(double val) { + if (std::isnan(val)) { + return "NaN"; + } + if (std::isinf(val)) { + return val < 0 ? "-Infinity" : "Infinity"; + } + return float_repr(val); +} + // recursively convert value to JSON string // TODO: avoid circular references static void value_to_json_internal(std::ostringstream & oss, const value & val, int curr_lvl, int indent, const std::string_view item_sep, const std::string_view key_sep) { @@ -1598,7 +1617,7 @@ static void value_to_json_internal(std::ostringstream & oss, const value & val, } else if (is_val(val)) { oss << val->as_int(); } else if (is_val(val)) { - oss << float_to_json(val->as_float()); + oss << float_json(val->as_float()); } else if (is_val(val)) { oss << "\""; for (char c : val->as_string().str()) { @@ -1645,7 +1664,8 @@ static void value_to_json_internal(std::ostringstream & oss, const value & val, size_t i = 0; for (const auto & pair : obj) { oss << indent_str() << (indent > 0 ? std::string(indent, ' ') : ""); - value_to_json_internal(oss, mk_val(pair.first->as_string().str()), curr_lvl + 1, indent, item_sep, key_sep); + const std::string key = is_val(pair.first) ? float_json(pair.first->as_float()) : pair.first->as_string().str(); + value_to_json_internal(oss, mk_val(key), curr_lvl + 1, indent, item_sep, key_sep); oss << key_sep; value_to_json_internal(oss, pair.second, curr_lvl + 1, indent, item_sep, key_sep); if (i < obj.size() - 1) { diff --git a/tests/test-jinja.cpp b/tests/test-jinja.cpp index 9b53e8cec3d4..7132c9e22483 100644 --- a/tests/test-jinja.cpp +++ b/tests/test-jinja.cpp @@ -1083,6 +1083,13 @@ static void test_literals(testing & t) { "1.0 1e+308 inf" ); + // json.dumps writes NaN and Infinity, str() writes nan and inf + test_template(t, "nan and inf", + "{{ [p|float, n|float, q|float]|tojson }} {{ {p|float: 1}|tojson }} {{ p|float|string }} {{ n|float|string }} {{ q|float|safe }}", + {{"p", "inf"}, {"n", "-inf"}, {"q", "nan"}}, + "[Infinity, -Infinity, NaN] {\"Infinity\": 1} inf -inf nan" + ); + test_template(t, "string|tojson", "{{ 'hello' | tojson }}", json::object(), From 08698c81ed1053b7be1ade00d1e127daf1ae9ba3 Mon Sep 17 00:00:00 2001 From: William Tan <1284324+Ninja3047@users.noreply.github.com> Date: Thu, 15 Oct 2026 05:00:57 -0400 Subject: [PATCH 018/423] jinja : improve context copy constructor performance (#30562) Assisted-by: Codex --- common/jinja/runtime.h | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/common/jinja/runtime.h b/common/jinja/runtime.h index bbd0c5caeacc..c01902a1a038 100644 --- a/common/jinja/runtime.h +++ b/common/jinja/runtime.h @@ -74,16 +74,12 @@ struct context { } ~context() = default; - context(const context & parent) : context() { + context(const context & parent) : + src(parent.src), + current_time(parent.current_time), + is_get_stats(parent.is_get_stats), // inherit variables (for example, when entering a new scope) - auto & pvar = parent.env->as_ordered_object(); - for (const auto & pair : pvar) { - set_val(pair.first, pair.second); - } - current_time = parent.current_time; - is_get_stats = parent.is_get_stats; - src = parent.src; - } + env(mk_val(*parent.env)) {} value get_val(const std::string & name) { value default_val = mk_val(name); From 3e4dfa8f39f904e681737d9f512a6e6338d848aa Mon Sep 17 00:00:00 2001 From: Karthik Vijayan Date: Thu, 15 Oct 2026 09:18:46 +0000 Subject: [PATCH 019/423] models : add maion-coder (#29778) * Add Support to maion-coder arch * iAdd Support to support Draft models * Address comments * Remove maion context shift function --------- Co-authored-by: Karthik Vijayan --- gguf-py/gguf/constants.py | 27 +++++++ src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/llama-kv-cache.cpp | 2 +- src/llama-model-saver.cpp | 3 +- src/llama-model.cpp | 3 + src/models/maion-coder.cpp | 149 +++++++++++++++++++++++++++++++++++++ src/models/models.h | 13 ++++ tests/test-llama-archs.cpp | 16 +++- 9 files changed, 211 insertions(+), 4 deletions(-) create mode 100644 src/models/maion-coder.cpp diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index f9baf687ee16..0c9872228a8a 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -684,6 +684,7 @@ class MODEL_ARCH(IntEnum): SPARK2_5 = auto() LLAMA_EMBED = auto() MAINCODER = auto() + MAION_CODER = auto() KIMI_LINEAR = auto() KIMI_K3 = auto() TALKIE = auto() @@ -1488,6 +1489,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.SPARK2_5: "spark2_5", MODEL_ARCH.LLAMA_EMBED: "llama-embed", MODEL_ARCH.MAINCODER: "maincoder", + MODEL_ARCH.MAION_CODER: "maion-coder", MODEL_ARCH.KIMI_LINEAR: "kimi-linear", MODEL_ARCH.KIMI_K3: "kimi-k3", MODEL_ARCH.TALKIE: "talkie", @@ -5742,6 +5744,31 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, ], + MODEL_ARCH.MAION_CODER: [ + MODEL_TENSOR.TOKEN_EMBD, + 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, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.MOE_LATENT_DOWN, + MODEL_TENSOR.MOE_LATENT_UP, + ], MODEL_ARCH.KIMI_LINEAR: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5744c5fb844d..d1dd4cbaeb67 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -154,6 +154,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_SPARK2_5, "spark2_5" }, { LLM_ARCH_LLAMA_EMBED, "llama-embed" }, { LLM_ARCH_MAINCODER, "maincoder" }, + { LLM_ARCH_MAION_CODER, "maion-coder" }, { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, { LLM_ARCH_KIMI_K3, "kimi-k3" }, { LLM_ARCH_GLM5_NEXT, "glm5-next" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 94771c2a1c32..4e8e53421af3 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -154,6 +154,7 @@ enum llm_arch { LLM_ARCH_SPARK2_5, LLM_ARCH_LLAMA_EMBED, LLM_ARCH_MAINCODER, + LLM_ARCH_MAION_CODER, LLM_ARCH_KIMI_LINEAR, LLM_ARCH_KIMI_K3, LLM_ARCH_GLM5_NEXT, diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 60a7e855497e..b7dfadef4638 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2009,7 +2009,7 @@ class llm_graph_input_k_shift : public llm_graph_input_i { void llm_graph_input_k_shift::set_input(const llama_ubatch * ubatch) { GGML_UNUSED(ubatch); - if (k_shift) { + if (k_shift && k_shift->buffer) { kv_self->set_input_k_shift(k_shift); } diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 47e36d3743ee..e011e24ea605 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -225,8 +225,9 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_BLOCK_COUNT, hparams.n_layer_all); add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); add_kv(LLM_KV_FEED_FORWARD_LENGTH, hparams.n_ff_arr, true); - add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp()); + add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, true); add_kv(LLM_KV_EXPERT_LATENT_LENGTH, hparams.n_expert_latent); + add_kv(LLM_KV_MOE_LATENT_SIZE, hparams.moe_latent_size); add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp); add_kv(LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH, hparams.n_ff_chexp); add_kv(LLM_KV_SWIGLU_CLAMP_EXP, std::vector( diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 923c2c0a2b6e..ded790b2b86d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -52,6 +52,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_llama_embed(params); case LLM_ARCH_MAINCODER: return new llama_model_maincoder(params); + case LLM_ARCH_MAION_CODER: + return new llama_model_maion_coder(params); case LLM_ARCH_TALKIE: return new llama_model_talkie(params); case LLM_ARCH_DECI: @@ -3556,6 +3558,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_GEMMA3: case LLM_ARCH_GEMMA3N: case LLM_ARCH_GEMMA4: + case LLM_ARCH_MAION_CODER: case LLM_ARCH_GEMMA4_ASSISTANT: case LLM_ARCH_GEMMA_EMBEDDING: case LLM_ARCH_GEMMA_EMBEDDING2: diff --git a/src/models/maion-coder.cpp b/src/models/maion-coder.cpp new file mode 100644 index 000000000000..bb3ebabc7ed3 --- /dev/null +++ b/src/models/maion-coder.cpp @@ -0,0 +1,149 @@ +#include "models.h" + +void llama_model_maion_coder::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all); + ml.get_key(LLM_KV_INTERLEAVE_MOE_LAYER_STEP, hparams.n_moe_layer_step); + ml.get_key(LLM_KV_MOE_LATENT_SIZE, hparams.moe_latent_size); + + if (hparams.n_moe_layer_step == 0 || hparams.moe_latent_size == 0 || hparams.n_expert == 0) { + throw std::runtime_error("maion-coder requires a positive MoE layer step, latent size and expert count"); + } + + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + uint32_t n_swa_layers = 0; + ml.get_arr_n(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, n_swa_layers); + if (hparams.n_swa == 0 || n_swa_layers != hparams.n_layer_all) { + throw std::runtime_error("maion-coder requires a positive sliding window and a sliding window pattern for each layer"); + } + ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl); + hparams.rope_pattern = hparams.is_swa_impl; + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; + hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + + type = LLM_TYPE_UNKNOWN; +} + +void llama_model_maion_coder::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t moe_n_embd = hparams.moe_latent_size; + + 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}, 0); + + for (int i = 0; i < n_layer; ++i) { + const bool is_moe_layer = (i + 1) % hparams.n_moe_layer_step == 0; + auto & layer = layers[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_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.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0); + + if (is_moe_layer) { + const int64_t n_ff_exp = hparams.n_ff_exp(i); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, 0); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, 0); + layer.ffn_latent_up = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP, "weight", i), {moe_n_embd, n_embd}, 0); + } else { + 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_maion_coder::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_maion_coder::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_k(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_v()); + GGML_ASSERT(n_embd_head == n_rot); + + 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) { + // export per-layer input hidden state so speculative drafts (e.g. dflash/eagle3) can read it + res->t_layer_inp[il] = inpL; + + ggml_tensor * inpSA = inpL; + ggml_tensor * cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, n_embd_head, n_head, n_head_kv, il); + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il); + Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il); + + // Full-attention layers do not use RoPE. + if (hparams.is_swa(il)) { + const float freq_base_l = model.get_rope_freq_base(cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, ext_factor, attn_factor, beta_fast, beta_slow); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, ext_factor, attn_factor, beta_fast, beta_slow); + } + + cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, 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); + } + + cur = build_norm(cur, model.layers[il].attn_post_norm, NULL, LLM_NORM_RMS, il); + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); + + if (model.layers[il].ffn_gate_inp == nullptr) { + cur = build_ffn(cur, + model.layers[il].ffn_up, NULL, NULL, + model.layers[il].ffn_gate, NULL, NULL, + model.layers[il].ffn_down, NULL, NULL, + NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); + } else { + ggml_tensor * router_logits = build_lora_mm(model.layers[il].ffn_gate_inp, cur); + ggml_tensor * inp_latent = build_lora_mm(model.layers[il].ffn_latent_down, cur); + cur = build_moe_ffn(inp_latent, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + nullptr, + n_expert, hparams.n_expert_used(il), + LLM_FFN_SILU, true, 0.0f, + LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, + il, router_logits); + cur = build_lora_mm(model.layers[il].ffn_latent_up, cur); + } + + cur = build_norm(cur, model.layers[il].ffn_post_norm, NULL, LLM_NORM_RMS, il); + cur = ggml_add(ctx0, cur, ffn_inp); + cur = build_cvec(cur, il); + inpL = cur; + } + + ggml_tensor * cur = build_norm(inpL, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + cur = build_lora_mm(model.output, cur, model.output_s); + cb(cur, "result_output", -1); + res->t_logits = cur; +} diff --git a/src/models/models.h b/src/models/models.h index a55983cfbc0f..88fd07d973cd 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -196,6 +196,19 @@ struct llama_model_maincoder : public llama_model_base { }; +struct llama_model_maion_coder : public llama_model_base { + llama_model_maion_coder(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_talkie : public llama_model_base { llama_model_talkie(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 1e8e915b151a..3dd8f0655ee9 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -135,6 +135,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { uint32_t n_layer = 2; if (arch == LLM_ARCH_LLAMA4) { n_layer = 4; // hparams.n_no_rope_layer_step is hard-coded to 4 + } else if (arch == LLM_ARCH_MAION_CODER) { + n_head = 4; + n_layer = 4; // dense and MoE layers, each with local and full attention } else if (arch == LLM_ARCH_GEMMA4) { n_embd = 128; n_head = 2; @@ -182,7 +185,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { uint32_t n_head_kv = n_head; if (arch == LLM_ARCH_QWEN3) { n_head_kv = 1; // MQA coverage - } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) { + } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE || arch == LLM_ARCH_MAION_CODER) { n_head_kv = 2; // GQA coverage } const uint32_t n_embd_head = n_embd / n_head; @@ -289,7 +292,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, uint32_t(8)); ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW, n_ctx/8); - if (arch == LLM_ARCH_GEMMA4) { + if (arch == LLM_ARCH_MAION_CODER) { + ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, std::vector{1, 1, 0, 0}); + } else if (arch == LLM_ARCH_GEMMA4) { ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, n_embd/2); ms.add_kv(LLM_KV_ATTENTION_SHARED_KV_LAYERS, uint32_t(0)); ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_SWA, n_embd_head); @@ -443,6 +448,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } } + if (arch == LLM_ARCH_MAION_CODER) { + ms.add_kv(LLM_KV_MOE_LATENT_SIZE, n_embd / 2); + ms.add_kv(LLM_KV_EXPERT_COUNT, uint32_t(4)); + ms.add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, std::vector{n_ff, n_ff / 2, n_ff, n_ff}); + } + ms.add_kv(LLM_KV_POSNET_EMBEDDING_LENGTH, n_embd); ms.add_kv(LLM_KV_POSNET_BLOCK_COUNT, n_layer); ms.add_kv(LLM_KV_CONVNEXT_EMBEDDING_LENGTH, n_embd); @@ -654,6 +665,7 @@ static bool check_causal_attn_toggle( static bool moe_mandatory(const llm_arch arch) { switch (arch) { case LLM_ARCH_LLAMA4: + case LLM_ARCH_MAION_CODER: case LLM_ARCH_COHERE2MOE: case LLM_ARCH_GROK: case LLM_ARCH_QWEN2MOE: From fd55cb98fe29d3df0054d3f3e6e6f15958ea72ab Mon Sep 17 00:00:00 2001 From: Aleksander Grygier Date: Thu, 15 Oct 2026 12:36:06 +0200 Subject: [PATCH 020/423] fix(ui): take the agentic turns default from DEFAULT_AGENTIC_CONFIG (#30582) The registry hard-coded 10 for AGENTIC_MAX_TURNS while DEFAULT_AGENTIC_CONFIG said 100. The registry seeds SETTING_CONFIG_DEFAULT, so the loop's `|| DEFAULT_AGENTIC_CONFIG.maxTurns` fallback never fires and users got 10 turns. Reference the shared constant, as the adjacent MCP timeout does. Assisted by: llama-ui (ggml-org/Qwen3.8-Flash-Next:IQ4_NL) --- .../ui/src/lib/constants/settings.constants.ts | 3 ++- .../unit/agentic-max-turns-default.test.ts | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) create mode 100644 tools/ui/tests/unit/agentic-max-turns-default.test.ts diff --git a/tools/ui/src/lib/constants/settings.constants.ts b/tools/ui/src/lib/constants/settings.constants.ts index f6596bfb6211..9d01f3ad0166 100644 --- a/tools/ui/src/lib/constants/settings.constants.ts +++ b/tools/ui/src/lib/constants/settings.constants.ts @@ -1,3 +1,4 @@ +import { DEFAULT_AGENTIC_CONFIG } from './agentic.constants'; import { CLI_FLAGS } from './cli-flags.constants'; import { DEFAULT_MCP_CONFIG } from './mcp.constants'; import { SETTINGS_KEYS } from './settings-keys.constants'; @@ -351,7 +352,7 @@ export const SETTINGS_REGISTRY: SettingsSectionEntry[] = [ icon: ListRestart, settings: [ { - defaultValue: 10, + defaultValue: DEFAULT_AGENTIC_CONFIG.maxTurns, help: 'Maximum number of tool execution cycles before stopping (prevents infinite loops).', isPositiveInteger: true, key: SETTINGS_KEYS.AGENTIC_MAX_TURNS, diff --git a/tools/ui/tests/unit/agentic-max-turns-default.test.ts b/tools/ui/tests/unit/agentic-max-turns-default.test.ts new file mode 100644 index 000000000000..91475231b991 --- /dev/null +++ b/tools/ui/tests/unit/agentic-max-turns-default.test.ts @@ -0,0 +1,18 @@ +import { DEFAULT_AGENTIC_CONFIG, SETTINGS_KEYS } from '$lib/constants'; +import { describe, expect, it } from 'vitest'; + +/** + * The registry `defaultValue` is the effective default: the store seeds its + * config from SETTING_CONFIG_DEFAULT, so the loop's + * `Number(setting) || DEFAULT_AGENTIC_CONFIG.maxTurns` fallback never fires. + * Once these disagreed (10 vs 100) and the loop silently stopped early. + */ +describe('AGENTIC_MAX_TURNS default value', () => { + it('sources the settings default from DEFAULT_AGENTIC_CONFIG.maxTurns', async () => { + const { SETTING_CONFIG_DEFAULT } = await import('$lib/constants'); + + expect(SETTING_CONFIG_DEFAULT[SETTINGS_KEYS.AGENTIC_MAX_TURNS]).toBe( + DEFAULT_AGENTIC_CONFIG.maxTurns + ); + }, 15000); +}); From 711780ffc75bbc5645173ed8647a7b3f328c052a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Thu, 15 Oct 2026 14:15:55 +0200 Subject: [PATCH 021/423] ci : bump ty to 0.0.88 (#30587) * bump ty to 0.0.88 * arange -> linspace.tolist() to pass expected type --- .github/workflows/python-type-check.yml | 2 +- scripts/server-bench.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/python-type-check.yml b/.github/workflows/python-type-check.yml index 5be177e9c504..8fb4ffd47a7c 100644 --- a/.github/workflows/python-type-check.yml +++ b/.github/workflows/python-type-check.yml @@ -35,7 +35,7 @@ jobs: uses: actions/setup-python@v6 with: python-version: "3.11" - pip-install: -r requirements/requirements-all.txt ty==0.0.84 + pip-install: -r requirements/requirements-all.txt ty==0.0.88 # - name: Type-check with Pyright # uses: jakebailey/pyright-action@v2 # with: diff --git a/scripts/server-bench.py b/scripts/server-bench.py index ce9737148760..f0984026c862 100755 --- a/scripts/server-bench.py +++ b/scripts/server-bench.py @@ -264,7 +264,7 @@ def benchmark( bin_max = np.ceil(token_t_last) + 1 plt.figure() - plt.hist(token_t, np.arange(0, bin_max)) + plt.hist(token_t, np.linspace(0, bin_max, bin_max, endpoint=False).tolist()) plt.xlim(0, bin_max + 1) plt.title(name or "") plt.xlabel("Time [s]") From 4ed1c8507e78c5dddc8e48d2822b87c2099e00d1 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 15 Oct 2026 15:38:20 +0300 Subject: [PATCH 022/423] ggml : bump version to 0.27.0 --- ggml/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index 3eacd4cd55ea..35252a83a5ea 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -4,7 +4,7 @@ project("ggml" C CXX ASM) ### GGML Version set(GGML_VERSION_MAJOR 0) -set(GGML_VERSION_MINOR 26) +set(GGML_VERSION_MINOR 27) set(GGML_VERSION_PATCH 0) set(GGML_VERSION_BASE "${GGML_VERSION_MAJOR}.${GGML_VERSION_MINOR}.${GGML_VERSION_PATCH}") From c0cb53ac4e34a0c16cbb4bdc963f0b5f3e60eacf Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 15 Oct 2026 16:16:53 +0300 Subject: [PATCH 023/423] sync : ggml --- scripts/sync-ggml.last | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/sync-ggml.last b/scripts/sync-ggml.last index 68997f378ffb..08b4084a6a91 100644 --- a/scripts/sync-ggml.last +++ b/scripts/sync-ggml.last @@ -1 +1 @@ -ffa4e8b80930029a35991f94e7c8a93cd67730ab +f1aa6a3465ffad40dadd390ac16b9deedd4e97f7 From b0d239164b26adcd317be1f0cffd2a1294f6ba60 Mon Sep 17 00:00:00 2001 From: Pascal Date: Thu, 15 Oct 2026 15:24:01 +0200 Subject: [PATCH 024/423] server: prefer empty slots and keep states overwritten by child tasks (#30567) Slot selection takes free slots by id or by last use only, so child tasks often land on a slot holding a parent state while empty slots are idle, and overwrite it without saving it to the prompt cache. Free slots are now ordered empty first, then least recently used, for both the LRU fallback and the child slots, and a child task saves the state it overwrites. With 4 grouped decision requests in flight and 12 states in rotation, prompt cache reads go from 5 to 36. [no release] --- tools/server/server-context.cpp | 40 ++++++++++++----------- tools/server/tests/unit/test_systemone.py | 38 +++++++++++++++++++++ 2 files changed, 59 insertions(+), 19 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d5b4a7029d4f..c1683d9ef61a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1763,25 +1763,14 @@ struct server_context_impl { } } - // find the slot that has been least recently used + // find an empty slot, or the one that has been least recently used if (ret == nullptr) { - int64_t t_last = -1; + const auto free_slots = get_free_slots(1, -1); - for (server_slot & slot : slots) { - // skip the slot if it is not available - if (slot.is_processing()) { - continue; - } + if (!free_slots.empty()) { + ret = free_slots[0]; - // select the current slot if the criteria match - if (!ret || slot.t_last_used <= t_last) { - t_last = slot.t_last_used; - ret = &slot; - } - } - - if (ret != nullptr) { - SLT_INF(*ret, "selected slot by LRU, t_last = %" PRId64 "\n", t_last); + SLT_INF(*ret, "selected slot by LRU, n_tokens = %d, t_last = %" PRId64 "\n", ret->prompt.n_tokens(), ret->t_last_used); update_cache = true; } @@ -2480,16 +2469,24 @@ struct server_context_impl { return true; } + // idle slots that lose the least by being reused: empty slots first, then the least recently used ones std::vector get_free_slots(size_t n_slots_needed, int exclude_id_slot) { std::vector free_slots; for (auto & slot : slots) { if (!slot.is_processing() && slot.id != exclude_id_slot) { free_slots.push_back(&slot); } - if (free_slots.size() >= n_slots_needed) { - break; - } } + + std::stable_sort(free_slots.begin(), free_slots.end(), [](const server_slot * a, const server_slot * b) { + if ((a->prompt.n_tokens() == 0) != (b->prompt.n_tokens() == 0)) { + return a->prompt.n_tokens() == 0; + } + return a->t_last_used < b->t_last_used; + }); + + free_slots.resize(std::min(free_slots.size(), n_slots_needed)); + return free_slots; } @@ -2519,6 +2516,11 @@ struct server_context_impl { size_t idx = 0; GGML_ASSERT(child_slots.size() == parent_task.child_tasks.size()); for (auto * slot : child_slots) { + // keep the state that the child task is about to overwrite + if (prompt_cache && slot->prompt_save(*prompt_cache)) { + prompt_cache->update(); + } + int id_child = parent_task.child_tasks[idx].id; if (!launch_slot_with_task(*slot, std::move(parent_task.child_tasks[idx]))) { SRV_ERR("failed to launch slot with child task, id_task = %d\n", id_child); diff --git a/tools/server/tests/unit/test_systemone.py b/tools/server/tests/unit/test_systemone.py index 61c2ab09897f..67e128e11e88 100644 --- a/tools/server/tests/unit/test_systemone.py +++ b/tools/server/tests/unit/test_systemone.py @@ -184,6 +184,44 @@ def test_systemone_prompt_cache(cache_ram, tmp_path): assert ("found better prompt" in log) == (cache_ram is not None) +def test_systemone_prompt_cache_child_slots(tmp_path): + global server + server = ServerPreset.tinyopenjev() + server.n_slots = len(TEST_QUESTIONS) + server.cache_ram = 100 + server.debug = True + server.log_path = str(tmp_path / "server.log") + server.start() + + states = [ + "My parcel shows delivered but it never arrived at my door.", + "The app crashes every time I open the settings page.", + "Please close my account and delete all my personal data.", + ] + angry = {"angry": TEST_QUESTIONS["angry"]} + + def decide(state, questions): + res = server.make_request("POST", "/v1/systemone", data={"state": state, "questions": questions}) + assert res.status_code == 200 + + # one state per slot, then a request whose child tasks take every slot + for state in states: + decide(state, angry) + decide(TEST_STATE, TEST_QUESTIONS) + + with open(server.log_path) as f: + n_restored = f.read().count("found better prompt with f_keep = 1.000") + + # every overwritten state comes back whole from the prompt cache + for state in states: + decide(state, angry) + + server.stop() + + with open(server.log_path) as f: + assert f.read().count("found better prompt with f_keep = 1.000") == n_restored + len(states) + + def test_systemone_images(): global server server = ServerPreset.tinyopenjev() From ea736c1d5c1d2cd0489abe788bc3a1859af5dc0b Mon Sep 17 00:00:00 2001 From: Pascal Date: Thu, 15 Oct 2026 15:24:19 +0200 Subject: [PATCH 025/423] server: save the draft state in the slot file (#30574) The slot file held the target context and the context checkpoints, but never the draft context, so a restored slot drafted with an empty or stale draft KV until the cache refilled. The file now ends with the live draft state, the full draft context and the speculative state, in the same appendix style as the checkpoints. A draft state that does not fit the draft context is dropped and the draft starts empty, while older files and older servers keep working. [no release] --- tools/server/server-context.cpp | 92 ++++++++++++++++++++++- tools/server/tests/unit/test_slot_save.py | 49 ++++++++++++ 2 files changed, 139 insertions(+), 2 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index c1683d9ef61a..c20167b90b62 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2736,6 +2736,87 @@ struct server_context_impl { return n_read; } + // the live draft state follows the checkpoints: the full draft context and the speculative state + // the llama state payload holds the target context only + static constexpr uint32_t SLOT_DRAFT_MAGIC = 0x54464453; // "SDFT" + static constexpr uint32_t SLOT_DRAFT_VERSION = 1; + + // returns false if the draft state could not be written completely + bool save_slot_draft(const std::string & filepath, const server_slot & slot, size_t & n_written) const { + n_written = 0; + std::vector data_dft; + std::vector data_spec; + if (ctx_dft != nullptr) { + data_dft.resize(llama_state_seq_get_size_ext(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_NONE)); + llama_state_seq_get_data_ext(ctx_dft, data_dft.data(), data_dft.size(), slot.id, LLAMA_STATE_SEQ_FLAGS_NONE); + } + common_speculative_get_state(spec.get(), slot.id, data_spec); + if (data_dft.empty() && data_spec.empty()) { + return true; + } + std::ofstream ofs(std::filesystem::u8path(filepath), std::ios::binary | std::ios::app); + if (!ofs) { + SLT_WRN(slot, "failed to append the draft state to '%s'\n", filepath.c_str()); + return false; + } + const uint32_t magic = SLOT_DRAFT_MAGIC; + const uint32_t version = SLOT_DRAFT_VERSION; + ckpt_write(ofs, &magic, sizeof(magic), n_written); + ckpt_write(ofs, &version, sizeof(version), n_written); + ckpt_write_buf(ofs, data_dft, n_written); + ckpt_write_buf(ofs, data_spec, n_written); + ofs.flush(); + if (!ofs) { + SLT_WRN(slot, "failed to append the draft state to '%s' - the draft state is incomplete\n", filepath.c_str()); + return false; + } + SLT_INF(slot, "appended the draft state (%.3f MiB) to '%s'\n", (float) n_written / 1024 / 1024, filepath.c_str()); + return true; + } + + // returns the number of bytes consumed, 0 if there is no usable draft state + // the draft starts empty unless the file holds a draft state that fits the draft context + size_t load_slot_draft(const std::string & filepath, size_t offset, server_slot & slot) const { + if (ctx_dft != nullptr) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, -1, -1); + } + std::ifstream ifs(std::filesystem::u8path(filepath), std::ios::binary | std::ios::ate); + const size_t file_size = ifs ? (size_t) ifs.tellg() : 0; + if (!ifs || file_size < offset || !ifs.seekg(offset)) { + return 0; + } + const size_t n_avail = file_size - offset; + size_t n_read = 0; + uint32_t magic = 0; + uint32_t version = 0; + if (!ckpt_read(ifs, &magic, sizeof(magic), n_read) || magic != SLOT_DRAFT_MAGIC) { + return 0; + } + std::vector data_dft; + std::vector data_spec; + if (!ckpt_read(ifs, &version, sizeof(version), n_read) || version != SLOT_DRAFT_VERSION || + !ckpt_read_buf(ifs, data_dft, n_avail, n_read) || + !ckpt_read_buf(ifs, data_spec, n_avail, n_read)) { + SLT_WRN(slot, "invalid draft state in '%s' - ignored\n", filepath.c_str()); + return 0; + } + if (!data_dft.empty()) { + if (ctx_dft == nullptr || + llama_state_seq_set_data_ext(ctx_dft, data_dft.data(), data_dft.size(), slot.id, LLAMA_STATE_SEQ_FLAGS_NONE) != data_dft.size()) { + if (ctx_dft != nullptr) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, -1, -1); + } + SLT_WRN(slot, "draft state in '%s' does not match the draft context - dropped\n", filepath.c_str()); + return n_read; + } + } + if (!data_spec.empty()) { + common_speculative_set_state(spec.get(), slot.id, data_spec); + } + SLT_INF(slot, "restored the draft state (%.3f MiB) from '%s'\n", (float) n_read / 1024 / 1024, filepath.c_str()); + return n_read; + } + // returns false to decline the task, it is offered again after the decode is done bool process_single_task(server_task && task, bool is_yielding) { // while yielding, an encode / decode is running and only reading the server state is safe @@ -2951,6 +3032,12 @@ struct server_context_impl { break; } + size_t nwrite_dft = 0; + if (!save_slot_draft(filepath, *slot, nwrite_dft)) { + send_error(task, "Unable to save slot: incomplete draft state", ERROR_TYPE_SERVER); + break; + } + const int64_t t_end = ggml_time_us(); const double t_save_ms = (t_end - t_start) / 1000.0; @@ -2960,7 +3047,7 @@ struct server_context_impl { res->filename = filename; res->is_save = true; res->n_tokens = slot->prompt.tokens.size(); - res->n_bytes = nwrite + nwrite_ckpt; + res->n_bytes = nwrite + nwrite_ckpt + nwrite_dft; res->t_ms = t_save_ms; queue_results.send(std::move(res)); } break; @@ -3018,6 +3105,7 @@ struct server_context_impl { // nread is the end offset of the llama state payload within the file const size_t nread_ckpt = load_slot_checkpoints(filepath, nread, *slot); + const size_t nread_dft = load_slot_draft(filepath, nread + nread_ckpt, *slot); const int64_t t_end = ggml_time_us(); const double t_restore_ms = (t_end - t_start) / 1000.0; @@ -3028,7 +3116,7 @@ struct server_context_impl { res->filename = filename; res->is_save = false; res->n_tokens = slot->prompt.tokens.size(); - res->n_bytes = nread + nread_ckpt; + res->n_bytes = nread + nread_ckpt + nread_dft; res->t_ms = t_restore_ms; queue_results.send(std::move(res)); } break; diff --git a/tools/server/tests/unit/test_slot_save.py b/tools/server/tests/unit/test_slot_save.py index 33bcede934cf..793f1a0a7aa9 100644 --- a/tools/server/tests/unit/test_slot_save.py +++ b/tools/server/tests/unit/test_slot_save.py @@ -132,6 +132,55 @@ def test_slot_restore_legacy_token_list(): assert res.body["timings"]["prompt_n"] == 6 # only the different part is processed +# the slot file carries the live draft state: a restored slot drafts like a slot that was never saved +def test_slot_restore_draft_state(): + global server + server.model_draft_hf_repo = server.model_hf_repo + server.spec_type = "draft-simple" + server.start() + + prompt = "What is the capital of France?" + turn = {} + for id_slot in [0, 1]: + res = server.make_request("POST", "/completion", data={ + "prompt": prompt, + "n_predict": 16, + "id_slot": id_slot, + "cache_prompt": True, + }) + assert res.status_code == 200 + turn[id_slot] = res.body["content"] + assert turn[0] == turn[1] + + res = server.make_request("POST", "/slots/0?action=save", data={ + "filename": "draft_slot0.bin", + }) + assert res.status_code == 200 + + res = server.make_request("POST", "/slots/0?action=erase") + assert res.status_code == 200 + + res = server.make_request("POST", "/slots/0?action=restore", data={ + "filename": "draft_slot0.bin", + }) + assert res.status_code == 200 + + timings = {} + for id_slot in [0, 1]: + res = server.make_request("POST", "/completion", data={ + "prompt": prompt + turn[id_slot] + " And", + "n_predict": 32, + "id_slot": id_slot, + "cache_prompt": True, + }) + assert res.status_code == 200 + timings[id_slot] = res.body["timings"] + + assert timings[0]["prompt_n"] == timings[1]["prompt_n"] + assert timings[0]["draft_n_accepted"] > 0 + assert timings[0]["draft_n"] == timings[1]["draft_n"] + assert timings[0]["draft_n_accepted"] == timings[1]["draft_n_accepted"] + def test_slot_erase(): global server From a6b91d8c5e0bd3849a4862c7561ea18d9483b986 Mon Sep 17 00:00:00 2001 From: Mikhail Podvitskii Date: Thu, 15 Oct 2026 15:24:47 +0200 Subject: [PATCH 026/423] model : Mellum - read MoE gating params from GGUF, unify RoPE (#30578) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * model : mellum - read MoE gating params from GGUF, unify RoPE - read/write expert_weights_norm and expert_gating_func instead of hardcoded values (defaults unchanged) - use per-layer rope params via get_rope_freq_base/scale - check n_expert once, outside the layer loop * Update conversion/mellum.py Co-authored-by: Sigbjørn Skjæret --------- Co-authored-by: Sigbjørn Skjæret [no release] --- conversion/mellum.py | 3 ++ src/models/mellum.cpp | 71 ++++++++++++++++++++----------------------- 2 files changed, 36 insertions(+), 38 deletions(-) diff --git a/conversion/mellum.py b/conversion/mellum.py index 1e50f92aeafd..6bcb81df8cfe 100644 --- a/conversion/mellum.py +++ b/conversion/mellum.py @@ -29,6 +29,9 @@ def set_gguf_parameters(self): self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in self.hparams["layer_types"]]) logger.info(f"gguf: sliding window pattern length = {len(self.hparams['layer_types'])}") + if (norm_topk_prob := self.hparams.get("norm_topk_prob")) is not None: + self.gguf_writer.add_expert_weights_norm(norm_topk_prob) + _experts: list[dict[str, Tensor]] | None = None def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: diff --git a/src/models/mellum.cpp b/src/models/mellum.cpp index 145597fc33ea..bc36dfa09cdb 100644 --- a/src/models/mellum.cpp +++ b/src/models/mellum.cpp @@ -5,13 +5,18 @@ void llama_model_mellum::load_arch_hparams(llama_model_loader & ml) { ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all); ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false); + hparams.expert_weights_norm = true; + hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX; + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); + if (hparams.n_swa > 0) { hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; load_swa_pattern(ml, 4); hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train; - hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train; + hparams.rope_freq_scale_train_swa = 1.0f; // SWA layers use plain RoPE, no YaRN ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false); } else { @@ -33,6 +38,13 @@ void llama_model_mellum::load_arch_tensors(llama_model_loader &) { output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + if (n_expert == 0) { + throw std::runtime_error("n_expert must be > 0 for Mellum"); + } + if (n_expert_used == 0) { + throw std::runtime_error("n_expert_used must be > 0 for Mellum"); + } + for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; @@ -48,13 +60,6 @@ void llama_model_mellum::load_arch_tensors(llama_model_loader &) { layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); - if (n_expert == 0) { - throw std::runtime_error("n_expert must be > 0 for Mellum"); - } - if (n_expert_used == 0) { - throw std::runtime_error("n_expert_used must be > 0 for Mellum"); - } - const int64_t n_ff_exp = hparams.n_ff_exp() ? hparams.n_ff_exp() : n_ff / n_expert_used; layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, 0); @@ -119,34 +124,24 @@ llama_model_mellum::graph::graph(const llama_model & model, const llm_grap const bool is_swa = hparams.is_swa(il); - if (is_swa) { - // For sliding window layers, use regular rope with no yarn rope scaling. - // This is achieved here by setting freq_scale and attn_factor to 1. - // We also set ext_factor to 0 to avoid a few unnecessary computations. - Qcur = ggml_rope_ext( - ctx0, Qcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, 1.0, - 0.0, 1.0, beta_fast, beta_slow - ); - - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, 1.0, - 0.0, 1.0, beta_fast, beta_slow - ); - } else { - Qcur = ggml_rope_ext( - ctx0, Qcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); - - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); - } + const float freq_base_l = model.get_rope_freq_base(cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + const float ext_factor_l = is_swa ? 0.0f : ext_factor; + const float attn_factor_l = is_swa ? 1.0f : attn_factor; + const float beta_fast_l = is_swa ? 32.0f : beta_fast; + const float beta_slow_l = is_swa ? 1.0f : beta_slow; + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l + ); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l + ); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); @@ -177,9 +172,9 @@ llama_model_mellum::graph::graph(const llama_model & model, const llm_grap model.layers[il].ffn_down_exps, nullptr, n_expert, n_expert_used, - LLM_FFN_SILU, true, + LLM_FFN_SILU, hparams.expert_weights_norm, hparams.expert_weights_scale, - LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, + (llama_expert_gating_func_type) hparams.expert_gating_func, il, nullptr, nullptr, model.layers[il].ffn_up_exps_s, From 61ad80bf93de0ac03a29f83f5f3e9f596bf40c2d Mon Sep 17 00:00:00 2001 From: Pascal Date: Thu, 15 Oct 2026 15:25:19 +0200 Subject: [PATCH 027/423] server: fix the off-by-two in the SWA reuse threshold (#30580) The first evaluated token attends to the positions above pos - n_swa, so a SWA state starting at pos - n_swa + 1 holds the whole window, but the reuse threshold asked for two more cells. A saved SWA state keeps only the cells inside its window, so a state read back from a slot file or from the RAM prompt cache was always rejected and the slot rolled back to a checkpoint. The SWA threshold now matches the window, while a state without SWA keeps the exact bound. The test harness gains override_kv, used to give the SWA test model a window smaller than the test context. [no release] --- tools/server/server-context.cpp | 8 ++- tools/server/tests/unit/test_slot_save.py | 86 +++++++++++++++++++++++ tools/server/tests/utils.py | 3 + 3 files changed, 95 insertions(+), 2 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index c20167b90b62..a18b7542c3af 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3784,8 +3784,12 @@ struct server_context_impl { // ref: https://github.com/ggml-org/llama.cpp/pull/24110 const bool has_new_tokens = (n_past < slot.task->n_tokens()); - // the largest pos_min required for a checkpoint to be useful - const auto pos_min_thold = std::max(0, pos_next - n_swa - (has_new_tokens ? 0 : 1)); + // the first evaluated token attends to the positions > pos_eval - n_swa, so a SWA state + // starting at pos_eval - n_swa + 1 is enough, while a state without SWA must end before pos_eval + const llama_pos pos_eval = pos_next - (has_new_tokens ? 0 : 1); + + // a state or a checkpoint is useful when its pos_min is below this threshold + const auto pos_min_thold = std::max(0, n_swa > 0 ? pos_eval - n_swa + 2 : pos_eval); if (n_past > 0 && n_past <= slot.prompt.n_tokens()) { const auto pos_min = llama_memory_seq_pos_min(llama_get_memory(ctx_tgt), slot.id); diff --git a/tools/server/tests/unit/test_slot_save.py b/tools/server/tests/unit/test_slot_save.py index 793f1a0a7aa9..488b63701d1b 100644 --- a/tools/server/tests/unit/test_slot_save.py +++ b/tools/server/tests/unit/test_slot_save.py @@ -838,3 +838,89 @@ def test_slot_restore_checkpoints_draft_kv_type_change(swa_server, ctkd_restore) }) assert res.status_code == 200 assert res.body["timings"]["prompt_n"] == n_live + + +# a SWA window smaller than the prompt, so a saved state holds only the cells inside the window +@pytest.fixture +def swa_window_server(tmp_path): + swa = ServerPreset.tinygemma3() + swa.override_kv = "gemma3.attention.sliding_window=int:64" + swa.slot_save_path = str(tmp_path) + swa.temperature = 0.0 + swa.n_predict = 8 + return swa + + +SWA_PROMPT = "The quick brown fox jumps over the lazy dog. " * 30 + + +def tokenize(server, text): + res = server.make_request("POST", "/tokenize", data={"content": text}) + assert res.status_code == 200 + return res.body["tokens"] + + +# the first turn and the exact token continuation of it, so the next turn extends the saved state +def swa_turns(server, id_slot): + res = server.make_request("POST", "/completion", data={ + "prompt": tokenize(server, SWA_PROMPT), + "id_slot": id_slot, + "cache_prompt": True, + "return_tokens": True, + }) + assert res.status_code == 200 + return tokenize(server, SWA_PROMPT) + res.body["tokens"] + tokenize(server, " The end.") + + +def swa_next(server, id_slot, prompt): + res = server.make_request("POST", "/completion", data={ + "prompt": prompt, + "id_slot": id_slot, + "cache_prompt": True, + }) + assert res.status_code == 200 + return res.body["timings"] + + +# a restored SWA state holds the window the next token needs, so it is reused as is +def test_slot_restore_swa_window(swa_window_server): + server = swa_window_server + server.cache_ram = 0 + server.start() + + prompt = {id_slot: swa_turns(server, id_slot) for id_slot in [0, 1]} + assert prompt[0] == prompt[1] + + res = server.make_request("POST", "/slots/0?action=save", data={"filename": "swa_slot0.bin"}) + assert res.status_code == 200 + res = server.make_request("POST", "/slots/0?action=erase") + assert res.status_code == 200 + res = server.make_request("POST", "/slots/0?action=restore", data={"filename": "swa_slot0.bin"}) + assert res.status_code == 200 + + timings = {id_slot: swa_next(server, id_slot, prompt[id_slot]) for id_slot in [0, 1]} + + assert timings[0]["cache_n"] == timings[1]["cache_n"] + assert timings[0]["prompt_n"] == timings[1]["prompt_n"] + + +# a SWA state loaded back from the RAM prompt cache is reused as is +def test_prompt_cache_swa_window(swa_window_server): + server = swa_window_server + server.n_slots = 1 + + def run(other_prompt_first): + server.start() + prompt = swa_turns(server, 0) + if other_prompt_first: + res = server.make_request("POST", "/completion", data={"prompt": "Once upon a time " * 60, "cache_prompt": True}) + assert res.status_code == 200 + timings = swa_next(server, 0, prompt) + server.stop() + return timings + + direct = run(False) + cached = run(True) + + assert cached["cache_n"] == direct["cache_n"] + assert cached["prompt_n"] == direct["prompt_n"] diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index b3207d3874aa..1c84f5121145 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -116,6 +116,7 @@ class ServerProcess: media_path: str | None = None sleep_idle_seconds: int | None = None cache_ram: int | None = None + override_kv: str | None = None no_cache_idle_slots: bool = False log_path: str | None = None ui_mcp_proxy: bool = False @@ -286,6 +287,8 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None: server_args.extend(["--sleep-idle-seconds", self.sleep_idle_seconds]) if self.cache_ram is not None: server_args.extend(["--cache-ram", self.cache_ram]) + if self.override_kv: + server_args.extend(["--override-kv", self.override_kv]) if self.no_cache_idle_slots: server_args.append("--no-cache-idle-slots") if self.ui_mcp_proxy: From a11b6530b541dd1698c7afe815d6f96de6e2b6cc Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 15 Oct 2026 16:28:09 +0300 Subject: [PATCH 028/423] scripts : add ci-status.sh to list active CI runs (#30586) Small helper around the GitHub Actions API: lists the workflow runs that are not finished yet (in_progress, queued, pending, waiting, requested), grouped by status. Supports selecting runs by status, by branch (server-side, to reduce API traffic) or by a regex on the branch name, and can cancel the selected runs via --kill. Assisted-by: pi:llama.cpp/Qwen3.8-Flash-Next --- scripts/ci-status.sh | 101 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100755 scripts/ci-status.sh diff --git a/scripts/ci-status.sh b/scripts/ci-status.sh new file mode 100755 index 000000000000..65d8eb5bb22d --- /dev/null +++ b/scripts/ci-status.sh @@ -0,0 +1,101 @@ +#!/usr/bin/env bash +# List GitHub Actions workflow runs that are not finished yet, grouped by status. + +set -euo pipefail + +usage() { + echo "usage: $(basename "$0") [-r owner/repo] [-s status]... [-b branch] [-w regex] [-t regex] [-f regex] [--kill [-y]]" + echo " -r, --repo repository (default: current)" + echo " -s, --status run status to include, repeatable (default: the not-finished ones)" + echo " -b, --branch branch to fetch runs for; exact match, handled by the API" + echo " -w, --workflow case-insensitive regex matched against the workflow name (client-side)" + echo " -t, --title case-insensitive regex matched against the run title (client-side)" + echo " -f, --filter case-insensitive regex matched against the branch name (client-side)" + echo " --kill cancel the selected runs (requires -b, -w, -t or -f)" + echo " -y, --yes skip the confirmation prompt of --kill" +} + +repo="" +branch="" +workflow="" +title="" +filter="" +kill=false +assume_yes=false +statuses=() +need_value() { echo "missing value for $1" >&2; exit 1; } +while [[ $# -gt 0 ]]; do + case "$1" in + -r|--repo) [[ $# -ge 2 ]] || need_value "$1"; repo="$2"; shift 2 ;; + -s|--status) [[ $# -ge 2 ]] || need_value "$1"; statuses+=("$2"); shift 2 ;; + -b|--branch) [[ $# -ge 2 ]] || need_value "$1"; branch="$2"; shift 2 ;; + -w|--workflow) [[ $# -ge 2 ]] || need_value "$1"; workflow="$2"; shift 2 ;; + -t|--title) [[ $# -ge 2 ]] || need_value "$1"; title="$2"; shift 2 ;; + -f|--filter) [[ $# -ge 2 ]] || need_value "$1"; filter="$2"; shift 2 ;; + --kill) kill=true; shift ;; + -y|--yes) assume_yes=true; shift ;; + -h|--help) usage; exit 0 ;; + *) echo "unknown argument: $1" >&2; exit 1 ;; + esac +done + +# one request per status: the API does not accept multiple statuses at once +[[ ${#statuses[@]} -gt 0 ]] || statuses=(in_progress queued pending waiting requested) + +[[ -n "$repo" ]] || repo=$(gh repo view --json nameWithOwner -q .nameWithOwner) + +query=("repos/$repo/actions/runs" -f per_page=100) +[[ -z "$branch" ]] || query+=(-f branch="$branch") + +selected=$( + for status in "${statuses[@]}"; do + gh api --paginate --method GET "${query[@]}" -f status="$status" -q '.workflow_runs[]' + done | jq -sr --arg branch "$branch" --arg workflow "$workflow" --arg title "$title" --arg filter "$filter" ' + unique_by(.id) + | map( + .age = ((now - (.run_started_at // .created_at | fromdate)) / 60 | floor) + | select( + ($branch == "" or .head_branch == $branch) and + ($workflow == "" or (.name | test($workflow; "i"))) and + ($title == "" or (.display_title | test($title; "i"))) and + ($filter == "" or (.head_branch | test($filter; "i"))) + ) + ) + | sort_by(.status, -.age)' +) + +tsv=$(jq -r '.[] | [ .status, (.age | tostring) + "m", .name, .head_branch, .display_title, .html_url ] | @tsv' <<< "$selected") + +if [[ -z "$tsv" ]]; then + echo "no active runs in $repo" + exit 0 +fi + +printf 'STATUS\tAGE\tWORKFLOW\tBRANCH\tTITLE\tURL\n%s\n' "$tsv" | column -t -s $'\t' +echo +printf '%s:' "$repo" +for status in "${statuses[@]}"; do + printf ' %s=%s' "$status" "$(awk -F'\t' -v s="$status" '$1 == s' <<< "$tsv" | wc -l | tr -d ' ')" +done +echo + +if [[ "$kill" == false ]]; then + exit 0 +fi + +if [[ -z "$branch$workflow$title$filter" ]]; then + echo "refusing to cancel everything: pass -b, -w, -t or -f to select the runs to cancel" >&2 + exit 1 +fi + +count=$(wc -l <<< "$tsv" | tr -d ' ') +if [[ "$assume_yes" == false ]]; then + printf 'cancel %s run(s)? [y/N] ' "$count" + read -r reply &2; exit 1; } + [[ "$reply" == [yY]* ]] || { echo "aborted"; exit 1; } +fi + +jq -r '.[] | [ .id, .status ] | @tsv' <<< "$selected" | +while IFS=$'\t' read -r id status; do + gh run cancel "$id" --repo "$repo" || echo "failed to cancel run $id (status: $status)" +done From cd0a5b6a180dff0048e5fd993098b21c0b2c91b8 Mon Sep 17 00:00:00 2001 From: Kartik Gulia Date: Thu, 15 Oct 2026 19:04:56 +0530 Subject: [PATCH 029/423] models : add Nemotron Super 3.5 VL and Nemotron Bolt 74B A9B support (#30520) * add support * remove unneeded bos fix * add video flag * fix flake error * remove video-specific budgeting * convert : fix lilicorr-dflash bid clash * spec : enforce lilicorr n_tokens condition * spec : fix lilicorr walk [no release] --- common/speculative.cpp | 2 +- conversion/__init__.py | 1 + conversion/base.py | 2 +- conversion/nemotron.py | 144 +++++++++++++++++++++++++--- gguf-py/gguf/constants.py | 1 + gguf-py/gguf/tensor_mapping.py | 14 +++ src/models/dflash.cpp | 4 +- src/models/nemotron-h.cpp | 24 ++++- tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-impl.h | 4 + tools/mtmd/clip-model.h | 2 +- tools/mtmd/clip.cpp | 19 ++++ tools/mtmd/models/models.h | 6 ++ tools/mtmd/models/nemotron-omni.cpp | 56 +++++++++++ tools/mtmd/mtmd.cpp | 7 ++ 15 files changed, 265 insertions(+), 22 deletions(-) create mode 100644 tools/mtmd/models/nemotron-omni.cpp diff --git a/common/speculative.cpp b/common/speculative.cpp index 701093308292..d8afd4095498 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1338,7 +1338,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { GGML_ASSERT(lattice && "DFlash2 selector produced no lattice"); int32_t predecessor = 0; - for (int32_t i = 1; i < n_block_tokens; ++i) { + for (int32_t i = 1; i < (is_dflash2 ? n_block_tokens : params.n_max + 1); ++i) { const float * row = lattice + (size_t) (beg + i) * n_embd_dec; const int32_t top_k = is_dflash2 ? selector_top_k : lilicorr_top_k; const float * scores = row + top_k + (size_t) predecessor * top_k; diff --git a/conversion/__init__.py b/conversion/__init__.py index d5b480a9bf41..8ce4cb7337d6 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -353,6 +353,7 @@ "MiniCPMV4_6ForConditionalGeneration": "minicpm", "MiniCPMV4_7ForConditionalGeneration": "minicpm", "Mistral3ForConditionalGeneration": "llava", + "NemotronH_Omni_Reasoning_V3": "nemotron", "NemotronH_Nano_VL_V2": "nemotron", "MuseGlimmerForConditionalGeneration": "muse_glimmer", "PaddleOCRVisionModel": "ernie", diff --git a/conversion/base.py b/conversion/base.py index efb7ea72e2cd..a5a80f27dad4 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -568,7 +568,7 @@ def dequant_fp8() -> None: tensors_to_remove.append(base_name + "_zero_point") else: raise NotImplementedError(f"Quant format {quant_format!r} for method {quant_method!r} is not yet supported") - elif quant_method == "modelopt": + elif quant_method in ("modelopt", "modelopt_mixed"): dequant_fp8() elif quant_method is not None: raise NotImplementedError(f"Quant method is not yet supported: {quant_method!r}") diff --git a/conversion/nemotron.py b/conversion/nemotron.py index 581f41a0230b..c900481943c2 100644 --- a/conversion/nemotron.py +++ b/conversion/nemotron.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Callable, Iterable, TYPE_CHECKING +from typing import Any, Callable, Iterable, Sequence, TYPE_CHECKING import numpy as np import torch @@ -89,19 +89,22 @@ def set_gguf_parameters(self): self.gguf_writer.add_lilicorr_vector_eps(float(dflash_config["lilicorr_vector_eps"])) self.gguf_writer.add_lilicorr_logit_scale(float(dflash_config["lilicorr_logit_scale"])) + def map_tensor_name(self, name: str, try_suffixes: Sequence[str] = (".weight", ".bias")) -> str: + if name.startswith("model.lilicorr.layers."): + parts = name.split(".", 4) + if not parts[4].startswith("attn.in_proj."): + parts[3] = str(self.block_count + int(parts[3])) + name = ".".join(parts) + elif name.startswith("model.lilicorr.feature_norm."): + feature_bid = self.block_count + self.lilicorr_geometry["num_layers"] + name = name.replace("model.lilicorr.feature_norm.", f"model.lilicorr.feature_mlp.{feature_bid}.norm.") + elif name.startswith("model.lilicorr.feature_mlp."): + feature_bid = self.block_count + self.lilicorr_geometry["num_layers"] + name = name.replace("model.lilicorr.feature_mlp.", f"model.lilicorr.feature_mlp.{feature_bid}.") + return super().map_tensor_name(name, try_suffixes) + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: if name.startswith("model.lilicorr."): - if name.startswith("model.lilicorr.layers."): - parts = name.split(".", 4) - if not parts[4].startswith("attn.in_proj."): - parts[3] = str(self.block_count + int(parts[3])) - name = ".".join(parts) - elif name.startswith("model.lilicorr.feature_norm."): - feature_bid = self.block_count + self.lilicorr_geometry["num_layers"] - name = name.replace("model.lilicorr.feature_norm.", f"model.lilicorr.feature_mlp.{feature_bid}.norm.") - elif name.startswith("model.lilicorr.feature_mlp."): - feature_bid = self.block_count + self.lilicorr_geometry["num_layers"] - name = name.replace("model.lilicorr.feature_mlp.", f"model.lilicorr.feature_mlp.{feature_bid}.") yield self.map_tensor_name(name), data_torch return yield from super().modify_tensors(data_torch, name, bid) @@ -251,6 +254,113 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield mapped_name, tensor +@ModelBase.register("NemotronH_Omni_Reasoning_V3") +class NemotronOmniVLModel(NemotronNanoV2VLModel): + _omni_tensor_prefixes = ( + "vision_model.embeddings.", + "vision_model.encoder.layer.", + "vision_projector.mlp1.", + ) + + def get_vision_config(self) -> dict[str, Any] | None: + vision_config = dict(self.global_config["vision_config"]) + vision_config["intermediate_size"] = int(vision_config["hidden_size"] * vision_config["mlp_ratio"]) + vision_config["image_size"] = 512 # warmup only; the full position table is resized at runtime + return vision_config + + def set_gguf_parameters(self): + preprocessor_config = self.preprocessor_config + + preprocessor_config["image_mean"] = preprocessor_config["norm_mean"] + preprocessor_config["image_std"] = preprocessor_config["norm_std"] + + MmprojModel.set_gguf_parameters(self) + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_OMNI) + self.gguf_writer.add_vision_attention_layernorm_eps(1e-6) + self.gguf_writer.add_vision_use_gelu(True) + downsample_ratio = self.global_config.get("downsample_ratio", 0.5) + self.gguf_writer.add_vision_projector_scale_factor(int(1.0 / downsample_ratio)) + patch_area = self.hparams["patch_size"] ** 2 + self.gguf_writer.add_vision_min_pixels(preprocessor_config["min_num_patches"] * patch_area) + self.gguf_writer.add_vision_max_pixels(preprocessor_config["max_num_patches"] * patch_area) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + + # Transformers can rename RADIO and projector tensors in Omni checkpoints. + # Accept both checkpoint layouts. + if name == "vision_model.summary_idxs": + # This buffer is only used by RadioModel.summary. The multimodal + # projector consumes RadioModel.features for both images and video. + return None + if name.startswith(cls._omni_tensor_prefixes): + return name, gen + + if "radio_model.model.patch_generator.video_embedder" in name or name.startswith("vision_projector.vision_final_layernorm."): + return MmprojModel.filter_tensors((name, gen)) + + return super().filter_tensors(item) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + if new_name in ("v.position_embd.weight", "v.class_embd"): + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.startswith(self._omni_tensor_prefixes): + if name == "vision_model.embeddings.video_patch_projection.weight": + data_torch = data_torch.reshape( + self.hparams["hidden_size"], + 6, + self.hparams["patch_size"], + self.hparams["patch_size"], + ) + yield gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_ENC_EMBD_PATCH] + ".weight.1", data_torch + return + + if name.startswith("vision_projector.mlp1."): + projector_mapping = { + "vision_projector.mlp1.norm.weight": "mlp1.0.weight", + "vision_projector.mlp1.linear1.weight": "mlp1.1.weight", + "vision_projector.mlp1.linear2.weight": "mlp1.3.weight", + } + name = projector_mapping[name] + elif name == "vision_model.embeddings.patch_projection.weight": + data_torch = data_torch.reshape( + self.hparams["hidden_size"], + 3, + self.hparams["patch_size"], + self.hparams["patch_size"], + ) + elif name == "vision_model.embeddings.position_embedding": + data_torch = data_torch.squeeze(0) + name += ".weight" + elif ".layer_scale" in name and name.endswith(".lambda1"): + name = name.removesuffix("lambda1") + "weight" + + yield from MmprojModel.modify_tensors(self, data_torch, name, bid) + return + + if "patch_generator.pos_embed" in name: + # Keep the full table for runtime interpolation. + yield from MmprojModel.modify_tensors(self, data_torch, name, bid) + return + + if "patch_generator.video_embedder" in name: + patch_size = self.hparams["patch_size"] + n_embd = self.hparams["hidden_size"] + temporal_patch_size = self.global_config["video_temporal_patch_size"] + if temporal_patch_size != 2: + raise ValueError("Current implementation only supports a temporal patch size of 2") + data_torch = data_torch.reshape(n_embd, temporal_patch_size * 3, patch_size, patch_size) + new_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_ENC_EMBD_PATCH] + ".weight.1" + yield new_name, data_torch + return + + yield from super().modify_tensors(data_torch, name, bid) + + @ModelBase.register("NemotronForCausalLM") @ModelBase.example("nvidia/Minitron-4B-Base") class NemotronModel(TextModel): @@ -382,6 +492,8 @@ def get_attn_layers(self): @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: name, gen = item + if name.startswith("language_model."): + name = name[len("language_model."):] if name.startswith("mtp."): # --no-mtp: drop the MTP head entirely if cls.no_mtp: @@ -531,6 +643,14 @@ def set_vocab(self): "mtp.layers.1.final_layernorm.weight": "model.layers.{bid}.shared_head.norm.weight", } + def _generate_nvfp4_tensors(self): + if self._mtp_bid is not None: + for name in list(self.model_tensors): + if name.startswith("mtp.layers.1.mixer"): + new_name = f"backbone.layers.{self._mtp_bid}.{name.split('.', 3)[3]}" + self.model_tensors[new_name] = self.model_tensors.pop(name) + super()._generate_nvfp4_tensors() + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: # mtp.layers.0: NextN input fusion + attention # mtp.layers.1: MoE + final head norm diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 0c9872228a8a..a39d046fe174 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -6265,6 +6265,7 @@ class VisionProjectorType: GLM5V = "glm5v" YOUTUVL = "youtuvl" NEMOTRON_V2_VL = "nemotron_v2_vl" + NEMOTRON_OMNI = "nemotron_omni" QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor POCKETTTS_SPKENC = "pockettts_spkenc" # audio: mimi encoder as voice-prompt encoder diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 5eab5b1fdcb2..63301cc114a9 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1735,6 +1735,7 @@ class TensorNameMap: "vision_model.class_embedding", # llama 4 "model.vision.patch_embedding.cls_embedding", # cogvlm "vision_model.radio_model.model.patch_generator.cls_token.token", # Nemotron Nano v2 VL + "vision_model.embeddings.cls_register_token", # Nemotron Omni "model.vision_model.embeddings.class_embedding", # Deepseek-OCR ), @@ -1756,6 +1757,7 @@ class TensorNameMap: "model.vision_model.embeddings.patch_embedding", # Deepseek-OCR CLIP "siglip2.vision_model.embeddings.patch_embedding", "vision_model.radio_model.model.patch_generator.embedder", # Nemotron Nano v2 VL + "vision_model.embeddings.patch_projection", # Nemotron Omni "model.vision_tower.patch_embedder.input_proj", # gemma4 "vision_tower.patch_embed.patchifier.proj", # dots.ocr "vision_encoder.patch_embed.proj", # dots3note @@ -1788,6 +1790,7 @@ class TensorNameMap: "model.vision.patch_embedding.position_embedding", # cogvlm "visual.embeddings.position_embedding", # glm4v "vision_model.radio_model.model.patch_generator.pos_embed", # Nemotron Nano v2 VL + "vision_model.embeddings.position_embedding", # Nemotron Omni "model.vision_tower.patch_embedder.position_embedding_table", # gemma4 "vision_model.positional_embedding", # Step3-VL "model.vision_embedder.pos_embedding", # gemma4 unified @@ -1818,6 +1821,7 @@ class TensorNameMap: ), MODEL_TENSOR.V_ENC_ATTN_Q: ( + "vision_model.encoder.layer.{bid}.attention.attention.query", # Nemotron Omni "model.vision_tower.vision_model.encoder.layers.{bid}.self_attn.q_proj", # Granite4Vision "vision_tower.vision_model.encoder.layers.{bid}.self_attn.q_proj", "model.vision_tower.encoder.layers.{bid}.self_attn.q_proj", # minicpmv4_6 @@ -1846,6 +1850,7 @@ class TensorNameMap: ), MODEL_TENSOR.V_ENC_ATTN_K: ( + "vision_model.encoder.layer.{bid}.attention.attention.key", # Nemotron Omni "model.vision_tower.vision_model.encoder.layers.{bid}.self_attn.k_proj", # Granite4Vision "vision_tower.vision_model.encoder.layers.{bid}.self_attn.k_proj", "model.vision_tower.encoder.layers.{bid}.self_attn.k_proj", # minicpmv4_6 @@ -1874,6 +1879,7 @@ class TensorNameMap: ), MODEL_TENSOR.V_ENC_ATTN_V: ( + "vision_model.encoder.layer.{bid}.attention.attention.value", # Nemotron Omni "model.vision_tower.vision_model.encoder.layers.{bid}.self_attn.v_proj", # Granite4Vision "vision_tower.vision_model.encoder.layers.{bid}.self_attn.v_proj", "model.vision_tower.encoder.layers.{bid}.self_attn.v_proj", # minicpmv4_6 @@ -1894,6 +1900,7 @@ class TensorNameMap: ), MODEL_TENSOR.V_ENC_INPUT_NORM: ( + "vision_model.encoder.layer.{bid}.norm1", # Nemotron Omni "vision.blocks.{bid}.norm1", # deepseek4v "model.vision_tower.vision_model.encoder.layers.{bid}.layer_norm1", # Granite4Vision "vision_tower.vision_model.encoder.layers.{bid}.layer_norm1", @@ -1920,6 +1927,7 @@ class TensorNameMap: ), MODEL_TENSOR.V_ENC_ATTN_O: ( + "vision_model.encoder.layer.{bid}.attention.output.dense", # Nemotron Omni "vision.blocks.{bid}.attn.wo", # deepseek4v "model.vision_tower.vision_model.encoder.layers.{bid}.self_attn.out_proj", # Granite4Vision "vision_tower.vision_model.encoder.layers.{bid}.self_attn.out_proj", @@ -1952,6 +1960,7 @@ class TensorNameMap: ), MODEL_TENSOR.V_ENC_POST_ATTN_NORM: ( + "vision_model.encoder.layer.{bid}.norm2", # Nemotron Omni "vision.blocks.{bid}.norm2", # deepseek4v "model.vision_tower.vision_model.encoder.layers.{bid}.layer_norm2", # Granite4Vision "vision_tower.vision_model.encoder.layers.{bid}.layer_norm2", @@ -1979,6 +1988,7 @@ class TensorNameMap: ), MODEL_TENSOR.V_ENC_FFN_UP: ( + "vision_model.encoder.layer.{bid}.mlp.fc1", # Nemotron Omni "vision.blocks.{bid}.mlp.w1_up", # deepseek4v (split from fused w1) "vision_encoder.blocks.{bid}.mlp.fc3", # dots3note "model.vision_tower.vision_model.encoder.layers.{bid}.mlp.fc1", # Granite4Vision @@ -2016,6 +2026,7 @@ class TensorNameMap: ), MODEL_TENSOR.V_ENC_FFN_DOWN: ( + "vision_model.encoder.layer.{bid}.mlp.fc2", # Nemotron Omni "vision.blocks.{bid}.mlp.w2", # deepseek4v "vision_encoder.blocks.{bid}.mlp.fc2", # dots3note "model.vision_tower.vision_model.encoder.layers.{bid}.mlp.fc2", # Granite4Vision @@ -2074,12 +2085,14 @@ class TensorNameMap: ), MODEL_TENSOR.V_LAYER_SCALE_1: ( + "vision_model.encoder.layer.{bid}.layer_scale1", # Nemotron Omni "vision_tower.vision_model.encoder.layers.{bid}.ls1", # InternVL "model.vision_tower.encoder.layer.{bid}.lambda_1", # Intern-S1 "vision_model.transformer.resblocks.{bid}.ls_1", # Step3-VL ), MODEL_TENSOR.V_LAYER_SCALE_2: ( + "vision_model.encoder.layer.{bid}.layer_scale2", # Nemotron Omni "vision_tower.vision_model.encoder.layers.{bid}.ls2", # InternVL "model.vision_tower.encoder.layer.{bid}.lambda_2", # Intern-S1 "vision_model.transformer.resblocks.{bid}.ls_2", # Step3-VL @@ -2114,6 +2127,7 @@ class TensorNameMap: "siglip2.vision_model.post_layernorm", "model.qwen2_model.model.model.norm", # Deepseek-OCR-2 qwen2 "model.vision_tower.ln_post", # muse-glimmer + "vision_projector.vision_final_layernorm", # Nemotron Omni Reasoning V3 ), MODEL_TENSOR.V_MM_POST_NORM: ( diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 982351dc860a..6f28d2f1db4c 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -930,7 +930,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra res->t_inp_tokens = inp_tokens; ggml_tensor * inp_anchor = nullptr; - if (model.lilicorr_pass_hidden_proj && ubatch.embd_state) { + if (model.lilicorr_pass_hidden_proj && ubatch.embd_state && n_tokens == (int64_t) hparams.dflash_block_size * ubatch.n_seqs_unq) { // LiliCorr sets inp_anchor to pass the context features to the lattice scorer auto inp = std::make_unique(n_embd, n_embd); inp->tokens = inp_tokens; @@ -1100,7 +1100,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra if (model.dflash_selector_hidden) { build_dflash2_selector(*this, model, inp_tokens); - } else if (model.lilicorr_pass_hidden_proj && ubatch.embd_state) { + } else if (model.lilicorr_pass_hidden_proj && ubatch.embd_state && n_tokens == (int64_t) hparams.dflash_block_size * ubatch.n_seqs_unq) { build_lilicorr_scores(*this, model, tok_embd, inp_anchor); } } diff --git a/src/models/nemotron-h.cpp b/src/models/nemotron-h.cpp index b14913f2a807..e9c61308feb1 100644 --- a/src/models/nemotron-h.cpp +++ b/src/models/nemotron-h.cpp @@ -109,21 +109,26 @@ void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) { layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED); } else { if (n_expert != 0) { + const auto * router_meta = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i).str().c_str()); + const int64_t n_expert_layer = router_meta ? router_meta->ne[1] : n_expert; + if (n_expert_layer < hparams.n_expert_used(i) || n_expert_layer > n_expert) { + throw std::runtime_error("Invalid per-layer Nemotron expert count"); + } // Use per-layer n_ff_exp; fall back to n_ff/n_expert_used if absent (existing GGUFs). const int64_t n_ff_exp_i = hparams.n_ff_exp(i) ? (int64_t)hparams.n_ff_exp(i) : hparams.n_ff(i) / (int64_t)hparams.n_expert_used(i); const int64_t n_ff_shexp = hparams.n_ff_shexp; - layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, trunk_flags); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, trunk_flags); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert_layer}, trunk_flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert_layer}, trunk_flags); // MoE branch layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, TENSOR_NOT_REQUIRED); layer.ffn_latent_up = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP, "weight", i), {moe_n_embd, n_embd}, TENSOR_NOT_REQUIRED); - layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp_i, moe_n_embd, n_expert}, trunk_flags); - layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp_i, n_expert}, trunk_flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp_i, moe_n_embd, n_expert_layer}, trunk_flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp_i, n_expert_layer}, trunk_flags); // Shared expert branch layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, trunk_flags); @@ -178,6 +183,15 @@ void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) { layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, mtp_flags); layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, mtp_flags); layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, mtp_flags); + + layer.ffn_down_exps_s = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "scale", i), {n_expert}, mtp_flags | TENSOR_NOT_REQUIRED); + layer.ffn_up_exps_s = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "scale", i), {n_expert}, mtp_flags | TENSOR_NOT_REQUIRED); + layer.ffn_down_exps_in_s = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "input_scale", i), {n_expert}, mtp_flags | TENSOR_NOT_REQUIRED); + layer.ffn_up_exps_in_s = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "input_scale", i), {n_expert}, mtp_flags | TENSOR_NOT_REQUIRED); + layer.ffn_down_shexp_s = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "scale", i), {1}, mtp_flags | TENSOR_NOT_REQUIRED); + layer.ffn_up_shexp_s = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "scale", i), {1}, mtp_flags | TENSOR_NOT_REQUIRED); + layer.ffn_down_shexp_in_s = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "input_scale", i), {1}, mtp_flags | TENSOR_NOT_REQUIRED); + layer.ffn_up_shexp_in_s = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "input_scale", i), {1}, mtp_flags | TENSOR_NOT_REQUIRED); } } @@ -304,7 +318,7 @@ ggml_tensor * llama_model_nemotron_h::graph::build_ffn_layer(ggml_tensor * cur, nullptr, // no gate model.layers[il].ffn_down_exps, model.layers[il].ffn_exp_probs_b, - n_expert, (int64_t)hparams.n_expert_used(il), + model.layers[il].ffn_up_exps->ne[2], (int64_t)hparams.n_expert_used(il), LLM_FFN_RELU_SQR, hparams.expert_weights_norm, hparams.expert_weights_scale, LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID, diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 20098937ce21..82f74bdc9665 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -46,6 +46,7 @@ add_library(mtmd models/kimivl.cpp models/kimik25.cpp models/nemotron-v2-vl.cpp + models/nemotron-omni.cpp models/muse-glimmer.cpp models/llama4.cpp models/llava.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 30cb3078916b..3d6c6f8c4801 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -498,6 +498,7 @@ enum projector_type { PROJECTOR_TYPE_YASA2, PROJECTOR_TYPE_KIMIK25, PROJECTOR_TYPE_NEMOTRON_V2_VL, + PROJECTOR_TYPE_NEMOTRON_OMNI, PROJECTOR_TYPE_HUNYUANVL, PROJECTOR_TYPE_PARAKEET, PROJECTOR_TYPE_EXAONE4_5, @@ -569,6 +570,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_YASA2, "yasa2"}, { PROJECTOR_TYPE_KIMIK25, "kimik25"}, { PROJECTOR_TYPE_NEMOTRON_V2_VL, "nemotron_v2_vl"}, + { PROJECTOR_TYPE_NEMOTRON_OMNI, "nemotron_omni"}, { PROJECTOR_TYPE_EXAONE4_5, "exaone4_5"}, { PROJECTOR_TYPE_HUNYUANVL, "hunyuanvl"}, { PROJECTOR_TYPE_MINICPMV4_6, "minicpmv4_6"}, @@ -904,11 +906,13 @@ static void clip_log_internal(enum ggml_log_level level, const char * format, .. struct clip_image_f32_batch { std::vector entries; bool is_audio = false; + bool is_video = false; clip_image_f32_batch clone() const { clip_image_f32_batch new_batch{ /* entries */ {}, /* is_audio */ is_audio, + /* is_video */ is_video, }; new_batch.entries.reserve(entries.size()); for (const auto & entry : entries) { diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 657671ded548..2c4ea398ccb7 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -566,7 +566,7 @@ struct clip_model { // embeddings ggml_tensor * class_embedding = nullptr; ggml_tensor * patch_embeddings_0 = nullptr; - ggml_tensor * patch_embeddings_1 = nullptr; // second Conv2D kernel when we decouple Conv3D along temporal dimension (Qwen2VL) + ggml_tensor * patch_embeddings_1 = nullptr; // additional patch kernel; interpretation is model-specific (Qwen2VL, Nemotron Omni) ggml_tensor * patch_bias = nullptr; ggml_tensor * position_embeddings = nullptr; ggml_tensor * norm_embd_w = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index fccb00c3aa70..4c1185a431b9 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1042,6 +1042,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_NEMOTRON_OMNI: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_LLAMA4: { builder = std::make_unique(ctx, img); @@ -1523,6 +1527,16 @@ struct clip_model_loader { { get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge, false); } break; + case PROJECTOR_TYPE_NEMOTRON_OMNI: + { + hparams.ffn_op = FFN_GELU_ERF; + log_ffn_op = "gelu_erf"; + hparams.image_resize_pad = PAD_NONE; + get_u32(KEY_PROJ_SCALE_FACTOR, hparams.n_merge); + get_u32(KEY_IMAGE_MIN_PIXELS, hparams.image_min_pixels); + get_u32(KEY_IMAGE_MAX_PIXELS, hparams.image_max_pixels); + hparams.warmup_image_size = hparams.image_size; + } break; case PROJECTOR_TYPE_PARAKEET: { get_u32(KEY_AUDIO_SUBSMPL_FACTOR, hparams.subsampling_factor); @@ -3269,6 +3283,7 @@ struct clip_model_loader { model.mm_3_b = get_tensor(string_format(TN_MVLM_PROJ_MLP, 3, "bias")); } break; case PROJECTOR_TYPE_NEMOTRON_V2_VL: + case PROJECTOR_TYPE_NEMOTRON_OMNI: { model.mm_0_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 0, "weight")); model.mm_1_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 1, "weight")); @@ -4320,6 +4335,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_COHERE2V: case PROJECTOR_TYPE_INTERNVL: case PROJECTOR_TYPE_NEMOTRON_V2_VL: + case PROJECTOR_TYPE_NEMOTRON_OMNI: case PROJECTOR_TYPE_LLAMA4: { // both X and Y are downscaled by the scale factor @@ -5428,6 +5444,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { case PROJECTOR_TYPE_COHERE2V: case PROJECTOR_TYPE_INTERNVL: case PROJECTOR_TYPE_NEMOTRON_V2_VL: + case PROJECTOR_TYPE_NEMOTRON_OMNI: case PROJECTOR_TYPE_QWEN2A: case PROJECTOR_TYPE_QWEN3A: case PROJECTOR_TYPE_GLMA: @@ -6149,6 +6166,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_3_w->ne[1]; // out_proj output dim case PROJECTOR_TYPE_INTERNVL: case PROJECTOR_TYPE_NEMOTRON_V2_VL: + case PROJECTOR_TYPE_NEMOTRON_OMNI: return ctx->model.mm_3_w->ne[1]; case PROJECTOR_TYPE_LLAMA4: return ctx->model.mm_model_proj->ne[1]; @@ -6225,6 +6243,7 @@ int clip_model_n_temporal_merge(const struct clip_ctx * ctx) { case PROJECTOR_TYPE_QWEN2VL: case PROJECTOR_TYPE_QWEN25VL: case PROJECTOR_TYPE_QWEN3VL: + case PROJECTOR_TYPE_NEMOTRON_OMNI: case PROJECTOR_TYPE_LING3VL: return 2; default: diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index adb5ede309dc..2f4676a54359 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -109,6 +109,12 @@ struct clip_graph_nemotron_v2_vl : clip_graph { ggml_cgraph * build() override; }; +struct clip_graph_nemotron_omni : clip_graph { + clip_graph_nemotron_omni(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; + ggml_tensor * build_inp_with_temporal_merge(); +}; + struct clip_graph_llama4 : clip_graph { clip_graph_llama4(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/models/nemotron-omni.cpp b/tools/mtmd/models/nemotron-omni.cpp new file mode 100644 index 000000000000..c4b1bff40621 --- /dev/null +++ b/tools/mtmd/models/nemotron-omni.cpp @@ -0,0 +1,56 @@ +#include "models.h" + +ggml_tensor * clip_graph_nemotron_omni::build_inp_with_temporal_merge() { + if (n_batch == 1 && !img_batch->is_video) { + return build_inp(); + } + + GGML_ASSERT(n_batch == 1 || n_batch == 2); + GGML_ASSERT(model.patch_embeddings_1 != nullptr); + GGML_ASSERT(model.patch_embeddings_1->ne[2] == 6); + + ggml_tensor * inp_raw = build_inp_raw(); + if (n_batch == 1) { + // pad the last video frame inside the graph + inp_raw = ggml_concat(ctx0, inp_raw, inp_raw, 2); + } else { + inp_raw = ggml_reshape_4d(ctx0, inp_raw, img.nx(), img.ny(), 6, 1); + } + + ggml_tensor * inp = ggml_conv_2d(ctx0, model.patch_embeddings_1, inp_raw, patch_size, patch_size, 0, 0, 1, 1); + inp = ggml_reshape_3d(ctx0, inp, n_patches, n_embd, 1); + return ggml_cont(ctx0, ggml_transpose(ctx0, inp)); +} + +ggml_cgraph * clip_graph_nemotron_omni::build() { + GGML_ASSERT(model.class_embedding != nullptr); + GGML_ASSERT(model.position_embeddings != nullptr); + + const int n_registers = model.class_embedding->ne[1]; + const int n_pos = n_patches + n_registers; + + ggml_tensor * inp = build_inp_with_temporal_merge(); + + inp = ggml_add(ctx0, inp, resize_position_embeddings(GGML_SCALE_MODE_BILINEAR)); + cb(inp, "inp_pos", -1); + + inp = ggml_concat(ctx0, model.class_embedding, inp, 1); + + ggml_tensor * cur = build_vit(inp, n_pos, NORM_TYPE_NORMAL, hparams.ffn_op, nullptr, nullptr); + + cur = ggml_view_2d(ctx0, cur, + n_embd, n_patches, + ggml_row_size(cur->type, n_embd), + n_registers * ggml_row_size(cur->type, n_embd)); + + cur = build_patch_merge_permute(cur, model.hparams.n_merge); + + { + cur = build_norm(cur, model.mm_0_w, nullptr, NORM_TYPE_RMS, 1e-5, -1); + cur = build_ffn(cur, model.mm_1_w, nullptr, nullptr, nullptr, model.mm_3_w, nullptr, FFN_RELU_SQR, -1); + } + + ggml_build_forward_expand(gf, cur); + + return gf; +} diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 2c2b069361d5..2fa02f83eb3f 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -898,6 +898,12 @@ struct mtmd_context { { image_preproc = std::make_unique(ctx_v); } break; + case PROJECTOR_TYPE_NEMOTRON_OMNI: + { + img_beg = ""; + img_end = ""; + image_preproc = std::make_unique(ctx_v); + } break; case PROJECTOR_TYPE_LFM2: { // multi-tile: @@ -1607,6 +1613,7 @@ struct mtmd_tokenizer { clip_image_f32_batch batch_f32; batch_f32.is_audio = false; + batch_f32.is_video = bitmaps[0]->mergeable; batch_f32.entries = std::move(preproc_out.entries); // do NOT use preproc_out from this point on, it's moved From 620d2225cc436ab0b001e3200e590e3b286ef9ab Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 15 Oct 2026 16:48:22 +0300 Subject: [PATCH 030/423] llama.cpp : bump version to 0.7.0 (#30603) --- CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 3086e81b1972..12f45558f5be 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -4,7 +4,7 @@ include(CheckIncludeFileCXX) ### llama.cpp version set(LLAMA_VERSION_MAJOR 0) -set(LLAMA_VERSION_MINOR 6) +set(LLAMA_VERSION_MINOR 7) set(LLAMA_VERSION_PATCH 0) set(LLAMA_VERSION_BASE "${LLAMA_VERSION_MAJOR}.${LLAMA_VERSION_MINOR}.${LLAMA_VERSION_PATCH}") From b76b13379d3b7d67ab0bc424516bb6fd2522bd8e Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 15 Oct 2026 20:32:52 +0300 Subject: [PATCH 031/423] pi : expand AI disclosure with hardware and stats (#30615) * pi : expand AI disclosure with hardware and stats Prefer PI_HARDWARE for the hardware field, spell out the disclosure block template (model link, harness, hardware, stats) and source [STATS] from ~/.pi/stats.sh. Record the hardware in the Assisted-by commit tag. Assisted-by: pi:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) * pi : drop square brackets from placeholders Brackets inside a markdown link text break the link on GitHub. Describe the disclosure bullets in words instead of embedding a bracketed template. Assisted-by: pi:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) --- .pi/gg/SYSTEM.md | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/.pi/gg/SYSTEM.md b/.pi/gg/SYSTEM.md index 8b8b8a1586fa..b086d72162bc 100644 --- a/.pi/gg/SYSTEM.md +++ b/.pi/gg/SYSTEM.md @@ -6,7 +6,8 @@ General: - PR and commit titles format: ` : `. Lookup recents for examples - Don't try to build or run the code unless you are explicitly asked to do so - Use the `gh` CLI tool when querying PRs, issues, or other GitHub resources -- When [MODEL] is needed, first try to get it from the `PI_MODEL_NAME` env var before asking the user +- When a model name is needed, first try to get it from the `PI_MODEL_NAME` env var before asking the user +- When a hardware description is needed, first try to get it from the `PI_HARDWARE` env var before asking the user - Never read the `AGENTS.md` file Coding: @@ -21,13 +22,20 @@ Pull requests (PRs): - Before opening a pull request, ask the user to confirm the description - Don't explicitly wrap lines in the PR description (each paragraph and bullet is a single line) - When creating a pull request, look for the repository's PR template and follow it -- For the AI usage disclosure section, write "YES. pi:llama.cpp/[MODEL]" -- If `PI_MODEL_NAME` env var is not set, ask the user to tell you what model was used and write it in place of [MODEL] +- For the AI usage disclosure section, write "YES", then these 4 bullets: + - Model: `<model>`, linked to its Hugging Face page - the link text is the plain model name + - Harness: `pi:llama.cpp` + - Hardware: `<hardware>` + - Stats: `<stats>` +- Run `~/.pi/stats.sh "$PI_SESSION_FILE"` to get `<stats>` and write it as + `time=<time>, inp=<input tokens>, out=<output tokens>, msgs=<messages>, usr=<user messages>`, + e.g. `time=12.3m, inp=120k, out=45k, msgs=87, usr=6` +- If `PI_MODEL_NAME` or `PI_HARDWARE` is not set, ask the user and use their answer for `<model>` or `<hardware>` - Always create the pull requests in draft mode - Never reply to review comments or post comments on issues/PRs without explicit permission from the user Commits: -- On every commit that you make, include a "Assisted-by: pi:llama.cpp/[MODEL]" tag +- On every commit that you make, include an `Assisted-by: pi:llama.cpp/<model> (<hardware>)` tag - Do not explicitly set the git author in commits - rely on the default git config - Always use `--no-gpg-sign` when committing - Never `git push` without explicit confirmation from the user From ba3c05f71da85cd0439e84ba42b0478b24d5db16 Mon Sep 17 00:00:00 2001 From: Pascal <admin@serveurperso.com> Date: Thu, 15 Oct 2026 23:27:12 +0200 Subject: [PATCH 032/423] model: add Muse Spark Open (text, vision and audio) (#30616) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * model: add Muse Spark Open (text, vision and audio) Muse Spark Open is a latent space MoE: the residual stream stays at 8192 while the 256 experts live in 4096, with a projection down and up around them, top-8 sigmoid routing with a correction bias. Attention interleaves three sliding layers at a 2048 window with rope for one global layer without rope, with a sigmoid output gate and a weightless QK-norm. Residuals are norm preserving, their gate is folded at conversion time. Vision reuses the Muse Glimmer tower. Audio is a causal convolutional pyramid with no attention and no spectrogram, four stages folding 16, 6, 5 and 4 samples into the channel dim, one embedding per 80 ms of 24 kHz. The quantizer keeps the latent projections and the attention gate at Q8_0, a little over 1% of the weights, since the whole residual stream passes through them and the low-bit types turned the output into a repeating loop. The model saver writes the post norm epsilon, and the arch is registered in test-llama-archs. Co-authored-by: Xuan Son Nguyen <son@huggingface.co> Co-authored-by: Pedro Cuenca <pedro@huggingface.co> * Muse Spark: drop the redundant forward expand * Muse Spark: drop the redundant forward expand and tighten the comments * Update src/llama-quant.cpp Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * Update src/llama-graph.cpp Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * Update conversion/muse_spark.py Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * Update conversion/muse_spark.py Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * Update conversion/muse_spark.py Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * Update conversion/muse_spark.py Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> --------- Co-authored-by: Xuan Son Nguyen <son@huggingface.co> Co-authored-by: Pedro Cuenca <pedro@huggingface.co> Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> --- common/chat.cpp | 6 + common/parsers/muse-spark.cpp | 146 +++++++++++++++ common/parsers/parsers.h | 1 + common/parsers/sources.cmake | 1 + conversion/__init__.py | 2 + conversion/base.py | 3 + conversion/muse_glimmer.py | 2 +- conversion/muse_spark.py | 221 +++++++++++++++++++++++ convert_hf_to_gguf_update.py | 2 + gguf-py/gguf/constants.py | 53 ++++++ gguf-py/gguf/gguf_writer.py | 12 ++ gguf-py/gguf/tensor_mapping.py | 11 ++ models/templates/muse-spark.jinja | 228 +++++++++++++++++++++++ src/llama-arch.cpp | 14 ++ src/llama-arch.h | 8 + src/llama-graph.cpp | 19 +- src/llama-graph.h | 8 +- src/llama-hparams.h | 3 + src/llama-model-saver.cpp | 1 + src/llama-model.cpp | 4 + src/llama-model.h | 11 ++ src/llama-quant.cpp | 21 +++ src/llama-vocab.cpp | 1 + src/models/models.h | 13 ++ src/models/muse-spark.cpp | 238 +++++++++++++++++++++++++ tests/test-chat.cpp | 44 +++++ tests/test-llama-archs.cpp | 14 +- tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-impl.h | 7 + tools/mtmd/clip-model.h | 26 +++ tools/mtmd/clip.cpp | 75 +++++++- tools/mtmd/models/models.h | 5 + tools/mtmd/models/muse-spark-audio.cpp | 66 +++++++ tools/mtmd/mtmd-audio.cpp | 20 +++ tools/mtmd/mtmd-audio.h | 6 + tools/mtmd/mtmd.cpp | 6 + 36 files changed, 1290 insertions(+), 9 deletions(-) create mode 100644 common/parsers/muse-spark.cpp create mode 100644 conversion/muse_spark.py create mode 100644 models/templates/muse-spark.jinja create mode 100644 src/models/muse-spark.cpp create mode 100644 tools/mtmd/models/muse-spark-audio.cpp diff --git a/common/chat.cpp b/common/chat.cpp index 4856713082ac..118072b6b514 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -1110,6 +1110,12 @@ static common_chat_params_init_fn common_chat_template_detect_params_init(const return common_chat_params_init_gpt_oss; } + // Muse Spark shares the Muse recipient protocol and adds a commentary recipient. + if (src.find("to=commentary") != std::string::npos && src.find("<|eom|>") != std::string::npos) { + LOG_DBG("Using specialized template: Muse Spark\n"); + return common_chat_params_init_muse_spark; + } + // Muse Glimmer format using " to=<recipient>" recipients and <|eom|>/<|eot|> message terminators. if (src.find("<atem:function_calls>") != std::string::npos && src.find("<|eom|>") != std::string::npos) { LOG_DBG("Using specialized template: Muse Glimmer\n"); diff --git a/common/parsers/muse-spark.cpp b/common/parsers/muse-spark.cpp new file mode 100644 index 000000000000..577090aa5cc9 --- /dev/null +++ b/common/parsers/muse-spark.cpp @@ -0,0 +1,146 @@ +#include "parsers.h" + +// An assistant turn is rendered as one or more messages, each +// "<|start|>assistant to=<recipient><|message|>{content}{END}" where END is +// <|eom|> (more messages follow) or <|eot|> (end of turn): +// - chain-of-thought: to=self, terminated by <|eom|> +// - user facing aside: to=commentary, terminated by <|eom|> +// - final answer: to=user, terminated by <|eot|> +// A commentary message carries content the user sees while the assistant keeps +// working, typically before or between tool calls. +// The generation prompt is just "<|start|>assistant"; the model emits its own +// " to=...<|message|>". +common_chat_params common_chat_params_init_muse_spark(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.generation_prompt = "<|start|>assistant"; + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + + data.preserved_tokens = { + "<|start|>", "<|message|>", "<|eom|>", "<|eot|>", + // ATEM tool-call markup emitted on " to=<tool>" turns. + "<atem:function_calls>", "<atem:invoke", "<atem:parameter", "</atem:parameter>", + "</atem:invoke>", "</atem:function_calls>", + }; + + data.message_delimiters = { + { COMMON_CHAT_ROLE_ASSISTANT, "<|start|>assistant" }, + { COMMON_CHAT_ROLE_USER, "<|start|>user" }, + { COMMON_CHAT_ROLE_SYSTEM, "<|start|>system" }, + { COMMON_CHAT_ROLE_TOOL, "<|start|>tool" }, + }; + + if (inputs.has_continuation()) { + const auto & msg = inputs.continue_msg; + + data.generation_prompt = "<|start|>assistant to=self<|message|>" + msg.reasoning_content; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += "<|eom|><|start|>assistant to=user<|message|>" + msg.render_content(); + } + + data.prompt.append(data.generation_prompt); + } + + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + // Constrained grammar whenever tools are offered. + auto include_grammar = has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE; + + data.parser = build_chat_peg_parser([&](common_chat_peg_builder & p) { + auto start = p.rule("start", p.literal("<|start|>assistant")); + + if (!extract_reasoning && !include_grammar) { + return start + p.content(p.rest()); + } + + if (extract_reasoning) { + p.rule("analysis", p.literal(" to=self<|message|>") + p.reasoning(p.until("<|eom|>")) + p.literal("<|eom|>")); + } else { + p.rule("analysis", p.literal(" to=self<|message|>") + p.content(p.until("<|eom|>")) + p.literal("<|eom|>")); + } + auto analysis = p.ref("analysis"); + + // Commentary is user facing content emitted mid turn, so it reads as + // content and always ends on <|eom|>. + auto preamble = p.rule("preamble", p.literal(" to=commentary<|message|>") + + p.content(p.until("<|eom|>")) + p.literal("<|eom|>")); + + auto aside = p.rule("aside", analysis | preamble); + + auto recipient = p.optional(p.literal(" to=user")); + auto final_msg = p.rule("final", recipient + p.literal("<|message|>") + + p.content(p.until_one_of({ "<|eot|>", "<|eom|>" }))); + + if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { + auto string_value = p.ac( + p.tool_arg_string_value(p.until("</atem:parameter>")) + p.tool_arg_close(p.literal("</atem:parameter>")), + "</atem:parameter>"); + + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) { + const auto & function = tool.at("function"); + const std::string name = function.at("name"); + + std::vector<common_peg_parser> arg_rules; + foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & prop, const common_chat_schema_document_ptr & doc) { + auto value_parser = p.eps(); + if (prop.schema->may_be_string()) { + value_parser = string_value; + } else { + value_parser = p.tool_arg_json_value( + p.schema(p.json(), "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index) + "-schema", doc, *prop.schema)) + + p.tool_arg_close(p.literal("</atem:parameter>")); + } + + arg_rules.push_back(p.tool_arg( + p.tool_arg_open(p.literal("<atem:parameter name=\"") + p.tool_arg_name(p.literal(prop.name)) + p.literal("\">")) + + value_parser)); + }); + + auto args = p.eps(); + if (!arg_rules.empty()) { + args = p.zero_or_more(p.choice(arg_rules) + p.space()); + } + + auto tool_parser = p.tool( + p.tool_open(p.literal(" to=") + p.until("<|message|>") + + p.literal("<|message|><atem:function_calls>") + p.space() + + p.literal("<atem:invoke name=\"") + p.tool_name(p.literal(name)) + p.literal("\">") + p.space()) + << p.tool_args(args) + << p.tool_close(p.literal("</atem:invoke>") + p.space() + p.literal("</atem:function_calls>"))); + + tool_choice |= p.rule("tool-" + std::to_string(tool_index), tool_parser); + }); + + auto tool_calls = inputs.parallel_tool_calls + ? p.trigger_rule("tool-call", tool_choice + p.zero_or_more(p.literal("<|eom|>") + start + tool_choice)) + : p.trigger_rule("tool-call", tool_choice); + + + if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) { + return p.zero_or_more(start + aside) + start + tool_calls; + } + auto trailing_calls = p.optional(p.literal("<|eom|>") + start + tool_calls); + return p.zero_or_more(start + aside) + start + (tool_calls | (final_msg + trailing_calls)); + } + + return p.zero_or_more(start + aside) + start + final_msg; + }); + + if (include_grammar) { + data.grammar_lazy = inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_REQUIRED; + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + data.parser.build_grammar(builder, data.grammar_lazy); + }); + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, + "<\\|start\\|>assistant( to=(?!self<\\|message\\|>)(?!commentary<\\|message\\|>)(?!user<\\|message\\|>)[^<]*?<\\|message\\|>)" }, + }; + } + + return data; +} diff --git a/common/parsers/parsers.h b/common/parsers/parsers.h index 8f0a97ce3446..fac69c096db1 100644 --- a/common/parsers/parsers.h +++ b/common/parsers/parsers.h @@ -79,6 +79,7 @@ common_chat_params common_chat_params_init_minimax_m3(const common_chat_template common_chat_params common_chat_params_init_ministral_3(const common_chat_template & tmpl, const autoparser::generation_params & inputs); common_chat_params common_chat_params_init_muse_glimmer(const common_chat_template & tmpl, const autoparser::generation_params & inputs); +common_chat_params common_chat_params_init_muse_spark(const common_chat_template & tmpl, const autoparser::generation_params & inputs); common_chat_params common_chat_params_init_qwen3_coder(const common_chat_template & tmpl, const autoparser::generation_params & inputs); diff --git a/common/parsers/sources.cmake b/common/parsers/sources.cmake index 152397074d22..5ee65b188a5b 100644 --- a/common/parsers/sources.cmake +++ b/common/parsers/sources.cmake @@ -19,6 +19,7 @@ set(LLAMA_CHAT_PARSERS_SOURCES ${CMAKE_CURRENT_LIST_DIR}/minimax-m3.cpp ${CMAKE_CURRENT_LIST_DIR}/ministral3.cpp ${CMAKE_CURRENT_LIST_DIR}/muse-glimmer.cpp + ${CMAKE_CURRENT_LIST_DIR}/muse-spark.cpp ${CMAKE_CURRENT_LIST_DIR}/qwen3-coder.cpp ${CMAKE_CURRENT_LIST_DIR}/translate-gemma.cpp ) diff --git a/conversion/__init__.py b/conversion/__init__.py index 8ce4cb7337d6..6c3914fbdfae 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -220,6 +220,7 @@ "OlmoeForCausalLM": "olmo", "MuseGlimmerAssistantModel": "muse_glimmer", "MuseGlimmerForConditionalGeneration": "muse_glimmer", + "MuseSparkForConditionalGeneration": "muse_spark", "OpenELMForCausalLM": "openelm", "OpenJevModel": "qwen", "OrionForCausalLM": "orion", @@ -356,6 +357,7 @@ "NemotronH_Omni_Reasoning_V3": "nemotron", "NemotronH_Nano_VL_V2": "nemotron", "MuseGlimmerForConditionalGeneration": "muse_glimmer", + "MuseSparkForConditionalGeneration": "muse_spark", "PaddleOCRVisionModel": "ernie", "Phi4ForCausalLMV": "phi", "Qwen2AudioForConditionalGeneration": "ultravox", diff --git a/conversion/base.py b/conversion/base.py index a5a80f27dad4..e3d265157add 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1869,6 +1869,9 @@ def get_vocab_base_pre(self, tokenizer) -> str: if chkhsh == "66b8d4e19ab16c3bfd89bce5d785fb7e0155e8648708a1f42077cb9fe002c273": # ref: https://huggingface.co/alvarobartt/grok-2-tokenizer res = "grok-2" + if chkhsh == "f5c965e4198da6e7525072d87d1ecf4c51d981fd486c7c4ec728ef9cec7e03fc": + # ref: https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B + res = "gpt-4o" if chkhsh == "b3d1dd861f1d4c5c0d2569ce36baf3f90fe8a102db3de50dd71ff860d91be3df": # ref: https://huggingface.co/aari1995/German_Semantic_V3 res = "jina-v2-de" diff --git a/conversion/muse_glimmer.py b/conversion/muse_glimmer.py index c86b33227366..9b63d3dba64a 100644 --- a/conversion/muse_glimmer.py +++ b/conversion/muse_glimmer.py @@ -94,7 +94,7 @@ def set_gguf_parameters(self): assert self.hparams_vision is not None c = self.hparams_vision # enriched vision_config from get_vision_config() - self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MUSE_GLIMMER) + self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.MUSE_GLIMMER) self.gguf_writer.add_vision_attention_layernorm_eps(float(c["layer_norm_eps"])) self.gguf_writer.add_vision_spatial_merge_size(int(c["merge_size"])) diff --git a/conversion/muse_spark.py b/conversion/muse_spark.py new file mode 100644 index 000000000000..150f22f95940 --- /dev/null +++ b/conversion/muse_spark.py @@ -0,0 +1,221 @@ +from __future__ import annotations + +from typing import Any, Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import ModelBase, MmprojModel, TextModel, gguf +from .muse_glimmer import MuseGlimmerVisionModel + + +# The reference clamps the squared stream coefficient with this floor. +RESIDUAL_GATE_EPS = 1e-3 + +# Block norms hold a delta around one; the output norm holds an absolute gain. +OFFSET_NORM_SUFFIXES = ( + "input_layernorm.weight", + "pre_feedforward_layernorm.weight", + "post_feedforward_layernorm.weight", + "pre_expert_norm.weight", + "post_expert_norm.weight", +) + + +def _unpermute_for_rope(tensor: "Tensor", n_heads: int) -> "Tensor": + """Invert transformers' `_permute_for_rope`: HF stores Q/K in rotate_half layout, + llama.cpp consumes the interleaved (NORM) layout. Vision Q/K carry a bias too.""" + if tensor.ndim == 1: + dim1, = tensor.shape + return tensor.view(n_heads, 2, dim1 // n_heads // 2).transpose(1, 2).reshape(dim1) + dim1, dim2 = tensor.shape + return tensor.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2) + + +@ModelBase.register("MuseSparkForConditionalGeneration") +@ModelBase.example("meta-models/Muse-Spark-Open-816B-A42B") +class MuseSparkModel(TextModel): + model_arch = gguf.MODEL_ARCH.MUSE_SPARK + + def set_vocab(self): + self._set_vocab_gpt2() + + from transformers import AutoTokenizer + tok = AutoTokenizer.from_pretrained(self.dir_model) + eot_id = tok.convert_tokens_to_ids("<|eot|>") # ty: ignore[unresolved-attribute] + if isinstance(eot_id, int) and eot_id >= 0: + self.gguf_writer.add_eot_token_id(eot_id) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + + # MetaP resolves to two scalars + # The output multiplier lands here and the attention scale rides on the synthesized QK-norm weights + self.gguf_writer.add_logit_scale(hparams["output_multiplier"]) + self.gguf_writer.add_final_logit_softcapping(hparams["final_logit_softcapping"]) + + self.gguf_writer.add_post_norm_eps(hparams["post_norm_eps"]) + + # Only the sliding layers carry rope, so the graph takes it from the pattern + sliding = [t == "sliding_attention" for t in hparams["layer_types"]] + theta = hparams["rope_parameters"]["rope_theta"] + assert sliding == [bool(t) for t in hparams["layer_rope_theta"]] + assert all(t in (0, theta) for t in hparams["layer_rope_theta"]) + assert all(t == "sparse" for t in hparams["mlp_layer_types"]) + + self.gguf_writer.add_sliding_window(hparams["sliding_window"]) + self.gguf_writer.add_rope_freq_base_swa(theta) + self.gguf_writer.add_sliding_window_pattern(sliding) + + # Experts live in a latent space narrower than the residual stream. + self.gguf_writer.add_moe_latent_size(hparams["moe_hidden_size"]) + self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"]) + self.gguf_writer.add_expert_weights_scale(hparams["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(True) + # the config carries no score function, the reference hardcodes sigmoid + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.endswith(OFFSET_NORM_SUFFIXES): + data_torch = data_torch + 1.0 + + if name.endswith(".self_attn.q_proj.weight"): + assert bid is not None + n_head = self.hparams["num_attention_heads"] + head_dim = self.hparams["head_dim"] + + yield self.map_tensor_name(name), _unpermute_for_rope(data_torch, n_head) + + # Synthesize QK-norm to carry the MetaP scale so the graph keeps the plain 1 / sqrt(head_dim) + ones = torch.ones(head_dim, dtype=torch.float32) + q_norm = ones * self.hparams["qk_scale_factor"] + yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q_NORM, bid), q_norm + yield self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_NORM, bid), ones + return + + if name.endswith(".self_attn.k_proj.weight"): + n_head_kv = self.hparams["num_key_value_heads"] + yield self.map_tensor_name(name), _unpermute_for_rope(data_torch, n_head_kv) + return + + # gate_up_proj fuses both projections over the output dim, gate first. + if name.endswith(".mlp.experts.gate_up_proj"): + assert bid is not None + n_ff = data_torch.shape[1] // 2 + yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid), data_torch[:, :n_ff, :] + yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_EXP, bid), data_torch[:, n_ff:, :] + return + + if name.endswith(".mlp.experts.down_proj"): + assert bid is not None + yield self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid), data_torch + return + + # The router scales its logits by the MetaP output multiplier before the sigmoid + # Fold it into the weight to keep a plain matmul + if name.endswith(".mlp.gate.weight"): + data_torch = data_torch.float() * self.hparams["output_multiplier"] + yield from super().modify_tensors(data_torch, name, bid) + return + + # Norm preserving residual gate: alpha scales the stream, beta the branch, with alpha^2 + beta^2 = 1 + # Both follow from the weight, so resolve them here + if name.endswith("_residual_gate.gate"): + assert bid is not None + attn = "post_attention" in name + + w = data_torch.float() / self.hparams["residual_gate_temperature"] + beta = torch.sigmoid(w) + alpha_sq = torch.sigmoid(-w) * (1.0 + beta) + alpha = torch.sqrt(torch.clamp(alpha_sq, min=RESIDUAL_GATE_EPS)) + + t_alpha = gguf.MODEL_TENSOR.ATTN_RESID_ALPHA if attn else gguf.MODEL_TENSOR.FFN_RESID_ALPHA + t_beta = gguf.MODEL_TENSOR.ATTN_RESID_BETA if attn else gguf.MODEL_TENSOR.FFN_RESID_BETA + yield self.format_tensor_name(t_alpha, bid), alpha + yield self.format_tensor_name(t_beta, bid), beta + return + + yield from super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("MuseSparkForConditionalGeneration") +@ModelBase.example("meta-models/Muse-Spark-Open-816B-A42B") +class MuseSparkVisionModel(MuseGlimmerVisionModel): + has_vision_encoder = True # same vision encoder as Muse Glimmer + has_audio_encoder = True + + def get_audio_config(self) -> dict[str, Any]: + c = self.global_config["audio_config"] + # the base class wants one width and one layer count: use the last stage width and the total number of conv blocks + return { + **c, + "hidden_size": c["hidden_sizes"][-1], + "intermediate_size": int(c["hidden_sizes"][-1] * c["mlp_ratio"]), + "num_attention_heads": 1, + "num_hidden_layers": sum(c["depths"]), + } + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + assert self.hparams_audio + a = self.hparams_audio + + self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.MUSE_SPARK_AUDIO) + # the input is the raw waveform, one channel takes the place of the mel bins + self.gguf_writer.add_audio_num_mel_bins(1) + self.gguf_writer.add_audio_attention_layernorm_eps(float(a["rms_norm_eps"])) + self.gguf_writer.add_audio_conv_kernel_size(int(a["conv_kernel_size"])) + self.gguf_writer.add_audio_hidden_sizes([int(x) for x in a["hidden_sizes"]]) + self.gguf_writer.add_audio_patch_sizes([int(x) for x in a["patch_sizes"]]) + self.gguf_writer.add_audio_depths([int(x) for x in a["depths"]]) + + @classmethod + def filter_tensors(cls, item): + name, gen = item + if name.startswith("model.audio_tower."): + return MmprojModel.filter_tensors((name, gen)) + return super().filter_tensors((name, gen)) + + _AUDIO_HEAD = { + "model.audio_tower.norm": gguf.MODEL_TENSOR.A_PRE_NORM, + "model.audio_tower.proj": gguf.MODEL_TENSOR.A_MMPROJ_FC, + "model.audio_tower.post_proj_norm": gguf.MODEL_TENSOR.A_POST_NORM, + } + + def _audio_tensor(self, data_torch, name, suffix): + assert self.hparams_audio + + stem = name[: -len(suffix) - 1] + + if stem in self._AUDIO_HEAD: + return self.format_tensor_name(self._AUDIO_HEAD[stem], suffix="." + suffix), data_torch + + parts = stem.split(".") + stage = int(parts[3]) + + if parts[4] == "patch_embedder": + key = (gguf.MODEL_TENSOR.A_ENC_PATCH_PROJ if parts[5] == "proj" + else gguf.MODEL_TENSOR.A_ENC_PATCH_NORM) + return self.format_tensor_name(key, bid=stage, suffix="." + suffix), data_torch + + # conv blocks are numbered per stage upstream and flat here + depths = [int(d) for d in self.hparams_audio["depths"]] + bid = sum(depths[:stage]) + int(parts[5]) + tail = ".".join(parts[6:]) + key = { + "conv": gguf.MODEL_TENSOR.A_ENC_CONV1D, + "norm": gguf.MODEL_TENSOR.A_ENC_CONV1D_NORM, + "mlp.fc1": gguf.MODEL_TENSOR.A_ENC_FFN_UP, + "mlp.fc2": gguf.MODEL_TENSOR.A_ENC_FFN_DOWN, + }[tail] + return self.format_tensor_name(key, bid=bid, suffix="." + suffix), data_torch + + def modify_tensors(self, data_torch, name, bid): + if name.startswith("model.audio_tower."): + yield self._audio_tensor(data_torch, name, name.rpartition(".")[2]) + return + yield from super().modify_tensors(data_torch, name, bid) diff --git a/convert_hf_to_gguf_update.py b/convert_hf_to_gguf_update.py index 0afbcc7e93a3..d49fb8f3ff3e 100755 --- a/convert_hf_to_gguf_update.py +++ b/convert_hf_to_gguf_update.py @@ -189,6 +189,8 @@ class TOKENIZER_TYPE(IntEnum): {"name": "qwen2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/Qwen/Qwen3-Embedding-0.6B", "chkhsh": "d4540891389ea895b53b399da6ac824becc30f2fba0e9ddbb98f92e55ca0e97c"}, {"name": "qwen35", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/openbmb/MiniCPM-V-4_6", "chkhsh": "1444df51289cfa8063b96f0e62b1125440111bc79a52003ea14b6eac7016fd5f"}, {"name": "grok-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/alvarobartt/grok-2-tokenizer", "chkhsh": "66b8d4e19ab16c3bfd89bce5d785fb7e0155e8648708a1f42077cb9fe002c273"}, + # muse-spark uses the l4_202k tokenizer, whose pre-tokenizer regex is gpt-4o's + {"name": "gpt-4o", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B", "chkhsh": "f5c965e4198da6e7525072d87d1ecf4c51d981fd486c7c4ec728ef9cec7e03fc"}, # jina-v2-de variants {"name": "jina-v2-de", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/aari1995/German_Semantic_V3", "chkhsh": "b3d1dd861f1d4c5c0d2569ce36baf3f90fe8a102db3de50dd71ff860d91be3df"}, {"name": "gpt-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/evilfreelancer/ruGPT3XL", "chkhsh": "0fe1cf6eda062318a1af7270f3331a85c539a01778ff948e24388e949c5282f4"}, diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a39d046fe174..a3a9c478b6e5 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -132,6 +132,7 @@ class LLM: EXPERTS_PER_GROUP = "{arch}.experts_per_group" MOE_EVERY_N_LAYERS = "{arch}.moe_every_n_layers" MOE_LATENT_SIZE = "{arch}.moe_latent_size" + POST_NORM_EPS = "{arch}.post_norm_epsilon" NEXTN_PREDICT_LAYERS = "{arch}.nextn_predict_layers" NUM_DEEPSTACK_LAYERS = "{arch}.n_deepstack_layers" DEEPSTACK_MAPPING = "{arch}.deepstack_mapping" @@ -466,6 +467,9 @@ class ClipAudio: SUBSAMPLING_FACTOR = "clip.audio.subsampling_factor" CHUNK_SIZE = "clip.audio.chunk_size" CONV_KERNEL_SIZE = "clip.audio.conv_kernel_size" + HIDDEN_SIZES = "clip.audio.hidden_sizes" # muse-spark, per stage + PATCH_SIZES = "clip.audio.patch_sizes" # muse-spark, per stage + DEPTHS = "clip.audio.depths" # muse-spark, per stage MAX_POS_EMB = "clip.audio.max_pos_emb" FEATURE_LAYERS = "clip.audio.feature_layer" # Granite Speech Plus RVQ_NUM_QUANTIZERS = "clip.audio.rvq.num_quantizers" @@ -611,6 +615,7 @@ class MODEL_ARCH(IntEnum): OLMO2 = auto() OLMOE = auto() MUSE_GLIMMER = auto() + MUSE_SPARK = auto() OPENELM = auto() ARCTIC = auto() DEEPSEEK = auto() @@ -743,6 +748,8 @@ class MODEL_TENSOR(IntEnum): ATTN_ROT_EMBD = auto() ATTN_SINKS = auto() ATTN_GATE = auto() + ATTN_RESID_ALPHA = auto() + ATTN_RESID_BETA = auto() FFN_GATE_INP = auto() FFN_GATE_INP_SHEXP = auto() FFN_NORM = auto() @@ -771,6 +778,10 @@ class MODEL_TENSOR(IntEnum): FFN_GATE_TID2EID = auto() MOE_LATENT_DOWN = auto() # nemotron 3 super MOE_LATENT_UP = auto() # nemotron 3 super + FFN_RESID_ALPHA = auto() + FFN_RESID_BETA = auto() + FFN_PRE_EXP_NORM = auto() + FFN_POST_EXP_NORM = auto() ATTN_Q_NORM = auto() ATTN_K_NORM = auto() LAYER_OUT_NORM = auto() @@ -1121,6 +1132,8 @@ class MODEL_TENSOR(IntEnum): A_ENC_EMBD_TO_LOGITS = auto() # lfm2 A_ENC_INP_PROJ = auto() # gemma4 A_ENC_CONV1D = auto() + A_ENC_PATCH_PROJ = auto() # muse-spark + A_ENC_PATCH_NORM = auto() # muse-spark A_ENC_CONV1D_NORM = auto() # gemma3n A_ENC_CONV2D = auto() A_ENC_CONV_OUT = auto() @@ -1415,6 +1428,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.OLMO2: "olmo2", MODEL_ARCH.OLMOE: "olmoe", MODEL_ARCH.MUSE_GLIMMER: "muse-glimmer", + MODEL_ARCH.MUSE_SPARK: "muse-spark", MODEL_ARCH.OPENELM: "openelm", MODEL_ARCH.ARCTIC: "arctic", MODEL_ARCH.DEEPSEEK: "deepseek", @@ -1544,6 +1558,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.ATTN_ROT_EMBD: "blk.{bid}.attn_rot_embd", MODEL_TENSOR.ATTN_SINKS: "blk.{bid}.attn_sinks", MODEL_TENSOR.ATTN_GATE: "blk.{bid}.attn_gate", + MODEL_TENSOR.ATTN_RESID_ALPHA: "blk.{bid}.attn_resid_alpha", + MODEL_TENSOR.ATTN_RESID_BETA: "blk.{bid}.attn_resid_beta", MODEL_TENSOR.ATTN_Q_NORM: "blk.{bid}.attn_q_norm", MODEL_TENSOR.ATTN_K_NORM: "blk.{bid}.attn_k_norm", MODEL_TENSOR.ATTN_OUT_NORM: "blk.{bid}.attn_output_norm", @@ -1576,6 +1592,10 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_GATE_TID2EID: "blk.{bid}.ffn_gate_tid2eid", MODEL_TENSOR.MOE_LATENT_DOWN: "blk.{bid}.ffn_latent_down", # nemotron 3 super MODEL_TENSOR.MOE_LATENT_UP: "blk.{bid}.ffn_latent_up", # nemotron 3 super + MODEL_TENSOR.FFN_RESID_ALPHA: "blk.{bid}.ffn_resid_alpha", + MODEL_TENSOR.FFN_RESID_BETA: "blk.{bid}.ffn_resid_beta", + MODEL_TENSOR.FFN_PRE_EXP_NORM: "blk.{bid}.ffn_pre_exp_norm", + MODEL_TENSOR.FFN_POST_EXP_NORM: "blk.{bid}.ffn_post_exp_norm", MODEL_TENSOR.LAYER_OUT_NORM: "blk.{bid}.layer_output_norm", MODEL_TENSOR.LAYER_OUT_SCALE: "blk.{bid}.layer_output_scale", MODEL_TENSOR.PER_LAYER_TOKEN_EMBD: "per_layer_token_embd", # gemma3n @@ -1921,6 +1941,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_ENC_EMBD_TO_LOGITS: "a.embd_to_logits", MODEL_TENSOR.A_ENC_INP_PROJ: "a.input_projection", MODEL_TENSOR.A_ENC_CONV1D: "a.conv1d.{bid}", + MODEL_TENSOR.A_ENC_PATCH_PROJ: "a.patch.{bid}", + MODEL_TENSOR.A_ENC_PATCH_NORM: "a.patch.{bid}.norm", MODEL_TENSOR.A_ENC_CONV2D: "a.conv2d.{bid}", MODEL_TENSOR.A_ENC_CONV_OUT: "a.conv_out", MODEL_TENSOR.A_ENC_CONV1D_NORM: "a.conv1d.{bid}.norm", @@ -2273,6 +2295,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_ENC_EMBD_TO_LOGITS, MODEL_TENSOR.A_ENC_INP_PROJ, MODEL_TENSOR.A_ENC_CONV1D, + MODEL_TENSOR.A_ENC_PATCH_PROJ, + MODEL_TENSOR.A_ENC_PATCH_NORM, MODEL_TENSOR.A_ENC_CONV2D, MODEL_TENSOR.A_ENC_CONV_OUT, MODEL_TENSOR.A_ENC_CONV1D_NORM, @@ -3979,6 +4003,34 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_PRE_NORM, MODEL_TENSOR.FFN_POST_NORM, ], + MODEL_ARCH.MUSE_SPARK: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_RESID_ALPHA, + MODEL_TENSOR.ATTN_RESID_BETA, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.FFN_RESID_ALPHA, + MODEL_TENSOR.FFN_RESID_BETA, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.MOE_LATENT_DOWN, + MODEL_TENSOR.FFN_PRE_EXP_NORM, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_POST_EXP_NORM, + MODEL_TENSOR.MOE_LATENT_UP, + ], MODEL_ARCH.OPENELM: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, @@ -6283,6 +6335,7 @@ class VisionProjectorType: COHERE2V = "cohere2v" D1OMNI_V = "d1omni_v" # lfm2 vision, without separator tokens D1OMNI_A = "d1omni_a" # lfm2a audio, with a residual block after the projector + MUSE_SPARK_AUDIO = "muse-spark-audio" # Items here are (block size, type size) diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index cc2bd8e7b3b8..a1b1f1263e66 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -924,6 +924,9 @@ def add_moe_every_n_layers(self, value: int) -> None: def add_moe_latent_size(self, value: int) -> None: self.add_uint32(Keys.LLM.MOE_LATENT_SIZE.format(arch=self.arch), value) + def add_post_norm_eps(self, value: float) -> None: + self.add_float32(Keys.LLM.POST_NORM_EPS.format(arch=self.arch), value) + def add_nextn_predict_layers(self, count: int) -> None: self.add_uint32(Keys.LLM.NEXTN_PREDICT_LAYERS.format(arch=self.arch), count) @@ -1591,6 +1594,15 @@ def add_audio_rvq_codebook_size(self, values: Sequence[int]) -> None: def add_audio_wa_pattern_mode(self, modes: Sequence[int]) -> None: self.add_array(Keys.ClipAudio.WA_PATTERN_MODE, modes) + def add_audio_hidden_sizes(self, values: Sequence[int]) -> None: + self.add_array(Keys.ClipAudio.HIDDEN_SIZES, values) + + def add_audio_patch_sizes(self, values: Sequence[int]) -> None: + self.add_array(Keys.ClipAudio.PATCH_SIZES, values) + + def add_audio_depths(self, values: Sequence[int]) -> None: + self.add_array(Keys.ClipAudio.DEPTHS, values) + def add_audio_window_size(self, value: int) -> None: self.add_uint32(Keys.ClipAudio.WINDOW_SIZE, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 63301cc114a9..ecea46463c37 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -530,6 +530,7 @@ class TensorNameMap: MODEL_TENSOR.FFN_EXP_PROBS_B: ( "model.layers.{bid}.mlp.gate.e_score_correction", # deepseek-v3 dots1 + "model.layers.{bid}.mlp.gate.e_score_correction_bias", # muse-spark "model.layers.{bid}.mlp.moe_statics.e_score_correction", # ernie4.5-moe "model.layers.{bid}.mlp.gate.expert_bias", # bailingmoe2 "model.layers.{bid}.mlp.expert_bias", # afmoe @@ -674,10 +675,20 @@ class TensorNameMap: MODEL_TENSOR.MOE_LATENT_DOWN: ( "backbone.layers.{bid}.mixer.fc1_latent_proj", # nemotron 3 super + "model.layers.{bid}.mlp.pre_expert_proj", # muse-spark ), MODEL_TENSOR.MOE_LATENT_UP: ( "backbone.layers.{bid}.mixer.fc2_latent_proj", # nemotron 3 super + "model.layers.{bid}.mlp.post_expert_proj", # muse-spark + ), + + MODEL_TENSOR.FFN_PRE_EXP_NORM: ( + "model.layers.{bid}.mlp.pre_expert_norm", # muse-spark + ), + + MODEL_TENSOR.FFN_POST_EXP_NORM: ( + "model.layers.{bid}.mlp.experts.post_expert_norm", # muse-spark ), # Feed-forward down diff --git a/models/templates/muse-spark.jinja b/models/templates/muse-spark.jinja new file mode 100644 index 000000000000..8fc6c8df225a --- /dev/null +++ b/models/templates/muse-spark.jinja @@ -0,0 +1,228 @@ +{%- set tools = tools | default(none) -%} +{%- set add_generation_prompt = add_generation_prompt | default(false) -%} +{%- set reasoning_effort = reasoning_effort | default(none) -%} +{%- set disable_reasoning = disable_reasoning | default(false) -%} +{%- set cutoff_date = cutoff_date | default('2026-01-04') -%} +{%- set current_date = current_date | default(strftime_now('%Y-%m-%d') if strftime_now is defined else none) -%} +{%- set reasoning_effort_map = {'minimal': 1, 'low': 16, 'medium': 64, 'high': 256} -%} + +{#- Neutralise control sequences in caller-supplied text so it cannot be + tokenized as structure. "<|" is broken because every special token has the + form "<|...|>"; the bare media placeholders are broken because the + multimodal processor rewrites them positionally and a literal one in text + would desynchronise placeholders from the attached media. -#} +{%- macro esc(value) -%} + {{- value | string + | replace('<|', '< |') + | replace('<image>', '< image>') + | replace('<video>', '< video>') + | replace('<audio>', '< audio>') -}} +{%- endmacro -%} + +{%- macro esc_header(value) -%} + {{- esc(value) | replace('\r', ' ') | replace('\n', ' ') -}} +{%- endmacro -%} + +{%- macro esc_attr(value) -%} + {{- esc_header(value) + | replace('&', '&') + | replace('"', '"') + | replace("'", ''') + | replace('<', '<') + | replace('>', '>') -}} +{%- endmacro -%} + +{%- macro render_content(content) -%} + {%- if content is string -%} + {{- esc(content | trim) -}} + {%- elif content is not none -%} + {%- for part in content -%} + {%- set part_type = part.get('type', '') -%} + {%- if part_type in ['image', 'input_image', 'image_url'] -%} + {{- '<|patch|>' -}} + {%- elif part_type in ['video', 'input_video', 'video_url'] -%} + {{- '<|video|>' -}} + {%- elif part_type in ['audio', 'input_audio', 'audio_url'] -%} + {{- '<|audio|>' -}} + {%- elif part_type in ['text', 'input_text', 'output_text'] -%} + {{- esc(part.get('text', '') | trim) -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} +{%- endmacro -%} + +{%- macro render_atem(tc) -%} + {%- set args = tc.function.arguments -%} + {%- if args is not mapping -%} + {{- raise_exception('This chat template requires tool_call.function.arguments to be a dict (mapping). The jinja sandbox has no JSON parser, so deserialize the arguments string before calling apply_chat_template.') -}} + {%- endif -%} + {{- '<atem:function_calls>\n<atem:invoke name="' -}}{{- esc_attr(tc.function.name) -}}{{- '">\n' -}} + {%- for k, v in args.items() -%} + {{- '<atem:parameter name="' -}}{{- esc_attr(k) -}}{{- '">' -}} + {%- if v is boolean -%} + {%- if v -%} + true + {%- else -%} + false + {%- endif -%} + {%- elif v is none -%} + null + {%- elif v is mapping or (v is iterable and v is not string) -%} + {{- esc(v | tojson) -}} + {%- else -%} + {{- esc(v) -}} + {%- endif -%} + {{- '</atem:parameter>\n' -}} + {%- endfor -%} + {{- '</atem:invoke>\n</atem:function_calls>' -}} +{%- endmacro -%} + +{%- macro render_tool_defs(tools) -%} + {{- 'In this environment you have access to a set of tools you can use to answer the user\'s question.\n\n' -}} + {{- 'You can invoke a function by writing a "<atem:function_calls>" block like the following as part of your reply to the user:\n' -}} + {{- '<atem:function_calls>\n<atem:invoke name="$FUNCTION_NAME">\n<atem:parameter name="$PARAMETER_NAME">$PARAMETER_VALUE</atem:parameter>\n...\n</atem:invoke>\n</atem:function_calls>\n\n' -}} + {{- 'String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions.\n' -}} + {{- 'Here are the functions available in JSONSchema format:\n' -}} + {{- '// Tool metadata\n{"name": "", "description": ""}\n' -}} + {{- '// Function schemas' -}} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {{- '\n{"name": ' -}}{{- fn.name | tojson -}} + {{- ', "description": ' -}}{{- (fn.description if fn.description is defined and fn.description else '') | tojson -}} + {{- ', "parameters": ' -}}{{- (fn.parameters if fn.parameters is defined and fn.parameters else {}) | tojson -}}{{- '}' -}} + {%- endfor -%} + {{- '\n\n' -}} + {{- 'Here are examples of how to call functions in the tool set. To make parallel tool calls, emit each call as its own message (one tool call per message) in consecutive messages, as shown below:\n' -}} + {{- '(If the tool namespace is not specified, invoke the function directly as `example_function_name_1` rather than `example_tool_name_1.example_function_name_1`)\n\n' -}} + {{- 'to=example_tool_name_1.example_function_name_1\n\n' -}} + {{- '<atem:function_calls>\n<atem:invoke name="example_tool_name_1.example_function_name_1">\n' -}} + {{- '<atem:parameter name="example_parameter_1">value_1</atem:parameter>\n' -}} + {{- '</atem:invoke>\n</atem:function_calls>\n\n' -}} + {{- 'to=example_tool_name_2.example_function_name_2\n\n' -}} + {{- '<atem:function_calls>\n<atem:invoke name="example_tool_name_2.example_function_name_2">\n' -}} + {{- '<atem:parameter name="example_parameter_2">This is the value for the second parameter\nthat can span\n"multiple" lines\n</atem:parameter>\n' -}} + {{- '<atem:parameter name="example_parameter_3">value_3</atem:parameter>\n' -}} + {{- '</atem:invoke>\n</atem:function_calls>' -}} +{%- endmacro -%} + +{%- macro render_valid_recipients(tools) -%} + {%- set rns = namespace(recipients=[]) -%} + {{- 'Use the appropriate recipient for each message:\n\n' -}} + {%- if not disable_reasoning -%} + {%- set rns.recipients = rns.recipients + ['"self"'] -%} + {{- '- "self": private reasoning and tool planning.\n' -}} + {%- endif -%} + {%- set rns.recipients = rns.recipients + ['"commentary"'] -%} + {{- '- "commentary": user-visible intermediate updates while the assistant will continue working, including messages sent to the user before or between tool calls.\n' -}} + {%- if tools -%} + {%- set first_tool = tools | first -%} + {%- set first_fn = (first_tool.function if first_tool.function is defined else first_tool).name -%} + {%- set rns.recipients = rns.recipients + ['"[func_name] (e.g., ' ~ esc_header(first_fn) ~ ')"'] -%} + {{- '- "tool" or "[tool name]": the name of the tool to call.\n' -}} + {%- endif -%} + {%- set rns.recipients = rns.recipients + ['"user"'] -%} + {{- '- "user": messages that end the assistant turn, such as a completed response or clarification question that waits for the user\'s reply; do not use for pre-tool updates or partial responses.\n' -}} + {{- '\n# Valid recipients: ' ~ rns.recipients | join(', ') ~ '.' -}} +{%- endmacro -%} + +{{- bos_token -}} +{%- set ns = namespace(instructions=none, rest=messages) -%} +{%- if messages and messages[0]['role'] in ['system', 'developer'] -%} + {%- set ns.instructions = render_content(messages[0]['content']) + | replace('Reasoning effort', 'Reasoning strength') + | replace('Reasoning Effort', 'Reasoning Strength') + | replace('reasoning effort', 'reasoning strength') + | replace('REASONING EFFORT', 'REASONING STRENGTH') -%} + {%- set ns.rest = messages[1:] -%} +{%- endif -%} +{%- set sys = namespace(parts=['You are a helpful AI assistant.', 'Knowledge cutoff: ' ~ cutoff_date ~ '.']) -%} +{%- if current_date -%} + {%- set sys.parts = sys.parts + ['Current date: ' ~ current_date ~ '.'] -%} +{%- endif -%} +{%- if ns.instructions -%} + {%- set sys.parts = sys.parts + [ns.instructions] -%} +{%- endif -%} +{%- if tools -%} + {%- set sys.parts = sys.parts + [render_tool_defs(tools) | trim] -%} +{%- endif -%} +{%- if not disable_reasoning and reasoning_effort in reasoning_effort_map and 'reasoning strength:' not in ((ns.instructions if ns.instructions else '') | lower) -%} + {%- set sys.parts = sys.parts + ['Reasoning strength: ' ~ reasoning_effort_map[reasoning_effort] ~ '.'] -%} +{%- endif -%} +{%- set sys.parts = sys.parts + [render_valid_recipients(tools) | trim] -%} +{{- '<|start|>system<|message|>' ~ sys.parts | join('\n\n') ~ '<|eot|>' -}} +{%- for message in ns.rest -%} + {%- set role = message['role'] -%} + {%- set end_token = '<|eom|>' if (not loop.last and ns.rest[loop.index0 + 1]['role'] == role) else '<|eot|>' -%} + {%- if role in ['system', 'developer'] -%} + {{- '<|start|>' + role + '<|message|>' -}} + {{- render_content(message['content']) -}} + {{- '<|eot|>' -}} + {%- elif role == 'user' -%} + {{- '<|start|>user<|message|>' -}} + {%- if message.get('render_timestamp') and message.get('timestamp') -%} + {%- set ts = message['timestamp'] -%} + {%- if ts is string -%} + {{- '[' ~ esc_header(ts) ~ '] ' -}} + {%- elif ts.strftime is defined -%} + {{- '[' ~ (ts.strftime('%a %Y-%m-%d %H:%M %Z') | trim) ~ '] ' -}} + {%- endif -%} + {%- endif -%} + {{- render_content(message['content']) -}} + {{- '<|eot|>' -}} + {%- elif role == 'tool' -%} + {%- set tname = message.get('name') -%} + {%- if not tname -%} + {%- set tcid = message.get('tool_call_id') -%} + {%- set tns = namespace(name=tcid if tcid else '') -%} + {%- for m in messages -%} + {%- if m.get('tool_calls') -%} + {%- for tc in m['tool_calls'] -%} + {%- if tcid is not none and tc.id is defined and tc.id == tcid -%} + {%- set tns.name = tc.function.name -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endfor -%} + {%- set tname = tns.name -%} + {%- endif -%} + {%- if not tname -%} + {{- raise_exception('tool message has no name and no matching tool_call_id') -}} + {%- endif -%} + {{- '<|start|>tool ' -}}{{- esc_header(tname) -}}{{- '<|message|><tool_output name="' -}}{{- esc_attr(tname) -}}{{- '">\n' -}} + {{- render_content(message['content']) -}} + {{- '\n</tool_output><|eot|>' -}} + {%- elif role == 'assistant' -%} + {%- set reasoning = message.get('reasoning_content') or message.get('reasoning') -%} + {%- if reasoning -%} + {{- '<|start|>assistant to=self<|message|>' -}}{{- esc(reasoning) -}}{{- '<|eom|>' -}} + {%- endif -%} + {%- set body = render_content(message.get('content')) -%} + {%- if message.get('tool_calls') -%} + {%- if body -%} + {{- '<|start|>assistant to=commentary<|message|>' + body + '<|eom|>' -}} + {%- endif -%} + {%- for tc in message['tool_calls'] -%} + {{- '<|start|>assistant to=' -}}{{- esc_header(tc.function.name) -}}{{- '<|message|>' -}} + {{- render_atem(tc) -}} + {{- (end_token if loop.last else '<|eom|>') -}} + {%- endfor -%} + {%- elif body or not reasoning -%} + {%- set recipient = message.get('recipient') or 'user' -%} + {%- if recipient == 'user' -%} + {%- set end_turn = true -%} + {%- elif recipient in ['self', 'commentary'] -%} + {%- set end_turn = false -%} + {%- elif message.get('end_turn') is not none -%} + {%- set end_turn = message.get('end_turn') -%} + {%- else -%} + {%- set end_turn = false -%} + {%- endif -%} + {{- '<|start|>assistant to=' -}}{{- esc_header(recipient) -}}{{- '<|message|>' -}} + {{- body -}} + {{- ('<|eot|>' if end_turn else '<|eom|>') -}} + {%- endif -%} + {%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- '<|start|>assistant' -}} +{%- endif -%} diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index d1dd4cbaeb67..73c6d51d9595 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -77,6 +77,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = { { LLM_ARCH_OLMO2, "olmo2" }, { LLM_ARCH_OLMOE, "olmoe" }, { LLM_ARCH_MUSE_GLIMMER, "muse-glimmer" }, + { LLM_ARCH_MUSE_SPARK, "muse-spark" }, { LLM_ARCH_OPENELM, "openelm" }, { LLM_ARCH_ARCTIC, "arctic" }, { LLM_ARCH_DEEPSEEK, "deepseek" }, @@ -228,6 +229,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = { { LLM_KV_EXPERTS_PER_GROUP, "%s.experts_per_group" }, { LLM_KV_MOE_EVERY_N_LAYERS, "%s.moe_every_n_layers" }, { LLM_KV_MOE_LATENT_SIZE, "%s.moe_latent_size" }, + { LLM_KV_POST_NORM_EPS, "%s.post_norm_epsilon" }, { LLM_KV_NEXTN_PREDICT_LAYERS, "%s.nextn_predict_layers" }, { LLM_KV_NUM_DEEPSTACK_LAYERS, "%s.n_deepstack_layers" }, { LLM_KV_DEEPSTACK_MAPPING, "%s.deepstack_mapping" }, @@ -500,6 +502,8 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = { { LLM_TENSOR_ATTN_Q_NORM, "blk.%d.attn_q_norm" }, { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, { LLM_TENSOR_ATTN_GATE, "blk.%d.attn_gate" }, + { LLM_TENSOR_ATTN_RESID_ALPHA, "blk.%d.attn_resid_alpha" }, + { LLM_TENSOR_ATTN_RESID_BETA, "blk.%d.attn_resid_beta" }, { LLM_TENSOR_FFN_POST_NORM, "blk.%d.post_ffw_norm" }, { LLM_TENSOR_FFN_POST_NORM_1, "blk.%d.post_ffw_norm_1" }, { LLM_TENSOR_FFN_POST_NORM_2, "blk.%d.post_ffw_norm_2" }, @@ -511,6 +515,10 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = { { LLM_TENSOR_FFN_EXP_PROBS_B_VL, "blk.%d.exp_probs_b_vl" }, { LLM_TENSOR_FFN_LATENT_DOWN, "blk.%d.ffn_latent_down" }, { LLM_TENSOR_FFN_LATENT_UP, "blk.%d.ffn_latent_up" }, + { LLM_TENSOR_FFN_PRE_EXP_NORM, "blk.%d.ffn_pre_exp_norm" }, + { LLM_TENSOR_FFN_POST_EXP_NORM, "blk.%d.ffn_post_exp_norm" }, + { LLM_TENSOR_FFN_RESID_ALPHA, "blk.%d.ffn_resid_alpha" }, + { LLM_TENSOR_FFN_RESID_BETA, "blk.%d.ffn_resid_beta" }, { LLM_TENSOR_ATTN_NORM_2, "blk.%d.attn_norm_2" }, { LLM_TENSOR_ATTN_QKV, "blk.%d.attn_qkv" }, { LLM_TENSOR_LAYER_OUT_NORM, "blk.%d.layer_output_norm" }, @@ -839,6 +847,8 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = { {LLM_TENSOR_ATTN_QKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_OUT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_ATTN_RESID_ALPHA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_ATTN_RESID_BETA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_FFN_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_FFN_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_FFN_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, @@ -1083,6 +1093,10 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = { // latent projections feed ggml_mul_mat, the buft probe must use MUL_MAT to keep them on GPU {LLM_TENSOR_FFN_LATENT_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_FFN_LATENT_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_PRE_EXP_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_POST_EXP_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_RESID_ALPHA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_RESID_BETA, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_MASKED_EMBD_CENTROIDS, {LLM_TENSOR_LAYER_INPUT, GGML_OP_NONE}}, {LLM_TENSOR_MASKED_EMBD_ORDERING, {LLM_TENSOR_LAYER_INPUT, GGML_OP_NONE}}, // eagle3 diff --git a/src/llama-arch.h b/src/llama-arch.h index 4e8e53421af3..4b607fec7100 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -82,6 +82,7 @@ enum llm_arch { LLM_ARCH_OLMO2, LLM_ARCH_OLMOE, LLM_ARCH_MUSE_GLIMMER, + LLM_ARCH_MUSE_SPARK, LLM_ARCH_OPENELM, LLM_ARCH_ARCTIC, LLM_ARCH_DEEPSEEK, @@ -233,6 +234,7 @@ enum llm_kv { LLM_KV_EXPERTS_PER_GROUP, LLM_KV_MOE_EVERY_N_LAYERS, LLM_KV_MOE_LATENT_SIZE, + LLM_KV_POST_NORM_EPS, LLM_KV_NEXTN_PREDICT_LAYERS, LLM_KV_NUM_DEEPSTACK_LAYERS, LLM_KV_DEEPSTACK_MAPPING, @@ -500,6 +502,8 @@ enum llm_tensor { LLM_TENSOR_ATTN_ROT_EMBD, LLM_TENSOR_ATTN_SINKS, LLM_TENSOR_ATTN_GATE, + LLM_TENSOR_ATTN_RESID_ALPHA, + LLM_TENSOR_ATTN_RESID_BETA, LLM_TENSOR_FFN_GATE_INP, LLM_TENSOR_FFN_GATE_INP_SHEXP, LLM_TENSOR_FFN_NORM, @@ -529,6 +533,10 @@ enum llm_tensor { LLM_TENSOR_FFN_EXP_PROBS_B_VL, LLM_TENSOR_FFN_LATENT_DOWN, LLM_TENSOR_FFN_LATENT_UP, + LLM_TENSOR_FFN_PRE_EXP_NORM, + LLM_TENSOR_FFN_POST_EXP_NORM, + LLM_TENSOR_FFN_RESID_ALPHA, + LLM_TENSOR_FFN_RESID_BETA, LLM_TENSOR_ATTN_Q_NORM, LLM_TENSOR_ATTN_K_NORM, LLM_TENSOR_LAYER_OUT_NORM, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 6cadda9716ff..e8e249b21a8d 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2093,7 +2093,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, - ggml_tensor * selected_experts_in) const { + ggml_tensor * selected_experts_in, + ggml_tensor * post_exp_norm, + float post_exp_norm_eps) const { return build_moe_ffn( cur, gate_inp, /* gate_inp_b */ nullptr, @@ -2114,7 +2116,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn( up_exps_s, gate_exps_s, down_exps_s, - selected_experts_in + selected_experts_in, + post_exp_norm, + post_exp_norm_eps ); } @@ -2142,7 +2146,9 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, - ggml_tensor * selected_experts_in) const { + ggml_tensor * selected_experts_in, + ggml_tensor * post_exp_norm, + float post_exp_norm_eps) const { const int64_t n_embd = cur->ne[0]; const int64_t n_tokens = cur->ne[1]; const bool weight_before_ffn = arch == LLM_ARCH_LLAMA4; // for llama4, we apply the sigmoid-ed weights before the FFN @@ -2448,6 +2454,13 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(experts, "ffn_moe_down_biased", il); } + // an expert output can be scaled by a norm weight, then normalized, before the router weights apply + if (post_exp_norm) { + experts = ggml_mul(ctx0, experts, post_exp_norm); + experts = ggml_rms_norm(ctx0, experts, post_exp_norm_eps); + cb(experts, "ffn_moe_post_exp_norm", il); + } + if (!weight_before_ffn) { experts = ggml_mul(ctx0, experts, weights); cb(experts, "ffn_moe_weighted", il); diff --git a/src/llama-graph.h b/src/llama-graph.h index e853a026955f..1a3698cad93f 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1178,7 +1178,9 @@ struct llm_graph_context { ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, ggml_tensor * down_exps_s = nullptr, - ggml_tensor * selected_experts_in = nullptr) const; + ggml_tensor * selected_experts_in = nullptr, + ggml_tensor * post_exp_norm = nullptr, + float post_exp_norm_eps = 1e-5f) const; ggml_tensor * build_moe_ffn( ggml_tensor * cur, @@ -1204,7 +1206,9 @@ struct llm_graph_context { ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, ggml_tensor * down_exps_s = nullptr, - ggml_tensor * selected_experts_in = nullptr) const; + ggml_tensor * selected_experts_in = nullptr, + ggml_tensor * post_exp_norm = nullptr, + float post_exp_norm_eps = 1e-5f) const; // the slots of the selected experts in the MoE cache, nullptr if the experts of the layer are not read from the cache ggml_tensor * build_moe_cache_slots( diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 922164a544b8..bc0d492152db 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -141,6 +141,9 @@ struct llama_hparams { float f_norm_rms_eps; float f_norm_group_eps; + // post norms use their own epsilon + float f_post_norm_eps = 0.0f; + float f_attn_logit_softcapping = 50.0f; float f_router_logit_softcapping = 30.0f; float f_final_logit_softcapping = 30.0f; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index e011e24ea605..46b1d2ae2744 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -279,6 +279,7 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, hparams.n_embd_head_v_full); add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + add_kv(LLM_KV_POST_NORM_EPS, hparams.f_post_norm_eps); add_kv(LLM_KV_ATTENTION_GROUPNORM_EPS, hparams.f_norm_group_eps); add_kv(LLM_KV_ATTENTION_GROUPNORM_GROUPS, hparams.n_norm_groups); add_kv(LLM_KV_ATTENTION_VALUE_EXPERT_COUNT, hparams.n_value_expert); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index ded790b2b86d..2660bf38a2f0 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -188,6 +188,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_olmoe(params); case LLM_ARCH_MUSE_GLIMMER: return new llama_model_muse_glimmer(params); + case LLM_ARCH_MUSE_SPARK: + return new llama_model_muse_spark(params); case LLM_ARCH_OPENELM: return new llama_model_openelm(params); case LLM_ARCH_GPTNEOX: @@ -991,6 +993,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_16x3_8B: return "16x3.8B"; case LLM_TYPE_10B_128x3_66B: return "10B+128x3.66B"; case LLM_TYPE_57B_A14B: return "57B.A14B"; + case LLM_TYPE_816B_A42B: return "816B.A42B"; case LLM_TYPE_17B_16E: return "17Bx16E (Scout)"; case LLM_TYPE_17B_128E: return "17Bx128E (Maverick)"; case LLM_TYPE_A13B: return "A13B"; @@ -3492,6 +3495,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_DEEPSEEK32: case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_MUSE_GLIMMER: + case LLM_ARCH_MUSE_SPARK: case LLM_ARCH_PLM: case LLM_ARCH_CHATGLM: case LLM_ARCH_GRANITE: diff --git a/src/llama-model.h b/src/llama-model.h index bcfbf1e2cb39..66f9129e68a1 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -114,6 +114,7 @@ enum llm_type { LLM_TYPE_16x3_8B, LLM_TYPE_10B_128x3_66B, LLM_TYPE_57B_A14B, + LLM_TYPE_816B_A42B, LLM_TYPE_17B_16E, // llama4 Scout LLM_TYPE_17B_128E, // llama4 Maverick LLM_TYPE_A13B, @@ -361,6 +362,16 @@ struct llama_layer { struct ggml_tensor * ffn_latent_down = nullptr; struct ggml_tensor * ffn_latent_up = nullptr; + // ff MoE norms around the experts, in latent space + struct ggml_tensor * ffn_pre_exp_norm = nullptr; + struct ggml_tensor * ffn_post_exp_norm = nullptr; + + // norm preserving residual coefficients + struct ggml_tensor * attn_resid_alpha = nullptr; + struct ggml_tensor * attn_resid_beta = nullptr; + struct ggml_tensor * ffn_resid_alpha = nullptr; + struct ggml_tensor * ffn_resid_beta = nullptr; + // ff shared expert (shexp) struct ggml_tensor * ffn_gate_inp_shexp = nullptr; struct ggml_tensor * ffn_gate_shexp = nullptr; diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 9c60667963b9..b8f96bcef30f 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -31,9 +31,11 @@ enum class tensor_category { ATTENTION_QKV, ATTENTION_KV_B, ATTENTION_OUTPUT, + ATTENTION_GATE, FFN_UP, FFN_GATE, FFN_DOWN, + FFN_LATENT, OUTPUT, OTHER }; @@ -147,6 +149,12 @@ static tensor_category tensor_get_category(const std::string & tensor_name) { if (tensor_name.find("ffn_gate") != std::string::npos) { return tensor_category::FFN_GATE; } + if (tensor_name.find("ffn_latent") != std::string::npos) { + return tensor_category::FFN_LATENT; + } + if (tensor_name.find("attn_gate.weight") != std::string::npos) { + return tensor_category::ATTENTION_GATE; + } if (tensor_name.find("ffn_down") != std::string::npos) { return tensor_category::FFN_DOWN; } @@ -490,6 +498,19 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type } } + // bumping to q8_0 is preferred since these tensors are not really big + // TODO: nemotron 3 super shares ffn_latent and afmoe, bailingmoe3, dots3note, hy-v4, kimi-k3 and laguna share attn_gate. + // Widening the guard is likely a win there too, but it changes their files, so it wants a perplexity comparison per architecture first. + if (arch == LLM_ARCH_MUSE_SPARK && + (category == tensor_category::FFN_LATENT || category == tensor_category::ATTENTION_GATE)) { + auto bits_per_weight = [](ggml_type t) { + return 8.0 * ggml_type_size(t) / ggml_blck_size(t); + }; + if (ggml_is_quantized(new_type) && bits_per_weight(new_type) < bits_per_weight(GGML_TYPE_Q8_0)) { + new_type = GGML_TYPE_Q8_0; + } + } + // for arches that share the same tensor between the token embeddings and the output, we quantize the token embeddings // with the quantization of the output tensor if (category == tensor_category::OUTPUT || (qs.has_tied_embeddings && category == tensor_category::TOKEN_EMBD)) { diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 9b5cc2e1c7b5..9fd5b9ef3f2a 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -3029,6 +3029,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { if (false || t.first == "<|eot_id|>" + || t.first == "<|eot|>" // muse-spark || t.first == "<|im_end|>" || t.first == "<|end|>" || t.first == "<|return|>" // o200k_harmony diff --git a/src/models/models.h b/src/models/models.h index 88fd07d973cd..51fdcfbb08c6 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1154,6 +1154,19 @@ struct llama_model_olmoe : public llama_model_base { }; +struct llama_model_muse_spark : public llama_model_base { + llama_model_muse_spark(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<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_muse_glimmer : public llama_model_base { llama_model_muse_glimmer(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/src/models/muse-spark.cpp b/src/models/muse-spark.cpp new file mode 100644 index 000000000000..5842ecb98110 --- /dev/null +++ b/src/models/muse-spark.cpp @@ -0,0 +1,238 @@ +#include "models.h" + +void llama_model_muse_spark::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_POST_NORM_EPS, hparams.f_post_norm_eps); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa); + ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping); + ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale); + ml.get_key(LLM_KV_MOE_LATENT_SIZE, hparams.moe_latent_size); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm); + ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all, false); + + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl); + + switch (hparams.n_layer()) { + case 62: type = LLM_TYPE_816B_A42B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_muse_spark::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t n_embd_moe = hparams.moe_latent_size; + const int64_t n_ff_exp = hparams.n_ff_exp(); + + 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}, 0); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + // Pre-attention norm. All block norms hold a delta folded at conversion time. + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + // Q/K/V/O projections. + 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); + + // QK-norm. Weights are synthesized at conversion to absorb the metap scale. + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); + + // Attention output gate: sigmoid(gate) * attn_out before o_proj. + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0); + + // Norm preserving residual: alpha scales the stream, beta the branch. + layer.attn_resid_alpha = create_tensor(tn(LLM_TENSOR_ATTN_RESID_ALPHA, "weight", i), {n_embd}, 0); + layer.attn_resid_beta = create_tensor(tn(LLM_TENSOR_ATTN_RESID_BETA, "weight", i), {n_embd}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0); + layer.ffn_resid_alpha = create_tensor(tn(LLM_TENSOR_FFN_RESID_ALPHA, "weight", i), {n_embd}, 0); + layer.ffn_resid_beta = create_tensor(tn(LLM_TENSOR_FFN_RESID_BETA, "weight", i), {n_embd}, 0); + + // Every layer routes. The router reads the full stream, the experts run + // in a latent space entered by ffn_latent_down and left by ffn_latent_up. + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + + layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, n_embd_moe}, 0); + layer.ffn_latent_up = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP, "weight", i), {n_embd_moe, n_embd}, 0); + + layer.ffn_pre_exp_norm = create_tensor(tn(LLM_TENSOR_FFN_PRE_EXP_NORM, "weight", i), {n_embd_moe}, 0); + layer.ffn_post_exp_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_EXP_NORM, "weight", i), {n_embd_moe}, 0); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd_moe, n_ff_exp, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd_moe, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd_moe, n_expert}, 0); + } +} + +llama_model_muse_spark::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_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + inpL = build_norm(inpL, nullptr, nullptr, LLM_NORM_RMS, -1); + cb(inpL, "embd_norm", -1); + + 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) { + res->t_layer_inp[il] = inpL; + + const float freq_base_l = model.get_rope_freq_base (cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + + ggml_tensor * inpSA = inpL; + + // RoPE runs on the SWA layers, NoPE on the global ones. + const bool use_rope = hparams.is_swa(il); + + cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + { + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, n_head_kv, il); + + ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, cur); + cb(gate, "attn_gate_proj", il); + + // attn_q_norm carries the metap scale, attn_k_norm is identity. + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il); + Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + cb(Kcur, "Kcur_normed", il); + + if (use_rope) { + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(Qcur, "Qcur_rope", il); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(Kcur, "Kcur_rope", il); + } + + cur = build_attn(inp_attn, + NULL, NULL, NULL, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "attn_out", il); + + // Weightless norm on the attention output, applied per head. + cur = ggml_reshape_3d(ctx0, cur, n_embd_head, n_head, n_tokens); + cur = ggml_rms_norm(ctx0, cur, hparams.f_norm_rms_eps); + cur = ggml_reshape_2d(ctx0, cur, n_embd_head * n_head, n_tokens); + cb(cur, "attn_o_norm", il); + + gate = ggml_sigmoid(ctx0, gate); + cb(gate, "attn_gate_sig", il); + cur = ggml_mul(ctx0, cur, gate); + cb(cur, "attn_gated", il); + + cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s); + cb(cur, "attn_o_proj", il); + } + + // Post-attention norm holds no weight. + cur = ggml_rms_norm(ctx0, cur, hparams.f_post_norm_eps); + cb(cur, "attn_post_norm", 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, + ggml_mul(ctx0, inpSA, model.layers[il].attn_resid_alpha), + ggml_mul(ctx0, cur, model.layers[il].attn_resid_beta)); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + // Enter the latent space, run the experts, leave it. + ggml_tensor * moe_inp = build_lora_mm(model.layers[il].ffn_latent_down, cur); + cb(moe_inp, "ffn_latent_down", il); + + moe_inp = build_norm(moe_inp, model.layers[il].ffn_pre_exp_norm, NULL, LLM_NORM_RMS, il); + cb(moe_inp, "ffn_pre_exp_norm", il); + + // The router reads the full stream, not the latent one. + ggml_tensor * router_logits = build_lora_mm(model.layers[il].ffn_gate_inp, cur); + cb(router_logits, "ffn_moe_logits", il); + + ggml_tensor * moe_out = build_moe_ffn(moe_inp, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + model.layers[il].ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + router_logits, + nullptr, nullptr, nullptr, nullptr, nullptr, + model.layers[il].ffn_post_exp_norm, hparams.f_post_norm_eps); + cb(moe_out, "ffn_moe_out", il); + + cur = build_lora_mm(model.layers[il].ffn_latent_up, moe_out); + cb(cur, "ffn_latent_up", il); + + // Post-FFN norm scales the input, then normalizes. + cur = ggml_mul(ctx0, cur, model.layers[il].ffn_post_norm); + cur = ggml_rms_norm(ctx0, cur, hparams.f_post_norm_eps); + cb(cur, "ffn_post_norm", il); + + cur = ggml_add(ctx0, + ggml_mul(ctx0, ffn_inp, model.layers[il].ffn_resid_alpha), + ggml_mul(ctx0, cur, model.layers[il].ffn_resid_beta)); + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + inpL = cur; + } + + cur = inpL; + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur, model.output_s); + cur = ggml_scale(ctx0, cur, hparams.f_logit_scale); + + if (hparams.f_final_logit_softcapping) { + cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping); + cur = ggml_tanh(ctx0, cur); + cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping); + } + + cb(cur, "result_output", -1); + res->t_logits = cur; +} + +std::unique_ptr<llm_graph_context> llama_model_muse_spark::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique<graph>(*this, params); +} diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 01b2220da4e4..49519eecab1d 100644 --- a/tests/test-chat.cpp +++ b/tests/test-chat.cpp @@ -6661,6 +6661,50 @@ static void test_template_output_peg_parsers(bool detailed_debug) { .run(); } + // Muse Spark format tests + { + auto tst = peg_tester("models/templates/muse-spark.jinja", detailed_debug); + + const std::string call_markup = + "<atem:function_calls>\n" + "<atem:invoke name=\"special_function\">\n" + "<atem:parameter name=\"arg1\">1</atem:parameter>\n" + "</atem:invoke>\n" + "</atem:function_calls>"; + + // A plain answer behaves as it does for Muse Glimmer + tst.test(" to=user<|message|>Hello, world!\nWhat's up?<|eot|>") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(message_assist) + .run(); + + // A commentary message is user facing content, not reasoning + tst.test(" to=commentary<|message|>Hello, world!\nWhat's up?<|eom|>" + "<|start|>assistant to=user<|message|>Done!<|eot|>") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect_content("Hello, world!\nWhat's up?Done!") + .run(); + + // Commentary announces the work, then the tool call follows in the same generation + tst.test(" to=commentary<|message|>Let me look that up.<|eom|>" + "<|start|>assistant to=special_function<|message|>" + + call_markup) + .tools({ special_function_tool }) + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect(message_with_content_and_tool_call("Let me look that up.", "special_function", + "{\"arg1\":1}")) + .run(); + + // Reasoning still reaches the reasoning channel with commentary in the picture + tst.test(" to=self<|message|>Thinking.<|eom|>" + "<|start|>assistant to=commentary<|message|>Working on it.<|eom|>" + "<|start|>assistant to=user<|message|>Hello!<|eot|>") + .reasoning_format(COMMON_REASONING_FORMAT_AUTO) + .expect_reasoning("Thinking.") + .expect_content("Working on it.Hello!") + .run(); + } + // GPT-OSS format tests { auto tst = peg_tester("models/templates/openai-gpt-oss-120b.jinja", detailed_debug); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 3dd8f0655ee9..2bfa1cf75bc1 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -185,7 +185,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { uint32_t n_head_kv = n_head; if (arch == LLM_ARCH_QWEN3) { n_head_kv = 1; // MQA coverage - } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE || arch == LLM_ARCH_MAION_CODER) { + } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_MUSE_SPARK || arch == LLM_ARCH_AFMOE || arch == LLM_ARCH_MAION_CODER) { n_head_kv = 2; // GQA coverage } const uint32_t n_embd_head = n_embd / n_head; @@ -303,7 +303,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { // SWA pattern: every 5th layer is full attention (matches E2B layer_types) ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5)); } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_SPARK2_5 || - arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_GRANITE_SWA || arch == LLM_ARCH_DOTS3NOTE || + arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_MUSE_SPARK || arch == LLM_ARCH_GRANITE_SWA || arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_MAPLE) { std::vector<uint32_t> pattern; pattern.reserve(n_layer); @@ -315,6 +315,15 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(2)); } + if (arch == LLM_ARCH_MUSE_SPARK) { + ms.add_kv(LLM_KV_POST_NORM_EPS, 1e-5f); + ms.add_kv(LLM_KV_ROPE_FREQ_BASE_SWA, 10000.0f); + ms.add_kv(LLM_KV_FINAL_LOGIT_SOFTCAPPING, 30.0f); + ms.add_kv(LLM_KV_MOE_LATENT_SIZE, n_embd/2); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); + } + // MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the // indexer head count is independent of the main attention head count. if (arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_GLM5_NEXT) { @@ -715,6 +724,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_MELLUM: case LLM_ARCH_LAGUNA: case LLM_ARCH_MAPLE: + case LLM_ARCH_MUSE_SPARK: return true; default: return false; diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 82f74bdc9665..edb5d3d89bbe 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -48,6 +48,7 @@ add_library(mtmd models/nemotron-v2-vl.cpp models/nemotron-omni.cpp models/muse-glimmer.cpp + models/muse-spark-audio.cpp models/llama4.cpp models/llava.cpp models/minicpmv.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 3d6c6f8c4801..ad291a5578bc 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -85,6 +85,9 @@ #define KEY_A_PROJ_STACK_FACTOR "clip.audio.projector.stack_factor" #define KEY_A_CHUNK_SIZE "clip.audio.chunk_size" #define KEY_A_CONV_KERNEL_SIZE "clip.audio.conv_kernel_size" +#define KEY_A_HIDDEN_SIZES "clip.audio.hidden_sizes" +#define KEY_A_PATCH_SIZES "clip.audio.patch_sizes" +#define KEY_A_DEPTHS "clip.audio.depths" #define KEY_A_MAX_POS_EMB "clip.audio.max_pos_emb" #define KEY_A_PROJ_WINDOW_SIZE "clip.audio.projector.window_size" #define KEY_A_PROJ_DOWNSAMPLE_RATE "clip.audio.projector.downsample_rate" @@ -369,6 +372,8 @@ #define TN_A_INP_PROJ "a.input_projection.%s" #define TN_A_CONV1D "a.conv1d.%d.%s" #define TN_A_CONV1D_NORM "a.conv1d.%d.norm.%s" +#define TN_A_PATCH "a.patch.%d.%s" +#define TN_A_PATCH_NORM "a.patch.%d.norm.%s" #define TN_A_OUT_PROJ "a.pre_encode.out.%s" #define TN_A_ATTN_PRE_NORM "%s.blk.%d.attn_pre_norm.%s" #define TN_A_ATTN_POST_NORM "%s.blk.%d.attn_post_norm.%s" @@ -515,6 +520,7 @@ enum projector_type { PROJECTOR_TYPE_POCKETTTS_GEN, PROJECTOR_TYPE_MUSE_GLIMMER, PROJECTOR_TYPE_COHERE2V, + PROJECTOR_TYPE_MUSE_SPARK_AUDIO, PROJECTOR_TYPE_UNKNOWN, }; @@ -587,6 +593,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"}, { PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"}, { PROJECTOR_TYPE_COHERE2V, "cohere2v"}, + { PROJECTOR_TYPE_MUSE_SPARK_AUDIO, "muse-spark-audio"}, }; static projector_type clip_projector_type_from_string(const std::string & str) { diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 2c4ea398ccb7..92ffe5e54fa2 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -134,6 +134,20 @@ struct clip_hparams { int32_t audio_proj_downsample_rate = 0; int32_t audio_proj_head_count = 0; + // audio pyramid (muse-spark): widths, folded samples and blocks per stage + std::vector<int32_t> audio_hidden_sizes; + std::vector<int32_t> audio_patch_sizes; + std::vector<int32_t> audio_depths; + + int32_t audio_samples_per_token() const { + int64_t n = 1; + for (int32_t p : audio_patch_sizes) { + n *= p; + GGML_ASSERT(n >= 1 && n <= (1 << 20)); + } + return (int32_t) n; + } + // audio-to-mel preprocessor params int32_t audio_chunk_len = -1; // in seconds int32_t audio_sample_rate = -1; @@ -283,6 +297,10 @@ struct clip_layer { ggml_tensor * ff_post_norm_w = nullptr; + // depthwise causal conv block (muse-spark audio), conv_norm_w is declared below + ggml_tensor * conv_w = nullptr; + ggml_tensor * conv_b = nullptr; + // layer scale (no bias) ggml_tensor * ls_1_w = nullptr; ggml_tensor * ls_2_w = nullptr; @@ -585,6 +603,14 @@ struct clip_model { std::vector<clip_layer> layers; + // one patch embedder per audio pyramid stage (muse-spark) + struct audio_patch { + ggml_tensor * proj_w = nullptr; + ggml_tensor * proj_b = nullptr; + ggml_tensor * norm_w = nullptr; + }; + std::vector<audio_patch> a_patches; + int32_t n_deepstack_layers = 0; // used by Qwen3-VL, calculated from clip_layer ggml_tensor * post_ln_w; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 4c1185a431b9..336c266b5762 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1059,6 +1059,10 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique<clip_graph_whisper_enc>(ctx, img); } break; + case PROJECTOR_TYPE_MUSE_SPARK_AUDIO: + { + builder = std::make_unique<clip_graph_muse_spark_audio>(ctx, img); + } break; case PROJECTOR_TYPE_KIMIVL: { builder = std::make_unique<clip_graph_kimivl>(ctx, img); @@ -1775,6 +1779,33 @@ struct clip_model_loader { hparams.set_limit_image_tokens(8, 576); hparams.set_warmup_n_tokens(16*16); } break; + case PROJECTOR_TYPE_MUSE_SPARK_AUDIO: + { + hparams.audio_sample_rate = 24000; // mono waveform, no frontend + get_u32(KEY_A_CONV_KERNEL_SIZE, hparams.audio_conv_kernel_size, true); + get_arr_int(KEY_A_HIDDEN_SIZES, hparams.audio_hidden_sizes, true); + get_arr_int(KEY_A_PATCH_SIZES, hparams.audio_patch_sizes, true); + get_arr_int(KEY_A_DEPTHS, hparams.audio_depths, true); + + // the graph indexes the three arrays by stage and model.layers by the running depth sum + const size_t n_stage = hparams.audio_patch_sizes.size(); + GGML_ASSERT(n_stage > 0); + GGML_ASSERT(hparams.audio_hidden_sizes.size() == n_stage); + GGML_ASSERT(hparams.audio_depths.size() == n_stage); + GGML_ASSERT(hparams.n_mel_bins == 1); + GGML_ASSERT(hparams.audio_conv_kernel_size >= 1); + + int64_t n_blocks = 0; + for (size_t i = 0; i < n_stage; ++i) { + GGML_ASSERT(hparams.audio_hidden_sizes[i] >= 1); + GGML_ASSERT(hparams.audio_depths[i] >= 0); + n_blocks += hparams.audio_depths[i]; + } + if (n_blocks != hparams.n_layer) { + throw std::runtime_error(string_format("%s: sum of audio depths (%" PRId64 ") does not match n_layer (%d)\n", __func__, n_blocks, hparams.n_layer)); + } + hparams.warmup_audio_size = hparams.audio_samples_per_token(); + } break; case PROJECTOR_TYPE_MUSE_GLIMMER: { hparams.n_merge = 2; // pixel-shuffle downsample after the ViT @@ -2187,7 +2218,9 @@ struct clip_model_loader { // GEMMA4UA is encoder-free: it uses n_mel_bins as a raw-waveform frame size (640) and has no FFT/filterbank, so the mel-range and FFT // checks below do not apply to it. // pocket-tts is encoder-free in the same sense: mimi convolves the raw waveform + // Muse Spark is encoder-free in the same sense: its conv pyramid reads the raw waveform, one channel takes the place of the mel bins const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA && + model.proj_type != PROJECTOR_TYPE_MUSE_SPARK_AUDIO && model.proj_type != PROJECTOR_TYPE_POCKETTTS_SPKENC; // Validate audio hparams loaded from GGUF metadata @@ -2348,7 +2381,8 @@ struct clip_model_loader { const bool has_standard_layers = ( model.proj_type != PROJECTOR_TYPE_GEMMA3NV && model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC && - model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN); + model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN && + model.proj_type != PROJECTOR_TYPE_MUSE_SPARK_AUDIO); // layers const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0; @@ -2619,6 +2653,38 @@ struct clip_model_loader { model.mm_merger_fc2_w = get_tensor(string_format(TN_MM_MERGER_FC2, "weight")); model.mm_merger_fc2_b = get_tensor(string_format(TN_MM_MERGER_FC2, "bias")); } break; + case PROJECTOR_TYPE_MUSE_SPARK_AUDIO: + { + model.layers.resize(hparams.n_layer); + model.a_patches.resize(hparams.audio_patch_sizes.size()); + for (size_t i = 0; i < model.a_patches.size(); ++i) { + auto & pe = model.a_patches[i]; + pe.proj_w = get_tensor(string_format(TN_A_PATCH, (int) i, "weight")); + pe.proj_b = get_tensor(string_format(TN_A_PATCH, (int) i, "bias")); + pe.norm_w = get_tensor(string_format(TN_A_PATCH_NORM, (int) i, "weight")); + if (pe.proj_w->ne[1] != hparams.audio_hidden_sizes[i]) { + throw std::runtime_error(string_format("%s: audio stage %zu projection has width %" PRId64 ", expected %d\n", __func__, i, pe.proj_w->ne[1], hparams.audio_hidden_sizes[i])); + } + } + + for (int il = 0; il < hparams.n_layer; ++il) { + auto & layer = model.layers[il]; + layer.conv_w = get_tensor(string_format(TN_CONV1D, il, "weight")); + layer.conv_b = get_tensor(string_format(TN_CONV1D, il, "bias")); + layer.conv_norm_w = get_tensor(string_format(TN_A_CONV1D_NORM, il, "weight")); + layer.ff_up_w = get_tensor(string_format(TN_FFN_UP, "a", il, "weight")); + layer.ff_up_b = get_tensor(string_format(TN_FFN_UP, "a", il, "bias")); + layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, "a", il, "weight")); + layer.ff_down_b = get_tensor(string_format(TN_FFN_DOWN, "a", il, "bias")); + } + + // pre_ln and post_ln come from the generic path above, where they are optional + if (!model.pre_ln_w || !model.post_ln_w) { + throw std::runtime_error(string_format("%s: missing pre_ln or post_ln weight\n", __func__)); + } + model.mm_fc_w = get_tensor(string_format(TN_MM_AUDIO_FC, "weight")); + model.mm_fc_b = get_tensor(string_format(TN_MM_AUDIO_FC, "bias")); + } break; case PROJECTOR_TYPE_MUSE_GLIMMER: { // 3-linear MLP: fc -> erf-GELU -> proj -> erf-GELU -> vision_proj (into LLM residual dim) @@ -4393,6 +4459,11 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { n_patches = n_patches_y * n_patches_x; } } break; + case PROJECTOR_TYPE_MUSE_SPARK_AUDIO: + { + const int64_t fold = ctx->model.hparams.audio_samples_per_token(); + n_patches = (int) ((img->nx() + fold - 1) / fold); + } break; case PROJECTOR_TYPE_VOXTRAL: case PROJECTOR_TYPE_ULTRAVOX: case PROJECTOR_TYPE_QWEN2A: @@ -5459,6 +5530,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { case PROJECTOR_TYPE_COGVLM: case PROJECTOR_TYPE_YASA2: case PROJECTOR_TYPE_GEMMA4UA: + case PROJECTOR_TYPE_MUSE_SPARK_AUDIO: case PROJECTOR_TYPE_QWEN3TTS_SPKENC: { // do nothing @@ -6157,6 +6229,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { case PROJECTOR_TYPE_GEMMA4UA: return ctx->model.mm_input_proj_w->ne[1]; case PROJECTOR_TYPE_IDEFICS3: + case PROJECTOR_TYPE_MUSE_SPARK_AUDIO: return ctx->model.mm_fc_w->ne[1]; case PROJECTOR_TYPE_ULTRAVOX: case PROJECTOR_TYPE_VOXTRAL: diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 2f4676a54359..7d7f5ac1daae 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -178,6 +178,11 @@ struct clip_graph_granite_speech : clip_graph { ggml_cgraph * build() override; }; +struct clip_graph_muse_spark_audio : clip_graph { + clip_graph_muse_spark_audio(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; +}; + struct clip_graph_gemma4a : clip_graph { clip_graph_gemma4a(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/models/muse-spark-audio.cpp b/tools/mtmd/models/muse-spark-audio.cpp new file mode 100644 index 000000000000..d5340d5f9faa --- /dev/null +++ b/tools/mtmd/models/muse-spark-audio.cpp @@ -0,0 +1,66 @@ +#include "models.h" + +// Muse Spark audio tower: a causal conv pyramid over the raw waveform, no attention and no spectrogram +// each stage folds `patch` samples into the channel dim, projects, then runs `depth` residual blocks (depthwise causal conv, RMS norm, GELU MLP) +ggml_cgraph * clip_graph_muse_spark_audio::build() { + const int kernel = hparams.audio_conv_kernel_size; + const int n_stage = (int) hparams.audio_patch_sizes.size(); + + // shape [1, n_samples], so the first stage folds the same way as the others + ggml_tensor * cur = build_inp_raw(1); + cur = ggml_reshape_2d(ctx0, cur, 1, img.nx()); + cb(cur, "inp_raw", -1); + + int bid = 0; + for (int stage = 0; stage < n_stage; ++stage) { + const int patch = hparams.audio_patch_sizes[stage]; + const auto & pe = model.a_patches[stage]; + + // fold `patch` positions into the channel dim + cur = ggml_reshape_2d(ctx0, cur, cur->ne[0] * patch, cur->ne[1] / patch); + cb(cur, "patch_folded", stage); + + GGML_ASSERT(pe.proj_w->ne[1] == hparams.audio_hidden_sizes[stage]); + cur = ggml_mul_mat(ctx0, pe.proj_w, cur); + cur = ggml_add(ctx0, cur, pe.proj_b); + cur = build_norm(cur, pe.norm_w, nullptr, NORM_TYPE_RMS, hparams.eps, stage); + cb(cur, "patch_embd", stage); + + for (int i = 0; i < hparams.audio_depths[stage]; ++i, ++bid) { + const auto & layer = model.layers[bid]; + ggml_tensor * residual = cur; + + // depthwise conv over time, pad on the left so it is causal + ggml_tensor * conv = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + conv = ggml_pad_ext(ctx0, conv, kernel - 1, 0, 0, 0, 0, 0, 0, 0); + conv = ggml_conv_1d_dw(ctx0, layer.conv_w, conv, 1, 0, 1); + conv = ggml_cont(ctx0, ggml_transpose(ctx0, conv)); + conv = ggml_add(ctx0, conv, layer.conv_b); + cb(conv, "conv", bid); + + cur = build_norm(conv, layer.conv_norm_w, nullptr, NORM_TYPE_RMS, hparams.eps, bid); + + cur = build_ffn(cur, + layer.ff_up_w, layer.ff_up_b, + nullptr, nullptr, + layer.ff_down_w, layer.ff_down_b, + FFN_GELU_ERF, bid); + cb(cur, "ffn_out", bid); + + cur = ggml_add(ctx0, cur, residual); + cb(cur, "block_out", bid); + } + } + + cur = build_norm(cur, model.pre_ln_w, nullptr, NORM_TYPE_RMS, hparams.eps, -1); + cb(cur, "enc_norm", -1); + + // straight into the text hidden size, no adapter + cur = ggml_mul_mat(ctx0, model.mm_fc_w, cur); + cur = ggml_add(ctx0, cur, model.mm_fc_b); + cur = build_norm(cur, model.post_ln_w, nullptr, NORM_TYPE_RMS, hparams.eps, -1); + cb(cur, "projected", -1); + + ggml_build_forward_expand(gf, cur); + return gf; +} diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index 1b4ea60bf961..bbeb57c80064 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1596,3 +1596,23 @@ bool mtmd_audio_preprocessor_pockettts::preprocess(const float * output.push_back(std::move(out)); return true; } + +// Muse Spark reads the raw waveform: pass the samples through, zero padded to a whole number of embeddings +bool mtmd_audio_preprocessor_muse_spark_audio::preprocess( + const float * samples, + size_t n_samples, + std::vector<mtmd_audio_mel> & output) const { + const int64_t fold = hparams.audio_samples_per_token(); + const int64_t n_padded = ((int64_t) n_samples + fold - 1) / fold * fold; + + mtmd_audio_mel mel; + mel.n_len = n_padded; + mel.n_len_org = (int64_t) n_samples; + mel.n_mel = 1; + mel.data.assign((size_t) n_padded, 0.0f); + std::copy(samples, samples + n_samples, mel.data.begin()); + + output.clear(); + output.push_back(std::move(mel)); + return true; +} diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index 4ca6cbd557e2..3b66c2758b05 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -70,6 +70,12 @@ struct mtmd_audio_preprocessor_whisper : mtmd_audio_preprocessor { mtmd_audio_cache cache; }; +struct mtmd_audio_preprocessor_muse_spark_audio : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_muse_spark_audio(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} + void initialize() override {} + bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) const override; +}; + struct mtmd_audio_preprocessor_conformer : mtmd_audio_preprocessor { mtmd_audio_preprocessor_conformer(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} mtmd_audio_preprocessor_conformer(const clip_hparams & hparams) : mtmd_audio_preprocessor(hparams) {} diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 2fa02f83eb3f..343d295719f3 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1009,6 +1009,12 @@ struct mtmd_context { aud_end = "<|audio_eos|>"; audio_preproc = std::make_unique<mtmd_audio_preprocessor_whisper>(ctx_a); } break; + case PROJECTOR_TYPE_MUSE_SPARK_AUDIO: + { + aud_beg = "<|soft_audio_start|>"; + aud_end = "<|soft_audio_end|>"; + audio_preproc = std::make_unique<mtmd_audio_preprocessor_muse_spark_audio>(ctx_a); + } break; case PROJECTOR_TYPE_QWEN3A: { aud_beg = "<|audio_start|>"; From 2f729cfa4ca6f55b8bacf67bb2ba92cb9548e291 Mon Sep 17 00:00:00 2001 From: Said Taghadouini <84044788+staghado@users.noreply.github.com> Date: Thu, 15 Oct 2026 23:53:02 +0200 Subject: [PATCH 033/423] mtmd: lightonocr does not require longest edge to always be 1540, native resolution is the default (#30600) Assisted-by: Codex --- tests/test-mtmd-impl.cpp | 21 +++++++++++++++++++++ tools/mtmd/mtmd-image.cpp | 14 +++++++++----- tools/mtmd/mtmd-image.h | 3 ++- 3 files changed, 32 insertions(+), 6 deletions(-) diff --git a/tests/test-mtmd-impl.cpp b/tests/test-mtmd-impl.cpp index 5321f2ed1eb9..dfa3e5e2c544 100644 --- a/tests/test-mtmd-impl.cpp +++ b/tests/test-mtmd-impl.cpp @@ -45,6 +45,27 @@ struct test_registry { // mtmd_image // +MAKE_TEST(test_image_preprocessor_longest_edge) { + clip_hparams hparams; + hparams.image_longest_edge = 1540; + hparams.patch_size = 14; + hparams.n_merge = 2; + + const std::vector<std::pair<clip_image_size, clip_image_size>> cases = { + { { 463, 463 }, { 476, 476 } }, + { { 463, 300 }, { 476, 308 } }, + { { 448, 280 }, { 448, 280 } }, + { { 2000, 1000 }, { 1540, 784 } }, + }; + + for (const auto & [input, expected] : cases) { + const clip_image_size actual = mtmd_image_preprocessor_longest_edge::get_target_size(hparams, input); + const std::string name = std::to_string(input.width) + "x" + std::to_string(input.height); + t.assert_equal(name + " width", expected.width, actual.width); + t.assert_equal(name + " height", expected.height, actual.height); + } +} + MAKE_TEST(test_image_preprocessor_lfm2) { clip_hparams hparams; hparams.patch_size = 16; diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index a34fd4548756..5debfe6f0dd8 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -863,20 +863,24 @@ mtmd_image_preproc_out mtmd_image_preprocessor_glm5v::preprocess(const clip_imag // mtmd_image_preprocessor_longest_edge // -mtmd_image_preproc_out mtmd_image_preprocessor_longest_edge::preprocess(const clip_image_u8 & img) const { +clip_image_size mtmd_image_preprocessor_longest_edge::get_target_size(const clip_hparams & hparams, const clip_image_size & original_size) { GGML_ASSERT(hparams.image_longest_edge > 0); - clip_image_u8 resized_image; - const clip_image_size original_size = img.get_size(); // the original pixtral model doesn't have n_merge const int cur_merge = hparams.n_merge == 0 ? 1 : hparams.n_merge; - const clip_image_size target_size = img_tool::calc_size_preserved_ratio( + return img_tool::calc_size_preserved_ratio( original_size, { /* align_size */ hparams.patch_size * cur_merge, /* min_pixels */ std::max(0, hparams.image_min_pixels), /* max_pixels */ std::max(0, hparams.image_max_pixels), - /* longest_edge */ hparams.image_longest_edge, + /* longest_edge */ std::min(hparams.image_longest_edge, + std::max(original_size.width, original_size.height)), }); +} + +mtmd_image_preproc_out mtmd_image_preprocessor_longest_edge::preprocess(const clip_image_u8 & img) const { + clip_image_u8 resized_image; + const clip_image_size target_size = get_target_size(hparams, img.get_size()); img_tool::resize(img, resized_image, target_size, hparams.image_resize_algo, hparams.image_resize_pad, diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index b318b2d4208d..1798ce51d4fb 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -140,9 +140,10 @@ struct mtmd_image_preprocessor_glm5v : mtmd_image_preprocessor { mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override; }; -// similar to mtmd_image_preprocessor_dyn_size, but resize the image to have longest edge equal to hparams.image_longest_edge, while preserving aspect ratio +// similar to mtmd_image_preprocessor_dyn_size, but cap the longest edge at hparams.image_longest_edge, while preserving aspect ratio struct mtmd_image_preprocessor_longest_edge : mtmd_image_preprocessor { mtmd_image_preprocessor_longest_edge(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} + static clip_image_size get_target_size(const clip_hparams & hparams, const clip_image_size & original_size); mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override; }; From e287bd872915dd359c3c4f17f66001ca173b0e10 Mon Sep 17 00:00:00 2001 From: Matt Rutkowski <mrutkows@us.ibm.com> Date: Thu, 15 Oct 2026 17:31:14 -0500 Subject: [PATCH 034/423] convert : fix regression in mapping of Granite speech tensors (#30628) Signed-off-by: Matt Rutkowski <mrutkows@us.ibm.com> --- gguf-py/gguf/tensor_mapping.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index ecea46463c37..01796e97879a 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2800,6 +2800,7 @@ class TensorNameMap: "conformer.layers.{bid}.conv.batch_norm", # lfm2 "conformer.layers.{bid}.lconv1d.pre_layer_norm", # gemma3n "sound_encoder.encoder.layers.{bid}.conv.norm", # parakeet + "encoder.layers.{bid}.conv.batch_norm", # granite_speech ), MODEL_TENSOR.A_ENC_CONV_NORM_MEAN: ( @@ -2808,7 +2809,6 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_CONV_NORM_VAR: ( "sound_encoder.encoder.layers.{bid}.conv.norm.running_var", # parakeet - "encoder.layers.{bid}.conv.batch_norm", # granite_speech ), MODEL_TENSOR.A_ENC_CONV_PW1: ( From 59d5fad7c111d30a547d690b69235ed9e05c035d Mon Sep 17 00:00:00 2001 From: Nigel Jones <nigel.l.jones+git@gmail.com> Date: Thu, 15 Oct 2026 23:41:01 +0100 Subject: [PATCH 035/423] convert : fix Granite Switch block_count for MultiSwitch checkpoints (#30621) * convert : fix Granite Switch block_count for MultiSwitch checkpoints Signed-off-by: Nigel Jones <jonesn@uk.ibm.com> * Update granite.py Co-authored-by: Gabe Goodhart <ghart@us.ibm.com> --------- Signed-off-by: Nigel Jones <jonesn@uk.ibm.com> Co-authored-by: Gabe Goodhart <ghart@us.ibm.com> --- conversion/granite.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index 796d37cca269..4cbeb7617892 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -239,8 +239,21 @@ class GraniteSwitchModel(GraniteMoeModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - # the weightless switch reserves one cache slot: one fewer block than num_hidden_layers - self.block_count = self.block_count - 1 + # The weightless switch layer(s) are counted in num_hidden_layers but carry no tensors. + # The count is not fixed, nor is it present in hparams, so we derive it here based on tensor block names. + layer_ids = { + int(m.group(1)) + for name in self.model_tensors + if (m := re.match(r"(?:model\.)?layers\.(\d+)\.", name)) + } + if not layer_ids: + raise ValueError("graniteswitch: no 'model.layers.N' tensors found in checkpoint") + n_weight_layers = max(layer_ids) + 1 + logger.info( + "gguf: (graniteswitch) num_hidden_layers %s counts %d weightless switch layer(s); using %d blocks", + self.hparams["num_hidden_layers"], self.block_count - n_weight_layers, n_weight_layers, + ) + self.block_count = n_weight_layers self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) self._n_adapters = int(self.hparams["num_adapters"]) From 2851b3bcbf8d993a2243ae751c0a7dfa8e944a40 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen <son@huggingface.co> Date: Fri, 16 Oct 2026 01:46:42 +0200 Subject: [PATCH 036/423] mtmd: add mtmd_audio_stream + support voxtral-realtime (#30629) * mtmd: add mtmd_audio_stream + support voxtral-realtime * address review --- conversion/__init__.py | 2 + conversion/mistral3.py | 10 + conversion/ultravox.py | 45 +++++ gguf-py/gguf/constants.py | 9 + gguf-py/gguf/gguf_writer.py | 6 + gguf-py/gguf/tensor_mapping.py | 14 ++ src/llama-arch.cpp | 5 + src/llama-arch.h | 3 + src/llama-ext.h | 7 + src/llama-graph.cpp | 13 ++ src/llama-graph.h | 1 + src/llama-hparams.h | 3 + src/llama-model.cpp | 61 ++++-- src/llama-model.h | 2 + src/models/mistral3.cpp | 62 +++++- src/models/models.h | 2 + tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-impl.h | 17 ++ tools/mtmd/clip.cpp | 129 +++++++++++++ tools/mtmd/clip.h | 10 + tools/mtmd/models/models.h | 7 + tools/mtmd/models/voxtral-rt.cpp | 111 +++++++++++ tools/mtmd/mtmd-audio.cpp | 25 +++ tools/mtmd/mtmd-audio.h | 3 + tools/mtmd/mtmd-cli.cpp | 101 ++++++++++ tools/mtmd/mtmd.cpp | 319 ++++++++++++++++++++++++++++++- tools/mtmd/mtmd.h | 50 +++++ 27 files changed, 998 insertions(+), 20 deletions(-) create mode 100644 tools/mtmd/models/voxtral-rt.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index 6c3914fbdfae..c2683aec471c 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -199,6 +199,7 @@ "MiniMaxM3SparseForConditionalGeneration": "minimax", "Ministral3ForCausalLM": "mistral3", "Mistral3ForConditionalGeneration": "mistral3", + "VoxtralRealtimeForConditionalGeneration": "mistral3", "MistralForCausalLM": "llama", "MixtralForCausalLM": "llama", "ModernBertDecisionModel": "bert", @@ -383,6 +384,7 @@ "UltravoxModel": "ultravox", "UnlimitedOCRForCausalLM": "deepseek", "VoxtralForConditionalGeneration": "ultravox", + "VoxtralRealtimeForConditionalGeneration": "ultravox", "YoutuVLForConditionalGeneration": "youtuvl", } diff --git a/conversion/mistral3.py b/conversion/mistral3.py index fee039b353c4..661475251aa8 100644 --- a/conversion/mistral3.py +++ b/conversion/mistral3.py @@ -66,3 +66,13 @@ def write_vocab(self): def write(self): self.impl.write() + + +@ModelBase.register("VoxtralRealtimeForConditionalGeneration") +@ModelBase.example("mistralai/Voxtral-Mini-4B-Realtime-2602") +class VoxtralRealtimeModel(LlamaModel): + model_arch = gguf.MODEL_ARCH.MISTRAL3 + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_realtime_delay_tokens(self.hparams["default_num_delay_tokens"]) diff --git a/conversion/ultravox.py b/conversion/ultravox.py index 62819e574d57..4a246daa1466 100644 --- a/conversion/ultravox.py +++ b/conversion/ultravox.py @@ -6,6 +6,7 @@ from torch import Tensor from .base import MmprojModel, ModelBase, TextModel, gguf +from .llama import LlamaModel @ModelBase.register("UltravoxModel") @@ -208,3 +209,47 @@ def tensor_force_quant(self, name, new_name, bid, n_dims): # Was trained in BF16, being safe, avoiding quantizing to FP16 return gguf.GGMLQuantizationType.F32 return super().tensor_force_quant(name, new_name, bid, n_dims) + + +@ModelBase.register("VoxtralRealtimeForConditionalGeneration") +@ModelBase.example("mistralai/Voxtral-Mini-4B-Realtime-2602") +class VoxtralRealtimeEncoderModel(MmprojModel): + has_vision_encoder = False + has_audio_encoder = True + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_audio is not None + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.VOXTRAL_RT) + self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"]) + self.gguf_writer.add_audio_head_dim(self.hparams_audio["head_dim"]) + self.gguf_writer.add_audio_attention_layernorm_eps(self.hparams_audio.get("rms_norm_eps", 1e-5)) + self.gguf_writer.add_audio_window_size(self.hparams_audio["sliding_window"]) + self.gguf_writer.add_audio_stack_factor(self.global_config["downsample_factor"]) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + if ".conv" in name and ".weight" in name: + return gguf.GGMLQuantizationType.F16 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + + # prevent clash naming with vision tensors + if name.startswith("multi_modal_projector"): + name = "audio." + name + + return super().filter_tensors((name, gen)) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if "conv1.bias" in name or "conv2.bias" in name: + data_torch = data_torch.unsqueeze(-1) + + # HF stores q/k in rotate-half layout, the encoder uses interleaved (normal) RoPE + if ".self_attn.q_proj." in name or ".self_attn.k_proj." in name: + assert self.hparams_audio is not None + n_head = self.hparams_audio["num_attention_heads"] + data_torch = LlamaModel.permute(data_torch, n_head, n_head) + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a3a9c478b6e5..92a69bd57369 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -175,6 +175,7 @@ class LLM: HAS_CONFIDENCE_HEAD = "{arch}.has_confidence_head" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" NORM_BEFORE_FC = "{arch}.norm_before_fc" + REALTIME_DELAY_TOKENS = "{arch}.realtime_delay_tokens" class Adapters: COUNT = "{arch}.adapters.count" @@ -481,6 +482,7 @@ class ClipAudio: class Attention: HEAD_COUNT = "clip.audio.attention.head_count" + HEAD_DIM = "clip.audio.attention.head_dim" # set when qkv width != n_embd LAYERNORM_EPS = "clip.audio.attention.layer_norm_epsilon" class Projector: @@ -758,6 +760,8 @@ class MODEL_TENSOR(IntEnum): FFN_POST_NORM = auto() FFN_POST_NORM_1 = auto() # gemma4 FFN_POST_NORM_2 = auto() # gemma4 + FFN_NORM_ADA_DOWN = auto() # voxtral realtime + FFN_NORM_ADA_UP = auto() # voxtral realtime FFN_GATE = auto() FFN_DOWN = auto() FFN_UP = auto() @@ -1572,6 +1576,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_PRE_NORM_2: "blk.{bid}.pre_ffw_norm_2", # gemma4 MODEL_TENSOR.FFN_POST_NORM_1: "blk.{bid}.post_ffw_norm_1", # gemma4 MODEL_TENSOR.FFN_POST_NORM_2: "blk.{bid}.post_ffw_norm_2", # gemma4 + MODEL_TENSOR.FFN_NORM_ADA_DOWN: "blk.{bid}.ffn_norm_ada_down", + MODEL_TENSOR.FFN_NORM_ADA_UP: "blk.{bid}.ffn_norm_ada_up", MODEL_TENSOR.FFN_GATE: "blk.{bid}.ffn_gate", MODEL_TENSOR.FFN_DOWN: "blk.{bid}.ffn_down", MODEL_TENSOR.FFN_UP: "blk.{bid}.ffn_up", @@ -5554,6 +5560,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.ATTN_ROT_EMBD, MODEL_TENSOR.FFN_GATE_INP, MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_NORM_ADA_DOWN, + MODEL_TENSOR.FFN_NORM_ADA_UP, MODEL_TENSOR.FFN_GATE, MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, @@ -6297,6 +6305,7 @@ class VisionProjectorType: GLMA = "glma" # audio QWEN25O = "qwen2.5o" # omni VOXTRAL = "voxtral" + VOXTRAL_RT = "voxtral_rt" MERALION = "meralion" # audio: Whisper + gated MLP adaptor LFM2 = "lfm2" KIMIVL = "kimivl" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index a1b1f1263e66..93ec95eb648e 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1138,6 +1138,9 @@ def add_attn_temperature_length(self, value: int) -> None: def add_attn_temperature_scale(self, value: float) -> None: self.add_float32(Keys.Attention.TEMPERATURE_SCALE.format(arch=self.arch), value) + def add_realtime_delay_tokens(self, value: int) -> None: + self.add_uint32(Keys.LLM.REALTIME_DELAY_TOKENS.format(arch=self.arch), value) + def add_pooling_type(self, value: PoolingType) -> None: self.add_uint32(Keys.LLM.POOLING_TYPE.format(arch=self.arch), value.value) @@ -1579,6 +1582,9 @@ def add_audio_block_count(self, value: int) -> None: def add_audio_head_count(self, value: int) -> None: self.add_uint32(Keys.ClipAudio.Attention.HEAD_COUNT, value) + def add_audio_head_dim(self, value: int) -> None: + self.add_uint32(Keys.ClipAudio.Attention.HEAD_DIM, value) + def add_audio_attention_layernorm_eps(self, value: float) -> None: self.add_float32(Keys.ClipAudio.Attention.LAYERNORM_EPS, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 01796e97879a..ab6abf9cb704 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -486,6 +486,14 @@ class TensorNameMap: ), # Post feed-forward norm + MODEL_TENSOR.FFN_NORM_ADA_DOWN: ( + "model.layers.{bid}.ada_rms_norm.linear1", # voxtral realtime + ), + + MODEL_TENSOR.FFN_NORM_ADA_UP: ( + "model.layers.{bid}.ada_rms_norm.linear2", # voxtral realtime + ), + MODEL_TENSOR.FFN_POST_NORM: ( "model.layers.{bid}.post_feedforward_layernorm", # gemma2 olmo2 "layers.{bid}.post_feedforward_layernorm", # embeddinggemma @@ -2491,6 +2499,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_CONV1D: ( "audio_tower.conv{bid}", # ultravox + "audio_tower.embedder.conv{bid}", # voxtral realtime "conformer.pre_encode.conv.{bid}", # lfm2 "model.audio_tower.subsample_conv_projection.conv_{bid}.conv", # gemma3n "conformer.subsample_conv_projection.layer{bid}.conv", # gemma4 @@ -2523,6 +2532,7 @@ class TensorNameMap: MODEL_TENSOR.A_POST_NORM: ( "audio_tower.layer_norm", # ultravox + "audio_tower.norm", # voxtral realtime "audio_encoder.dots_encoder.speech_encoder.layer_norm", # dots3note "audio_tower.ln_post", # qwen2omni "encoder.layer_norm", # mimo-audio-tokenizer @@ -2594,6 +2604,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_OUTPUT: ( "audio_tower.layers.{bid}.self_attn.out_proj", # ultravox + "audio_tower.layers.{bid}.self_attn.o_proj", # voxtral realtime "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.out_proj", # dots3note "conformer.layers.{bid}.self_attn.linear_out", # lfm2 "conformer.layers.{bid}.attention.post", # gemma3n @@ -2633,6 +2644,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_FFN_UP: ( "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc1_up", # dots3note (split from fc1 in conversion code) "audio_tower.layers.{bid}.fc1", # ultravox + "audio_tower.layers.{bid}.mlp.up_proj", # voxtral realtime "conformer.layers.{bid}.feed_forward1.linear1", # lfm2 "conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_1", # gemma4 @@ -2643,11 +2655,13 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_FFN_GATE: ( "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc1_gate", # dots3note (split from fc1 in conversion code) + "audio_tower.layers.{bid}.mlp.gate_proj", # voxtral realtime ), MODEL_TENSOR.A_ENC_FFN_DOWN: ( "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc2", # dots3note "audio_tower.layers.{bid}.fc2", # ultravox + "audio_tower.layers.{bid}.mlp.down_proj", # voxtral realtime "conformer.layers.{bid}.feed_forward1.linear2", # lfm2 "conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_2", # gemma4 diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 73c6d51d9595..1476a079a7cd 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -291,6 +291,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = { { LLM_KV_ATTENTION_VALUE_SCALE, "%s.attention.value_scale" }, { LLM_KV_ATTENTION_TEMPERATURE_LENGTH, "%s.attention.temperature_length" }, { LLM_KV_ATTENTION_TEMPERATURE_SCALE, "%s.attention.temperature_scale" }, + { LLM_KV_REALTIME_DELAY_TOKENS, "%s.realtime_delay_tokens" }, { LLM_KV_ATTENTION_KEY_LENGTH_MLA, "%s.attention.key_length_mla" }, { LLM_KV_ATTENTION_VALUE_LENGTH_MLA, "%s.attention.value_length_mla" }, { LLM_KV_ATTENTION_KEY_LENGTH_SWA, "%s.attention.key_length_swa" }, @@ -507,6 +508,8 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = { { LLM_TENSOR_FFN_POST_NORM, "blk.%d.post_ffw_norm" }, { LLM_TENSOR_FFN_POST_NORM_1, "blk.%d.post_ffw_norm_1" }, { LLM_TENSOR_FFN_POST_NORM_2, "blk.%d.post_ffw_norm_2" }, + { LLM_TENSOR_FFN_NORM_ADA_DOWN, "blk.%d.ffn_norm_ada_down" }, + { LLM_TENSOR_FFN_NORM_ADA_UP, "blk.%d.ffn_norm_ada_up" }, { LLM_TENSOR_FFN_PRE_NORM_2, "blk.%d.pre_ffw_norm_2" }, { LLM_TENSOR_FFN_GATE_SHEXP, "blk.%d.ffn_gate_shexp" }, { LLM_TENSOR_FFN_UP_SHEXP, "blk.%d.ffn_up_shexp" }, @@ -991,6 +994,8 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = { {LLM_TENSOR_FFN_PRE_NORM_2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_FFN_POST_NORM_1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_FFN_POST_NORM_2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_NORM_ADA_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_NORM_ADA_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_FFN_POST_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_FFN_NORM_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_ATTN_Q_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index 4b607fec7100..0c842d312db1 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -266,6 +266,7 @@ enum llm_kv { LLM_KV_HRM_H_CYCLES, LLM_KV_HRM_L_CYCLES, LLM_KV_HRM_PREFIX_LM, + LLM_KV_REALTIME_DELAY_TOKENS, LLM_KV_ATTENTION_HEAD_COUNT, LLM_KV_ATTENTION_HEAD_COUNT_KV, @@ -510,6 +511,8 @@ enum llm_tensor { LLM_TENSOR_FFN_POST_NORM, LLM_TENSOR_FFN_POST_NORM_1, LLM_TENSOR_FFN_POST_NORM_2, + LLM_TENSOR_FFN_NORM_ADA_DOWN, + LLM_TENSOR_FFN_NORM_ADA_UP, LLM_TENSOR_FFN_PRE_NORM_2, LLM_TENSOR_FFN_GATE, LLM_TENSOR_FFN_DOWN, diff --git a/src/llama-ext.h b/src/llama-ext.h index b8c5e8033e89..8b3932bd7507 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -154,3 +154,10 @@ LLAMA_API uint32_t llama_model_target_layer_ids_n(const struct llama_mode // if out is nullptr, returns the number of tokens without writing to out // caller must allocate enough memory for out before calling LLAMA_API uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out); + +// same as above, but only the rows of the given tokens, in order (n_tokens * n_embd elements) +// returns the number of elements written or 0 on error +LLAMA_API uint32_t llama_model_get_tok_embd_rows(const struct llama_model * model, const llama_token * tokens, uint32_t n_tokens, float * out); + +// realtime models: number of tokens the output lags behind the audio input, used by mtmd; 0 if not applicable +LLAMA_API uint32_t llama_model_n_tokens_delay(const struct llama_model * model); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index e8e249b21a8d..80be4bb57373 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1991,6 +1991,19 @@ ggml_tensor * llm_graph_context::build_ffn( cb(cur, "ffn_act", il); } } break; + case LLM_FFN_GELU_ERF: + if (gate && type_gate == LLM_FFN_PAR) { + cur = ggml_geglu_erf_split(ctx0, cur, tmp); + cb(cur, "ffn_geglu_erf", il); + type_gate = LLM_FFN_SEQ; + } else { + cur = ggml_gelu_erf(ctx0, cur); + cb(cur, "ffn_gelu_erf", il); + if (act_scales != NULL) { + cur = ggml_div(ctx0, cur, act_scales); + cb(cur, "ffn_act", il); + } + } break; case LLM_FFN_RELU: if (gate && type_gate == LLM_FFN_PAR) { cur = ggml_reglu_split(ctx0, cur, tmp); diff --git a/src/llama-graph.h b/src/llama-graph.h index 1a3698cad93f..816fa07cc7c4 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -59,6 +59,7 @@ enum llm_ffn_op_type : int { LLM_FFN_NONE = 0, // sentinel: unset; archs must assign before use LLM_FFN_SILU, LLM_FFN_GELU, + LLM_FFN_GELU_ERF, LLM_FFN_RELU, LLM_FFN_RELU_SQR, LLM_FFN_SWIGLU, diff --git a/src/llama-hparams.h b/src/llama-hparams.h index bc0d492152db..3e025368b91f 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -281,6 +281,9 @@ struct llama_hparams { float f_attn_temp_scale = 0.0f; float f_attn_temp_offset = 0.0f; // offset position index + // realtime models: number of tokens the output lags behind the audio input + uint32_t n_tokens_delay = 0; + // gemma3n altup uint32_t n_altup = 4; // altup_num_inputs uint32_t i_altup_act = 0; // altup_active_idx diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 2660bf38a2f0..5d48d702c5b8 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -3906,26 +3906,19 @@ uint32_t llama_model_target_layer_ids_n(const struct llama_model * model) { return (uint32_t) model->target_layer_ids.size(); } -uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out) { - if (model->vocab.n_tokens() == 0 || model->tok_embd == nullptr) { - return 0; - } - - const ggml_tensor * tensor = model->tok_embd; - const size_t nelements = ggml_nelements(tensor); - GGML_ASSERT(nelements <= UINT32_MAX); // for the return type - - if (out == nullptr) { - return (uint32_t) nelements; - } +// read rows [row0, row0 + n_rows) of a 2D tensor as F32 +static void llama_tensor_rows_to_f32(const ggml_tensor * tensor, int64_t row0, int64_t n_rows, float * out) { + const size_t nelements = (size_t) n_rows * tensor->ne[0]; + const size_t offset = (size_t) row0 * tensor->nb[1]; + const size_t nbytes = ggml_row_size(tensor->type, tensor->ne[0]) * n_rows; if (tensor->type == GGML_TYPE_F32) { - ggml_backend_tensor_get(tensor, out, 0, nelements * sizeof(float)); - return (uint32_t) nelements; + ggml_backend_tensor_get(tensor, out, offset, nbytes); + return; } - std::vector<uint8_t> buf(ggml_nbytes(tensor)); - ggml_backend_tensor_get(tensor, buf.data(), 0, buf.size()); + std::vector<uint8_t> buf(nbytes); + ggml_backend_tensor_get(tensor, buf.data(), offset, nbytes); const ggml_type_traits * traits = ggml_get_type_traits(tensor->type); if (tensor->type == GGML_TYPE_F16) { @@ -3937,6 +3930,42 @@ uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out) } else { GGML_ABORT("unsupported tensor type for dequantization: %s", ggml_type_name(tensor->type)); } +} + +uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out) { + if (model->vocab.n_tokens() == 0 || model->tok_embd == nullptr) { + return 0; + } + + const ggml_tensor * tensor = model->tok_embd; + const size_t nelements = ggml_nelements(tensor); + GGML_ASSERT(nelements <= UINT32_MAX); // for the return type + + if (out == nullptr) { + return (uint32_t) nelements; + } + + llama_tensor_rows_to_f32(tensor, 0, tensor->ne[1], out); return (uint32_t) nelements; } + +uint32_t llama_model_n_tokens_delay(const struct llama_model * model) { + return model->hparams.n_tokens_delay; +} + +uint32_t llama_model_get_tok_embd_rows(const struct llama_model * model, const llama_token * tokens, uint32_t n_tokens, float * out) { + if (model->vocab.n_tokens() == 0 || model->tok_embd == nullptr || tokens == nullptr || out == nullptr) { + return 0; + } + + const ggml_tensor * tensor = model->tok_embd; + for (uint32_t i = 0; i < n_tokens; i++) { + if (tokens[i] < 0 || tokens[i] >= tensor->ne[1]) { + return 0; + } + llama_tensor_rows_to_f32(tensor, tokens[i], 1, out + (size_t) i * tensor->ne[0]); + } + + return n_tokens * (uint32_t) tensor->ne[0]; +} diff --git a/src/llama-model.h b/src/llama-model.h index 66f9129e68a1..6335dc2bc318 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -327,6 +327,8 @@ struct llama_layer { struct ggml_tensor * ffn_post_norm_1 = nullptr; // gemma4 struct ggml_tensor * ffn_post_norm_2 = nullptr; // gemma4 struct ggml_tensor * ffn_pre_norm_2 = nullptr; // gemma4 + struct ggml_tensor * ffn_norm_ada_down = nullptr; // voxtral realtime + struct ggml_tensor * ffn_norm_ada_up = nullptr; // voxtral realtime struct ggml_tensor * layer_out_norm = nullptr; struct ggml_tensor * layer_out_norm_b = nullptr; struct ggml_tensor * ffn_norm_exps = nullptr; diff --git a/src/models/mistral3.cpp b/src/models/mistral3.cpp index 484e1b698ac5..af058d558d93 100644 --- a/src/models/mistral3.cpp +++ b/src/models/mistral3.cpp @@ -8,6 +8,8 @@ void llama_model_mistral3::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ROPE_SCALING_YARN_BETA_SLOW, hparams.yarn_beta_slow, false); ml.get_key(LLM_KV_ROPE_SCALING_YARN_LOG_MUL, hparams.rope_yarn_log_mul, false); + ml.get_key(LLM_KV_REALTIME_DELAY_TOKENS, hparams.n_tokens_delay, false); + hparams.f_attn_temp_offset = 0.0f; // TODO: maybe add n_attn_temp_floor_scale as a separate KV? @@ -26,7 +28,7 @@ void llama_model_mistral3::load_arch_hparams(llama_model_loader & ml) { } } -void llama_model_mistral3::load_arch_tensors(llama_model_loader &) { +void llama_model_mistral3::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); @@ -53,6 +55,13 @@ void llama_model_mistral3::load_arch_tensors(llama_model_loader &) { layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + // optional ada rms norm (voxtral realtime), the hidden size is taken from the tensor itself + if (const ggml_tensor * meta = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_NORM_ADA_DOWN, "weight", i).str().c_str())) { + const int64_t n_ada = meta->ne[1]; + layer.ffn_norm_ada_down = create_tensor(tn(LLM_TENSOR_FFN_NORM_ADA_DOWN, "weight", i), {n_embd, n_ada}, 0); + layer.ffn_norm_ada_up = create_tensor(tn(LLM_TENSOR_FFN_NORM_ADA_UP, "weight", i), {n_ada, n_embd}, 0); + } + if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_LONGROPE) { layer.rope_long = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_LONG, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); layer.rope_short = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); @@ -86,6 +95,34 @@ void llama_model_mistral3::load_arch_tensors(llama_model_loader &) { } } +// sinusoidal embedding of the output delay, constant for the whole session +class llama_model_mistral3::input_t_cond : public llm_graph_input_i { +public: + input_t_cond(uint32_t n_delay, int64_t n_embd) : data(n_embd) { + const int64_t half = n_embd / 2; + for (int64_t i = 0; i < half; i++) { + const double freq = exp(-log(10000.0) * i / half) * n_delay; + data[i] = cos(freq); + data[half + i] = sin(freq); + } + } + + void set_input(const llama_ubatch * ubatch) override { + GGML_UNUSED(ubatch); + ggml_backend_tensor_set(t_cond, data.data(), 0, ggml_nbytes(t_cond)); + } + + // the input does not depend on the batch + bool can_reuse(const llm_graph_params & params) override { + GGML_UNUSED(params); + return true; + } + + ggml_tensor * t_cond = nullptr; // F32 [n_embd] + + std::vector<float> data; +}; + std::unique_ptr<llm_graph_context> llama_model_mistral3::build_arch_graph(const llm_graph_params & params) const { return std::make_unique<graph>(*this, params); } @@ -112,6 +149,17 @@ llama_model_mistral3::graph::graph(const llama_model & model, const llm_graph_pa auto * inp_attn = build_attn_inp_kv(); + // (optional) time conditioning of the ada rms norm + ggml_tensor * t_cond = nullptr; + if (model.layers[0].ffn_norm_ada_down) { + auto inp = std::make_unique<input_t_cond>(hparams.n_tokens_delay, n_embd); + t_cond = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_embd); + ggml_set_input(t_cond); + cb(t_cond, "inp_t_cond", -1); + inp->t_cond = t_cond; + res->add_input(std::move(inp)); + } + const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; ggml_tensor * inp_out_ids = build_inp_out_ids(); @@ -176,6 +224,18 @@ llama_model_mistral3::graph::graph(const llama_model & model, const llm_graph_pa LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); + if (t_cond) { + // ada rms norm: scale the normalized input by (1 + mlp(t_cond)) + ggml_tensor * ada = build_ffn(t_cond, + model.layers[il].ffn_norm_ada_down, NULL, NULL, + NULL, NULL, NULL, + model.layers[il].ffn_norm_ada_up, NULL, NULL, + NULL, + LLM_FFN_GELU_ERF, LLM_FFN_SEQ, il); + cur = ggml_add(ctx0, cur, ggml_mul(ctx0, cur, ada)); + cb(cur, "ffn_norm_ada", il); + } + cur = build_ffn(cur, model.layers[il].ffn_up, model.layers[il].ffn_up_b, model.layers[il].ffn_up_s, model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, model.layers[il].ffn_gate_s, diff --git a/src/models/models.h b/src/models/models.h index 51fdcfbb08c6..94a7896172f3 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2744,6 +2744,8 @@ struct llama_model_mistral3 : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; + class input_t_cond; // voxtral realtime: time conditioning of the ada rms norm + struct graph : public llm_graph_context { graph(const llama_model & model, const llm_graph_params & params); }; diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index edb5d3d89bbe..398b73904f5c 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -69,6 +69,7 @@ add_library(mtmd models/step3vl.cpp models/siglip.cpp models/whisper-enc.cpp + models/voxtral-rt.cpp models/deepseekocr.cpp models/deepseekocr2.cpp models/mobilenetv5.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index ad291a5578bc..10eea6f74619 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -1,6 +1,7 @@ #pragma once #include "ggml.h" +#include "ggml-cpp.h" #include "gguf.h" #include "clip.h" @@ -446,6 +447,20 @@ // align x to upper multiple of n #define CLIP_ALIGN(x, n) ((((x) + (n) - 1) / (n)) * (n)) +// sliding window KV cache of a causal audio encoder, see clip_enc_state_init() +struct clip_enc_state { + int32_t n_ctx = 0; // ring slots: attention window + frames per call + int32_t n_past = 0; // absolute position of the next frame + + std::vector<int32_t> slot_pos; // absolute position held by each slot, -1 if empty + std::vector<float> stem_tail; // last conv stem output frame, empty on a cold start + + ggml_context_ptr ctx_data; + ggml_backend_buffer_ptr buf; + std::vector<ggml_tensor *> k; // per layer [n_embd_kv, n_ctx] + std::vector<ggml_tensor *> v; +}; + // forward declaration // TODO: improve this later struct clip_ctx; @@ -480,6 +495,7 @@ enum projector_type { PROJECTOR_TYPE_GLMA, PROJECTOR_TYPE_QWEN25O, // will be replaced by QWEN2A or QWEN25VL depending on clip_ctx PROJECTOR_TYPE_VOXTRAL, + PROJECTOR_TYPE_VOXTRAL_RT, PROJECTOR_TYPE_MERALION, PROJECTOR_TYPE_MUSIC_FLAMINGO, PROJECTOR_TYPE_LFM2, @@ -553,6 +569,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_GLMA, "glma"}, { PROJECTOR_TYPE_QWEN25O, "qwen2.5o"}, { PROJECTOR_TYPE_VOXTRAL, "voxtral"}, + { PROJECTOR_TYPE_VOXTRAL_RT, "voxtral_rt"}, { PROJECTOR_TYPE_MERALION, "meralion"}, { PROJECTOR_TYPE_MUSIC_FLAMINGO, "musicflamingo"}, { PROJECTOR_TYPE_LFM2, "lfm2"}, diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 336c266b5762..793b8cd1a631 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1050,6 +1050,10 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique<clip_graph_llama4>(ctx, img); } break; + case PROJECTOR_TYPE_VOXTRAL_RT: + { + builder = std::make_unique<clip_graph_voxtral_rt>(ctx, img, params ? params->enc_state : nullptr); + } break; case PROJECTOR_TYPE_ULTRAVOX: case PROJECTOR_TYPE_VOXTRAL: case PROJECTOR_TYPE_QWEN2A: @@ -1923,6 +1927,22 @@ struct clip_model_loader { hparams.audio_window_len = 400; hparams.audio_hop_len = 160; } break; + case PROJECTOR_TYPE_VOXTRAL_RT: + { + get_u32(KEY_A_PROJ_STACK_FACTOR, hparams.proj_stack_factor); + get_u32(KEY_A_ATTN_WINDOW_SIZE, hparams.attn_window_size); + hparams.rope_theta = 1000000.0f; // params.json + hparams.ffn_op = FFN_SILU; + log_ffn_op = "silu"; + + // audio preprocessing params + hparams.audio_sample_rate = 16000; + hparams.audio_n_fft = 400; + hparams.audio_window_len = 400; + hparams.audio_hop_len = 160; + // 2 context frames + 8 s of new frames, see clip_enc_state_init() + hparams.warmup_audio_size = 2 + 2*400; + } break; case PROJECTOR_TYPE_MIMO_AUDIO: { get_u32(KEY_A_RVQ_NUM_QUANTIZERS, hparams.rvq_num_quantizers, false); @@ -3320,6 +3340,7 @@ struct clip_model_loader { } } break; case PROJECTOR_TYPE_VOXTRAL: + case PROJECTOR_TYPE_VOXTRAL_RT: { model.conv1d_1_w = get_tensor(string_format(TN_CONV1D, 1, "weight")); model.conv1d_1_b = get_tensor(string_format(TN_CONV1D, 1, "bias")); @@ -4459,6 +4480,11 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { n_patches = n_patches_y * n_patches_x; } } break; + case PROJECTOR_TYPE_VOXTRAL_RT: + { + // 2 context frames, conv stride 2, then stacked frames + n_patches = (img->nx() - 2) / 2 / ctx->model.hparams.proj_stack_factor; + } break; case PROJECTOR_TYPE_MUSE_SPARK_AUDIO: { const int64_t fold = ctx->model.hparams.audio_samples_per_token(); @@ -4656,6 +4682,56 @@ static std::vector<c2w_state_slot> list_gen_state_slots(const clip_hparams & hpa } } +clip_enc_state * clip_enc_state_init(clip_ctx * ctx, int32_t n_frames_max) { + const auto & hparams = ctx->model.hparams; + if (ctx->proj_type() != PROJECTOR_TYPE_VOXTRAL_RT || hparams.attn_window_size <= 0) { + LOG_ERR("%s: the encoder does not support streaming input\n", __func__); + return nullptr; + } + if (n_frames_max <= 0) { + LOG_ERR("%s: invalid n_frames_max %d\n", __func__, n_frames_max); + return nullptr; + } + + const int n_layer = hparams.n_layer; + const int d_head = hparams.n_embd_head > 0 ? hparams.n_embd_head : hparams.n_embd / hparams.n_head; + const int n_embd_kv = hparams.n_head_kv * d_head; + + auto st = std::make_unique<clip_enc_state>(); + st->n_ctx = hparams.attn_window_size + n_frames_max; + st->slot_pos.assign(st->n_ctx, -1); + + ggml_init_params params = { + /*.mem_size =*/ ggml_tensor_overhead() * 2 * n_layer, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + st->ctx_data.reset(ggml_init(params)); + for (int il = 0; il < n_layer; il++) { + ggml_tensor * k = ggml_new_tensor_2d(st->ctx_data.get(), GGML_TYPE_F16, n_embd_kv, st->n_ctx); + ggml_tensor * v = ggml_new_tensor_2d(st->ctx_data.get(), GGML_TYPE_F16, n_embd_kv, st->n_ctx); + ggml_format_name(k, "enc_state_k_%d", il); + ggml_format_name(v, "enc_state_v_%d", il); + st->k.push_back(k); + st->v.push_back(v); + } + st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), ctx->backend)); + if (!st->buf) { + LOG_ERR("%s: failed to allocate the encoder state\n", __func__); + return nullptr; + } + // empty slots are masked out, but they must not hold NaN + ggml_backend_buffer_clear(st->buf.get(), 0); + + LOG_INF("%s: encoder state: %d slots, %.2f MiB\n", __func__, st->n_ctx, ggml_backend_buffer_get_size(st->buf.get()) / 1024.0 / 1024.0); + + return st.release(); +} + +void clip_enc_state_free(clip_enc_state * state) { + delete state; +} + bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { const clip_image_f32_batch & imgs = *params->imgs; int n_batch_cur = imgs.entries.size(); @@ -4839,6 +4915,44 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // set input per projector switch (ctx->model.proj_type) { + case PROJECTOR_TYPE_VOXTRAL_RT: + { + clip_enc_state * st = params->enc_state; + if (!st) { + LOG_ERR("%s: this encoder needs a clip_enc_state\n", __func__); + return false; + } + const int n_window = hparams.attn_window_size; + const int n_new = (imgs.entries[0].nx() - 2) / 2; + GGML_ASSERT(n_new <= st->n_ctx - n_window); + + std::vector<int32_t> pos(n_new); + std::vector<int64_t> idx(n_new); + std::vector<int32_t> slot_pos = st->slot_pos; // including the frames written by this call + for (int i = 0; i < n_new; i++) { + pos[i] = st->n_past + i; + idx[i] = pos[i] % st->n_ctx; + slot_pos[idx[i]] = pos[i]; + } + set_input_i32("inp_pos", pos); + ggml_backend_tensor_set(get_inp_tensor("inp_kv_idx"), idx.data(), 0, idx.size() * sizeof(int64_t)); + + // each frame sees itself and the n_window - 1 frames before it + std::vector<float> mask((size_t) st->n_ctx * n_new, -INFINITY); + for (int i = 0; i < n_new; i++) { + for (int s = 0; s < st->n_ctx; s++) { + const int32_t p = slot_pos[s]; + if (p >= 0 && p <= pos[i] && p > pos[i] - n_window) { + mask[(size_t) i * st->n_ctx + s] = 0.0f; + } + } + } + set_input_attn_mask(mask); + + std::vector<float> tail = st->stem_tail; + tail.resize(hparams.n_embd, 0.0f); + set_input_f32("inp_stem_tail", tail); + } break; case PROJECTOR_TYPE_MUSE_GLIMMER: { const int grid_w = pos_w; // image_size_width / patch_size @@ -6099,6 +6213,20 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { out_audio.resize((size_t) n_frames * hop); } } + if (params->enc_state != nullptr) { + clip_enc_state * st = params->enc_state; + ggml_tensor * tail = ggml_graph_get_tensor(gf, "out_stem_tail"); + GGML_ASSERT(tail != nullptr); + st->stem_tail.resize(ggml_nelements(tail)); + ggml_backend_tensor_get(tail, st->stem_tail.data(), 0, ggml_nbytes(tail)); + + const int n_new = (imgs.entries[0].nx() - 2) / 2; + for (int i = 0; i < n_new; i++) { + st->slot_pos[(st->n_past + i) % st->n_ctx] = st->n_past + i; + } + st->n_past += n_new; + } + if (params->state_out != nullptr) { auto & state_out = *params->state_out; size_t total = 0; @@ -6233,6 +6361,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_fc_w->ne[1]; case PROJECTOR_TYPE_ULTRAVOX: case PROJECTOR_TYPE_VOXTRAL: + case PROJECTOR_TYPE_VOXTRAL_RT: case PROJECTOR_TYPE_MUSIC_FLAMINGO: return ctx->model.mm_2_w->ne[1]; case PROJECTOR_TYPE_MERALION: diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index 9e12702cebfb..19db4760825d 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -13,6 +13,7 @@ #define MTMD_INTERNAL_HEADER struct clip_ctx; +struct clip_enc_state; struct clip_image_size { int width; @@ -119,9 +120,18 @@ struct clip_encode_params { std::vector<float> * out_audio = nullptr; // decoded PCM samples, F32 const std::vector<uint8_t> * state_in = nullptr; // state from previous call, null or wrong size means cold start std::vector<uint8_t> * state_out = nullptr; // state for the next call + + // streaming audio input, updated in place by clip_encode() + struct clip_enc_state * enc_state = nullptr; }; bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params); +// persistent state of a causal audio encoder, for streaming input (one per stream) +// it holds the sliding window of past keys/values, so each clip_encode() call only takes the new frames +// n_frames_max is the maximum number of encoder frames per call +struct clip_enc_state * clip_enc_state_init(struct clip_ctx * ctx, int32_t n_frames_max); +void clip_enc_state_free(struct clip_enc_state * state); + bool clip_is_llava(const struct clip_ctx * ctx); // note for contributor: this clip_is_(model) pattern is deprecated // do NOT add new functions like this diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 7d7f5ac1daae..294448605972 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -155,6 +155,13 @@ struct clip_graph_whisper_enc : clip_graph { ggml_cgraph * build() override; }; +// causal whisper-like encoder with a sliding window KV cache (voxtral realtime) +struct clip_graph_voxtral_rt : clip_graph { + clip_enc_state * state; // nullptr during warmup, the cache is then a graph input + clip_graph_voxtral_rt(clip_ctx * ctx, const clip_image_f32 & img, clip_enc_state * state) : clip_graph(ctx, img), state(state) {} + ggml_cgraph * build() override; +}; + struct clip_graph_deepseekocr : clip_graph { clip_graph_deepseekocr(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/models/voxtral-rt.cpp b/tools/mtmd/models/voxtral-rt.cpp new file mode 100644 index 000000000000..4a01512d21b5 --- /dev/null +++ b/tools/mtmd/models/voxtral-rt.cpp @@ -0,0 +1,111 @@ +#include "models.h" + +// the input holds 2 context mel frames followed by the new frames, so the causal conv stem +// sees the same left context as on the full spectrogram +ggml_cgraph * clip_graph_voxtral_rt::build() { + const int n_frames_in = img.nx(); + const int n_new = (n_frames_in - 2) / 2; // encoder frames, the conv stem has stride 2 + const int n_window = hparams.attn_window_size; + const int stack = hparams.proj_stack_factor; + const int n_embd_kv = n_head_kv * d_head; + const int n_ctx = state ? state->n_ctx : n_window + n_new; + + GGML_ASSERT(n_window > 0 && stack > 0); + GGML_ASSERT(n_new > 0 && n_new % stack == 0); + GGML_ASSERT(n_new <= n_ctx - n_window); + + ggml_tensor * inp = build_inp_raw(1); // [n_frames_in, n_mel] + + ggml_tensor * stem_tail = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, n_embd); + ggml_set_name(stem_tail, "inp_stem_tail"); + ggml_set_input(stem_tail); + + ggml_tensor * pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_new); + ggml_set_name(pos, "inp_pos"); + ggml_set_input(pos); + + ggml_tensor * kv_idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new); + ggml_set_name(kv_idx, "inp_kv_idx"); + ggml_set_input(kv_idx); + + ggml_tensor * kq_mask = build_inp_attn_mask(n_ctx, n_new); + + // causal conv stem + ggml_tensor * cur = ggml_conv_1d(ctx0, model.conv1d_1_w, inp, 1, 0, 1); + cur = ggml_add(ctx0, cur, model.conv1d_1_b); + cur = ggml_gelu_erf(ctx0, cur); // [2*n_new, n_embd] + + // the last frame is the left context of the second conv on the next call + ggml_tensor * tail = ggml_view_2d(ctx0, cur, 1, n_embd, cur->nb[1], (2*n_new - 1) * ggml_element_size(cur)); + tail = ggml_cont(ctx0, tail); + ggml_set_name(tail, "out_stem_tail"); + ggml_set_output(tail); + ggml_build_forward_expand(gf, tail); + + cur = ggml_concat(ctx0, stem_tail, cur, 0); // [2*n_new + 1, n_embd] + cur = ggml_conv_1d(ctx0, model.conv1d_2_w, cur, 2, 0, 1); + cur = ggml_add(ctx0, cur, model.conv1d_2_b); + cur = ggml_gelu_erf(ctx0, cur); // [n_new, n_embd] + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [n_embd, n_new] + cb(cur, "after_conv1d", -1); + + for (int il = 0; il < n_layer; il++) { + const auto & layer = model.layers[il]; + + ggml_tensor * x = build_norm(cur, layer.ln_1_w, nullptr, NORM_TYPE_RMS, eps, il); + + ggml_tensor * q = ggml_add(ctx0, build_mm(layer.q_w, x), layer.q_b); + ggml_tensor * k = build_mm(layer.k_w, x); + ggml_tensor * v = ggml_add(ctx0, build_mm(layer.v_w, x), layer.v_b); + + q = ggml_reshape_3d(ctx0, q, d_head, n_head, n_new); + k = ggml_reshape_3d(ctx0, k, d_head, n_head_kv, n_new); + + // interleaved rope on absolute frame positions + q = ggml_rope_ext(ctx0, q, pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + k = ggml_rope_ext(ctx0, k, pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + cb(q, "q", il); + cb(k, "k", il); + cb(v, "v", il); + + ggml_tensor * k_cache; + ggml_tensor * v_cache; + if (state) { + k_cache = state->k[il]; + v_cache = state->v[il]; + } else { + k_cache = ggml_new_tensor_2d(ctx0, GGML_TYPE_F16, n_embd_kv, n_ctx); + v_cache = ggml_new_tensor_2d(ctx0, GGML_TYPE_F16, n_embd_kv, n_ctx); + ggml_set_input(k_cache); + ggml_set_input(v_cache); + } + + // write the new frames into the ring, the mask selects the window out of the whole ring + ggml_build_forward_expand(gf, ggml_set_rows(ctx0, k_cache, ggml_reshape_2d(ctx0, k, n_embd_kv, n_new), kv_idx)); + ggml_build_forward_expand(gf, ggml_set_rows(ctx0, v_cache, v, kv_idx)); + + ggml_tensor * k_all = ggml_reshape_3d(ctx0, k_cache, d_head, n_head_kv, n_ctx); + ggml_tensor * v_all = ggml_reshape_3d(ctx0, v_cache, d_head, n_head_kv, n_ctx); + + x = build_attn(layer.o_w, layer.o_b, q, k_all, v_all, kq_mask, kq_scale, il); + cur = ggml_add(ctx0, cur, x); + + x = build_norm(cur, layer.ln_2_w, nullptr, NORM_TYPE_RMS, eps, il); + x = build_ffn(x, layer.ff_up_w, nullptr, layer.ff_gate_w, nullptr, layer.ff_down_w, layer.ff_down_b, hparams.ffn_op, il); + cur = ggml_add(ctx0, cur, x); + cb(cur, "layer_out", il); + } + + cur = build_norm(cur, model.post_ln_w, nullptr, NORM_TYPE_RMS, eps, -1); + cb(cur, "after_transformer", -1); + + cur = build_stack(cur, stack, n_embd); + cur = build_mm(model.mm_1_w, cur); + cur = ggml_gelu_erf(ctx0, cur); + cur = build_mm(model.mm_2_w, cur); + cb(cur, "projected", -1); + + ggml_build_forward_expand(gf, cur); + + return gf; +} diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index bbeb57c80064..55847545855a 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -366,6 +366,31 @@ static void log_mel_spectrogram_worker_thread(int ith, } } +void mtmd_audio_log_mel_frame(const mtmd_audio_cache & cache, const float * samples, int n_fft, int n_mel, float * out) { + const int n_fft_bins = n_fft / 2 + 1; + GGML_ASSERT((int) cache.hann_window.size() == n_fft); + GGML_ASSERT(cache.filters.n_mel == n_mel && cache.filters.n_fft == n_fft); + + std::vector<float> fft_in(n_fft * 2, 0.0f); + std::vector<float> fft_out(n_fft * 2 * 2 * 2); + for (int j = 0; j < n_fft; j++) { + fft_in[j] = cache.hann_window[j] * samples[j]; + } + fft(cache, fft_in.data(), n_fft, fft_out.data()); + + for (int j = 0; j < n_fft_bins; j++) { + fft_out[j] = fft_out[2 * j + 0] * fft_out[2 * j + 0] + fft_out[2 * j + 1] * fft_out[2 * j + 1]; + } + + for (int j = 0; j < n_mel; j++) { + double sum = 0.0; + for (int k = 0; k < n_fft_bins; k++) { + sum += fft_out[k] * cache.filters.data[(size_t) j * n_fft_bins + k]; + } + out[j] = log10(std::max(sum, 1e-10)); + } +} + // ref: https://github.com/openai/whisper/blob/main/whisper/audio.py#L110-L157 static bool log_mel_spectrogram( const float * samples, diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index 3b66c2758b05..26ace2143911 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -50,6 +50,9 @@ struct mtmd_audio_cache { ); }; +// compute one log10-mel frame from n_fft windowed samples, cache must be initialized for n_fft +void mtmd_audio_log_mel_frame(const mtmd_audio_cache & cache, const float * samples, int n_fft, int n_mel, float * out); + struct mtmd_audio_preprocessor { const clip_hparams & hparams; diff --git a/tools/mtmd/mtmd-cli.cpp b/tools/mtmd/mtmd-cli.cpp index 99ae5373c6b9..1d7d65ef8691 100644 --- a/tools/mtmd/mtmd-cli.cpp +++ b/tools/mtmd/mtmd-cli.cpp @@ -257,6 +257,99 @@ static int generate_response(mtmd_cli_context & ctx, int n_predict) { return 0; } +// feed an audio file to a realtime model chunk by chunk, printing the transcript as it is produced +static int run_audio_stream(mtmd_cli_context & ctx, const std::string & fname, int n_predict) { + auto res = mtmd_helper_bitmap_init_from_file(ctx.ctx_vision.get(), fname.c_str(), false, ctx.init_opt); + if (!res.bitmap) { + return 1; // error is already printed by libmtmd + } + mtmd::bitmap bmp(res.bitmap); + if (!mtmd_bitmap_is_audio(bmp.ptr.get())) { + LOG_ERR("%s: %s is not an audio file\n", __func__, fname.c_str()); + return 1; + } + const float * samples = (const float *) bmp.data(); + const size_t n_samples = bmp.n_bytes() / sizeof(float); + const int rate = mtmd_get_audio_sample_rate(ctx.ctx_vision.get()); + + const auto sparams = mtmd_audio_stream_params_default(); + mtmd::audio_stream_ptr stream(mtmd_audio_stream_init(ctx.ctx_vision.get(), sparams)); + if (!stream) { + return 1; + } + + const size_t n_chunk = (size_t) rate * sparams.step_ms / 1000; + const int n_embd = llama_model_n_embd_inp(ctx.model); + + LOG("Transcribing %.1f s of audio, fed in chunks of %d ms\n\n", (double) n_samples / rate, sparams.step_ms); + + size_t n_fed = 0; + int n_gen = 0; + llama_token tok = LLAMA_TOKEN_NULL; + const int64_t t_start = ggml_time_us(); + int64_t t_enc = 0; // time spent in the stream (mel + encoder) + int64_t t_dec = 0; // time spent in the decoder + while (!g_is_interrupted) { + const int64_t t0 = ggml_time_us(); + const int32_t n = mtmd_audio_stream_step(stream.get(), tok); + t_enc += ggml_time_us() - t0; + if (n < 0) { + if (n == -2) { + LOG_ERR("%s: stream error\n", __func__); + return 1; + } + break; // end of stream + } + if (n == 0) { + if (n_fed < n_samples) { + const size_t cnt = std::min(n_chunk, n_samples - n_fed); + mtmd_audio_stream_push(stream.get(), samples + n_fed, cnt); + n_fed += cnt; + } else if (mtmd_audio_stream_flush(stream.get(), true) != 0) { + break; // already ended + } + continue; + } + + const int64_t t1 = ggml_time_us(); + const float * rows = mtmd_audio_stream_get_output_embd(stream.get()); + ctx.batch.clear(); + for (int32_t i = 0; i < n; i++) { + const llama_pos pos[GGML_MROPE_SECTIONS] = { ctx.n_past, ctx.n_past, ctx.n_past, 0 }; + ctx.n_past++; + ctx.batch.add_embd({ rows + (size_t) i * n_embd, 1, (size_t) n_embd }, pos, 0, i == n - 1); + } + if (llama_process(ctx.lctx, LLAMA_PROCESS_TYPE_DECODE, ctx.batch.get())) { + LOG_ERR("%s: failed to decode\n", __func__); + return 1; + } + + tok = common_sampler_sample(ctx.smpl, ctx.lctx, -1); + common_sampler_accept(ctx.smpl, tok, true); + t_dec += ggml_time_us() - t1; + if (llama_vocab_is_eog(ctx.vocab, tok)) { + break; + } + // the model emits pad and word-boundary tokens between words + if (!llama_vocab_is_control(ctx.vocab, tok)) { + LOG("%s", common_token_to_piece(ctx.lctx, tok).c_str()); + fflush(stdout); + } + if (++n_gen >= n_predict) { + break; + } + } + + const double t_total = (ggml_time_us() - t_start) / 1e6; + LOG("\n\n%s: %d tokens, %.2f s audio in %.2f s (%.2fx realtime)\n", __func__, + n_gen, (double) n_samples / rate, t_total, (double) n_samples / rate / t_total); + LOG("%s: encoder %.2f s, decoder %.2f s, %.1f ms per 80 ms step\n", __func__, + t_enc / 1e6, t_dec / 1e6, (t_enc + t_dec) / 1e3 / std::max(n_gen, 1)); + llama_perf_context_print(ctx.lctx); + + return 0; +} + static std::string chat_add_and_format(mtmd_cli_context & ctx, common_chat_msg & new_msg) { LOG_DBG("chat_add_and_format: new_msg.role='%s', new_msg.content='%s'\n", new_msg.role.c_str(), new_msg.content.c_str()); @@ -472,6 +565,14 @@ int main(int argc, char ** argv) { if (g_is_interrupted) return 130; + if (mtmd_support_audio_stream(ctx.ctx_vision.get())) { + if (params.image.size() != 1) { + LOG_ERR("This model transcribes audio in realtime, pass exactly one --audio file\n"); + return 1; + } + return run_audio_stream(ctx, params.image[0], n_predict); + } + auto eval_system_prompt_if_present = [&] { if (params.system_prompt.empty()) { return 0; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 343d295719f3..26d5ee4105c7 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -7,6 +7,7 @@ #include "debug/mtmd-debug.h" #include "llama.h" +#include "../src/llama-ext.h" // fix problem with std::min and std::max #if defined(_WIN32) @@ -533,6 +534,7 @@ struct mtmd_context { std::string media_marker; const int n_embd_text = -1; // -1 means llm context not provided, skip checking this const llama_vocab * vocab = nullptr; // can be nullptr if text_model is not provided + const llama_model * text_model = nullptr; // for the token embeddings of mtmd_audio_stream mtmd_pos_type pos_type; // these are not token, but strings used to mark the beginning and end of image/audio embeddings @@ -578,6 +580,7 @@ struct mtmd_context { media_marker (ctx_params.media_marker), n_embd_text (text_model ? llama_model_n_embd_inp(text_model) : -1), vocab (text_model ? llama_model_get_vocab(text_model) : nullptr), + text_model (text_model), batch_max_tokens(ctx_params.batch_max_tokens) { if (ctx_params.image_marker != nullptr) { @@ -1033,6 +1036,10 @@ struct mtmd_context { aud_beg = "<sound>"; audio_preproc = std::make_unique<mtmd_audio_preprocessor_whisper>(ctx_a); } break; + case PROJECTOR_TYPE_VOXTRAL_RT: + { + // audio and text share positions, only usable via mtmd_audio_stream + } break; case PROJECTOR_TYPE_ULTRAVOX: case PROJECTOR_TYPE_GLMA: case PROJECTOR_TYPE_MERALION: @@ -1093,8 +1100,9 @@ struct mtmd_context { } // initialize audio preprocessor - GGML_ASSERT(audio_preproc != nullptr); - audio_preproc->initialize(); + if (audio_preproc) { + audio_preproc->initialize(); + } } // get clip ctx based on chunk type @@ -1690,6 +1698,10 @@ struct mtmd_tokenizer { samples = (const float *)buf.data(); n_samples = buf.size() / sizeof(float); } + if (!ctx->audio_preproc) { + LOG_ERR("%s: this model takes audio input only via mtmd_audio_stream\n", __func__); + return 2; + } bool ok = ctx->audio_preproc->preprocess(samples, n_samples, mel_spec_chunks); if (!ok) { LOG_ERR("Unable to preprocess audio\n"); @@ -2919,7 +2931,7 @@ void mtmd_debug_preprocess_image(mtmd_context * ctx, const std::vector<uint8_t> } void mtmd_debug_preprocess_audio(mtmd_context * ctx, const std::vector<float> & samples) { - if (!ctx->ctx_a) { + if (!ctx->ctx_a || !ctx->audio_preproc) { LOG_ERR("%s: model does not support audio input\n", __func__); return; } @@ -2982,3 +2994,304 @@ mtmd_memory_usage mtmd_get_memory_usage(const char * mmproj_fname, return {{}, -1, false}; } } + +// +// realtime audio input +// + +struct mtmd_audio_stream { + mtmd_context * ctx; + clip_ctx * ctx_clip; + clip_enc_state * enc_state = nullptr; + + // model constants + int n_mel; + int n_fft; + int hop; + int n_embd; // decoder embedding size + int frames_per_tok; // mel frames per decoder row + int samples_per_tok; + int n_frames_max; // encoder frames per clip_encode() call + int n_step_mel; // mel frames per encoder call + int n_right_pad; // silence tokens added by flush() + float log_mel_max; // global_log_mel_max in params.json + + std::vector<llama_token> prefix; // text tokens paired with the first rows + + mtmd_audio_cache cache; + + // samples not yet turned into mel frames, buf[0] is at absolute index buf_pos + std::vector<float> buf; + int64_t buf_pos = 0; + int64_t n_samples = 0; // total samples pushed, including padding + int64_t n_frames = 0; // mel frames computed so far + + std::vector<float> mel; // pending mel frames, frame-major + std::vector<float> mel_ctx; // last 2 mel frames, the left context of the conv stem + std::vector<float> embd; // pending audio embeddings + size_t embd_off = 0; // rows of embd already consumed + std::vector<float> out; // rows of the last step + std::vector<float> tok_embd; + + int64_t n_rows = 0; // rows produced so far + bool pending_flush = false; + bool ended = false; + + ~mtmd_audio_stream() { + clip_enc_state_free(enc_state); + } +}; + +bool mtmd_support_audio_stream(const mtmd_context * ctx) { + return ctx->ctx_a && ctx->text_model && clip_get_projector_type(ctx->ctx_a) == PROJECTOR_TYPE_VOXTRAL_RT; +} + +mtmd_audio_stream_params mtmd_audio_stream_params_default(void) { + mtmd_audio_stream_params params; + params.step_ms = 320; + return params; +} + +static void mtmd_audio_stream_add_silence(mtmd_audio_stream * s, int64_t n) { + s->buf.insert(s->buf.end(), n, 0.0f); + s->n_samples += n; +} + +mtmd_audio_stream * mtmd_audio_stream_init(mtmd_context * ctx, mtmd_audio_stream_params params) { + if (!mtmd_support_audio_stream(ctx)) { + LOG_ERR("%s: the model does not support streaming audio input\n", __func__); + return nullptr; + } + const clip_hparams * hparams = clip_get_hparams(ctx->ctx_a); + const llama_vocab * vocab = ctx->vocab; + + auto s = std::make_unique<mtmd_audio_stream>(); + s->ctx = ctx; + s->ctx_clip = ctx->ctx_a; + s->n_mel = hparams->n_mel_bins; + s->n_fft = hparams->audio_n_fft; + s->hop = hparams->audio_hop_len; + s->n_embd = clip_n_mmproj_embd(ctx->ctx_a); + if (s->n_embd != ctx->n_embd_text) { + LOG_ERR("%s: mmproj embedding size %d does not match the text model (%d)\n", __func__, s->n_embd, ctx->n_embd_text); + return nullptr; + } + + // one decoder row per 80 ms of audio: the conv stem has stride 2, then frames are stacked + s->frames_per_tok = 2 * hparams->proj_stack_factor; + s->samples_per_tok = s->frames_per_tok * s->hop; + s->n_frames_max = (hparams->warmup_audio_size - 2) / 2; + s->log_mel_max = 1.5f; + + // the encoder works in whole decoder rows + int n_step = (int) ((int64_t) params.step_ms * hparams->audio_sample_rate / 1000 / s->hop); + n_step = std::max(n_step, s->frames_per_tok); + n_step = CLIP_ALIGN(n_step, s->frames_per_tok); + s->n_step_mel = std::min(n_step, 2 * s->n_frames_max); + + // prefix: BOS then one pad token per silence token and per delay token + // the output lags behind the audio by n_delay rows, the text model is conditioned on the same value + const int n_delay = llama_model_n_tokens_delay(ctx->text_model); + const int n_left_pad = 32; + s->n_right_pad = n_delay + 1 + 10; + { + const std::string pad_str = "[STREAMING_PAD]"; + std::vector<llama_token> pad(8); + const int n_pad = llama_tokenize(vocab, pad_str.c_str(), pad_str.size(), pad.data(), pad.size(), false, true); + if (n_pad != 1) { + LOG_ERR("%s: cannot find the %s token\n", __func__, pad_str.c_str()); + return nullptr; + } + s->prefix.push_back(llama_vocab_bos(vocab)); + s->prefix.insert(s->prefix.end(), n_left_pad + n_delay, pad[0]); + } + + s->cache.fill_sin_cos_table(s->n_fft); + s->cache.fill_hann_window(hparams->audio_window_len, true); + s->cache.fill_mel_filterbank_matrix(s->n_mel, s->n_fft, hparams->audio_sample_rate); + + s->enc_state = clip_enc_state_init(s->ctx_clip, s->n_frames_max); + if (!s->enc_state) { + return nullptr; + } + + // the first mel frame is centered on the first sample, so its window starts before the audio + s->buf.assign(s->n_fft / 2, 0.0f); + s->buf_pos = -s->n_fft / 2; + s->mel_ctx.assign(2 * s->n_mel, 0.0f); + + // silence before the audio + mtmd_audio_stream_add_silence(s.get(), (int64_t) n_left_pad * s->samples_per_tok); + + return s.release(); +} + +void mtmd_audio_stream_free(mtmd_audio_stream * stream) { + delete stream; +} + +int32_t mtmd_audio_stream_push(mtmd_audio_stream * s, const float * samples, size_t n_samples) { + if (s->ended) { + return 1; + } + s->buf.insert(s->buf.end(), samples, samples + n_samples); + s->n_samples += n_samples; + return 0; +} + +int32_t mtmd_audio_stream_flush(mtmd_audio_stream * s, bool end) { + if (s->ended) { + return 1; + } + // align to a whole row, then add the silence the model expects after the audio + const int64_t n_align = (s->samples_per_tok - s->n_samples % s->samples_per_tok) % s->samples_per_tok; + mtmd_audio_stream_add_silence(s, n_align + (int64_t) s->n_right_pad * s->samples_per_tok); + s->pending_flush = true; + s->ended = end; + return 0; +} + +// mel frame k is centered on sample k * hop, with force the samples past the end are taken as silence +static void mtmd_audio_stream_compute_mel(mtmd_audio_stream * s, bool force) { + const int half = s->n_fft / 2; + std::vector<float> window(s->n_fft); + std::vector<float> frame(s->n_mel); + + while (true) { + const int64_t start = s->n_frames * s->hop - half; + if (start + s->n_fft > s->n_samples && !(force && start + half < s->n_samples)) { + break; + } + for (int j = 0; j < s->n_fft; j++) { + const int64_t idx = start + j - s->buf_pos; + window[j] = idx < (int64_t) s->buf.size() ? s->buf[idx] : 0.0f; + } + mtmd_audio_log_mel_frame(s->cache, window.data(), s->n_fft, s->n_mel, frame.data()); + for (int m = 0; m < s->n_mel; m++) { + const float v = std::max(frame[m], s->log_mel_max - 8.0f); + s->mel.push_back((v + 4.0f) / 4.0f); + } + s->n_frames++; + } + + // drop the samples no frame will need again + const int64_t keep_from = s->n_frames * s->hop - half; + if (keep_from > s->buf_pos) { + s->buf.erase(s->buf.begin(), s->buf.begin() + (keep_from - s->buf_pos)); + s->buf_pos = keep_from; + } +} + +static bool mtmd_audio_stream_encode(mtmd_audio_stream * s, bool force) { + while (true) { + const int n_avail = (int) (s->mel.size() / s->n_mel); + int n = n_avail - n_avail % s->frames_per_tok; + if (!force) { + if (n_avail < s->n_step_mel) { + break; + } + n = std::min(n, s->n_step_mel); + } + n = std::min(n, 2 * s->n_frames_max); + if (n <= 0) { + break; + } + + // 2 context frames, then the new frames, stored mel-major + const int n_in = 2 + n; + std::vector<float> data((size_t) s->n_mel * n_in); + for (int m = 0; m < s->n_mel; m++) { + data[(size_t) m * n_in + 0] = s->mel_ctx[m]; + data[(size_t) m * n_in + 1] = s->mel_ctx[s->n_mel + m]; + for (int t = 0; t < n; t++) { + data[(size_t) m * n_in + 2 + t] = s->mel[(size_t) t * s->n_mel + m]; + } + } + + clip_image_f32 img; + img.set_size({n_in, s->n_mel}, false, /* is_audio */ true); + img.cpy_buf(data); + clip_image_f32_batch batch; + batch.is_audio = true; + batch.entries.push_back(std::move(img)); + + std::vector<float> res((size_t) s->n_embd * (n / s->frames_per_tok)); + clip_encode_params params; + params.n_threads = s->ctx->n_threads; + params.imgs = &batch; + params.out_embd = &res; + params.enc_state = s->enc_state; + if (!clip_encode(s->ctx_clip, ¶ms)) { + LOG_ERR("%s: failed to encode audio\n", __func__); + return false; + } + s->embd.insert(s->embd.end(), res.begin(), res.end()); + + s->mel_ctx.assign(s->mel.begin() + (size_t) (n - 2) * s->n_mel, s->mel.begin() + (size_t) n * s->n_mel); + s->mel.erase(s->mel.begin(), s->mel.begin() + (size_t) n * s->n_mel); + } + return true; +} + +// out rows = audio embeddings + token embeddings +static bool mtmd_audio_stream_add_rows(mtmd_audio_stream * s, const llama_token * tokens, size_t n) { + s->tok_embd.resize(n * s->n_embd); + if (llama_model_get_tok_embd_rows(s->ctx->text_model, tokens, n, s->tok_embd.data()) == 0) { + LOG_ERR("%s: invalid token\n", __func__); + return false; + } + const float * audio = s->embd.data() + s->embd_off * s->n_embd; + for (size_t i = 0; i < n * s->n_embd; i++) { + s->out.push_back(audio[i] + s->tok_embd[i]); + } + s->embd_off += n; + s->n_rows += n; + return true; +} + +int32_t mtmd_audio_stream_step(mtmd_audio_stream * s, llama_token last_token) { + const bool force = s->pending_flush; + s->pending_flush = false; + + mtmd_audio_stream_compute_mel(s, force); + if (!mtmd_audio_stream_encode(s, force)) { + return -2; + } + + // drop the consumed rows once in a while + if (s->embd_off >= 1024) { + s->embd.erase(s->embd.begin(), s->embd.begin() + s->embd_off * s->n_embd); + s->embd_off = 0; + } + + const int64_t n_avail = (int64_t) (s->embd.size() / s->n_embd) - s->embd_off; + s->out.clear(); + + if (s->n_rows == 0) { + // the prefix rows go out together, the caller samples from the last one + const int64_t n_prefix = (int64_t) s->prefix.size(); + if (n_avail < n_prefix) { + return s->ended ? -1 : 0; + } + if (!mtmd_audio_stream_add_rows(s, s->prefix.data(), s->prefix.size())) { + return -2; + } + return (int32_t) n_prefix; + } + + if (n_avail == 0) { + return s->ended ? -1 : 0; + } + if (last_token == LLAMA_TOKEN_NULL) { + LOG_ERR("%s: the token sampled from the previous row is required\n", __func__); + return -2; + } + if (!mtmd_audio_stream_add_rows(s, &last_token, 1)) { + return -2; + } + return 1; +} + +const float * mtmd_audio_stream_get_output_embd(const mtmd_audio_stream * s) { + return s->out.data(); +} diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index bcd76a9d716a..df24c56568f7 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -363,6 +363,51 @@ struct mtmd_caps { }; MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname); +///////////////////////////////////////// +// EXPERIMENTAL API for realtime audio input, subjected to breaking changes +// +// a stream turns incoming PCM into decoder-ready embedding rows, one row per decoder position +// it keeps the encoder state between calls, so audio can be fed as it arrives +// how text tokens are mixed into the rows is a model detail hidden inside the stream +// the text model given to mtmd_init_from_file() is required + +// whether the current model supports mtmd_audio_stream +MTMD_API bool mtmd_support_audio_stream(const mtmd_context * ctx); + +struct mtmd_audio_stream; +typedef struct mtmd_audio_stream mtmd_audio_stream; + +struct mtmd_audio_stream_params { + int32_t step_ms; // audio encoded per encoder call, rounded up to the model step (default: 320) +}; +MTMD_API struct mtmd_audio_stream_params mtmd_audio_stream_params_default(void); + +// returns nullptr on failure +MTMD_API mtmd_audio_stream * mtmd_audio_stream_init(mtmd_context * ctx, struct mtmd_audio_stream_params params); +MTMD_API void mtmd_audio_stream_free(mtmd_audio_stream * stream); + +// samples are PCM F32 at mtmd_get_audio_sample_rate() +// returns 0 on success, 1 if the stream has ended +MTMD_API int32_t mtmd_audio_stream_push(mtmd_audio_stream * stream, const float * samples, size_t n_samples); + +// pad the stream so that all pending rows can be produced +// end = false: the stream stays open (use it on a pause, to get a partial transcript) +// end = true: no more push() is accepted +// returns 0 on success, 1 if the stream has ended +MTMD_API int32_t mtmd_audio_stream_flush(mtmd_audio_stream * stream, bool end); + +// produce the next decoder rows +// last_token is the token sampled from the previous row, LLAMA_TOKEN_NULL before the first sample +// return value: +// > 0 number of rows ready, read them with mtmd_audio_stream_get_output_embd() +// 0 more audio is needed +// -1 end of stream, nothing more will be produced +// -2 error +MTMD_API int32_t mtmd_audio_stream_step(mtmd_audio_stream * stream, llama_token last_token); + +// rows of the last step, the reading size (in bytes) is n_rows * llama_model_n_embd_inp(model) * sizeof(float) +MTMD_API const float * mtmd_audio_stream_get_output_embd(const mtmd_audio_stream * stream); + ///////////////////////////////////////// // EXPERIMENTAL API for audio generation, subjected to breaking changes @@ -493,6 +538,11 @@ struct mtmd_batch_deleter { }; using batch_ptr = std::unique_ptr<mtmd_batch, mtmd_batch_deleter>; +struct mtmd_audio_stream_deleter { + void operator()(mtmd_audio_stream * val) { mtmd_audio_stream_free(val); } +}; +using audio_stream_ptr = std::unique_ptr<mtmd_audio_stream, mtmd_audio_stream_deleter>; + struct bitmap { bitmap_ptr ptr; bitmap() : ptr(nullptr) {} From 43d57fb8ca78d1a39da69069cc865c9cab497b2f Mon Sep 17 00:00:00 2001 From: Neo Zhang <zhang.jianyu@outlook.com> Date: Fri, 16 Oct 2026 09:58:07 +0800 Subject: [PATCH 037/423] [SYCL] Enable ssl in build script and CI (#30456) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * enable OPENSSL=ON as default in build script, update it in CI and dockerfile * change to LLAMA_BUILD_BORINGSSL on windows * refactor * rm Download UI build * Apply suggestion from @CISC Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * Apply suggestion from @CISC Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> * Apply suggestion from @CISC Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> --------- Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co> --- .devops/intel.Dockerfile | 2 +- .github/workflows/build-sycl.yml | 7 ++- .github/workflows/release.yml | 7 ++- examples/sycl/build.sh | 19 ++++++-- examples/sycl/win-build-sycl.bat | 75 ++++++++++++++++++++++++-------- 5 files changed, 86 insertions(+), 24 deletions(-) diff --git a/.devops/intel.Dockerfile b/.devops/intel.Dockerfile index ced0d641a889..e20971d128d9 100644 --- a/.devops/intel.Dockerfile +++ b/.devops/intel.Dockerfile @@ -107,7 +107,7 @@ RUN mkdir /tmp/neo/ && cd /tmp/neo/ \ && dpkg --install *.deb RUN apt-get update \ - && apt-get install -y libgomp1 curl ffmpeg \ + && apt-get install -y libgomp1 curl ffmpeg openssl libssl-dev \ && apt autoremove -y \ && apt clean -y \ && rm -rf /tmp/* /var/tmp/* \ diff --git a/.github/workflows/build-sycl.yml b/.github/workflows/build-sycl.yml index 54380f40f50b..ecfab4e52403 100644 --- a/.github/workflows/build-sycl.yml +++ b/.github/workflows/build-sycl.yml @@ -86,6 +86,12 @@ jobs: wget -q "https://github.com/oneapi-src/level-zero/releases/download/v${LEVEL_ZERO_VERSION}/libze-dev_${LEVEL_ZERO_VERSION}%2B${LEVEL_ZERO_UBUNTU_VERSION}_amd64.deb" -O libze-dev.deb sudo apt-get install -y ./libze1.deb ./libze-dev.deb + - name: Dependencies + id: depends + run: | + sudo apt-get update + sudo apt-get install -y libssl-dev + - name: ccache uses: ggml-org/ccache-action@v1.2.24 with: @@ -111,7 +117,6 @@ jobs: -DGGML_SYCL=ON \ -DCMAKE_C_COMPILER=icx \ -DCMAKE_CXX_COMPILER=icpx \ - -DLLAMA_OPENSSL=OFF \ -DGGML_NATIVE=OFF \ -DGGML_SYCL_F16=${{ matrix.fp16 }} \ -DGGML_SYCL_SUPPORT_LEVEL_ZERO_API=ON \ diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 035ea958e26d..0a732a914659 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1453,6 +1453,12 @@ jobs: name: llama-ui.zip path: tools/ui/dist + - name: Dependencies + id: depends + run: | + sudo apt-get update + sudo apt-get install -y libssl-dev + - name: ccache uses: ggml-org/ccache-action@v1.2.24 with: @@ -1471,7 +1477,6 @@ jobs: -DCMAKE_CXX_COMPILER=icpx \ -DCMAKE_INSTALL_RPATH='$ORIGIN' \ -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ - -DLLAMA_OPENSSL=OFF \ -DGGML_NATIVE=OFF \ -DGGML_SYCL_F16=${{ matrix.fp16 }} time cmake --build build --config Release -j $(nproc) diff --git a/examples/sycl/build.sh b/examples/sycl/build.sh index bec96886b19d..dc9e44a8ea28 100755 --- a/examples/sycl/build.sh +++ b/examples/sycl/build.sh @@ -4,15 +4,18 @@ # SPDX-License-Identifier: MIT print_usage() { - echo "Usage: ./build.sh [fp32|fp16] [--help]" + echo "Usage: ./build.sh [fp32|fp16] [ssl|nossl] [--help]" echo "" echo "Options:" echo " fp32 Build with FP32 precision (default)" echo " fp16 Build with FP16 precision (faster for long-prompt inference)" + echo " ssl Build with OpenSSL (default)" + echo " nossl Build without OpenSSL" echo " --help Print this help message" } PRECISION=fp32 +SSL=ON for arg in "$@"; do case "$arg" in @@ -23,6 +26,12 @@ for arg in "$@"; do fp32|fp16) PRECISION="$arg" ;; + ssl) + SSL=ON + ;; + nossl) + SSL=OFF + ;; *) echo "Error: unknown option '$arg'" print_usage @@ -37,10 +46,13 @@ source /opt/intel/oneapi/setvars.sh if [ "$PRECISION" = "fp16" ]; then #for FP16 - cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DGGML_SYCL_F16=ON -DLLAMA_OPENSSL=ON # faster for long-prompt inference + echo "cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DGGML_SYCL_F16=ON -DLLAMA_OPENSSL=$SSL # faster for long-prompt inference +" + cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DGGML_SYCL_F16=ON -DLLAMA_OPENSSL="$SSL" # faster for long-prompt inference else #for FP32 - cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DLLAMA_OPENSSL=ON + echo "cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DLLAMA_OPENSSL=$SSL" + cmake .. -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DLLAMA_OPENSSL="$SSL" fi #build example/main @@ -50,4 +62,5 @@ fi #cmake --build . --config Release --target llama-bench #build all binary +echo "cmake --build . --config Release -j$((($(nproc)+1)/2)) -v" cmake --build . --config Release -j$((($(nproc)+1)/2)) -v diff --git a/examples/sycl/win-build-sycl.bat b/examples/sycl/win-build-sycl.bat index 9a82edbefe62..4b3906a6371f 100644 --- a/examples/sycl/win-build-sycl.bat +++ b/examples/sycl/win-build-sycl.bat @@ -3,23 +3,62 @@ :: Copyright (C) 2024 Intel Corporation :: SPDX-License-Identifier: MIT -IF /I "%1"=="--help" ( - echo Usage: win-build-sycl.bat [fp32^|fp16] [--help] - echo. - echo Options: - echo fp32 Build with FP32 precision ^(default^) - echo fp16 Build with FP16 precision ^(faster for long-prompt inference^) - echo --help Print this help message - exit /B 0 -) +SET "PRECISION=fp32" +SET "SSL=ON" +SET "OPENSSL_OPTION=" -SET PRECISION=%1 -IF "%PRECISION%"=="" SET PRECISION=fp32 -IF /I NOT "%PRECISION%"=="fp32" IF /I NOT "%PRECISION%"=="fp16" ( - echo Error: invalid value '%PRECISION%'. Use 'fp32' or 'fp16'. - echo Usage: win-build-sycl.bat [fp32^|fp16] [--help] - exit /B 1 -) +:PARSE_ARGS +IF "%~1"=="" GOTO ARGS_DONE +IF /I "%~1"=="--help" GOTO HELP +IF /I "%~1"=="fp32" GOTO ARG_FP32 +IF /I "%~1"=="fp16" GOTO ARG_FP16 +IF /I "%~1"=="ssl" GOTO ARG_SSL +IF /I "%~1"=="nossl" GOTO ARG_NOSSL +GOTO INVALID_ARG + +:ARG_FP32 +SET "PRECISION=fp32" +GOTO NEXT_ARG + +:ARG_FP16 +SET "PRECISION=fp16" +GOTO NEXT_ARG + +:ARG_SSL +SET "SSL=ON" +SET "OPENSSL_OPTION=" +GOTO NEXT_ARG + +:ARG_NOSSL +SET "SSL=OFF" +SET "OPENSSL_OPTION=-DLLAMA_OPENSSL=OFF" +GOTO NEXT_ARG + +:NEXT_ARG +SHIFT +GOTO PARSE_ARGS + +:HELP +SET "ARG_ERROR=0" +GOTO PRINT_USAGE + +:INVALID_ARG +SET "ARG_ERROR=1" +echo Error: unknown option '%~1' + +:PRINT_USAGE +echo Usage: win-build-sycl.bat [fp32^|fp16] [ssl^|nossl] [--help] +echo. +echo Options: +echo fp32 Build with FP32 precision ^(default^) +echo fp16 Build with FP16 precision ^(faster for long-prompt inference^) +echo ssl Build with BoringSSL ^(default^) +echo nossl Build without BoringSSL +echo --help Print this help message +IF "%ARG_ERROR%"=="1" exit /B 1 +exit /B 0 + +:ARGS_DONE IF not exist build (mkdir build) cd build @@ -31,10 +70,10 @@ if %errorlevel% neq 0 goto ERROR IF /I "%PRECISION%"=="fp16" ( :: for FP16 :: faster for long-prompt inference - cmake -G "MinGW Makefiles" .. -DLLAMA_OPENSSL=OFF -DGGML_SYCL=ON -DCMAKE_CXX_COMPILER=icx -DBUILD_SHARED_LIBS=ON -DCMAKE_BUILD_TYPE=Release -DGGML_SYCL_F16=ON + cmake -G "MinGW Makefiles" .. -DLLAMA_BUILD_BORINGSSL=%SSL% %OPENSSL_OPTION% -DGGML_SYCL=ON -DCMAKE_CXX_COMPILER=icx -DBUILD_SHARED_LIBS=ON -DCMAKE_BUILD_TYPE=Release -DGGML_SYCL_F16=ON ) ELSE ( :: for FP32 - cmake -G "Ninja" .. -DLLAMA_OPENSSL=OFF -DGGML_SYCL=ON -DCMAKE_C_COMPILER=cl -DCMAKE_CXX_COMPILER=icx -DBUILD_SHARED_LIBS=ON -DCMAKE_BUILD_TYPE=Release + cmake -G "Ninja" .. -DLLAMA_BUILD_BORINGSSL=%SSL% %OPENSSL_OPTION% -DGGML_SYCL=ON -DCMAKE_C_COMPILER=cl -DCMAKE_CXX_COMPILER=icx -DBUILD_SHARED_LIBS=ON -DCMAKE_BUILD_TYPE=Release ) if %errorlevel% neq 0 goto ERROR From e707790c73590c2ff5303045b714afb0c5e22035 Mon Sep 17 00:00:00 2001 From: Frosty40 <newjordan@gmail.com> Date: Thu, 15 Oct 2026 20:59:38 -0500 Subject: [PATCH 038/423] sycl: support reordered Q8_0 MoE weights (#29337) Preserve the existing MoE accumulation order. Assisted-by: Codex Assisted-by: Angel --- ggml/src/ggml-sycl/ggml-sycl.cpp | 49 ++++++++++++++++++++++++++++++-- ggml/src/ggml-sycl/mmvq.cpp | 29 +++++++++++++++---- ggml/src/ggml-sycl/vecdotq.hpp | 7 +++-- 3 files changed, 74 insertions(+), 11 deletions(-) diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 6759af0edc40..1cfc84126412 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -4442,6 +4442,47 @@ static bool reorder_qw_q8_0(uint8_t * data_device, const int ncols, const int nr return true; } +// Reorder each Q8_0 expert slice into a self-contained [qs][d] SoA layout. +static bool reorder_qw_q8_0_moe(uint8_t * data_device, size_t expert_bytes, int64_t n_expert, dpct::queue_ptr stream) { + GGML_ASSERT(expert_bytes % sizeof(block_q8_0) == 0); + const int blocks_per_expert = (int) (expert_bytes / sizeof(block_q8_0)); + const size_t total_bytes = expert_bytes * (size_t) n_expert; + + sycl_reorder_temp_buffer tmp(stream, total_bytes); + if (!tmp) { + GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, total_bytes); + return false; + } + uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr); + + sycl::event copy_event; + SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, total_bytes))); + if (!g_ggml_sycl_use_async_mem_op) { + copy_event.wait(); + } + + const int total_blocks = blocks_per_expert * (int) n_expert; + auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) { + const int gb = gb_; + const int e = gb / blocks_per_expert; + const int ib = gb % blocks_per_expert; + const block_q8_0 * x = (const block_q8_0 *) (tmp_buf + (size_t) e * expert_bytes); + uint8_t * base = data_device + (size_t) e * expert_bytes; + + auto * qs_ptr = base; + auto * d_ptr = (sycl::half *) (qs_ptr + (size_t) blocks_per_expert * QK8_0); + + for (int j = 0; j < QK8_0; ++j) { + *((int8_t *) qs_ptr + ib * QK8_0 + j) = x[ib].qs[j]; + } + d_ptr[ib] = x[ib].d; + }); + if (!g_ggml_sycl_use_async_mem_op) { + reorder_event.wait_and_throw(); + } + return true; +} + static bool reorder_qw_q4_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) { GGML_ASSERT(size % sizeof(block_q4_K) == 0); GGML_ASSERT(offset % sizeof(block_q4_K) == 0); @@ -4966,6 +5007,8 @@ static bool reorder_qw(const ggml_tensor * src0, dpct::queue_ptr stream) { if (src0->ne[2] > 1) { GGML_ASSERT((size_t) size == (size_t) src0->ne[2] * src0->nb[2]); switch (src0->type) { + case GGML_TYPE_Q8_0: + return reorder_qw_q8_0_moe(data_device, src0->nb[2], src0->ne[2], stream); case GGML_TYPE_Q4_K: return reorder_qw_q4_k_moe(data_device, src0->nb[2], src0->ne[2], stream); case GGML_TYPE_Q5_K: @@ -5051,8 +5094,8 @@ static void opt_for_reorder_id(ggml_backend_sycl_context * ctx, const ggml_tenso if (!g_ggml_sycl_enable_optimize || !ctx->opt_feature.reorder) { return; } - if (src0->type != GGML_TYPE_Q4_K && src0->type != GGML_TYPE_Q5_K && src0->type != GGML_TYPE_Q6_K && - src0->type != GGML_TYPE_MXFP4) { + if (src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_Q4_K && + src0->type != GGML_TYPE_Q5_K && src0->type != GGML_TYPE_Q6_K && src0->type != GGML_TYPE_MXFP4) { return; } // The MXFP4 reorder kernels use 8-byte vector loads, so every expert slice must stay aligned. @@ -5438,7 +5481,7 @@ static bool ggml_sycl_mul_mat_id_mmvq_fused( const int n_experts_used = (int) n_ids_per_group; const int nrows = (int) src0->ne[1]; - // Lazily reorder the (Q4_K) expert weights into a per-expert SoA layout, then run the reorder + // Lazily reorder supported expert weights into a per-expert SoA layout, then run the reorder // GEMV. Placed after the bail checks so a non-dispatchable op does not pay the reorder cost. opt_for_reorder_id(&ctx, src0); const ggml_tensor_extra_gpu * src0_extra = diff --git a/ggml/src/ggml-sycl/mmvq.cpp b/ggml/src/ggml-sycl/mmvq.cpp index eade8949fbb1..e9057777cb4f 100644 --- a/ggml/src/ggml-sycl/mmvq.cpp +++ b/ggml/src/ggml-sycl/mmvq.cpp @@ -3202,9 +3202,11 @@ static void mul_mat_vec_q_moe_reorder( const auto sg = item_ct1.get_sub_group(); + constexpr int vdr_mmvq = reorder_vec_dot_q_sycl::gtype == GGML_TYPE_Q8_0 ? + VDR_Q8_0_Q8_1_MMVQ : block_traits::vdr_mmvq; const int blocks_per_row = ncols / block_traits::qk; - constexpr int blocks_per_subgroup = ceil_div(block_traits::vdr_mmvq * WARP_SIZE, block_traits::qi); - constexpr int block_elements_per_subgroup = block_traits::qi / block_traits::vdr_mmvq; + constexpr int blocks_per_subgroup = ceil_div(vdr_mmvq * WARP_SIZE, block_traits::qi); + constexpr int block_elements_per_subgroup = block_traits::qi / vdr_mmvq; const int nblocks = nrows * (ncols / block_traits::qk); static_assert(blocks_per_subgroup > 0); @@ -3223,12 +3225,24 @@ static void mul_mat_vec_q_moe_reorder( #pragma unroll for (int elem = 0; elem < block_elements_per_subgroup; elem += WARP_SIZE) { - const int iqs = elem + block_traits::vdr_mmvq * (sg.get_local_linear_id() % block_elements_per_subgroup); - partial_sum += reorder_vec_dot_q_sycl()(vx, bx_offset, d_offset, q8_1_quant_ptr, q8_1_ds_ptr, iqs); + const int iqs = elem + vdr_mmvq * (sg.get_local_linear_id() % block_elements_per_subgroup); + if constexpr (reorder_vec_dot_q_sycl::gtype == GGML_TYPE_Q8_0) { + partial_sum += reorder_vec_dot_q_sycl()(vx, bx_offset, d_offset, q8_1_quant_ptr, q8_1_ds_ptr, iqs, vdr_mmvq); + } else { + partial_sum += reorder_vec_dot_q_sycl()(vx, bx_offset, d_offset, q8_1_quant_ptr, q8_1_ds_ptr, iqs); + } } } - auto sum = sycl::reduce_over_group(sg, partial_sum, std::plus<>()); + float sum = partial_sum; + if constexpr (reorder_vec_dot_q_sycl::gtype == GGML_TYPE_Q8_0) { + // Keep the accumulation order of the non-reordered MoE kernel. + for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { + sum += dpct::permute_sub_group_by_xor(sg, sum, mask); + } + } else { + sum = sycl::reduce_over_group(sg, partial_sum, std::plus<>()); + } if (sg.leader()) { dst[row] = sum; } @@ -3269,6 +3283,11 @@ bool ggml_sycl_mul_mat_vec_q_id_reorder( size_t src1_row_stride, dpct::queue_ptr stream) { switch (src0_type) { + case GGML_TYPE_Q8_0: + launch_mul_mat_vec_q_moe_reorder<reorder_vec_dot_q_sycl<GGML_TYPE_Q8_0>>( + vx_base, vy, ids_dev, dst_base, ncols, nrows, n_experts_used, + expert_weight_stride, dst_row_stride, src1_row_stride, stream); + return true; case GGML_TYPE_Q4_K: launch_mul_mat_vec_q_moe_reorder<reorder_vec_dot_q_sycl<GGML_TYPE_Q4_K>>( vx_base, vy, ids_dev, dst_base, ncols, nrows, n_experts_used, diff --git a/ggml/src/ggml-sycl/vecdotq.hpp b/ggml/src/ggml-sycl/vecdotq.hpp index d1b6092bea9f..57360a004b64 100644 --- a/ggml/src/ggml-sycl/vecdotq.hpp +++ b/ggml/src/ggml-sycl/vecdotq.hpp @@ -483,7 +483,8 @@ template <> struct reorder_vec_dot_q_sycl<GGML_TYPE_Q8_0> { __dpct_inline__ float operator()(const void * __restrict__ vbq, const std::pair<int, int> ibx_offset, const std::pair<int, int> d_offset, const int8_t * q8_1_quant_ptr, - const sycl::half2 * q8_1_ds, const int & iqs) { + const sycl::half2 * q8_1_ds, const int & iqs, + const int vdr = q8_0_traits::vdr_mmvq) { const uint8_t * base = static_cast<const uint8_t *>(vbq); const int8_t * qs = reinterpret_cast<const int8_t *>(base + ibx_offset.first); const ggml_half d = *reinterpret_cast<const ggml_half *>(base + d_offset.first); @@ -492,14 +493,14 @@ template <> struct reorder_vec_dot_q_sycl<GGML_TYPE_Q8_0> { int u[q8_0_traits::vdr_mmvq]; #pragma unroll - for (size_t i = 0; i < q8_0_traits::vdr_mmvq; ++i) { + for (int i = 0; i < vdr; ++i) { v[i] = get_int_from_int8(qs, iqs + i); u[i] = get_int_from_int8_aligned(q8_1_quant_ptr, iqs + i); } int sumi = 0; #pragma unroll - for (size_t i = 0; i < q8_0_traits::vdr_mmvq; ++i) { + for (int i = 0; i < vdr; ++i) { sumi = dpct::dp4a(v[i], u[i], sumi); } From b6334e37e6ef286fc5091a41b6561b2113d915fb Mon Sep 17 00:00:00 2001 From: Gaurav Garg <gaugarg@nvidia.com> Date: Fri, 16 Oct 2026 08:20:26 +0530 Subject: [PATCH 039/423] Limit max CUDA pinned host buffer size to 4 GB (#30449) --- ggml/src/ggml-cuda/ggml-cuda.cu | 8 +++++++- src/llama-model.cpp | 3 ++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 7c45d0d1fb23..6fa22289f331 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1348,6 +1348,12 @@ static ggml_backend_buffer_t ggml_backend_cuda_host_buffer_type_alloc_buffer(ggm return buffer; } +static size_t ggml_backend_cuda_host_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { + GGML_UNUSED(buft); + // very large pinned allocations can fail (e.g. on Windows) and fall back to pageable memory, smaller ones succeed + return 4ull*1024*1024*1024; +} + ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type() { static struct ggml_backend_buffer_type ggml_backend_cuda_buffer_type_host = { /* .iface = */ { @@ -1355,7 +1361,7 @@ ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type() { /* .alloc_buffer = */ ggml_backend_cuda_host_buffer_type_alloc_buffer, /* .alloc_buffer_n = */ NULL, /* .get_alignment = */ ggml_backend_cpu_buffer_type()->iface.get_alignment, - /* .get_max_size = */ NULL, // defaults to SIZE_MAX + /* .get_max_size = */ ggml_backend_cuda_host_buffer_type_get_max_size, /* .get_alloc_size = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size, /* .get_alloc_size_n = */ NULL, /* .is_host = */ ggml_backend_cpu_buffer_type()->iface.is_host, diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 5d48d702c5b8..3da32333b10a 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1950,7 +1950,8 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { if (buf == nullptr) { throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); } - if (use_mlock && ggml_backend_buffer_is_host(buf)) { + // a multi_buffer has no base to lock, its parts are usually pinned memory (host buffer types with a max size) + if (use_mlock && ggml_backend_buffer_is_host(buf) && ggml_backend_buffer_get_base(buf) != nullptr) { pimpl->mlock_bufs.emplace_back(new llama_mlock); auto & mlock_buf = pimpl->mlock_bufs.back(); mlock_buf->init (ggml_backend_buffer_get_base(buf)); From 6afa9a3134dc4833ab58e625464aa67e5fc7eea5 Mon Sep 17 00:00:00 2001 From: Aman Gupta <amangupta052@gmail.com> Date: Fri, 16 Oct 2026 08:44:36 +0530 Subject: [PATCH 040/423] llama : share the k-pool input views between layers (#30613) The scheduler uploads each view of a host input on its own, so the per-layer reshape of new_pool_idxs (and of sel_mask in glm5-next) was copied once per layer. --- src/models/glm5-next.cpp | 9 +++++++-- src/models/qwen4exp.cpp | 5 ++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index a1b722950730..90870bd9675d 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -271,8 +271,10 @@ class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i { ggml_tensor * pool_mask = nullptr; // F32/F16 [n_pool, n_tokens] ggml_tensor * tail_idxs = nullptr; // I32 [kpool - 1, n_tokens] ggml_tensor * sel_mask = nullptr; // F32 [n_sel, 1, 1, n_tokens] 0 for live selection slots, -inf for dead ones + ggml_tensor * sel_mask_2d = nullptr; // F32 [n_sel, n_tokens] view of sel_mask, shared by the layers // n_new is never below 1, see build_inp_kpool ggml_tensor * new_pool_idxs = nullptr; // I32 [kpool, n_new] members of the pools completed this ubatch + ggml_tensor * new_pool_flat = nullptr; // I32 [kpool*n_new] view of new_pool_idxs, shared by the layers ggml_tensor * new_pool_rep = nullptr; // I64 [n_new] cell to write each new pooled key into const llama_memory_hybrid_idx_context * mctx; @@ -326,12 +328,15 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui ggml_set_input(inp->sel_mask); // set_input_kpool always fills the mask, so it stays allocated in every graph ggml_build_forward_expand(gf, inp->sel_mask); + // one view for all layers: each view of a host input is a separate upload + inp->sel_mask_2d = ggml_reshape_2d(ctx0, inp->sel_mask, n_sel, n_tokens); } inp->n_new = n_new; inp->new_pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, n_new); ggml_set_input(inp->new_pool_idxs); + inp->new_pool_flat = ggml_reshape_1d(ctx0, inp->new_pool_idxs, kpool*n_new); // one scatter row per new pool, each a distinct rep row (see kpool_build_state) inp->new_pool_rep = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new); ggml_set_input(inp->new_pool_rep); @@ -917,7 +922,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( const int64_t n_kv = mctx_lid->get_n_kv(); // Pool the entries completed by this ubatch. The last one is a dummy when the ubatch completes none. - ggml_tensor * rows = kpool_cache.gather_key_gate(ggml_reshape_1d(ctx0, inp_kpool->new_pool_idxs, kpool*n_new)); + ggml_tensor * rows = kpool_cache.gather_key_gate(inp_kpool->new_pool_flat); rows = ggml_reshape_3d(ctx0, rows, 2*n_embd_indexer, kpool, n_new); ggml_tensor * pk = ggml_view_3d(ctx0, rows, n_embd_indexer, kpool, n_new, rows->nb[1], rows->nb[2], 0); @@ -1005,7 +1010,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( // Live slots (visible pools, real tail cells) address disjoint cells. Each dead slot writes its own dump row // n_kv + slot, so the scatter indices of a token are unique: idx = dump + live*(idx - dump), live = exp(mask). GGML_ASSERT(inp_kpool->sel_mask->ne[0] == n_sel && inp_kpool->sel_mask->ne[3] == n_tokens); - ggml_tensor * live = ggml_exp(ctx0, ggml_reshape_2d(ctx0, inp_kpool->sel_mask, n_sel, n_tokens)); + ggml_tensor * live = ggml_exp(ctx0, inp_kpool->sel_mask_2d); ggml_tensor * dump = ggml_arange(ctx0, (float) n_kv, (float) (n_kv + n_sel), 1.0f); ggml_tensor * idx_f = ggml_cast(ctx0, sel_idx, GGML_TYPE_F32); idx_f = ggml_add(ctx0, ggml_mul(ctx0, ggml_sub(ctx0, idx_f, dump), live), dump); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index e055abdf55d2..77430e7afdc4 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -682,6 +682,7 @@ class llama_model_qwen4exp::llm_graph_input_kpool : public llm_graph_input_i { ggml_tensor * pool_mask = nullptr; // F16 [n_pool, n_tokens] ggml_tensor * tail_idxs = nullptr; // I32 [kpool - 1, n_tokens] ggml_tensor * new_pool_idxs = nullptr; // I32 [kpool, n_new] members of the blocks to re-pool this ubatch + ggml_tensor * new_pool_flat = nullptr; // I32 [kpool*n_new] view of new_pool_idxs, shared by the QSA layers ggml_tensor * new_pool_rep = nullptr; // I64 [n_new] cell to write each new pooled key into ggml_tensor * new_pool_pos = nullptr; // I32 [4*n_new] M-RoPE position of each new block's first member @@ -724,6 +725,8 @@ llama_model_qwen4exp::llm_graph_input_kpool * llama_model_qwen4exp::graph::build inp->new_pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, inp->n_new); ggml_set_input(inp->new_pool_idxs); + // one view for all layers: each view of a host input is a separate upload + inp->new_pool_flat = ggml_reshape_1d(ctx0, inp->new_pool_idxs, kpool*inp->n_new); // one scatter row per new pool, each a distinct rep row (see kpool_build_state) inp->new_pool_rep = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, inp->n_new); ggml_set_input(inp->new_pool_rep); @@ -765,7 +768,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_qsa_sel( auto kpool_cache = mctx_hyb->get_kpool_access(ctx0, il, idx_dim); // pool only the blocks this ubatch completes or regroups - ggml_tensor * rows = kpool_cache.gather_key_gate(ggml_reshape_1d(ctx0, inp_kpool->new_pool_idxs, kpool*n_new)); + ggml_tensor * rows = kpool_cache.gather_key_gate(inp_kpool->new_pool_flat); rows = ggml_reshape_3d(ctx0, rows, idx_dim, kpool, n_new); // mean over the members; kpool is small, so summing slices beats a transpose plus sum_rows From b797b9ca816bbced1e71c29c8b28cb35b103b624 Mon Sep 17 00:00:00 2001 From: Max Krasnyansky <maxk@qti.qualcomm.com> Date: Thu, 15 Oct 2026 20:30:11 -0700 Subject: [PATCH 041/423] hexagon: tensor repacking rework and support for tiled and flat quantized and fp weights (#30635) * hex-repack: Remove need_repack hack and streamline the rules -- breaks GET_ROWS tests for K-Quants, will fix in the next round Add support for GET_ROWS with tiled Q5_K Add support for raw k-quant tensors * hex-get-rows: vectorize raw k-quant ukernels * hex-binary: refactor dma-loop thread functions (no need for macros, etc) * hex-quant: move/cleanup flat and tiled dequant ukernels * hex-repack: add support for repacking/tiling F16/F32/BF16 2D weights * hex-get-rows: remove redundant cast * hex-repack: proper filtering for BF16 that were leaking through during the tests * hex-repack: streamling tiling rules (broken MM for now) * hex-repack: fix mmid after tiling updates * hex-repack: rename repacked to tiled in all flags/funcs/kernels/vars * hex-mm: fix valid_rows calc in various kernels * hex-mm: remove unused hvx_mm_nx_2d_flat * hex-mm: cleanup and optimize hvx-tiled vec_dot_f16 * hex-mm: another pass at hvx-tiled f16 matmul * hex-get-rows: optimize small src1 * hex-get-rows: one more pass at recovering get-rows perf regression * hex-repack: export GGML_HEXAGON_RT_REPACK via run.py * hex-mm-tiled: fix accuracy regression with hmx-tiled-f16 matmul * hex-mm: fold tiled_f16 checks into the general flow * hex-get-rows: remove dead checks for buffer_is_hexagon and tiling during supports_op * hex-quant: proper tail handling in new dequant ukernels * hex-get-rows: remove unused code and cleanup support checks * hex-mm: cleanup ne1 padding/tiling checks * hex-mm: cleanup fusion checks for tiled tensors * hex-repack: add a little can_repack helper to reduce duplication * hex-repack: never repack host buffers * hex-repack: factor out repack functions into separate .h and .cpp * hex-tensor: cleanup tensor ifaces for tiled/padded sizes * hex-repack: update k-quant tiling comments and cleanup #defines * hex-dcache: factor out l2/dcache dirty tracking and flushing * hex-tensor: update docs and remove #defs duplication * hex-tensor: remove unused/incorrect get_data interface * hex-get-rows: cleanup vtcm layout build * tests: disable huge MM tests * hex-dcache: add missing header --- docs/backend/snapdragon/developer.md | 28 +- ggml/src/ggml-hexagon/CMakeLists.txt | 2 + ggml/src/ggml-hexagon/ggml-hexagon.cpp | 2232 +++-------------- ggml/src/ggml-hexagon/htp-opnode.h | 8 +- ggml/src/ggml-hexagon/htp-repack.cpp | 1000 ++++++++ ggml/src/ggml-hexagon/htp-repack.h | 64 + ggml/src/ggml-hexagon/htp/CMakeLists.txt | 2 +- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 8 +- ggml/src/ggml-hexagon/htp/get-rows-ops.c | 499 ++-- ggml/src/ggml-hexagon/htp/get-rows-ops.h | 83 +- .../ggml-hexagon/htp/hmx-mm-kernels-tiled.h | 64 +- ggml/src/ggml-hexagon/htp/htp-ctx.h | 9 +- .../htp/{htp-tensor.c => htp-dcache.c} | 8 +- ggml/src/ggml-hexagon/htp/htp-dcache.h | 20 + ggml/src/ggml-hexagon/htp/htp-ops.h | 13 +- ggml/src/ggml-hexagon/htp/htp-tensor.h | 105 +- .../ggml-hexagon/htp/hvx-mm-kernels-float.h | 205 +- .../ggml-hexagon/htp/hvx-mm-kernels-tiled.h | 219 +- ggml/src/ggml-hexagon/htp/hvx-quant.h | 457 +++- ggml/src/ggml-hexagon/htp/main.c | 8 +- ggml/src/ggml-hexagon/htp/matmul-ops.c | 1732 ++++++------- ggml/src/ggml-hexagon/htp/matmul-ops.h | 205 +- scripts/snapdragon/run.py | 3 + tests/test-backend-ops.cpp | 3 + 24 files changed, 3539 insertions(+), 3438 deletions(-) create mode 100644 ggml/src/ggml-hexagon/htp-repack.cpp create mode 100644 ggml/src/ggml-hexagon/htp-repack.h rename ggml/src/ggml-hexagon/htp/{htp-tensor.c => htp-dcache.c} (97%) create mode 100644 ggml/src/ggml-hexagon/htp/htp-dcache.h diff --git a/docs/backend/snapdragon/developer.md b/docs/backend/snapdragon/developer.md index 3608250dc729..e088390e8111 100644 --- a/docs/backend/snapdragon/developer.md +++ b/docs/backend/snapdragon/developer.md @@ -134,6 +134,18 @@ Writing high-performance operators for Hexagon requires following specific guide - **Distinction**: `octx->src[N]` pointers *can* be NULL by design and must be checked when optional. Examples include attention masks, optional bias or weights in fused kernels, and frequency factors. +### `supports_op` Independence from Buffer and Repack State + +- `supports_op` (and op-specific checkers such as `supported_mul_mat()`, `supported_get_rows()`) runs during GGML graph scheduling **before** buffers are allocated and before tensor data is populated. +- **Never check buffer usage or repacking state in `supports_op`**: + - `t->buffer` may be NULL or uninitialized. + - `t->data` may be NULL. + - `buffer->usage` is unknown at scheduling time. + - `is_tiled` or `extra->flags & GGML_HEXAGON_TENSOR_TILED` is not set yet. +- Operator support decisions must depend strictly on static tensor properties: operand types (`t->type`), shapes (`t->ne`), strides (`t->nb`), and operation parameters. +- Repacking (converting weights to tiled layout) is an internal optimization applied during `set_tensor()` or rt-repack. Non-weight tensors or test harness tensors with the same type remain in flat layout. +- Therefore, any supported type (such as BF16, F16, quantized types) must be supported both when repacked into tiles and when stored flat in memory. + ### Multiline Macro Formatting - Keep trailing backslashes in multiline `#define` macros cleanly aligned to a consistent column. @@ -190,7 +202,7 @@ sessions. Common partitioning logic is factored into reusable inline helpers in [`htp-tensor.h`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h): -1. [`htp_tensor_mdev_rows_per_chunk`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L71): +1. [`htp_tensor_mdev_rows_per_chunk`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L104): Determines the minimum number of rows per chunk so that the chunk byte size is a multiple of 128 bytes: ``` @@ -201,7 +213,7 @@ Common partitioning logic is factored into reusable inline helpers in Returns `false` if the tensor cannot be safely row-partitioned (such as unaligned base pointer, permuted layout, or non-128-byte aligned outer strides). -2. [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L98): +2. [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L131): Calculates the per-device work range `struct htp_tensor_mdev_range { uint32_t start; uint32_t count; }` given `total_units`, `units_per_chunk`, `mdev_idx`, `mdev_count`, and the precomputed `mdev_count_div`. Handles chunk distribution across devices, assigns remainder units to the last device, and automatically triggers @@ -240,11 +252,11 @@ if (nrows == 0) { For flat element-wise operations that write directly to DDR: - Partition total linear elements N = ne0 * ne1 * ne2 * ne3 in 128-byte cache line chunks (`elems_per_line = (elem_size == 4) ? 32 : 64`). - Requires strict 1D contiguity: - [`htp_tensor_is_contiguous(dst, elem_size)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L32) + [`htp_tensor_is_contiguous(dst, elem_size)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L65) and 128-byte aligned destination pointer - [`htp_tensor_mdev_data_aligned(dst)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L51). + [`htp_tensor_mdev_data_aligned(dst)`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L84). - If contiguous and aligned, pass `elems_per_line` to - [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L98); + [`htp_tensor_mdev_partition`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h#L131); otherwise pass 0 to trigger Device 0 fallback. ### Single-Device Fallback (Device 0) @@ -312,7 +324,7 @@ Multi-device execution synchronizes worker sessions through atomic fence slots a (Output Slice 0) (Output Slice 1) | | Tensor Cache Flush Tensor Cache Flush - (htp_tensor_flush_all) (htp_tensor_flush_all) + (htp_dcache_flush_dirty_tensors) (htp_dcache_flush_dirty_tensors) | | Post-Op/Batch Barrier ---------------------- Post-Op/Batch Barrier (htp_mdev_group_barrier) (htp_mdev_group_barrier) @@ -350,8 +362,8 @@ Multi-device execution synchronizes worker sessions through atomic fence slots a ### Tensor Cache Flush and Pipeline Completion - In the kernel, ensure all pushed DMA operations have been popped in strict FIFO order to drain the queue. -- Use [`htp_tensor_flush_all()`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h) to flush specific dirty tensors back to DDR: - - [`htp_tensor_flush_all()`](../../../ggml/src/ggml-hexagon/htp/htp-tensor.h) flushes modified tensor address ranges, or the +- Use [`htp_dcache_flush_dirty_tensors()`](../../../ggml/src/ggml-hexagon/htp/htp-dcache.h) to flush specific dirty tensors back to DDR: + - [`htp_dcache_flush_dirty_tensors()`](../../../ggml/src/ggml-hexagon/htp/htp-dcache.h) flushes modified tensor address ranges, or the full D-cache when their total size exceeds the flush threshold, ensuring peer devices and the host CPU observe consistent data in DDR. - Never signal completion before all DMA transfers are drained and dirty tensor flushes have completed. diff --git a/ggml/src/ggml-hexagon/CMakeLists.txt b/ggml/src/ggml-hexagon/CMakeLists.txt index 20c720263b73..277a4c863ad0 100644 --- a/ggml/src/ggml-hexagon/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/CMakeLists.txt @@ -50,6 +50,8 @@ endif() set(TARGET_NAME ggml-hexagon) ggml_add_backend_library(${TARGET_NAME} ggml-hexagon.cpp + htp-repack.cpp + htp-repack.h htp-drv.cpp htp-drv.h libdl.h diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 08c15ff786a9..54c474122408 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -77,6 +77,7 @@ #include "htp/im2col-ops.h" #include "htp_iface.h" #include "htp-drv.h" +#include "htp-repack.h" using intvec = std::vector<int>; using uintvec = std::vector<unsigned int>; @@ -164,7 +165,8 @@ static int opt_opbatch = 1280; // max number of ops in a batch static int opt_opqueue = 8; // max number of pending batches static int opt_optrace = 0; // trace buffer size per thread (0 means default) static int opt_oppoll = 0; // polling for batch completions -static int opt_opfusion = 1; // enable/disable op fusion +static int opt_opfusion = 1; // enable/disable op fusion +static int opt_rt_repack = 1; // runtime repacking: 0 = disabled, 1 = quantized only (default), 2 = always enum ggml_hexagon_fusion_flags { GGML_HEXAGON_FUSE_ALLREDUCE_ADD = (1 << 1), // 2 @@ -309,42 +311,12 @@ static void ggml_hexagon_dump_trace_events(const std::string & sess_name, const } enum ggml_hexagon_tensor_flags { - GGML_HEXAGON_TENSOR_REPACK = (1 << 0), + GGML_HEXAGON_TENSOR_TILED = (1 << 0), GGML_HEXAGON_TENSOR_WEIGHT = (1 << 1), GGML_HEXAGON_TENSOR_FENCE = (1 << 2), GGML_HEXAGON_TENSOR_FUSEABLE = (1 << 3), }; -static inline bool ggml_hexagon_is_repack_type(enum ggml_type type) { - return type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q4_1 || - type == GGML_TYPE_Q8_0 || type == GGML_TYPE_IQ4_NL || - type == GGML_TYPE_MXFP4 || type == GGML_TYPE_Q6_K || - type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K || - type == GGML_TYPE_Q3_K || type == GGML_TYPE_Q2_K; -} - -// Size of one repacked row in the DSP tiled layout. The Q3_K..Q6_K tiles store uncompressed scales/mins, -// so they are larger than the ggml blocks. For the other repack types (Q2_K included) the tile has the same size. -static inline size_t ggml_hexagon_tiled_row_size(enum ggml_type type, int64_t ne0) { - if (type == GGML_TYPE_Q6_K) { - return (size_t) (ne0 / 32) * (HTP_MM_WEIGHT_TILE_SIZE_Q6_K / 32); - } - if (type == GGML_TYPE_Q4_K) { - return (size_t) (ne0 / 32) * (HTP_MM_WEIGHT_TILE_SIZE_Q4_1 / 32); - } - if (type == GGML_TYPE_Q5_K) { - return (size_t) (ne0 / 32) * (HTP_MM_WEIGHT_TILE_SIZE_Q5_K / 32); - } - if (type == GGML_TYPE_Q3_K) { - return (size_t) (ne0 / 32) * (HTP_MM_WEIGHT_TILE_SIZE_Q3_K / 32); - } - return ggml_row_size(type, ne0); -} - -static inline bool ggml_hexagon_is_hmx_weight_type(enum ggml_type type) { - return type == GGML_TYPE_F16 || type == GGML_TYPE_F32 || ggml_hexagon_is_repack_type(type); -} - struct ggml_hexagon_session; static void ggml_hexagon_precompute_matmul_params( @@ -519,6 +491,32 @@ static inline bool ggml_hexagon_tensor_is_fuseable(const struct ggml_tensor * t) return (extra->flags & GGML_HEXAGON_TENSOR_FUSEABLE) != 0; } +static inline bool ggml_hexagon_can_rt_repack(const struct ggml_tensor * t) { + if (opt_rt_repack == 0 || !htp_is_tiled_weight_tensor(t)) { + return false; + } + if (t->buffer && ggml_backend_buft_is_host(t->buffer->buft)) { + return false; + } + return opt_rt_repack > 1 || ggml_is_quantized(t->type); +} + +static inline bool ggml_hexagon_tensor_is_tiled(const struct ggml_tensor * t) { + const auto * extra = (const ggml_hexagon_tensor_extra *) t->extra; + if (extra && (extra->flags & GGML_HEXAGON_TENSOR_TILED)) { + return true; + } + return ggml_hexagon_can_rt_repack(t); +} + +static inline bool ggml_hexagon_needs_repack(const struct ggml_tensor * t) { + const auto * extra = (const ggml_hexagon_tensor_extra *) t->extra; + if (extra && (extra->flags & GGML_HEXAGON_TENSOR_TILED)) { + return false; + } + return ggml_hexagon_can_rt_repack(t); +} + static inline bool ggml_hexagon_tensors_overlap(const struct ggml_tensor * a, const struct ggml_tensor * b) { const uintptr_t a0 = (uintptr_t) a->data; const uintptr_t b0 = (uintptr_t) b->data; @@ -600,8 +598,6 @@ struct ggml_hexagon_session { uint64_t cached_uid = 0; std::vector<htp_opnode> cached_nodes; - mutable std::unordered_set<const ggml_tensor *> needs_repack; - ggml_hexagon_mdev_group mdev; ggml_backend_dev_t dev = nullptr; ggml_backend_hexagon_device_context * dev_ctx = nullptr; @@ -850,1617 +846,54 @@ struct ggml_hexagon_fence_buffer : public ggml_hexagon_shared_buffer { if (ptr) { memset(ptr, 0, (size_t) n_slots * GGML_HEXAGON_FENCE_SLOT_SIZE); } - return ptr; - } - - void free_slot(void * ptr, uint32_t n_slots = 1) { - if (!ptr) return; - uint32_t slot = ((uint8_t *) ptr - base()) / GGML_HEXAGON_FENCE_SLOT_SIZE; - for (uint32_t i = 0; i < n_slots; i++) { - free_slots.push_back(slot + i); - } - } -}; - -inline uint8_t * ggml_hexagon_session::alloc_fence(uint32_t n_slots) { - uint8_t * ptr = fence_buf->alloc_slot(n_slots); - GGML_ASSERT(ptr); - return ptr; -} - -inline void ggml_hexagon_session::free_fence(void * ptr, uint32_t n_slots) { - if (fence_buf) { - fence_buf->free_slot(ptr, n_slots); - } -} - -static void ggml_backend_hexagon_buffer_free_buffer(ggml_backend_buffer_t buffer) { - auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(buffer->context); - sbuf->sess->unclone_buffer(sbuf); - delete sbuf; -} - -static void * ggml_backend_hexagon_buffer_get_base(ggml_backend_buffer_t buffer) { - auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(buffer->context); - return sbuf->base(); -} - -static enum ggml_status ggml_backend_hexagon_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { - auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(buffer->context); - auto sess = sbuf->sess; - - HEX_VERBOSE("ggml-hex: %s init-tensor %s : base %p data %p nbytes %zu\n", sess->c_name(), - tensor->name, (void *) sbuf->base(), tensor->data, ggml_nbytes(tensor)); - - auto extra = new ggml_hexagon_tensor_extra(); - sbuf->tensor_extra.push_back(extra); - - tensor->extra = extra; - if (ggml_hexagon_is_repack_type(tensor->type)) { - if (sess->needs_repack.count(tensor)) { - extra->flags |= GGML_HEXAGON_TENSOR_REPACK; - sess->needs_repack.erase(tensor); - } - } - - return GGML_STATUS_SUCCESS; -} - -// ** Repack helpers for tiled quantized weights - -static void unpack_q4_0_quants(uint8_t * qs, const block_q4_0 * x, unsigned int bi) { - static const int qk = QK4_0; - - for (unsigned int i = 0; i < qk / 2; ++i) { - const int x0 = (x->qs[i] & 0x0F); - const int x1 = (x->qs[i] >> 4); - qs[bi * qk + i + 0] = x0; - qs[bi * qk + i + qk / 2] = x1; - } -} - -static void pack_q4_0_quants(block_q4_0 * x, const uint8_t * qs, unsigned int bi) { - static const int qk = QK4_0; - - for (unsigned int i = 0; i < qk / 2; ++i) { - const uint8_t x0 = qs[bi * qk + i + 0]; - const uint8_t x1 = qs[bi * qk + i + qk / 2]; - x->qs[i] = x0 | (x1 << 4); - } -} - -static void unpack_q4_1_quants(uint8_t * qs, const block_q4_1 * x, unsigned int bi) { - static const int qk = QK4_1; - - for (unsigned int i = 0; i < qk / 2; ++i) { - const int x0 = (x->qs[i] & 0x0F); - const int x1 = (x->qs[i] >> 4); - qs[bi * qk + i + 0] = x0; - qs[bi * qk + i + qk / 2] = x1; - } -} - -static void pack_q4_1_quants(block_q4_1 * x, const uint8_t * qs, unsigned int bi) { - static const int qk = QK4_1; - - for (unsigned int i = 0; i < qk / 2; ++i) { - const uint8_t x0 = qs[bi * qk + i + 0]; - const uint8_t x1 = qs[bi * qk + i + qk / 2]; - x->qs[i] = x0 | (x1 << 4); - } -} - -static void unpack_mxfp4_quants(uint8_t * qs, const block_mxfp4 * x, unsigned int bi) { - static const int qk = QK_MXFP4; - - for (unsigned int i = 0; i < qk / 2; ++i) { - const int x0 = (x->qs[i] & 0x0F); - const int x1 = (x->qs[i] >> 4); - qs[bi * qk + i + 0] = x0; - qs[bi * qk + i + qk / 2] = x1; - } -} - -static void pack_mxfp4_quants(block_mxfp4 * x, const uint8_t * qs, unsigned int bi) { - static const int qk = QK_MXFP4; - - for (unsigned int i = 0; i < qk / 2; ++i) { - const uint8_t x0 = qs[bi * qk + i + 0]; - const uint8_t x1 = qs[bi * qk + i + qk / 2]; - x->qs[i] = x0 | (x1 << 4); - } -} - -// repack q4_0 data into q4_0_tiled tensor -static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { - const block_q4_0 * src_matrix = (const block_q4_0 *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - int n_col_tiles = ne1_padded / 32; - int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_0; - const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - - size_t slice_size = ne1 * ggml_row_size(t->type, ne0); - int64_t start_slice = offset / slice_size; - int64_t end_slice = (offset + size + slice_size - 1) / slice_size; - if (end_slice > ne2 * ne3) { - end_slice = ne2 * ne3; - } - - for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { - const block_q4_0 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - unpack_q4_0_quants(tile_quants[row], &src_slice[r * (ne0 / 32) + kt], 0); - } else { - memset(tile_quants[row], 8, 32); - } - } - - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; - } - } - - ggml_half * scale_dst = (ggml_half *)(tile_dst + 512); - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * (ne0 / 32) + kt].d : 0; - } - } - } - } -} - -// repack q4_0_tiled tensor into q4_0 data -static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t offset, size_t size) { - block_q4_0 * dst_matrix = (block_q4_0 *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - int n_col_tiles = ne1_padded / 32; - int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_0; - const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - - size_t slice_size = ne1 * ggml_row_size(t->type, ne0); - size_t row_size_bytes = ggml_row_size(t->type, ne0); - int64_t start_slice = offset / slice_size; - int64_t end_slice = (offset + size + slice_size - 1) / slice_size; - if (end_slice > ne2 * ne3) { - end_slice = ne2 * ne3; - } - - for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { - size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); - size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); - size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; - size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; - - int64_t start_row = slice_offset_start / row_size_bytes; - int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; - end_row = (std::min)(end_row, ne1); - - int start_ct = start_row / 32; - int end_ct = (end_row + 31) / 32; - end_ct = (std::min)(end_ct, n_col_tiles); - - block_q4_0 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_q4_0); - const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; - - for (int ct = start_ct; ct < end_ct; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - uint8_t val = tile_src[cp * 32 + row]; - tile_quants[row][2 * cp + 0] = val & 0x0F; - tile_quants[row][2 * cp + 1] = val >> 4; - } - } - - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r >= start_row && r < end_row && kt < ne0 / 32) { - pack_q4_0_quants(&dst_slice[(r - start_row) * (ne0 / 32) + kt], tile_quants[row], 0); - } - } - - const ggml_half * scale_src = (const ggml_half *)(tile_src + 512); - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r >= start_row && r < end_row && kt < ne0 / 32) { - dst_slice[(r - start_row) * (ne0 / 32) + kt].d = scale_src[row]; - } - } - } - } - } -} - -// repack q4_1 data into q4_1_tiled tensor -static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { - const block_q4_1 * src_matrix = (const block_q4_1 *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - int n_col_tiles = ne1_padded / 32; - int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_1; - const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - - size_t slice_size = ne1 * ggml_row_size(t->type, ne0); - int64_t start_slice = offset / slice_size; - int64_t end_slice = (offset + size + slice_size - 1) / slice_size; - if (end_slice > ne2 * ne3) { - end_slice = ne2 * ne3; - } - - for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { - const block_q4_1 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - unpack_q4_1_quants(tile_quants[row], &src_slice[r * (ne0 / 32) + kt], 0); - } else { - memset(tile_quants[row], 0, 32); - } - } - - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; - } - } - - ggml_half * scale_dst = (ggml_half *)(tile_dst + 512); - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - scale_dst[2 * row + 0] = src_slice[r * (ne0 / 32) + kt].d; - scale_dst[2 * row + 1] = src_slice[r * (ne0 / 32) + kt].m; - } else { - scale_dst[2 * row + 0] = 0; - scale_dst[2 * row + 1] = 0; - } - } - } - } - } -} - -// repack q4_1_tiled tensor into q4_1 data -static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t offset, size_t size) { - block_q4_1 * dst_matrix = (block_q4_1 *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - int n_col_tiles = ne1_padded / 32; - int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_1; - const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - - size_t slice_size = ne1 * ggml_row_size(t->type, ne0); - size_t row_size_bytes = ggml_row_size(t->type, ne0); - int64_t start_slice = offset / slice_size; - int64_t end_slice = (offset + size + slice_size - 1) / slice_size; - if (end_slice > ne2 * ne3) { - end_slice = ne2 * ne3; - } - - for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { - size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); - size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); - size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; - size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; - - int64_t start_row = slice_offset_start / row_size_bytes; - int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; - end_row = (std::min)(end_row, ne1); - - int start_ct = start_row / 32; - int end_ct = (end_row + 31) / 32; - end_ct = (std::min)(end_ct, n_col_tiles); - - block_q4_1 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_q4_1); - const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; - - for (int ct = start_ct; ct < end_ct; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - uint8_t val = tile_src[cp * 32 + row]; - tile_quants[row][2 * cp + 0] = val & 0x0F; - tile_quants[row][2 * cp + 1] = val >> 4; - } - } - - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r >= start_row && r < end_row && kt < ne0 / 32) { - pack_q4_1_quants(&dst_slice[(r - start_row) * (ne0 / 32) + kt], tile_quants[row], 0); - } - } - - const ggml_half * scale_src = (const ggml_half *)(tile_src + 512); - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r >= start_row && r < end_row && kt < ne0 / 32) { - dst_slice[(r - start_row) * (ne0 / 32) + kt].d = scale_src[2 * row]; - dst_slice[(r - start_row) * (ne0 / 32) + kt].m = scale_src[2 * row + 1]; - } - } - } - } - } -} - -// repack q8_0 data into q8_0_tiled tensor -static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { - const block_q8_0 * src_matrix = (const block_q8_0 *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - int n_col_tiles = ne1_padded / 32; - int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q8_0; - const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - - size_t slice_size = ne1 * ggml_row_size(t->type, ne0); - int64_t start_slice = offset / slice_size; - int64_t end_slice = (offset + size + slice_size - 1) / slice_size; - if (end_slice > ne2 * ne3) { - end_slice = ne2 * ne3; - } - - for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { - const block_q8_0 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - - for (int cp = 0; cp < 16; cp++) { - int col0 = cp * 2; - int col1 = col0 + 1; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - const block_q8_0 * b = (r < ne1 && kt < ne0 / 32) ? &src_slice[r * (ne0 / 32) + kt] : NULL; - tile_dst[cp * 64 + 2 * row + 0] = b ? b->qs[col0] : 0; - tile_dst[cp * 64 + 2 * row + 1] = b ? b->qs[col1] : 0; - } - } - - ggml_half * scale_dst = (ggml_half *)(tile_dst + 1024); - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * (ne0 / 32) + kt].d : 0; - } - } - } - } -} - -// repack q8_0_tiled tensor into q8_0 data -static void repack_tiled_q8_0(void * data, const ggml_tensor * t, size_t offset, size_t size) { - block_q8_0 * dst_matrix = (block_q8_0 *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - int n_col_tiles = ne1_padded / 32; - int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q8_0; - const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - - size_t slice_size = ne1 * ggml_row_size(t->type, ne0); - size_t row_size_bytes = ggml_row_size(t->type, ne0); - int64_t start_slice = offset / slice_size; - int64_t end_slice = (offset + size + slice_size - 1) / slice_size; - if (end_slice > ne2 * ne3) { - end_slice = ne2 * ne3; - } - - for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { - size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); - size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); - size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; - size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; - - int64_t start_row = slice_offset_start / row_size_bytes; - int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; - end_row = (std::min)(end_row, ne1); - - int start_ct = start_row / 32; - int end_ct = (end_row + 31) / 32; - end_ct = (std::min)(end_ct, n_col_tiles); - - block_q8_0 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_q8_0); - const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; - - for (int ct = start_ct; ct < end_ct; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - for (int cp = 0; cp < 16; cp++) { - int col0 = cp * 2; - int col1 = col0 + 1; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r >= start_row && r < end_row && kt < ne0 / 32) { - block_q8_0 & b = dst_slice[(r - start_row) * (ne0 / 32) + kt]; - b.qs[col0] = tile_src[cp * 64 + 2 * row + 0]; - b.qs[col1] = tile_src[cp * 64 + 2 * row + 1]; - } - } - } - - const ggml_half * scale_src = (const ggml_half *)(tile_src + 1024); - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r >= start_row && r < end_row && kt < ne0 / 32) { - dst_slice[(r - start_row) * (ne0 / 32) + kt].d = scale_src[row]; - } - } - } - } - } -} - -// repack mxfp4 data into mxfp4_tiled tensor -static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { - const block_mxfp4 * src_matrix = (const block_mxfp4 *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - int n_col_tiles = ne1_padded / 32; - int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_MXFP4; - const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - - size_t slice_size = ne1 * ggml_row_size(t->type, ne0); - int64_t start_slice = offset / slice_size; - int64_t end_slice = (offset + size + slice_size - 1) / slice_size; - if (end_slice > ne2 * ne3) { - end_slice = ne2 * ne3; - } - - for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { - const block_mxfp4 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - unpack_mxfp4_quants(tile_quants[row], &src_slice[r * (ne0 / 32) + kt], 0); - } else { - memset(tile_quants[row], 0, 32); - } - } - - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; - } - } - - uint8_t * scale_dst = tile_dst + 512; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * (ne0 / 32) + kt].e : 0; - } - } - } - } -} - -// repack mxfp4_tiled tensor into mxfp4 data -static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t offset, size_t size) { - block_mxfp4 * dst_matrix = (block_mxfp4 *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - int n_col_tiles = ne1_padded / 32; - int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_MXFP4; - const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - - size_t slice_size = ne1 * ggml_row_size(t->type, ne0); - size_t row_size_bytes = ggml_row_size(t->type, ne0); - int64_t start_slice = offset / slice_size; - int64_t end_slice = (offset + size + slice_size - 1) / slice_size; - if (end_slice > ne2 * ne3) { - end_slice = ne2 * ne3; - } - - for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { - size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); - size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); - size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; - size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; - - int64_t start_row = slice_offset_start / row_size_bytes; - int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; - end_row = (std::min)(end_row, ne1); - - int start_ct = start_row / 32; - int end_ct = (end_row + 31) / 32; - end_ct = (std::min)(end_ct, n_col_tiles); - - block_mxfp4 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_mxfp4); - const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; - - for (int ct = start_ct; ct < end_ct; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - uint8_t val = tile_src[cp * 32 + row]; - tile_quants[row][2 * cp + 0] = val & 0x0F; - tile_quants[row][2 * cp + 1] = val >> 4; - } - } - - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r >= start_row && r < end_row && kt < ne0 / 32) { - pack_mxfp4_quants(&dst_slice[(r - start_row) * (ne0 / 32) + kt], tile_quants[row], 0); - } - } - - const uint8_t * scale_src = tile_src + 512; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r >= start_row && r < end_row && kt < ne0 / 32) { - dst_slice[(r - start_row) * (ne0 / 32) + kt].e = scale_src[row]; - } - } - } - } - } -} - -// unsigned 6-bit value (0..63) of element e of a Q6_K block, same bit layout as dequantize_row_q6_K -static inline uint8_t q6_K_get_quant(const block_q6_K * b, int e) { - const int c = e / 128; - const int w = e % 128; - const int g = w / 32; - const int l = w % 32; - const uint8_t * ql = b->ql + c * 64; - const uint8_t * qh = b->qh + c * 32; - uint8_t lo, hi; - switch (g) { - case 0: lo = ql[l] & 0xF; hi = (qh[l] >> 0) & 3; break; - case 1: lo = ql[l + 32] & 0xF; hi = (qh[l] >> 2) & 3; break; - case 2: lo = ql[l] >> 4; hi = (qh[l] >> 4) & 3; break; - default: lo = ql[l + 32] >> 4; hi = (qh[l] >> 6) & 3; break; - } - return (uint8_t) (lo | (hi << 4)); -} - -// tile layout: see HTP_MM_WEIGHT_TILE_SIZE_Q6_K in htp/matmul-ops.h -static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { - GGML_ASSERT(offset == 0); - - const block_q6_K * src_matrix = (const block_q6_K *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - GGML_ASSERT(ne0 % QK_K == 0); - - const int n_col_tiles = ne1_padded / 32; - const int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q6_K; - const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; - - const int64_t sb_per_row = ne0 / QK_K; - - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_q6_K * src_slice = src_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - memset(matrix_dst, 0, matrix_size); // padding rows and the OR-ed nibbles below need zeroed tiles - - for (int64_t r = 0; r < ne1; r++) { - const int ct = (int) (r / 32); - const int row = (int) (r % 32); - const block_q6_K * src_row = src_slice + r * sb_per_row; - - for (int kt = 0; kt < n_k_tiles; kt++) { - const int kt_local = kt % 8; // k-tile within the super-block - const block_q6_K * b = &src_row[kt / 8]; - const float d = GGML_FP16_TO_FP32(b->d); - - uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size; - uint8_t * lo_pl = tile; - uint8_t * hi_pl = tile + 512; - ggml_half * sc_pl = (ggml_half *) (tile + 768); - - for (int lk = 0; lk < 32; lk++) { - const uint8_t q6 = q6_K_get_quant(b, kt_local * 32 + lk); - const int g = lk >> 2; - const int pos = row * 4 + (lk & 3); - lo_pl[(g >> 1) * 128 + pos] |= (uint8_t) ((q6 & 0xF) << ((g & 1) * 4)); - hi_pl[(g >> 2) * 128 + pos] |= (uint8_t) ((q6 >> 4) << ((g & 3) * 2)); - } - for (int sub = 0; sub < 2; sub++) { - sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) b->scales[kt_local * 2 + sub]); - } - } - } - } - } - - GGML_UNUSED(size); -} - -// Reverse of repack_q6_K_tiled. Unpacks quants losslessly and normalizes sub-block scales. Read-back only. -static void repack_tiled_q6_K(void * data, const ggml_tensor * t, size_t offset, size_t size) { - GGML_ASSERT(offset == 0); - - block_q6_K * dst_matrix = (block_q6_K *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - GGML_ASSERT(ne0 % QK_K == 0); - - const int n_col_tiles = ne1_padded / 32; - const int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q6_K; - const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; - - const int64_t sb_per_row = ne0 / QK_K; - - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_q6_K * dst_slice = dst_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int64_t r = 0; r < ne1; r++) { - const int ct = (int) (r / 32); - const int row = (int) (r % 32); - block_q6_K * dst_row = dst_slice + r * sb_per_row; - - for (int64_t sb = 0; sb < sb_per_row; sb++) { - block_q6_K * b = &dst_row[sb]; - memset(b, 0, sizeof(block_q6_K)); - - float sub_scales[16]; - for (int kt_local = 0; kt_local < 8; kt_local++) { - const int kt = sb * 8 + kt_local; - const uint8_t * tile = matrix_src + ((size_t) ct * n_k_tiles + kt) * tile_size; - const uint8_t * lo_pl = tile; - const uint8_t * hi_pl = tile + 512; - const ggml_half * sc_pl = (const ggml_half *) (tile + 768); - - const int c = kt_local / 4; - const int g = kt_local % 4; - uint8_t * ql = b->ql + c * 64; - uint8_t * qh = b->qh + c * 32; - - for (int lk = 0; lk < 32; lk++) { - const int g_tile = lk >> 2; - const int pos = row * 4 + (lk & 3); - const uint8_t lo = (lo_pl[(g_tile >> 1) * 128 + pos] >> ((g_tile & 1) * 4)) & 0xF; - const uint8_t hi = (hi_pl[(g_tile >> 2) * 128 + pos] >> ((g_tile & 3) * 2)) & 3; - - switch (g) { - case 0: - ql[lk] |= lo; - qh[lk] |= (hi << 0); - break; - case 1: - ql[lk + 32] |= lo; - qh[lk] |= (hi << 2); - break; - case 2: - ql[lk] |= (lo << 4); - qh[lk] |= (hi << 4); - break; - default: - ql[lk + 32] |= (lo << 4); - qh[lk] |= (hi << 6); - break; - } - } - - for (int sub = 0; sub < 2; sub++) { - sub_scales[kt_local * 2 + sub] = GGML_FP16_TO_FP32(sc_pl[sub * 32 + row]); - } - } - - float max_abs_scale = 0.0f; - for (int s = 0; s < 16; s++) { - float abs_scale = fabsf(sub_scales[s]); - if (abs_scale > max_abs_scale) { - max_abs_scale = abs_scale; - } - } - - if (max_abs_scale == 0.0f) { - b->d = GGML_FP32_TO_FP16(0.0f); - memset(b->scales, 0, sizeof(b->scales)); - } else { - float d_flt = max_abs_scale / 127.0f; - b->d = GGML_FP32_TO_FP16(d_flt); - float d_actual = GGML_FP16_TO_FP32(b->d); - float inv_d = (d_actual != 0.0f) ? (1.0f / d_actual) : 0.0f; - for (int s = 0; s < 16; s++) { - int sc = (int) roundf(sub_scales[s] * inv_d); - b->scales[s] = (int8_t) (std::max)(-128, (std::min)(127, sc)); - } - } - } - } - } - } - - GGML_UNUSED(size); -} - -// low 2 bits (0..3) of element e of a Q2_K or Q3_K block, same bit layout as dequantize_row_q2_K / q3_K -static inline uint8_t q2_3_K_get_low2(const uint8_t * qs, int e) { - const int c = e / 128; - const int j = (e % 128) / 32; - const int l = e % 32; - return (qs[c * 32 + l] >> (2 * j)) & 3; -} - -// hmask bit of element e of a Q3_K block, same bit layout as dequantize_row_q3_K -static inline bool q3_K_get_hbit(const block_q3_K * b, int e) { - const int c = e / 128; - const int j = (e % 128) / 32; - const int l = e % 32; - return (b->hmask[l] >> (c * 4 + j)) & 1; -} - -// signed 6-bit scale j (-32..31) of a Q3_K block, same packing as quantize_row_q3_K_ref -static inline int q3_K_get_scale(const uint8_t * scales, int j) { - const int lo = (j < 8) ? (scales[j] & 0xF) : (scales[j - 8] >> 4); - const int hi = (scales[8 + j % 4] >> (2 * (j / 4))) & 3; - return (lo | (hi << 4)) - 32; -} - -// read-back: find fp16 d and l[j] in [lmin, lmax] with fp16(d * l[j]) == prod[j] for all j, false if none -static bool hexagon_recover_k_scales(const ggml_half * prod, int n, int lmin, int lmax, ggml_half * d_out, int * l_out) { - int jmax = 0; - for (int j = 1; j < n; j++) { - if (fabsf(GGML_FP16_TO_FP32(prod[j])) > fabsf(GGML_FP16_TO_FP32(prod[jmax]))) { - jmax = j; - } - } - const float pmax = GGML_FP16_TO_FP32(prod[jmax]); - if (pmax == 0.0f) { - *d_out = GGML_FP32_TO_FP16(0.0f); - for (int j = 0; j < n; j++) { - l_out[j] = 0; - } - return true; - } - - // the quantizers put the largest scale at or near the range end, so try large |l| first - const int lext = (std::max)(-lmin, lmax); - for (int a = lext; a >= 1; a--) { - for (int sign : { 1, -1 }) { - const int lj = sign * a; - if (lj < lmin || lj > lmax) { - continue; - } - const ggml_half d0 = GGML_FP32_TO_FP16(pmax / (float) lj); - for (int ulp : { 0, -1, 1 }) { - ggml_half d = d0; - uint16_t bits; - memcpy(&bits, &d, sizeof(bits)); - bits = (uint16_t) (bits + ulp); - memcpy(&d, &bits, sizeof(bits)); - - const float df = GGML_FP16_TO_FP32(d); - if (!std::isfinite(df) || df == 0.0f) { - continue; - } - bool ok = true; - for (int j = 0; j < n && ok; j++) { - const int l = (int) roundf(GGML_FP16_TO_FP32(prod[j]) / df); - const ggml_half p = GGML_FP32_TO_FP16(df * (float) l); - ok = l >= lmin && l <= lmax && memcmp(&p, &prod[j], sizeof(p)) == 0; - l_out[j] = l; - } - if (ok) { - *d_out = d; - return true; - } - } - } - } - return false; -} - -// tile layout: see HTP_MM_WEIGHT_TILE_SIZE_Q3_K in htp/matmul-ops.h -static void repack_q3_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { - GGML_ASSERT(offset == 0); - - const block_q3_K * src_matrix = (const block_q3_K *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - GGML_ASSERT(ne0 % QK_K == 0); - - const int n_col_tiles = ne1_padded / 32; - const int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q3_K; - const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; - - const int64_t sb_per_row = ne0 / QK_K; - - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_q3_K * src_slice = src_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - memset(matrix_dst, 0, matrix_size); // padding rows and the OR-ed bits below need zeroed tiles - - for (int64_t r = 0; r < ne1; r++) { - const int ct = (int) (r / 32); - const int row = (int) (r % 32); - const block_q3_K * src_row = src_slice + r * sb_per_row; - - for (int kt = 0; kt < n_k_tiles; kt++) { - const int kt_local = kt % 8; // k-tile within the super-block - const block_q3_K * b = &src_row[kt / 8]; - const float d = GGML_FP16_TO_FP32(b->d); - - uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size; - uint8_t * lo_pl = tile; - uint8_t * neg_pl = tile + 256; - ggml_half * sc_pl = (ggml_half *) (tile + 384); - - for (int lk = 0; lk < 32; lk++) { - const int e = kt_local * 32 + lk; - const int g = lk >> 2; - const int pos = row * 4 + (lk & 3); - lo_pl[(g >> 2) * 128 + pos] |= (uint8_t) (q2_3_K_get_low2(b->qs, e) << ((g & 3) * 2)); - if (!q3_K_get_hbit(b, e)) { - neg_pl[pos] |= (uint8_t) (1 << g); - } - } - for (int sub = 0; sub < 2; sub++) { - sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) q3_K_get_scale(b->scales, kt_local * 2 + sub)); - } - } - } - } - } - - GGML_UNUSED(size); -} - -// Reverse of repack_q3_K_tiled. Unpacks quants losslessly and normalizes sub-block scales. Read-back only. -static void repack_tiled_q3_K(void * data, const ggml_tensor * t, size_t offset, size_t size) { - GGML_ASSERT(offset == 0); - - block_q3_K * dst_matrix = (block_q3_K *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - GGML_ASSERT(ne0 % QK_K == 0); - - const int n_col_tiles = ne1_padded / 32; - const int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q3_K; - const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; - - const int64_t sb_per_row = ne0 / QK_K; - - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_q3_K * dst_slice = dst_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int64_t r = 0; r < ne1; r++) { - const int ct = (int) (r / 32); - const int row = (int) (r % 32); - block_q3_K * dst_row = dst_slice + r * sb_per_row; - - for (int64_t sb = 0; sb < sb_per_row; sb++) { - block_q3_K * b = &dst_row[sb]; - memset(b, 0, sizeof(block_q3_K)); - - ggml_half sub_scales[16]; - for (int kt_local = 0; kt_local < 8; kt_local++) { - const int kt = sb * 8 + kt_local; - const uint8_t * tile = matrix_src + ((size_t) ct * n_k_tiles + kt) * tile_size; - const uint8_t * lo_pl = tile; - const uint8_t * neg_pl = tile + 256; - const ggml_half * sc_pl = (const ggml_half *) (tile + 384); - - for (int lk = 0; lk < 32; lk++) { - const int e = kt_local * 32 + lk; - const int g = lk >> 2; - const int pos = row * 4 + (lk & 3); - const uint8_t lo = (lo_pl[(g >> 2) * 128 + pos] >> ((g & 3) * 2)) & 3; - - const int c = e / 128; - const int j = (e % 128) / 32; - const int l = e % 32; - b->qs[c * 32 + l] |= (uint8_t) (lo << (2 * j)); - if (!((neg_pl[pos] >> g) & 1)) { - b->hmask[l] |= (uint8_t) (1 << (c * 4 + j)); - } - } - - for (int sub = 0; sub < 2; sub++) { - sub_scales[kt_local * 2 + sub] = sc_pl[sub * 32 + row]; - } - } - - int ls[16]; - if (!hexagon_recover_k_scales(sub_scales, 16, -32, 31, &b->d, ls)) { - // no exact match: same scale choice as quantize_row_q3_K_ref - float max_scale = 0.0f; - for (int s = 0; s < 16; s++) { - if (fabsf(GGML_FP16_TO_FP32(sub_scales[s])) > fabsf(max_scale)) { - max_scale = GGML_FP16_TO_FP32(sub_scales[s]); - } - } - b->d = GGML_FP32_TO_FP16(-max_scale / 32.0f); - const float d_actual = GGML_FP16_TO_FP32(b->d); - const float inv_d = (d_actual != 0.0f) ? (1.0f / d_actual) : 0.0f; - for (int s = 0; s < 16; s++) { - ls[s] = (std::max)(-32, (std::min)(31, (int) roundf(GGML_FP16_TO_FP32(sub_scales[s]) * inv_d))); - } - } - - for (int s = 0; s < 16; s++) { - const int l = ls[s] + 32; - if (s < 8) { - b->scales[s] = l & 0xF; - } else { - b->scales[s - 8] |= (uint8_t) ((l & 0xF) << 4); - } - b->scales[s % 4 + 8] |= (uint8_t) ((l >> 4) << (2 * (s / 4))); - } - } - } - } - } - - GGML_UNUSED(size); -} - -// tile layout: see HTP_MM_WEIGHT_TILE_SIZE_Q2_K in htp/matmul-ops.h -static void repack_q2_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { - GGML_ASSERT(offset == 0); - - const block_q2_K * src_matrix = (const block_q2_K *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - GGML_ASSERT(ne0 % QK_K == 0); - - const int n_col_tiles = ne1_padded / 32; - const int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q2_K; - const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; - - const int64_t sb_per_row = ne0 / QK_K; - - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_q2_K * src_slice = src_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - memset(matrix_dst, 0, matrix_size); // padding rows and the OR-ed bits below need zeroed tiles - - for (int64_t r = 0; r < ne1; r++) { - const int ct = (int) (r / 32); - const int row = (int) (r % 32); - const block_q2_K * src_row = src_slice + r * sb_per_row; - - for (int kt = 0; kt < n_k_tiles; kt++) { - const int kt_local = kt % 8; // k-tile within the super-block - const block_q2_K * b = &src_row[kt / 8]; - - uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size; - uint8_t * lo_pl = tile; - uint8_t * sc_pl = tile + 256; - ggml_half * dd_pl = (ggml_half *) (tile + 320); - - for (int lk = 0; lk < 32; lk++) { - const int g = lk >> 2; - const int pos = row * 4 + (lk & 3); - lo_pl[(g >> 2) * 128 + pos] |= (uint8_t) (q2_3_K_get_low2(b->qs, kt_local * 32 + lk) << ((g & 3) * 2)); - } - for (int sub = 0; sub < 2; sub++) { - sc_pl[sub * 32 + row] = b->scales[kt_local * 2 + sub]; - } - if (row / 8 == kt_local % 4) { - dd_pl[row % 8] = kt_local < 4 ? b->d : (ggml_half) (b->dmin ^ 0x8000); // -dmin - } - } - } - } - } - - GGML_UNUSED(size); -} - -// Reverse of repack_q2_K_tiled. Lossless. Read-back only. -static void repack_tiled_q2_K(void * data, const ggml_tensor * t, size_t offset, size_t size) { - GGML_ASSERT(offset == 0); - - block_q2_K * dst_matrix = (block_q2_K *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - GGML_ASSERT(ne0 % QK_K == 0); - - const int n_col_tiles = ne1_padded / 32; - const int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q2_K; - const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; - - const int64_t sb_per_row = ne0 / QK_K; - - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_q2_K * dst_slice = dst_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int64_t r = 0; r < ne1; r++) { - const int ct = (int) (r / 32); - const int row = (int) (r % 32); - block_q2_K * dst_row = dst_slice + r * sb_per_row; - - for (int64_t sb = 0; sb < sb_per_row; sb++) { - block_q2_K * b = &dst_row[sb]; - memset(b, 0, sizeof(block_q2_K)); - - for (int kt_local = 0; kt_local < 8; kt_local++) { - const int kt = sb * 8 + kt_local; - const uint8_t * tile = matrix_src + ((size_t) ct * n_k_tiles + kt) * tile_size; - const uint8_t * lo_pl = tile; - const uint8_t * sc_pl = tile + 256; - const ggml_half * dd_pl = (const ggml_half *) (tile + 320); - - for (int lk = 0; lk < 32; lk++) { - const int e = kt_local * 32 + lk; - const int g = lk >> 2; - const int pos = row * 4 + (lk & 3); - const uint8_t lo = (lo_pl[(g >> 2) * 128 + pos] >> ((g & 3) * 2)) & 3; - b->qs[(e / 128) * 32 + e % 32] |= (uint8_t) (lo << (2 * ((e % 128) / 32))); - } - - for (int sub = 0; sub < 2; sub++) { - b->scales[kt_local * 2 + sub] = sc_pl[sub * 32 + row]; - } - if (row / 8 == kt_local % 4) { - if (kt_local < 4) { - b->d = dd_pl[row % 8]; - } else { - b->dmin = (ggml_half) (dd_pl[row % 8] ^ 0x8000); - } - } - } - } - } - } - } - - GGML_UNUSED(size); -} - -static inline void get_scale_min_k4(int j, const uint8_t * q, uint8_t * d, uint8_t * m) { - if (j < 4) { - *d = q[j] & 63; - *m = q[j + 4] & 63; - } else { - *d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4); - *m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); - } -} - -// tile layout: see HTP_MM_WEIGHT_TILE_SIZE_Q4_1 in htp/matmul-ops.h -static void repack_q4_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { - GGML_ASSERT(offset == 0); - - const block_q4_K * src_matrix = (const block_q4_K *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - GGML_ASSERT(ne0 % QK_K == 0); - - const int n_col_tiles = ne1_padded / 32; - const int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_1; - const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; - - const int64_t sb_per_row = ne0 / QK_K; - - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_q4_K * src_slice = src_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - memset(matrix_dst, 0, matrix_size); - - for (int64_t r = 0; r < ne1; r++) { - const int ct = (int) (r / 32); - const int row = (int) (r % 32); - const block_q4_K * src_row = src_slice + r * sb_per_row; - - for (int kt = 0; kt < n_k_tiles; kt++) { - const int kt_local = kt % 8; - const block_q4_K * b = &src_row[kt / 8]; - const float d = GGML_FP16_TO_FP32(b->d); - const float dmin = GGML_FP16_TO_FP32(b->dmin); - - uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size; - - uint8_t sc, m; - get_scale_min_k4(kt_local, b->scales, &sc, &m); - - const float D = d * (float) sc; - const float M = -dmin * (float) m; - - const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32; - const int shift = (kt_local & 1) ? 4 : 0; - - for (int cp = 0; cp < 16; cp++) { - const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F; - const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F; - tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0); - } - - ggml_half * scale_dst = (ggml_half *) (tile_dst + 512); - scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(D); - scale_dst[2 * row + 1] = GGML_FP32_TO_FP16(M); - } - } - } - } - - GGML_UNUSED(size); -} - -// Reverse of repack_q4_K_tiled. Unpacks quants and normalizes scales/mins. Read-back only. -static void repack_tiled_q4_K(void * data, const ggml_tensor * t, size_t offset, size_t size) { - GGML_ASSERT(offset == 0); - - block_q4_K * dst_matrix = (block_q4_K *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - GGML_ASSERT(ne0 % QK_K == 0); - - const int n_col_tiles = ne1_padded / 32; - const int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_1; - const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; - - const int64_t sb_per_row = ne0 / QK_K; - - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_q4_K * dst_slice = dst_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int64_t r = 0; r < ne1; r++) { - const int ct = (int) (r / 32); - const int row = (int) (r % 32); - block_q4_K * dst_row = dst_slice + r * sb_per_row; - - for (int64_t sb = 0; sb < sb_per_row; sb++) { - block_q4_K * b = &dst_row[sb]; - memset(b, 0, sizeof(block_q4_K)); - - float sub_scales[8]; - float sub_mins[8]; - - for (int kt_local = 0; kt_local < 8; kt_local++) { - const int kt = sb * 8 + kt_local; - const uint8_t * tile_src = matrix_src + ((size_t) ct * n_k_tiles + kt) * tile_size; - const ggml_half * scale_src = (const ggml_half *) (tile_src + 512); - - uint8_t * qs_sub = b->qs + (kt_local / 2) * 32; - const int shift = (kt_local & 1) ? 4 : 0; - - for (int cp = 0; cp < 16; cp++) { - const uint8_t val = tile_src[cp * 32 + row]; - const uint8_t q0 = val & 0x0F; - const uint8_t q1 = val >> 4; - qs_sub[2 * cp + 0] |= (uint8_t) (q0 << shift); - qs_sub[2 * cp + 1] |= (uint8_t) (q1 << shift); - } - - const float D = GGML_FP16_TO_FP32(scale_src[2 * row + 0]); - const float M = GGML_FP16_TO_FP32(scale_src[2 * row + 1]); - sub_scales[kt_local] = (D > 0.0f) ? D : 0.0f; - sub_mins[kt_local] = (-M > 0.0f) ? -M : 0.0f; - } - - float max_scale = 0.0f; - float max_min = 0.0f; - for (int j = 0; j < 8; j++) { - if (sub_scales[j] > max_scale) max_scale = sub_scales[j]; - if (sub_mins[j] > max_min) max_min = sub_mins[j]; - } - - float inv_scale = 0.0f; - if (max_scale > 0.0f) { - b->d = GGML_FP32_TO_FP16(max_scale / 63.0f); - const float d_actual = GGML_FP16_TO_FP32(b->d); - inv_scale = (d_actual > 0.0f) ? (1.0f / d_actual) : 0.0f; - } else { - b->d = GGML_FP32_TO_FP16(0.0f); - } - - float inv_min = 0.0f; - if (max_min > 0.0f) { - b->dmin = GGML_FP32_TO_FP16(max_min / 63.0f); - const float dmin_actual = GGML_FP16_TO_FP32(b->dmin); - inv_min = (dmin_actual > 0.0f) ? (1.0f / dmin_actual) : 0.0f; - } else { - b->dmin = GGML_FP32_TO_FP16(0.0f); - } - - for (int j = 0; j < 8; j++) { - uint8_t ls = (uint8_t) roundf(inv_scale * sub_scales[j]); - uint8_t lm = (uint8_t) roundf(inv_min * sub_mins[j]); - ls = (std::min)((uint8_t) 63, ls); - lm = (std::min)((uint8_t) 63, lm); - if (j < 4) { - b->scales[j] = ls; - b->scales[j + 4] = lm; - } else { - b->scales[j + 4] = (ls & 0xF) | ((lm & 0xF) << 4); - b->scales[j - 4] |= ((ls >> 4) << 6); - b->scales[j - 0] |= ((lm >> 4) << 6); - } - } - } - } - } - } - - GGML_UNUSED(size); -} - -// tile layout: see HTP_MM_WEIGHT_TILE_SIZE_Q5_K in htp/matmul-ops.h -static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { - GGML_ASSERT(offset == 0); - - const block_q5_K * src_matrix = (const block_q5_K *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - GGML_ASSERT(ne0 % QK_K == 0); - - const int n_col_tiles = ne1_padded / 32; - const int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q5_K; - const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; - - const int64_t sb_per_row = ne0 / QK_K; - - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_q5_K * src_slice = src_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - memset(matrix_dst, 0, matrix_size); - - for (int64_t r = 0; r < ne1; r++) { - const int ct = (int) (r / 32); - const int row = (int) (r % 32); - const block_q5_K * src_row = src_slice + r * sb_per_row; - - for (int kt = 0; kt < n_k_tiles; kt++) { - const int kt_local = kt % 8; - const block_q5_K * b = &src_row[kt / 8]; - const float d = GGML_FP16_TO_FP32(b->d); - const float dmin = GGML_FP16_TO_FP32(b->dmin); - - uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size; - uint8_t * plane = tile_dst + 640; - - uint8_t sc, m; - get_scale_min_k4(kt_local, b->scales, &sc, &m); - - const float D = d * (float) sc; - const float M = -dmin * (float) m; - - const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32; - const int shift = (kt_local & 1) ? 4 : 0; - const uint8_t hbit = (uint8_t) (1 << kt_local); - - for (int cp = 0; cp < 16; cp++) { - const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F; - const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F; - tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0); - - const int i = cp / 4; - const int lane = (cp % 4) * 32 + row; - if (b->qh[2 * cp + 0] & hbit) { - plane[lane] |= (uint8_t) (1 << (2 * i)); - } - if (b->qh[2 * cp + 1] & hbit) { - plane[lane] |= (uint8_t) (1 << (2 * i + 1)); - } - } - - ggml_half * scale_dst = (ggml_half *) (tile_dst + 512); - scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(D); - scale_dst[2 * row + 1] = GGML_FP32_TO_FP16(M); - } - } - } - } - - GGML_UNUSED(size); -} - -// Reverse of repack_q5_K_tiled. Unpacks quants losslessly and normalizes scales/mins. Read-back only. -static void repack_tiled_q5_K(void * data, const ggml_tensor * t, size_t offset, size_t size) { - GGML_ASSERT(offset == 0); - - block_q5_K * dst_matrix = (block_q5_K *) data; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - int64_t ne2 = t->ne[2]; - int64_t ne3 = t->ne[3]; - int64_t ne0_padded = hex_round_up(ne0, 32); - int64_t ne1_padded = hex_round_up(ne1, 32); - - GGML_ASSERT(ne0 % QK_K == 0); - - const int n_col_tiles = ne1_padded / 32; - const int n_k_tiles = ne0_padded / 32; - const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q5_K; - const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; - - const int64_t sb_per_row = ne0 / QK_K; - - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_q5_K * dst_slice = dst_matrix + (i3 * ne2 + i2) * (ne1 * sb_per_row); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int64_t r = 0; r < ne1; r++) { - const int ct = (int) (r / 32); - const int row = (int) (r % 32); - block_q5_K * dst_row = dst_slice + r * sb_per_row; - - for (int64_t sb = 0; sb < sb_per_row; sb++) { - block_q5_K * b = &dst_row[sb]; - memset(b, 0, sizeof(block_q5_K)); - - float sub_scales[8]; - float sub_mins[8]; - - for (int kt_local = 0; kt_local < 8; kt_local++) { - const int kt = sb * 8 + kt_local; - const uint8_t * tile_src = matrix_src + ((size_t) ct * n_k_tiles + kt) * tile_size; - const uint8_t * plane = tile_src + 640; - const ggml_half * scale_src = (const ggml_half *) (tile_src + 512); - - uint8_t * qs_sub = b->qs + (kt_local / 2) * 32; - const int shift = (kt_local & 1) ? 4 : 0; - const uint8_t hbit = (uint8_t) (1 << kt_local); - - for (int cp = 0; cp < 16; cp++) { - const uint8_t val = tile_src[cp * 32 + row]; - const uint8_t q0 = val & 0x0F; - const uint8_t q1 = val >> 4; - qs_sub[2 * cp + 0] |= (uint8_t) (q0 << shift); - qs_sub[2 * cp + 1] |= (uint8_t) (q1 << shift); - - const int i = cp / 4; - const int lane = (cp % 4) * 32 + row; - if (plane[lane] & (1 << (2 * i))) { - b->qh[2 * cp + 0] |= hbit; - } - if (plane[lane] & (1 << (2 * i + 1))) { - b->qh[2 * cp + 1] |= hbit; - } - } - - const float D = GGML_FP16_TO_FP32(scale_src[2 * row + 0]); - const float M = GGML_FP16_TO_FP32(scale_src[2 * row + 1]); - sub_scales[kt_local] = (D > 0.0f) ? D : 0.0f; - sub_mins[kt_local] = (-M > 0.0f) ? -M : 0.0f; - } - - float max_scale = 0.0f; - float max_min = 0.0f; - for (int j = 0; j < 8; j++) { - if (sub_scales[j] > max_scale) max_scale = sub_scales[j]; - if (sub_mins[j] > max_min) max_min = sub_mins[j]; - } - - float inv_scale = 0.0f; - if (max_scale > 0.0f) { - b->d = GGML_FP32_TO_FP16(max_scale / 63.0f); - const float d_actual = GGML_FP16_TO_FP32(b->d); - inv_scale = (d_actual > 0.0f) ? (1.0f / d_actual) : 0.0f; - } else { - b->d = GGML_FP32_TO_FP16(0.0f); - } - - float inv_min = 0.0f; - if (max_min > 0.0f) { - b->dmin = GGML_FP32_TO_FP16(max_min / 63.0f); - const float dmin_actual = GGML_FP16_TO_FP32(b->dmin); - inv_min = (dmin_actual > 0.0f) ? (1.0f / dmin_actual) : 0.0f; - } else { - b->dmin = GGML_FP32_TO_FP16(0.0f); - } + return ptr; + } - for (int j = 0; j < 8; j++) { - uint8_t ls = (uint8_t) roundf(inv_scale * sub_scales[j]); - uint8_t lm = (uint8_t) roundf(inv_min * sub_mins[j]); - ls = (std::min)((uint8_t) 63, ls); - lm = (std::min)((uint8_t) 63, lm); - if (j < 4) { - b->scales[j] = ls; - b->scales[j + 4] = lm; - } else { - b->scales[j + 4] = (ls & 0xF) | ((lm & 0xF) << 4); - b->scales[j - 4] |= ((ls >> 4) << 6); - b->scales[j - 0] |= ((lm >> 4) << 6); - } - } - } - } + void free_slot(void * ptr, uint32_t n_slots = 1) { + if (!ptr) return; + uint32_t slot = ((uint8_t *) ptr - base()) / GGML_HEXAGON_FENCE_SLOT_SIZE; + for (uint32_t i = 0; i < n_slots; i++) { + free_slots.push_back(slot + i); } } +}; - GGML_UNUSED(size); +inline uint8_t * ggml_hexagon_session::alloc_fence(uint32_t n_slots) { + uint8_t * ptr = fence_buf->alloc_slot(n_slots); + GGML_ASSERT(ptr); + return ptr; } -static void repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t size) { - switch (tensor->type) { - case GGML_TYPE_Q4_0: - repack_q4_0_tiled(tensor, data, 0, size); - break; - - case GGML_TYPE_Q4_1: - repack_q4_1_tiled(tensor, data, 0, size); - break; - - case GGML_TYPE_Q4_K: - repack_q4_K_tiled(tensor, data, 0, size); - break; - - case GGML_TYPE_Q8_0: - repack_q8_0_tiled(tensor, data, 0, size); - break; +inline void ggml_hexagon_session::free_fence(void * ptr, uint32_t n_slots) { + if (fence_buf) { + fence_buf->free_slot(ptr, n_slots); + } +} - case GGML_TYPE_IQ4_NL: - repack_q4_0_tiled(tensor, data, 0, size); - break; +static void ggml_backend_hexagon_buffer_free_buffer(ggml_backend_buffer_t buffer) { + auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(buffer->context); + sbuf->sess->unclone_buffer(sbuf); + delete sbuf; +} - case GGML_TYPE_MXFP4: - repack_mxfp4_tiled(tensor, data, 0, size); - break; +static void * ggml_backend_hexagon_buffer_get_base(ggml_backend_buffer_t buffer) { + auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(buffer->context); + return sbuf->base(); +} - case GGML_TYPE_Q5_K: - repack_q5_K_tiled(tensor, data, 0, size); - break; +static enum ggml_status ggml_backend_hexagon_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + auto sbuf = static_cast<ggml_hexagon_shared_buffer *>(buffer->context); + auto sess = sbuf->sess; - case GGML_TYPE_Q6_K: - repack_q6_K_tiled(tensor, data, 0, size); - break; + HEX_VERBOSE("ggml-hex: %s init-tensor %s : base %p data %p nbytes %zu\n", sess->c_name(), + tensor->name, (void *) sbuf->base(), tensor->data, ggml_nbytes(tensor)); - case GGML_TYPE_Q3_K: - repack_q3_K_tiled(tensor, data, 0, size); - break; + auto extra = new ggml_hexagon_tensor_extra(); + sbuf->tensor_extra.push_back(extra); - case GGML_TYPE_Q2_K: - repack_q2_K_tiled(tensor, data, 0, size); - break; + tensor->extra = extra; - default: - break; - } + return GGML_STATUS_SUCCESS; } static void ggml_backend_hexagon_buffer_set_tensor(ggml_backend_buffer_t buffer, @@ -2474,21 +907,21 @@ static void ggml_backend_hexagon_buffer_set_tensor(ggml_backend_buffer_t buffer, if (ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { extra->flags |= GGML_HEXAGON_TENSOR_WEIGHT; - if (ggml_hexagon_is_repack_type(tensor->type)) { - extra->flags |= GGML_HEXAGON_TENSOR_REPACK; + if (htp_is_tiled_weight_tensor(tensor)) { + extra->flags |= GGML_HEXAGON_TENSOR_TILED; } } HEX_VERBOSE("ggml-hex: %s set-tensor %s : data %p offset %zu size %zu usage %d flags 0x%x\n", sess->c_name(), tensor->name, data, offset, size, (int) buffer->usage, extra->flags); - if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + if ((extra->flags & GGML_HEXAGON_TENSOR_TILED) == 0) { memcpy((char *) tensor->data + offset, data, size); return; } if (offset == 0 && size == ggml_nbytes(tensor) && extra->shadow_buf.empty()) { - repack_tensor_tiled(tensor, data, size); + htp_repack_tensor_tiled(tensor, data, size); return; } @@ -2499,7 +932,7 @@ static void ggml_backend_hexagon_buffer_set_tensor(ggml_backend_buffer_t buffer, extra->shadow_size += size; if (extra->shadow_size >= ggml_nbytes(tensor)) { - repack_tensor_tiled(tensor, extra->shadow_buf.data(), extra->shadow_buf.size()); + htp_repack_tensor_tiled(tensor, extra->shadow_buf.data(), extra->shadow_buf.size()); extra->shadow_buf.clear(); extra->shadow_buf.shrink_to_fit(); extra->shadow_size = 0; @@ -2518,76 +951,13 @@ static void ggml_backend_hexagon_buffer_get_tensor(ggml_backend_buffer_t buffer, HEX_VERBOSE("ggml-hex: %s get-tensor %s : data %p offset %zu size %zu usage %d flags 0x%x\n", sess->c_name(), tensor->name, data, offset, size, (int) buffer->usage, extra->flags); - if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + if ((extra->flags & GGML_HEXAGON_TENSOR_TILED) == 0) { memcpy(data, (const char *) tensor->data + offset, size); return; } - switch (tensor->type) { - case GGML_TYPE_Q4_0: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q4_0(data, tensor, offset, size); - break; - - case GGML_TYPE_Q4_1: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q4_1(data, tensor, offset, size); - break; - - case GGML_TYPE_Q4_K: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q4_K(data, tensor, offset, size); - break; - - case GGML_TYPE_Q5_K: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q5_K(data, tensor, offset, size); - break; - - case GGML_TYPE_Q8_0: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q8_0(data, tensor, offset, size); - break; - - case GGML_TYPE_IQ4_NL: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q4_0(data, tensor, offset, size); - break; - - case GGML_TYPE_MXFP4: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_mxfp4(data, tensor, offset, size); - break; - - case GGML_TYPE_Q6_K: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q6_K(data, tensor, offset, size); - break; - - case GGML_TYPE_Q3_K: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q3_K(data, tensor, offset, size); - break; - - case GGML_TYPE_Q2_K: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q2_K(data, tensor, offset, size); - break; - - default: - memcpy(data, (const char *) tensor->data + offset, size); - break; - } + GGML_ASSERT(offset == 0); + htp_unpack_tensor_tiled(data, tensor, offset, size); } static bool ggml_backend_hexagon_buffer_cpy_tensor(ggml_backend_buffer_t buffer, @@ -2615,15 +985,15 @@ static void ggml_backend_hexagon_buffer_set_tensor_2d(ggml_backend_buffer_t buff if (ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { extra->flags |= GGML_HEXAGON_TENSOR_WEIGHT; - if (ggml_hexagon_is_repack_type(tensor->type)) { - extra->flags |= GGML_HEXAGON_TENSOR_REPACK; + if (htp_is_tiled_weight_tensor(tensor)) { + extra->flags |= GGML_HEXAGON_TENSOR_TILED; } } HEX_VERBOSE("ggml-hex: %s set-tensor-2d %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d flags 0x%x\n", sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, (int) buffer->usage, extra->flags); - if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + if ((extra->flags & GGML_HEXAGON_TENSOR_TILED) == 0) { for (size_t i = 0; i < n_copies; i++) { memcpy((uint8_t *) tensor->data + offset + i * stride_tensor, (const uint8_t *) data + i * stride_data, size); } @@ -2639,7 +1009,7 @@ static void ggml_backend_hexagon_buffer_set_tensor_2d(ggml_backend_buffer_t buff extra->shadow_size += n_copies * size; if (extra->shadow_size >= ggml_nbytes(tensor)) { - repack_tensor_tiled(tensor, extra->shadow_buf.data(), extra->shadow_buf.size()); + htp_repack_tensor_tiled(tensor, extra->shadow_buf.data(), extra->shadow_buf.size()); extra->shadow_buf.clear(); extra->shadow_buf.shrink_to_fit(); extra->shadow_size = 0; @@ -2661,7 +1031,7 @@ static void ggml_backend_hexagon_buffer_get_tensor_2d(ggml_backend_buffer_t buff HEX_VERBOSE("ggml-hex: %s get-tensor-2d %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d\n", sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, (int) buffer->usage); - if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + if ((extra->flags & GGML_HEXAGON_TENSOR_TILED) == 0) { for (size_t i = 0; i < n_copies; i++) { memcpy((uint8_t *)data + i * stride_data, (const uint8_t *)tensor->data + offset + i * stride_tensor, size); } @@ -2679,52 +1049,7 @@ static void ggml_backend_hexagon_buffer_get_tensor_2d(ggml_backend_buffer_t buff GGML_ASSERT((offset + temp_size) <= ggml_nbytes(tensor)); std::vector<uint8_t> temp_buf(temp_size); - - switch (tensor->type) { - case GGML_TYPE_Q4_0: - repack_tiled_q4_0(temp_buf.data(), tensor, offset, temp_size); - break; - - case GGML_TYPE_Q4_1: - repack_tiled_q4_1(temp_buf.data(), tensor, offset, temp_size); - break; - - case GGML_TYPE_Q4_K: - repack_tiled_q4_K(temp_buf.data(), tensor, offset, temp_size); - break; - - case GGML_TYPE_Q5_K: - repack_tiled_q5_K(temp_buf.data(), tensor, offset, temp_size); - break; - - case GGML_TYPE_Q8_0: - repack_tiled_q8_0(temp_buf.data(), tensor, offset, temp_size); - break; - - case GGML_TYPE_IQ4_NL: - repack_tiled_q4_0(temp_buf.data(), tensor, offset, temp_size); - break; - - case GGML_TYPE_MXFP4: - repack_tiled_mxfp4(temp_buf.data(), tensor, offset, temp_size); - break; - - case GGML_TYPE_Q6_K: - repack_tiled_q6_K(temp_buf.data(), tensor, offset, temp_size); - break; - - case GGML_TYPE_Q3_K: - repack_tiled_q3_K(temp_buf.data(), tensor, offset, temp_size); - break; - - case GGML_TYPE_Q2_K: - repack_tiled_q2_K(temp_buf.data(), tensor, offset, temp_size); - break; - - default: - memcpy(temp_buf.data(), (const uint8_t *) tensor->data + offset, temp_size); - break; - } + htp_unpack_tensor_tiled(temp_buf.data(), tensor, offset, temp_size); for (size_t i = 0; i < n_copies; i++) { memcpy((uint8_t *) data + i * stride_data, temp_buf.data() + i * stride_tensor, size); @@ -2834,16 +1159,15 @@ static size_t ggml_backend_hexagon_buffer_type_get_alignment(ggml_backend_buffer } static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const struct ggml_tensor * t) { - if (ggml_hexagon_is_repack_type(t->type)) { + if (!ggml_backend_buft_is_host(buft) && htp_is_tiled_weight_tensor(t)) { int64_t ne0 = hex_round_up(t->ne[0], 32); int64_t ne1 = hex_round_up(t->ne[1], 32); int64_t ne2 = t->ne[2]; int64_t ne3 = t->ne[3]; - return ggml_hexagon_tiled_row_size(t->type, ne0) * ne1 * ne2 * ne3; + size_t tiled_size = htp_tiled_row_size(t->type, ne0) * ne1 * ne2 * ne3; + return std::max(tiled_size, ggml_nbytes(t)); } return ggml_nbytes(t); - - GGML_UNUSED(buft); } static size_t parse_size(const char * str, size_t default_unit = 1024 * 1024) { @@ -3196,20 +1520,19 @@ struct ggml_hexagon_opbatch { bool same_shape(const htp_tensor * h, const ggml_tensor * t) const { auto extra = (ggml_hexagon_tensor_extra *) t->extra; - int64_t ne0 = t->ne[0]; - int64_t ne1 = t->ne[1]; - const bool is_repack = (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0; - if (is_repack) { - ne0 = hex_round_up(ne0, 32); - ne1 = hex_round_up(ne1, 32); - } - int64_t nb1 = is_repack ? (int64_t) ggml_hexagon_tiled_row_size(t->type, ne0) : t->nb[1]; - int64_t nb2 = is_repack ? nb1 * ne1 : t->nb[2]; - int64_t nb3 = is_repack ? nb2 * t->ne[2] : t->nb[3]; + const bool is_tiled = (extra->flags & GGML_HEXAGON_TENSOR_TILED) != 0; + const int64_t ne0_padded = hex_round_up(t->ne[0], 32); + const int64_t ne1_padded = hex_round_up(t->ne[1], 32); + + int64_t nb1 = is_tiled ? (int64_t) htp_tiled_row_size(t->type, ne0_padded) : t->nb[1]; + int64_t nb2 = is_tiled ? nb1 * ne1_padded : t->nb[2]; + int64_t nb3 = is_tiled ? nb2 * t->ne[2] : t->nb[3]; + int64_t nb0 = (is_tiled && (t->type == GGML_TYPE_F32 || t->type == GGML_TYPE_BF16)) ? sizeof(ggml_fp16_t) : t->nb[0]; + enum ggml_type expected_type = (is_tiled && (t->type == GGML_TYPE_F32 || t->type == GGML_TYPE_BF16)) ? GGML_TYPE_F16 : t->type; - return (h->type == t->type) && - (h->ne[0] == ne0) && (h->ne[1] == ne1) && (h->ne[2] == t->ne[2]) && (h->ne[3] == t->ne[3]) && - (h->nb[0] == t->nb[0]) && (h->nb[1] == nb1) && (h->nb[2] == nb2) && (h->nb[3] == nb3); + return (h->type == expected_type) && + (h->ne[0] == t->ne[0]) && (h->ne[1] == t->ne[1]) && (h->ne[2] == t->ne[2]) && (h->ne[3] == t->ne[3]) && + (h->nb[0] == nb0) && (h->nb[1] == nb1) && (h->nb[2] == nb2) && (h->nb[3] == nb3); } // add tensor and return its index @@ -3243,23 +1566,28 @@ struct ggml_hexagon_opbatch { h.ti = ti; h.data = t_offset; h.type = t->type; - - const bool is_repack = (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0; - if (is_repack) { - h.ne[0] = hex_round_up(t->ne[0], 32); - h.ne[1] = hex_round_up(t->ne[1], 32); - h.ne[2] = t->ne[2]; - h.ne[3] = t->ne[3]; - - h.nb[0] = t->nb[0]; - h.nb[1] = ggml_hexagon_tiled_row_size(t->type, h.ne[0]); - h.nb[2] = h.nb[1] * h.ne[1]; + h.ne[0] = t->ne[0]; + h.ne[1] = t->ne[1]; + h.ne[2] = t->ne[2]; + h.ne[3] = t->ne[3]; + + const bool is_tiled = (extra->flags & GGML_HEXAGON_TENSOR_TILED) != 0; + if (is_tiled) { + if (t->type == GGML_TYPE_F32 || t->type == GGML_TYPE_BF16) { + h.type = GGML_TYPE_F16; + h.nb[0] = sizeof(ggml_fp16_t); + } else { + h.nb[0] = t->nb[0]; + } + const int64_t ne0_padded = hex_round_up(t->ne[0], 32); + const int64_t ne1_padded = hex_round_up(t->ne[1], 32); + h.nb[1] = htp_tiled_row_size(t->type, ne0_padded); + h.nb[2] = h.nb[1] * ne1_padded; h.nb[3] = h.nb[2] * h.ne[2]; h.size = h.nb[3] * h.ne[3]; t_size = h.size; } else { h.size = t_size; - h.ne[0] = t->ne[0]; h.ne[1] = t->ne[1]; h.ne[2] = t->ne[2]; h.ne[3] = t->ne[3]; h.nb[0] = t->nb[0]; h.nb[1] = t->nb[1]; h.nb[2] = t->nb[2]; h.nb[3] = t->nb[3]; } @@ -3267,8 +1595,8 @@ struct ggml_hexagon_opbatch { if ((extra->flags & GGML_HEXAGON_TENSOR_WEIGHT) != 0) { h.flags |= HTP_TENSOR_WEIGHT; } - if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0) { - h.flags |= HTP_TENSOR_REPACK; + if ((extra->flags & GGML_HEXAGON_TENSOR_TILED) != 0) { + h.flags |= HTP_TENSOR_TILED; } if ((extra->flags & GGML_HEXAGON_TENSOR_FENCE) != 0) { h.flags |= HTP_TENSOR_FENCE; @@ -3666,6 +1994,9 @@ struct ggml_hexagon_opbatch { if (src2->type != GGML_TYPE_F32) return false; + // the fused kernel writes the shape of the MUL_MAT output: the ADD must not broadcast it + if (!ggml_are_same_shape(mm_out, node.dst())) return false; + const struct htp_mm_kernel_params * orig_kparams = (const struct htp_mm_kernel_params *) last_node.kernel_params; struct htp_mm_kernel_params kparams; ggml_hexagon_precompute_fused_matmul_add_params(sess, src0, src1, src2, node.dst(), &kparams); @@ -3753,7 +2084,8 @@ struct ggml_hexagon_opbatch { const ggml_tensor * w0 = last_node.inputs[0]; const ggml_tensor * x = last_node.inputs[curr_n]; - if (x_in != x || w_in->type != w0->type || w_in->ne[0] != w0->ne[0]) { + if (x_in != x || w_in->type != w0->type || w_in->ne[0] != w0->ne[0] || + ggml_hexagon_tensor_is_tiled(w_in) != ggml_hexagon_tensor_is_tiled(w0)) { return false; } if (!last_node.fused.empty() && (mm_is_hmx_eligible(last_node.fused[0]) != mm_is_hmx_eligible(node.node))) { @@ -3909,7 +2241,8 @@ struct ggml_hexagon_opbatch { const ggml_tensor * x = last_node.inputs[curr_n]; const ggml_tensor * ids = last_node.inputs[curr_n + 1]; - if (x_in != x || ids_in != ids || w_in->type != w0->type || w_in->ne[0] != w0->ne[0] || w_in->ne[2] != w0->ne[2]) { + if (x_in != x || ids_in != ids || w_in->type != w0->type || w_in->ne[0] != w0->ne[0] || w_in->ne[2] != w0->ne[2] || + ggml_hexagon_tensor_is_tiled(w_in) != ggml_hexagon_tensor_is_tiled(w0)) { return false; } if (!last_node.fused.empty() && (mm_is_hmx_eligible(last_node.fused[0]) != mm_is_hmx_eligible(node.node))) { @@ -5433,7 +3766,6 @@ static bool ggml_hexagon_matmul_is_hmx_eligible( const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst, - int ne01_padded, bool is_matmul_id, bool is_batched ) { @@ -5442,17 +3774,18 @@ static bool ggml_hexagon_matmul_is_hmx_eligible( } const int ne00 = src0->ne[0]; + const int ne01 = src0->ne[1]; const int ne11 = src1->ne[1]; const int ne12 = src1->ne[2]; const int wtype = src0->type; // HMX weight tiles accept non-32-aligned N for non-matmul_id. - if (ne01_padded % 32 != 0 && is_matmul_id) { + if (ne01 % 32 != 0 && is_matmul_id) { return false; } // HMX supports F16, F32, and repack quantized types. - if (!ggml_hexagon_is_hmx_weight_type((ggml_type) wtype)) { + if (!htp_is_hmx_weight_type((ggml_type) wtype)) { return false; } @@ -5461,8 +3794,12 @@ static bool ggml_hexagon_matmul_is_hmx_eligible( return false; } - // Quantized HMX kernels only handle flat 2D matmul (or matmul_id wrapping flat 2D matmuls). - if (!is_matmul_id && is_batched && wtype != GGML_TYPE_F16) { + // Quantized and tiled HMX kernels only handle flat 2D matmul (or matmul_id wrapping flat 2D matmuls). + const bool is_tiled = ggml_hexagon_tensor_is_tiled(src0); + if (ggml_is_quantized((ggml_type) wtype) && !is_tiled) { + return false; + } + if (!is_matmul_id && is_batched && is_tiled) { return false; } @@ -5501,7 +3838,8 @@ static bool ggml_hexagon_precompute_hmx_mm_params( size_t vtcm_budget, struct htp_mm_kernel_params * kparams ) { - const int aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype); + const bool is_tiled = ggml_hexagon_tensor_is_tiled(src0); + const int aligned_tile_size = is_tiled ? htp_tensor_get_tile_padded_size(wtype) : 0; const int n_threads = (int)sess->n_threads; const int ne10 = act->ne[0]; @@ -5528,7 +3866,7 @@ static bool ggml_hexagon_precompute_hmx_mm_params( bool use_grouped = false; int act_threads_selected = 0; - if (is_batched_val && wtype == GGML_TYPE_F16 && group_size > 1) { + if (is_batched_val && wtype == GGML_TYPE_F16 && !is_tiled && group_size > 1) { // Try grouped path first if (htp_mm_hmx_solve_batched_params(wtype, ne00_padded, ne01_padded, m_for_solver, group_size, n_threads, pipeline, src2_size, vtcm_budget, &m_chunk, &n_chunk, &act_threads_selected, &vtcm_size)) { use_grouped = true; @@ -5559,9 +3897,9 @@ static bool ggml_hexagon_precompute_hmx_mm_params( kparams->n_chunk = n_chunk; kparams->n_threads = n_threads; kparams->n_act_threads = act_threads_selected; - kparams->tile_size = htp_mm_get_weight_tile_size(wtype); + kparams->tile_size = is_tiled ? htp_tensor_get_tile_size(wtype) : 0; kparams->aligned_tile_size = aligned_tile_size; - kparams->act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + kparams->act_row_size = htp_mm_tiled_act_row_size(wtype, ne10); kparams->vtcm_size = vtcm_size; kparams->vtcm_src0_size = 0; kparams->div_n_act_threads = init_fastdiv_values(act_threads_selected); @@ -5571,13 +3909,11 @@ static bool ggml_hexagon_precompute_hmx_mm_params( kparams->vtcm_dst_size = 0; if (is_batched && !is_matmul_id) { - kparams->kernel_type = HTP_MM_KERNEL_HMX_F16_BATCHED; + kparams->kernel_type = HTP_MM_KERNEL_HMX_BATCHED; } else { kparams->kernel_type = HTP_MM_KERNEL_HMX_2D; } return true; - - GGML_UNUSED(src0); } static void ggml_hexagon_precompute_hvx_mm_params( @@ -5586,6 +3922,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( const struct ggml_tensor * src1, const struct ggml_tensor * dst, int wtype, + bool is_tiled, int ne02, int ne03, int ne10, @@ -5608,19 +3945,19 @@ static void ggml_hexagon_precompute_hvx_mm_params( GGML_UNUSED(ne02); GGML_UNUSED(ne03); - const bool is_quant = (wtype != GGML_TYPE_F16 && wtype != GGML_TYPE_F32); const int src1_nrows = ne11 * ne12 * ne13; - if (is_quant) { - // Quantized HVX - kparams->tile_size = htp_mm_get_weight_tile_size(wtype); - kparams->aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype); + if (is_tiled) { + // Tiled HVX (quantized weights or tiled F16/F32/BF16) + kparams->tile_size = htp_tensor_get_tile_size(wtype); + kparams->aligned_tile_size = htp_tensor_get_tile_padded_size(wtype); const bool k_align = (ne10 % 32 == 0); if (is_matmul_id) { - kparams->kernel_type = (src1_nrows < (int) sess->n_threads) ? HTP_MM_KERNEL_HVX_QUANT_BLOCK : HTP_MM_KERNEL_HVX_QUANT_ROW; - kparams->act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + kparams->kernel_type = (src1_nrows < (int) sess->n_threads) + ? HTP_MM_KERNEL_HVX_QUANT_BLOCK : HTP_MM_KERNEL_HVX_QUANT_ROW; + kparams->act_row_size = htp_mm_tiled_act_row_size(wtype, ne10); struct htp_mm_hvx_vtcm_layout L; uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16; @@ -5645,13 +3982,11 @@ static void ggml_hexagon_precompute_hvx_mm_params( kparams->vtcm_act_size = L.act_bytes; kparams->vtcm_bias_size = 0; kparams->vtcm_dst_size = L.dst_bytes; - goto done_quant; + return; } else { bool try_tiled = (k_align && opt_mm_select >= 1); if (try_tiled) { - kparams->act_row_size = htp_mm_weight_has_offset(wtype) - ? htp_mm_q8_1_tiled_row_size(ne10) - : htp_mm_q8_0_tiled_row_size(ne10); + kparams->act_row_size = htp_mm_tiled_act_row_size(wtype, ne10); if (src1_nrows < (int) sess->n_threads) { kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_BLOCK; } else { @@ -5684,17 +4019,18 @@ static void ggml_hexagon_precompute_hvx_mm_params( kparams->vtcm_act_size = L.act_bytes; kparams->vtcm_bias_size = L.bias_bytes; kparams->vtcm_dst_size = L.dst_bytes; - goto done_quant; + return; } } kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; return; } - - done_quant:; + } else if (is_matmul_id) { + kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; + return; } else if (wtype == GGML_TYPE_F16) { - // F16 HVX + // Flat F16 HVX struct htp_mm_hvx_vtcm_layout L; uint32_t m_chunk = 0; if (htp_mm_hvx_solve_vtcm_params( @@ -5715,8 +4051,8 @@ static void ggml_hexagon_precompute_hvx_mm_params( kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; return; - } else { - // F32 HVX + } else if (wtype == GGML_TYPE_F32) { + // Flat F32 HVX struct htp_mm_hvx_vtcm_layout L; uint32_t m_chunk = 0; if (htp_mm_hvx_solve_vtcm_params( @@ -5735,6 +4071,31 @@ static void ggml_hexagon_precompute_hvx_mm_params( return; } + kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; + return; + } else if (wtype == GGML_TYPE_BF16) { + // Flat BF16 HVX + struct htp_mm_hvx_vtcm_layout L; + uint32_t m_chunk = 0; + if (htp_mm_hvx_solve_vtcm_params( + HTP_MM_KERNEL_HVX_BF16_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads, + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, vtcm_budget, + &L, &m_chunk)) { + kparams->kernel_type = HTP_MM_KERNEL_HVX_BF16_F32_VTCM; + kparams->m_chunk = (m_chunk < (uint32_t) src1_nrows) ? m_chunk : 0; + kparams->act_row_size = hex_round_up(ne10 * 4, 128); + kparams->vtcm_size = L.total_bytes; + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_act_size = L.act_bytes; + kparams->vtcm_bias_size = L.bias_bytes; + kparams->vtcm_dst_size = L.dst_bytes; + kparams->n_prefetch = 16; + return; + } + + kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; + return; + } else { kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; return; } @@ -5761,14 +4122,13 @@ static void ggml_hexagon_precompute_matmul_params_impl( const int ne12 = src1->ne[2]; const int ne13 = src1->ne[3]; - const int wtype = src0->type; - const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype); - const int ne00_padded = is_repack ? hex_round_up(ne00, 32) : ne00; - const int ne01_padded = is_repack ? hex_round_up(ne01, 32) : ne01; + const bool is_tiled = ggml_hexagon_tensor_is_tiled(src0); + const int wtype = (is_tiled && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_BF16)) ? GGML_TYPE_F16 : src0->type; + const int ne00_padded = is_tiled ? hex_round_up(ne00, 32) : ne00; const int ne11_padded = hex_round_up(ne11, 32); // VTCM has to hold whole 32-row weight tiles, so size for the rounded-up N // even when the tensor itself is ragged. - const int ne01_tiled = hex_round_up(ne01_padded, 32); + const int ne01_tiled = hex_round_up(ne01, 32); const bool is_matmul_id = (dst->op == GGML_OP_MUL_MAT_ID); const bool is_batched = (ne02 * ne03 > 1 || ne12 * ne13 > 1); @@ -5777,14 +4137,14 @@ static void ggml_hexagon_precompute_matmul_params_impl( // Check HMX eligibility and try precomputing HMX parameters bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 2); - if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, src1, dst, ne01_padded, is_matmul_id, is_batched)) { + if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, src1, dst, is_matmul_id, is_batched)) { if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, src1, dst, wtype, ne00_padded, ne01_tiled, ne02, ne11, ne12, ne11_padded, is_matmul_id, is_batched, src2_size, vtcm_budget, kparams)) { goto finalize; } } // Fallback to HVX parameter computation - ggml_hexagon_precompute_hvx_mm_params(sess, src0, src1, dst, wtype, ne02, ne03, ne10, ne11, ne12, ne13, is_matmul_id, src2_row_size, vtcm_budget, kparams); + ggml_hexagon_precompute_hvx_mm_params(sess, src0, src1, dst, wtype, is_tiled, ne02, ne03, ne10, ne11, ne12, ne13, is_matmul_id, src2_row_size, vtcm_budget, kparams); finalize: kparams->div_ne12_ne1 = init_fastdiv_values(ne12 * ne11); @@ -6170,29 +4530,24 @@ static void ggml_hexagon_precompute_get_rows_params( const uint32_t ne12 = src1->ne[2]; const uint32_t nr = ne10 * ne11 * ne12; - const ggml_tensor * src0_base = src0->view_src ? src0->view_src : src0; - const auto * extra = src0_base->buffer && ggml_backend_buffer_is_hexagon(src0_base->buffer) ? - (const ggml_hexagon_tensor_extra *) src0_base->extra : nullptr; - const bool tiled = src0->type == GGML_TYPE_Q4_0 || (extra && (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0) || - sess->needs_repack.count(src0_base) || sess->needs_repack.count(src0); + const bool tiled = ggml_hexagon_tensor_is_tiled(src0); + const int wtype = (tiled && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_BF16)) ? GGML_TYPE_F16 : src0->type; - if (src0->type == dst->type) { - kparams->kernel_type = HTP_GET_ROWS_KERNEL_SAMETYPE; - } else if (tiled) { + if (tiled) { kparams->kernel_type = HTP_GET_ROWS_KERNEL_TILED; + } else if (src0->type == dst->type) { + kparams->kernel_type = HTP_GET_ROWS_KERNEL_SAMETYPE; } else { kparams->kernel_type = HTP_GET_ROWS_KERNEL_FLAT; } - const uint32_t chunks_per_row = 1; - const uint32_t chunk_size = ne00; const uint32_t total_tasks = nr; kparams->n_threads = (std::min)((uint32_t)sess->n_threads, total_tasks); struct htp_get_rows_vtcm_layout vtcm_layout = {}; while (kparams->n_threads > 0) { - htp_get_rows_vtcm_layout_build(&vtcm_layout, kparams->kernel_type, src0->type, ne00, kparams->n_threads); + htp_get_rows_vtcm_layout_build(&vtcm_layout, kparams->kernel_type, wtype, ne00, total_tasks, src1->type, kparams->n_threads); if (vtcm_layout.total_bytes <= sess->vtcm_size) { break; } @@ -6200,19 +4555,15 @@ static void ggml_hexagon_precompute_get_rows_params( } if (kparams->n_threads == 0 && total_tasks > 0) { - htp_get_rows_vtcm_layout_build(&vtcm_layout, kparams->kernel_type, src0->type, ne00, 1); + htp_get_rows_vtcm_layout_build(&vtcm_layout, kparams->kernel_type, wtype, ne00, total_tasks, src1->type, 1); } kparams->vtcm_size = (total_tasks == 0) ? 0 : vtcm_layout.total_bytes; kparams->tasks_per_thread = kparams->n_threads > 0 ? (total_tasks + kparams->n_threads - 1) / kparams->n_threads : 0; - - kparams->chunks_per_row = chunks_per_row; - kparams->chunk_size = chunk_size; kparams->total_tasks = total_tasks; kparams->div_ne10 = ne10 > 0 ? init_fastdiv_values(ne10) : fastdiv_values{0, 0}; kparams->div_ne10_ne11 = (ne10 * ne11) > 0 ? init_fastdiv_values(ne10 * ne11) : fastdiv_values{0, 0}; - kparams->div_chunks_per_row = chunks_per_row > 0 ? init_fastdiv_values(chunks_per_row) : fastdiv_values{0, 0}; kparams->div_ne02 = ne02 > 0 ? init_fastdiv_values(ne02) : fastdiv_values{0, 0}; kparams->div_ne03 = ne03 > 0 ? init_fastdiv_values(ne03) : fastdiv_values{0, 0}; } @@ -6920,49 +5271,44 @@ static void ggml_hexagon_precompute_fused_mmnx_params( const int ne12 = act->ne[2]; const int ne13 = act->ne[3]; - const int wtype = src0->type; - const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype); - const int ne00_padded = is_repack ? hex_round_up(ne00, 32) : ne00; - const int ne01_padded = is_repack ? hex_round_up(ne01, 32) : ne01; + const bool is_tiled = ggml_hexagon_tensor_is_tiled(src0); + const int wtype = (is_tiled && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_BF16)) ? GGML_TYPE_F16 : src0->type; + const int ne00_padded = is_tiled ? hex_round_up(ne00, 32) : ne00; const int ne11_padded = hex_round_up(ne11, 32); - const int ne01_tiled = hex_round_up(ne01_padded, 32); + const int ne01_tiled = hex_round_up(ne01, 32); const size_t vtcm_budget = sess->vtcm_size; const bool is_batched = (ne02 * ne03 > 1 || ne12 * ne13 > 1); bool hmx_enabled = (sess->n_hmx > 0) && (opt_mm_select >= 2); - if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, act, nullptr, ne01_padded, false, is_batched)) { + if (hmx_enabled && ggml_hexagon_matmul_is_hmx_eligible(src0, act, nullptr, false, is_batched)) { if (ggml_hexagon_precompute_hmx_mm_params(sess, src0, act, nullptr, wtype, ne00_padded, ne01_tiled, ne02, ne11, ne12, ne11_padded, false, is_batched, 0, vtcm_budget, kparams)) { kparams->n_weights = n_weights; goto finalize; } } - if (!is_repack) { + if (!is_tiled || ne10 % 32 != 0) { kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; return; } { const int act_nrows = ne11 * ne12 * ne13; - const size_t act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + const size_t act_row_size = htp_mm_tiled_act_row_size(wtype, ne10); const size_t src0_row_size = src0->nb[1]; - uint32_t best_n_prefetch = 16; - - if (is_repack) { - const uint32_t max_prefetch = (act_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16; - best_n_prefetch = 2; - for (uint32_t d = max_prefetch; d >= 2; d /= 2) { - struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, act_nrows, sess->n_threads, - 0, src0_row_size, act_row_size, 0, d, false, true - ); - if (L.total_bytes <= sess->vtcm_size) { - best_n_prefetch = d; - break; - } + const uint32_t max_prefetch = (act_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16; + uint32_t best_n_prefetch = 2; + for (uint32_t d = max_prefetch; d >= 2; d /= 2) { + struct htp_mm_hvx_vtcm_layout L; + htp_mm_hvx_vtcm_layout_build( + &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, act_nrows, sess->n_threads, + 0, src0_row_size, act_row_size, 0, d, false, true + ); + if (L.total_bytes <= sess->vtcm_size) { + best_n_prefetch = d; + break; } } @@ -6976,15 +5322,17 @@ static void ggml_hexagon_precompute_fused_mmnx_params( ); if (try_tiled && L.total_bytes <= sess->vtcm_size) { - kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW; - kparams->act_row_size = act_row_size; - kparams->vtcm_src0_size = L.src0_bytes; - kparams->vtcm_act_size = L.act_bytes; - kparams->vtcm_bias_size = 0; - kparams->vtcm_dst_size = L.dst_bytes; - kparams->vtcm_size = L.total_bytes; - kparams->n_prefetch = best_n_prefetch; - kparams->n_weights = n_weights; + kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW; + kparams->tile_size = htp_tensor_get_tile_size(wtype); + kparams->aligned_tile_size = htp_tensor_get_tile_padded_size(wtype); + kparams->act_row_size = act_row_size; + kparams->vtcm_src0_size = L.src0_bytes; + kparams->vtcm_act_size = L.act_bytes; + kparams->vtcm_bias_size = 0; + kparams->vtcm_dst_size = L.dst_bytes; + kparams->vtcm_size = L.total_bytes; + kparams->n_prefetch = best_n_prefetch; + kparams->n_weights = n_weights; } else { kparams->kernel_type = HTP_MM_KERNEL_UNSUPPORTED; return; @@ -7011,16 +5359,6 @@ static void ggml_hexagon_precompute_fused_mmidnx_params( kparams->n_weights = n_weights; } -static bool ggml_hexagon_tensor_is_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) { - return t && t->buffer && ggml_backend_buft_is_host(t->buffer->buft); - GGML_UNUSED(sess); -} - -static bool ggml_hexagon_tensor_is_non_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) { - return t && t->buffer && !ggml_backend_buft_is_host(t->buffer->buft); - GGML_UNUSED(sess); -} - static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * sess, const struct ggml_tensor * dst) { const struct ggml_tensor * src0 = dst->src[0]; const struct ggml_tensor * src1 = dst->src[1]; @@ -7044,7 +5382,7 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s case GGML_TYPE_Q6_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q2_K: - if (!ggml_is_contiguous(src0) || ggml_is_permuted(src0)) { + if (!ggml_is_contiguous(src0) || ggml_is_permuted(src0) || ggml_is_view(src0)) { return false; } @@ -7059,10 +5397,6 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s if (src1->ne[2] % src0->ne[2] != 0 || src1->ne[3] % src0->ne[3] != 0) { return false; } - - if (!src0->buffer) { - sess->needs_repack.insert(src0); - } break; case GGML_TYPE_F16: @@ -7092,6 +5426,18 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s } break; + case GGML_TYPE_BF16: + if (src0->nb[1] < src0->nb[0]) { + return false; + } + if (src1->ne[2] < src0->ne[2] || src1->ne[3] < src0->ne[3]) { + return false; + } + if (src1->ne[2] % src0->ne[2] != 0 || src1->ne[3] % src0->ne[3] != 0) { + return false; + } + break; + default: return false; } @@ -7131,7 +5477,7 @@ static bool ggml_hexagon_supported_mul_mat_id(const struct ggml_hexagon_session case GGML_TYPE_Q6_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q2_K: - if (!ggml_is_contiguous(src0) || ggml_is_permuted(src0)) { + if (!ggml_is_contiguous(src0) || ggml_is_permuted(src0) || ggml_is_view(src0)) { return false; } @@ -7139,9 +5485,16 @@ static bool ggml_hexagon_supported_mul_mat_id(const struct ggml_hexagon_session src0->type == GGML_TYPE_Q3_K || src0->type == GGML_TYPE_Q2_K) ? QK_K : 32)) { return false; } + break; - if (!src0->buffer) { - sess->needs_repack.insert(src0); + case GGML_TYPE_F16: + case GGML_TYPE_F32: + case GGML_TYPE_BF16: + if (!ggml_is_contiguous(src0) || ggml_is_permuted(src0)) { + return false; + } + if (src0->ne[0] % 32 != 0) { + return false; } break; @@ -7564,40 +5917,12 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session * const struct ggml_tensor * src1 = op->src[1]; // indices const struct ggml_tensor * dst = op; - if ((src0->type == GGML_TYPE_Q4_0 || src0->type == GGML_TYPE_Q4_K || src0->type == GGML_TYPE_Q6_K) && src0->view_src) { - return false; - } - - const ggml_tensor * src0_base = src0->view_src ? src0->view_src : src0; - bool is_repacked = false; - if (src0_base->buffer && ggml_backend_buffer_is_hexagon(src0_base->buffer) && src0_base->extra) { - const auto * extra = (const ggml_hexagon_tensor_extra *) src0_base->extra; - is_repacked = (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0; - if (is_repacked && src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q4_K && src0->type != GGML_TYPE_Q6_K && src0->type != GGML_TYPE_Q8_0) { - return false; - } - } - is_repacked = is_repacked || sess->needs_repack.count(src0_base) || sess->needs_repack.count(src0); - - // View offsets use the raw quantized layout and cannot address a tiled allocation. - if (src0->view_src && is_repacked) { - return false; - } - - if ((src0->type == GGML_TYPE_Q4_0 || src0->type == GGML_TYPE_Q4_K || src0->type == GGML_TYPE_Q6_K) && src0->buffer && ggml_backend_buffer_get_size(src0->buffer) != 0 && !is_repacked) { - return false; - } - - if (src0->type != dst->type && src0->ne[0] < 32) { + if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16 && src0->type != GGML_TYPE_BF16 && + src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q4_K && src0->type != GGML_TYPE_Q5_K && src0->type != GGML_TYPE_Q6_K && src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_I32) { return false; } - if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16 && - src0->type != GGML_TYPE_Q4_0 && src0->type != GGML_TYPE_Q4_K && src0->type != GGML_TYPE_Q6_K && src0->type != GGML_TYPE_Q8_0 && src0->type != GGML_TYPE_I32) { - return false; - } - - if ((src0->type == GGML_TYPE_Q4_K || src0->type == GGML_TYPE_Q6_K) && (!ggml_is_contiguous(src0) || ggml_is_permuted(src0) || src0->ne[0] % QK_K)) { + if (!ggml_is_contiguous_rows(src0) || !ggml_is_contiguous_rows(src1) || !ggml_is_contiguous_rows(dst)) { return false; } @@ -7626,11 +5951,6 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session * return false; } - // Tiled quantized weights have no raw fallback. Mark only accepted tensors for repacking. - if ((src0->type == GGML_TYPE_Q4_0 || src0->type == GGML_TYPE_Q4_K || src0->type == GGML_TYPE_Q6_K) && !src0->buffer) { - sess->needs_repack.insert(src0); - } - return true; GGML_UNUSED(sess); @@ -8436,20 +6756,13 @@ static bool mm_is_hmx_eligible(const ggml_tensor * t) { if (ggml_hexagon_matmul_can_collapse(src0, src1, t)) { const ggml_tensor src1_c = ggml_hexagon_tensor_collapse_rows(src1); - const int wtype = src0->type; - const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype); - const int ne01_padded = is_repack ? hex_round_up(src0->ne[1], 32) : src0->ne[1]; - return ggml_hexagon_matmul_is_hmx_eligible(src0, &src1_c, t, ne01_padded, false, false); + return ggml_hexagon_matmul_is_hmx_eligible(src0, &src1_c, t, false, false); } - const int wtype = src0->type; - const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype); const bool is_matmul_id = (t->op == GGML_OP_MUL_MAT_ID); const bool is_batched = (src0->ne[2] * src0->ne[3] > 1 || src1->ne[2] * src1->ne[3] > 1); - const int ne01_padded = is_repack ? hex_round_up(src0->ne[1], 32) : src0->ne[1]; - - return ggml_hexagon_matmul_is_hmx_eligible(src0, src1, t, ne01_padded, is_matmul_id, is_batched); + return ggml_hexagon_matmul_is_hmx_eligible(src0, src1, t, is_matmul_id, is_batched); } static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struct htp_mm_kernel_params * kparams) { @@ -8457,8 +6770,8 @@ static bool is_supported_mul_mat_nx_kernel(const ggml_tensor * src0, const struc return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D; } - if (!ggml_hexagon_is_repack_type(src0->type) || src0->type == GGML_TYPE_Q6_K) { - return false; // Q6_K has no fused HVX kernel + if (!ggml_hexagon_tensor_is_tiled(src0)) { + return false; } return kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW; @@ -8469,7 +6782,7 @@ static bool is_supported_mul_mat_id_nx_kernel(const ggml_tensor * src0, const st return kparams->kernel_type == HTP_MM_KERNEL_HMX_2D; } - if (!ggml_hexagon_is_repack_type(src0->type)) { + if (!ggml_hexagon_tensor_is_tiled(src0)) { return false; } @@ -8484,13 +6797,13 @@ static bool is_mergeable_mul_mat(const ggml_tensor * t) { if (src0->ne[2] != 1 || src0->ne[3] != 1) return false; if (mm_is_hmx_eligible(t)) { - return ggml_hexagon_is_hmx_weight_type(src0->type); + return htp_is_hmx_weight_type(src0->type); } - // HVX path requires F32 activations and repacked weights (except Q6_K) + // HVX path requires F32 activations and tiled weights if (src1->type != GGML_TYPE_F32) return false; - return ggml_hexagon_is_repack_type(src0->type) && src0->type != GGML_TYPE_Q6_K; + return ggml_hexagon_tensor_is_tiled(src0); } static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2) { @@ -8506,6 +6819,9 @@ static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor if (n1->src[0]->type != n2->src[0]->type) { return false; } + if (ggml_hexagon_tensor_is_tiled(n1->src[0]) != ggml_hexagon_tensor_is_tiled(n2->src[0])) { + return false; + } if (mm_is_hmx_eligible(n1) != mm_is_hmx_eligible(n2)) { return false; } @@ -8516,7 +6832,11 @@ static bool is_mergeable_mul_mat_id(const ggml_tensor * t) { if (t->op != GGML_OP_MUL_MAT_ID) return false; const ggml_tensor * src0 = t->src[0]; - return ggml_hexagon_is_repack_type(src0->type); + if (mm_is_hmx_eligible(t)) { + return htp_is_hmx_weight_type(src0->type); + } + + return ggml_hexagon_tensor_is_tiled(src0); } static bool is_mergeable_mul_mat_id_pair(const ggml_tensor * n1, const ggml_tensor * n2) { @@ -8538,6 +6858,9 @@ static bool is_mergeable_mul_mat_id_pair(const ggml_tensor * n1, const ggml_tens if (n1->src[0]->type != n2->src[0]->type) { return false; } + if (ggml_hexagon_tensor_is_tiled(n1->src[0]) != ggml_hexagon_tensor_is_tiled(n2->src[0])) { + return false; + } if (mm_is_hmx_eligible(n1) != mm_is_hmx_eligible(n2)) { return false; } @@ -8588,6 +6911,26 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg htp_opnode node(HTP_OP_INVALID, n); node.opcode = op_remap_to_htp(n); + if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID || node.opcode == HTP_OP_GET_ROWS) { + ggml_tensor * src0 = node.node->src[0]; + if (ggml_hexagon_needs_repack(src0)) { + if (opt_verbose) { + char dims_buf[64]; + char strides_buf[64]; + htp_opformat fmt; + fmt.format_tensor_dims(dims_buf, sizeof(dims_buf), src0); + fmt.format_tensor_strides(strides_buf, sizeof(strides_buf), src0); + GGML_LOG_DEBUG("ggml-hex: %s rt-repack %s|%s|%s|%s|size %zu\n", + sess->c_name(), src0->name, dims_buf, ggml_type_name(src0->type), strides_buf, ggml_nbytes(src0)); + } + std::vector<uint8_t> raw(ggml_nbytes(src0)); + memcpy(raw.data(), src0->data, raw.size()); + htp_repack_tensor_tiled(src0, raw.data(), raw.size()); + auto * extra = (ggml_hexagon_tensor_extra *) src0->extra; + extra->flags |= GGML_HEXAGON_TENSOR_TILED; + } + } + if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID) { ggml_hexagon_precompute_matmul_params(sess, node.node->src[0], node.node->src[1], node.node, @@ -9378,6 +7721,9 @@ static bool ggml_hexagon_supported_roll(const struct ggml_hexagon_session * sess return true; } +// IMPORTANT: supports_op runs before buffer allocation and before tensor data is populated. +// Never check buffer usage, buffer properties, or repacking state (is_tiled) here. +// Support must depend only on static tensor properties (shape, type, strides). static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { auto dev_ctx = static_cast<ggml_backend_hexagon_device_context *>(dev->context); auto sess = dev_ctx->session(); @@ -9971,12 +8317,15 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { "please update hexagon_type to match ggml_type"); static_assert((unsigned int) HTP_TYPE_Q2_K == (unsigned int) GGML_TYPE_Q2_K, "please update hexagon_type to match ggml_type"); + static_assert((unsigned int) HTP_TYPE_BF16 == (unsigned int) GGML_TYPE_BF16, + "please update hexagon_type to match ggml_type"); const char * str_verbose = getenv("GGML_HEXAGON_VERBOSE"); const char * str_opbatch = getenv("GGML_HEXAGON_OPBATCH"); const char * str_opqueue = getenv("GGML_HEXAGON_OPQUEUE"); const char * str_oppoll = getenv("GGML_HEXAGON_OPPOLL"); const char * str_opfusion = getenv("GGML_HEXAGON_OPFUSION"); + const char * str_rt_repack = getenv("GGML_HEXAGON_RT_REPACK"); const char * str_opfilter = getenv("GGML_HEXAGON_OPFILTER"); const char * str_profile = getenv("GGML_HEXAGON_PROFILE"); const char * str_etm = getenv("GGML_HEXAGON_ETM"); @@ -10034,6 +8383,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { opt_optrace = str_optrace ? strtoul(str_optrace, NULL, 0) : (opt_opbatch * 256); opt_oppoll = str_oppoll ? strtoul(str_oppoll, NULL, 0) : opt_oppoll; opt_opfusion = str_opfusion ? atoi(str_opfusion) : opt_opfusion; + opt_rt_repack = str_rt_repack ? atoi(str_rt_repack) : opt_rt_repack; opt_profile = str_profile ? atoi(str_profile) : 0; opt_etm = str_etm ? atoi(str_etm) : 0; if (str_power_mode) { diff --git a/ggml/src/ggml-hexagon/htp-opnode.h b/ggml/src/ggml-hexagon/htp-opnode.h index 8ff38545c74e..d62ce1d000b5 100644 --- a/ggml/src/ggml-hexagon/htp-opnode.h +++ b/ggml/src/ggml-hexagon/htp-opnode.h @@ -326,10 +326,12 @@ struct htp_opformat { const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params; const char * path = "unknown"; int32_t type = kparams->kernel_type; - if (type == HTP_MM_KERNEL_HMX_2D || type == HTP_MM_KERNEL_HMX_F16_BATCHED) { - path = "hmx-tiled"; + if (type == HTP_MM_KERNEL_HMX_2D || type == HTP_MM_KERNEL_HMX_BATCHED) { + path = (kparams->tile_size > 0) ? "hmx-tiled" : "hmx-flat"; } else if (type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || type == HTP_MM_KERNEL_HVX_F32_F32_VTCM || - type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { + type == HTP_MM_KERNEL_HVX_BF16_F32_VTCM) { + path = "hvx-flat"; + } else if (type == HTP_MM_KERNEL_HVX_QUANT_ROW || type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { path = "hvx-tiled"; } snprintf(str, max_size, "%s vtcm %d", path, (int) kparams->vtcm_size); diff --git a/ggml/src/ggml-hexagon/htp-repack.cpp b/ggml/src/ggml-hexagon/htp-repack.cpp new file mode 100644 index 000000000000..38beabe82a68 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp-repack.cpp @@ -0,0 +1,1000 @@ +#include "htp-repack.h" + +#define GGML_COMMON_IMPL_CPP +#include "ggml-common.h" +#include "ggml-impl.h" +#include "ggml-quants.h" + +#include <algorithm> +#include <cmath> +#include <cstring> +#include <initializer_list> + +// --- 4-bit block-32 nibble helpers (shared by Q4_0, IQ4_NL, Q4_1, MXFP4) --- + +static inline void unpack_nibbles_32(uint8_t * dst, const uint8_t * qs) { + for (unsigned int i = 0; i < 16; ++i) { + dst[i + 0] = qs[i] & 0x0F; + dst[i + 16] = qs[i] >> 4; + } +} + +static inline void pack_nibbles_32(uint8_t * qs, const uint8_t * src) { + for (unsigned int i = 0; i < 16; ++i) { + qs[i] = src[i + 0] | (src[i + 16] << 4); + } +} + +static inline void transpose_nibbles_to_tile(uint8_t * tile_dst, const uint8_t tile_quants[32][32]) { + for (int cp = 0; cp < 16; cp++) { + for (int row = 0; row < 32; row++) { + tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; + } + } +} + +static inline void transpose_tile_to_nibbles(uint8_t tile_quants[32][32], const uint8_t * tile_src) { + for (int cp = 0; cp < 16; cp++) { + for (int row = 0; row < 32; row++) { + const uint8_t val = tile_src[cp * 32 + row]; + tile_quants[row][2 * cp + 0] = val & 0x0F; + tile_quants[row][2 * cp + 1] = val >> 4; + } + } +} + +// --- Generic 32x32 tile grid walkers for block-32 and float types --- + +template <typename SrcElemT, int ELEMS_PER_KTILE, typename PackTileFn> +static void repack_to_tiled_grid_32(ggml_tensor * t, const void * data, size_t offset, size_t size, + size_t tile_size, PackTileFn && pack_tile) { + const SrcElemT * src_matrix = (const SrcElemT *) data; + const int64_t ne0 = t->ne[0]; + const int64_t ne1 = t->ne[1]; + const int64_t ne2 = t->ne[2]; + const int64_t ne3 = t->ne[3]; + const int n_col_tiles = hex_round_up(ne1, 32) / 32; + const int n_k_tiles = hex_round_up(ne0, 32) / 32; + const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; + + const size_t slice_size = (size_t) ne1 * ggml_row_size(t->type, ne0); + const int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + + const int64_t row_stride = (ne0 / (32 / ELEMS_PER_KTILE)); + + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + const SrcElemT * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * row_stride); + uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; + + for (int ct = 0; ct < n_col_tiles; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + uint8_t * tile_dst = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size; + pack_tile(tile_dst, src_slice, ne0, ne1, row_stride, ct, kt); + } + } + } +} + +template <typename DstElemT, int ELEMS_PER_KTILE, typename UnpackTileFn> +static void unpack_from_tiled_grid_32(void * data, const ggml_tensor * t, size_t offset, size_t size, + size_t tile_size, UnpackTileFn && unpack_tile) { + uint8_t * dst_bytes = (uint8_t *) data; + const int64_t ne0 = t->ne[0]; + const int64_t ne1 = t->ne[1]; + const int64_t ne2 = t->ne[2]; + const int64_t ne3 = t->ne[3]; + const int n_col_tiles = hex_round_up(ne1, 32) / 32; + const int n_k_tiles = hex_round_up(ne0, 32) / 32; + const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; + + const size_t row_size_bytes = ggml_row_size(t->type, ne0); + const size_t slice_size = (size_t) ne1 * row_size_bytes; + const int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + + const int64_t row_stride = (ne0 / (32 / ELEMS_PER_KTILE)); + + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + const size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); + const size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); + const size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; + const size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; + + const int64_t start_row = slice_offset_start / row_size_bytes; + const int64_t end_row = (std::min)((int64_t) ((slice_offset_end + row_size_bytes - 1) / row_size_bytes), ne1); + + const int start_ct = (int) (start_row / 32); + const int end_ct = (std::min)((int) ((end_row + 31) / 32), n_col_tiles); + + DstElemT * dst_slice = (DstElemT *) (dst_bytes + (cur_start_byte - offset)); + const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; + + for (int ct = start_ct; ct < end_ct; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + const uint8_t * tile_src = matrix_src + ((size_t) ct * n_k_tiles + kt) * tile_size; + unpack_tile(dst_slice, tile_src, ne0, start_row, end_row, row_stride, ct, kt); + } + } + } +} + +// --- Block-32 Quantized Types (Q4_0, IQ4_NL, Q4_1, Q8_0, MXFP4) --- + +static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { + repack_to_tiled_grid_32<block_q4_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_0, + [](uint8_t * tile_dst, const block_q4_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) { + uint8_t tile_quants[32][32]; + ggml_half * scale_dst = (ggml_half *) (tile_dst + 512); + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + if (r < ne1 && kt < ne0 / 32) { + const block_q4_0 & b = src_slice[r * row_stride + kt]; + unpack_nibbles_32(tile_quants[row], b.qs); + scale_dst[row] = b.d; + } else { + memset(tile_quants[row], 8, 32); + scale_dst[row] = 0; + } + } + transpose_nibbles_to_tile(tile_dst, tile_quants); + }); +} + +static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t offset, size_t size) { + unpack_from_tiled_grid_32<block_q4_0, 1>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q4_0, + [](block_q4_0 * dst_slice, const uint8_t * tile_src, int64_t ne0, int64_t start_row, int64_t end_row, int64_t row_stride, int ct, int kt) { + uint8_t tile_quants[32][32]; + transpose_tile_to_nibbles(tile_quants, tile_src); + const ggml_half * scale_src = (const ggml_half *) (tile_src + 512); + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + block_q4_0 & b = dst_slice[(r - start_row) * row_stride + kt]; + pack_nibbles_32(b.qs, tile_quants[row]); + b.d = scale_src[row]; + } + } + }); +} + +static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { + repack_to_tiled_grid_32<block_q4_1, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_1, + [](uint8_t * tile_dst, const block_q4_1 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) { + uint8_t tile_quants[32][32]; + ggml_half * scale_dst = (ggml_half *) (tile_dst + 512); + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + if (r < ne1 && kt < ne0 / 32) { + const block_q4_1 & b = src_slice[r * row_stride + kt]; + unpack_nibbles_32(tile_quants[row], b.qs); + scale_dst[2 * row + 0] = b.d; + scale_dst[2 * row + 1] = b.m; + } else { + memset(tile_quants[row], 0, 32); + scale_dst[2 * row + 0] = 0; + scale_dst[2 * row + 1] = 0; + } + } + transpose_nibbles_to_tile(tile_dst, tile_quants); + }); +} + +static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t offset, size_t size) { + unpack_from_tiled_grid_32<block_q4_1, 1>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q4_1, + [](block_q4_1 * dst_slice, const uint8_t * tile_src, int64_t ne0, int64_t start_row, int64_t end_row, int64_t row_stride, int ct, int kt) { + uint8_t tile_quants[32][32]; + transpose_tile_to_nibbles(tile_quants, tile_src); + const ggml_half * scale_src = (const ggml_half *) (tile_src + 512); + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + block_q4_1 & b = dst_slice[(r - start_row) * row_stride + kt]; + pack_nibbles_32(b.qs, tile_quants[row]); + b.d = scale_src[2 * row + 0]; + b.m = scale_src[2 * row + 1]; + } + } + }); +} + +static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { + repack_to_tiled_grid_32<block_mxfp4, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_MXFP4, + [](uint8_t * tile_dst, const block_mxfp4 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) { + uint8_t tile_quants[32][32]; + uint8_t * scale_dst = tile_dst + 512; + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + if (r < ne1 && kt < ne0 / 32) { + const block_mxfp4 & b = src_slice[r * row_stride + kt]; + unpack_nibbles_32(tile_quants[row], b.qs); + scale_dst[row] = b.e; + } else { + memset(tile_quants[row], 0, 32); + scale_dst[row] = 0; + } + } + transpose_nibbles_to_tile(tile_dst, tile_quants); + }); +} + +static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t offset, size_t size) { + unpack_from_tiled_grid_32<block_mxfp4, 1>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_MXFP4, + [](block_mxfp4 * dst_slice, const uint8_t * tile_src, int64_t ne0, int64_t start_row, int64_t end_row, int64_t row_stride, int ct, int kt) { + uint8_t tile_quants[32][32]; + transpose_tile_to_nibbles(tile_quants, tile_src); + const uint8_t * scale_src = tile_src + 512; + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + block_mxfp4 & b = dst_slice[(r - start_row) * row_stride + kt]; + pack_nibbles_32(b.qs, tile_quants[row]); + b.e = scale_src[row]; + } + } + }); +} + +static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { + repack_to_tiled_grid_32<block_q8_0, 1>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q8_0, + [](uint8_t * tile_dst, const block_q8_0 * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) { + for (int cp = 0; cp < 16; cp++) { + const int col0 = cp * 2; + const int col1 = col0 + 1; + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + const block_q8_0 * b = (r < ne1 && kt < ne0 / 32) ? &src_slice[r * row_stride + kt] : nullptr; + tile_dst[cp * 64 + 2 * row + 0] = b ? b->qs[col0] : 0; + tile_dst[cp * 64 + 2 * row + 1] = b ? b->qs[col1] : 0; + } + } + ggml_half * scale_dst = (ggml_half *) (tile_dst + 1024); + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * row_stride + kt].d : 0; + } + }); +} + +static void repack_tiled_q8_0(void * data, const ggml_tensor * t, size_t offset, size_t size) { + unpack_from_tiled_grid_32<block_q8_0, 1>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q8_0, + [](block_q8_0 * dst_slice, const uint8_t * tile_src, int64_t ne0, int64_t start_row, int64_t end_row, int64_t row_stride, int ct, int kt) { + for (int cp = 0; cp < 16; cp++) { + const int col0 = cp * 2; + const int col1 = col0 + 1; + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + block_q8_0 & b = dst_slice[(r - start_row) * row_stride + kt]; + b.qs[col0] = tile_src[cp * 64 + 2 * row + 0]; + b.qs[col1] = tile_src[cp * 64 + 2 * row + 1]; + } + } + } + const ggml_half * scale_src = (const ggml_half *) (tile_src + 1024); + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + dst_slice[(r - start_row) * row_stride + kt].d = scale_src[row]; + } + } + }); +} + +// --- Float Types (F16, F32, BF16 -> f16_tiled HMX layout) --- + +template <typename SrcT, typename ToFp16Fn> +static void repack_float_to_f16_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size, ToFp16Fn && to_fp16) { + repack_to_tiled_grid_32<SrcT, 32>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_F16, + [&](uint8_t * tile_bytes, const SrcT * src_slice, int64_t ne0, int64_t ne1, int64_t row_stride, int ct, int kt) { + uint16_t * tile_dst = (uint16_t *) tile_bytes; + for (int cp = 0; cp < 16; cp++) { + const int64_t c0 = kt * 32 + cp * 2; + const int64_t c1 = c0 + 1; + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + const SrcT * r_ptr = (r < ne1) ? &src_slice[r * row_stride] : nullptr; + tile_dst[cp * 64 + row * 2 + 0] = (r_ptr && c0 < ne0) ? to_fp16(r_ptr[c0]) : 0; + tile_dst[cp * 64 + row * 2 + 1] = (r_ptr && c1 < ne0) ? to_fp16(r_ptr[c1]) : 0; + } + } + }); +} + +template <typename DstT, typename FromFp16Fn> +static void unpack_f16_tiled_to_float(void * data, const ggml_tensor * t, size_t offset, size_t size, FromFp16Fn && from_fp16) { + unpack_from_tiled_grid_32<DstT, 32>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_F16, + [&](DstT * dst_slice, const uint8_t * tile_bytes, int64_t ne0, int64_t start_row, int64_t end_row, int64_t row_stride, int ct, int kt) { + const uint16_t * tile_src = (const uint16_t *) tile_bytes; + for (int cp = 0; cp < 16; cp++) { + const int64_t c0 = kt * 32 + cp * 2; + const int64_t c1 = c0 + 1; + for (int row = 0; row < 32; row++) { + const int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row) { + DstT * r_ptr = dst_slice + (r - start_row) * row_stride; + if (c0 < ne0) r_ptr[c0] = from_fp16(tile_src[cp * 64 + row * 2 + 0]); + if (c1 < ne0) r_ptr[c1] = from_fp16(tile_src[cp * 64 + row * 2 + 1]); + } + } + } + }); +} + +// --- Generic K-Quant (super-block 256) grid walkers --- + +template <typename BlockT, typename PackKtFn> +static void repack_kquant_to_tiled_grid(ggml_tensor * t, const void * data, size_t offset, size_t size, + size_t tile_size, PackKtFn && pack_kt) { + GGML_ASSERT(offset == 0); + GGML_UNUSED(size); + + const BlockT * src_matrix = (const BlockT *) data; + const int64_t ne0 = t->ne[0]; + const int64_t ne1 = t->ne[1]; + const int64_t ne2 = t->ne[2]; + const int64_t ne3 = t->ne[3]; + GGML_ASSERT(ne0 % QK_K == 0); + + const int n_col_tiles = hex_round_up(ne1, 32) / 32; + const int n_k_tiles = hex_round_up(ne0, 32) / 32; + const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; + const int64_t sb_per_row = ne0 / QK_K; + + for (int64_t slice_idx = 0; slice_idx < ne2 * ne3; slice_idx++) { + const BlockT * src_slice = src_matrix + slice_idx * (ne1 * sb_per_row); + uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; + + // Padding rows and OR-ed bit planes require zeroed tiles + memset(matrix_dst, 0, matrix_size); + + for (int64_t r = 0; r < ne1; r++) { + const int ct = (int) (r / 32); + const int row = (int) (r % 32); + const BlockT * src_row = src_slice + r * sb_per_row; + + for (int kt = 0; kt < n_k_tiles; kt++) { + uint8_t * tile = matrix_dst + ((size_t) ct * n_k_tiles + kt) * tile_size; + pack_kt(tile, &src_row[kt / 8], row, kt % 8); + } + } + } +} + +template <typename BlockT, typename UnpackSbFn> +static void unpack_kquant_from_tiled_grid(void * data, const ggml_tensor * t, size_t offset, size_t size, + size_t tile_size, UnpackSbFn && unpack_sb) { + uint8_t * dst_bytes = (uint8_t *) data; + const int64_t ne0 = t->ne[0]; + const int64_t ne1 = t->ne[1]; + const int64_t ne2 = t->ne[2]; + const int64_t ne3 = t->ne[3]; + GGML_ASSERT(ne0 % QK_K == 0); + + const int n_col_tiles = hex_round_up(ne1, 32) / 32; + const int n_k_tiles = hex_round_up(ne0, 32) / 32; + const size_t matrix_size = (size_t) n_col_tiles * n_k_tiles * tile_size; + const int64_t sb_per_row = ne0 / QK_K; + + const size_t row_size_bytes = ggml_row_size(t->type, ne0); + const size_t slice_size = (size_t) ne1 * row_size_bytes; + const int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + const size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); + const size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); + const size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; + const size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; + + const int64_t start_row = slice_offset_start / row_size_bytes; + const int64_t end_row = (std::min)((int64_t) ((slice_offset_end + row_size_bytes - 1) / row_size_bytes), ne1); + + BlockT * dst_slice = (BlockT *) (dst_bytes + (cur_start_byte - offset)); + const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; + + for (int64_t r = start_row; r < end_row; r++) { + const int ct = (int) (r / 32); + const int row = (int) (r % 32); + BlockT * dst_row = dst_slice + (r - start_row) * sb_per_row; + + for (int64_t sb = 0; sb < sb_per_row; sb++) { + BlockT * b = &dst_row[sb]; + memset(b, 0, sizeof(BlockT)); + const uint8_t * sb_tiles = matrix_src + ((size_t) ct * n_k_tiles + sb * 8) * tile_size; + unpack_sb(b, sb_tiles, tile_size, row); + } + } + } +} + +// --- K-Quant Helpers & Per-Type Implementations (Q2_K, Q3_K, Q4_K, Q5_K, Q6_K) --- + +// unsigned 6-bit value (0..63) of element e of a Q6_K block, same bit layout as dequantize_row_q6_K +static inline uint8_t q6_K_get_quant(const block_q6_K * b, int e) { + const int c = e / 128; + const int w = e % 128; + const int g = w / 32; + const int l = w % 32; + const uint8_t * ql = b->ql + c * 64; + const uint8_t * qh = b->qh + c * 32; + uint8_t lo, hi; + switch (g) { + case 0: lo = ql[l] & 0xF; hi = (qh[l] >> 0) & 3; break; + case 1: lo = ql[l + 32] & 0xF; hi = (qh[l] >> 2) & 3; break; + case 2: lo = ql[l] >> 4; hi = (qh[l] >> 4) & 3; break; + default: lo = ql[l + 32] >> 4; hi = (qh[l] >> 6) & 3; break; + } + return (uint8_t) (lo | (hi << 4)); +} + +static void repack_q6_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { + repack_kquant_to_tiled_grid<block_q6_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q6_K, + [](uint8_t * tile, const block_q6_K * b, int row, int kt_local) { + const float d = GGML_FP16_TO_FP32(b->d); + uint8_t * lo_pl = tile; + uint8_t * hi_pl = tile + 512; + ggml_half * sc_pl = (ggml_half *) (tile + 768); + + for (int lk = 0; lk < 32; lk++) { + const uint8_t q6 = q6_K_get_quant(b, kt_local * 32 + lk); + const int g = lk >> 2; + const int pos = row * 4 + (lk & 3); + lo_pl[(g >> 1) * 128 + pos] |= (uint8_t) ((q6 & 0xF) << ((g & 1) * 4)); + hi_pl[(g >> 2) * 128 + pos] |= (uint8_t) ((q6 >> 4) << ((g & 3) * 2)); + } + for (int sub = 0; sub < 2; sub++) { + sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) b->scales[kt_local * 2 + sub]); + } + }); +} + +static void repack_tiled_q6_K(void * data, const ggml_tensor * t, size_t offset, size_t size) { + unpack_kquant_from_tiled_grid<block_q6_K>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q6_K, + [](block_q6_K * b, const uint8_t * sb_tiles, size_t tile_size, int row) { + float sub_scales[16]; + for (int kt_local = 0; kt_local < 8; kt_local++) { + const uint8_t * tile = sb_tiles + kt_local * tile_size; + const uint8_t * lo_pl = tile; + const uint8_t * hi_pl = tile + 512; + const ggml_half * sc_pl = (const ggml_half *) (tile + 768); + + const int c = kt_local / 4; + const int g = kt_local % 4; + uint8_t * ql = b->ql + c * 64; + uint8_t * qh = b->qh + c * 32; + + for (int lk = 0; lk < 32; lk++) { + const int g_tile = lk >> 2; + const int pos = row * 4 + (lk & 3); + const uint8_t lo = (lo_pl[(g_tile >> 1) * 128 + pos] >> ((g_tile & 1) * 4)) & 0xF; + const uint8_t hi = (hi_pl[(g_tile >> 2) * 128 + pos] >> ((g_tile & 3) * 2)) & 3; + + switch (g) { + case 0: ql[lk] |= lo; qh[lk] |= (hi << 0); break; + case 1: ql[lk + 32] |= lo; qh[lk] |= (hi << 2); break; + case 2: ql[lk] |= (lo << 4); qh[lk] |= (hi << 4); break; + default: ql[lk + 32] |= (lo << 4); qh[lk] |= (hi << 6); break; + } + } + + for (int sub = 0; sub < 2; sub++) { + sub_scales[kt_local * 2 + sub] = GGML_FP16_TO_FP32(sc_pl[sub * 32 + row]); + } + } + + float max_abs_scale = 0.0f; + for (int s = 0; s < 16; s++) { + max_abs_scale = (std::max)(max_abs_scale, fabsf(sub_scales[s])); + } + + if (max_abs_scale == 0.0f) { + b->d = GGML_FP32_TO_FP16(0.0f); + } else { + b->d = GGML_FP32_TO_FP16(max_abs_scale / 127.0f); + const float d_actual = GGML_FP16_TO_FP32(b->d); + const float inv_d = (d_actual != 0.0f) ? (1.0f / d_actual) : 0.0f; + for (int s = 0; s < 16; s++) { + const int sc = (int) roundf(sub_scales[s] * inv_d); + b->scales[s] = (int8_t) (std::max)(-128, (std::min)(127, sc)); + } + } + }); +} + +// low 2 bits (0..3) of element e of a Q2_K or Q3_K block, same bit layout as dequantize_row_q2_K / q3_K +static inline uint8_t q2_3_K_get_low2(const uint8_t * qs, int e) { + const int c = e / 128; + const int j = (e % 128) / 32; + const int l = e % 32; + return (qs[c * 32 + l] >> (2 * j)) & 3; +} + +// hmask bit of element e of a Q3_K block, same bit layout as dequantize_row_q3_K +static inline bool q3_K_get_hbit(const block_q3_K * b, int e) { + const int c = e / 128; + const int j = (e % 128) / 32; + const int l = e % 32; + return (b->hmask[l] >> (c * 4 + j)) & 1; +} + +// signed 6-bit scale j (-32..31) of a Q3_K block, same packing as quantize_row_q3_K_ref +static inline int q3_K_get_scale(const uint8_t * scales, int j) { + const int lo = (j < 8) ? (scales[j] & 0xF) : (scales[j - 8] >> 4); + const int hi = (scales[8 + j % 4] >> (2 * (j / 4))) & 3; + return (lo | (hi << 4)) - 32; +} + +// read-back: find fp16 d and l[j] in [lmin, lmax] with fp16(d * l[j]) == prod[j] for all j, false if none +static bool hexagon_recover_k_scales(const ggml_half * prod, int n, int lmin, int lmax, ggml_half * d_out, int * l_out) { + int jmax = 0; + for (int j = 1; j < n; j++) { + if (fabsf(GGML_FP16_TO_FP32(prod[j])) > fabsf(GGML_FP16_TO_FP32(prod[jmax]))) { + jmax = j; + } + } + const float pmax = GGML_FP16_TO_FP32(prod[jmax]); + if (pmax == 0.0f) { + *d_out = GGML_FP32_TO_FP16(0.0f); + for (int j = 0; j < n; j++) { + l_out[j] = 0; + } + return true; + } + + // the quantizers put the largest scale at or near the range end, so try large |l| first + const int lext = (std::max)(-lmin, lmax); + for (int a = lext; a >= 1; a--) { + for (int sign : { 1, -1 }) { + const int lj = sign * a; + if (lj < lmin || lj > lmax) { + continue; + } + const ggml_half d0 = GGML_FP32_TO_FP16(pmax / (float) lj); + for (int ulp : { 0, -1, 1 }) { + ggml_half d = d0; + uint16_t bits; + memcpy(&bits, &d, sizeof(bits)); + bits = (uint16_t) (bits + ulp); + memcpy(&d, &bits, sizeof(bits)); + + const float df = GGML_FP16_TO_FP32(d); + if (!std::isfinite(df) || df == 0.0f) { + continue; + } + bool ok = true; + for (int j = 0; j < n && ok; j++) { + const int l = (int) roundf(GGML_FP16_TO_FP32(prod[j]) / df); + const ggml_half p = GGML_FP32_TO_FP16(df * (float) l); + ok = l >= lmin && l <= lmax && memcmp(&p, &prod[j], sizeof(p)) == 0; + l_out[j] = l; + } + if (ok) { + *d_out = d; + return true; + } + } + } + } + return false; +} + +static void repack_q3_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { + repack_kquant_to_tiled_grid<block_q3_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q3_K, + [](uint8_t * tile, const block_q3_K * b, int row, int kt_local) { + const float d = GGML_FP16_TO_FP32(b->d); + uint8_t * lo_pl = tile; + uint8_t * neg_pl = tile + 256; + ggml_half * sc_pl = (ggml_half *) (tile + 384); + + for (int lk = 0; lk < 32; lk++) { + const int e = kt_local * 32 + lk; + const int g = lk >> 2; + const int pos = row * 4 + (lk & 3); + lo_pl[(g >> 2) * 128 + pos] |= (uint8_t) (q2_3_K_get_low2(b->qs, e) << ((g & 3) * 2)); + if (!q3_K_get_hbit(b, e)) { + neg_pl[pos] |= (uint8_t) (1 << g); + } + } + for (int sub = 0; sub < 2; sub++) { + sc_pl[sub * 32 + row] = GGML_FP32_TO_FP16(d * (float) q3_K_get_scale(b->scales, kt_local * 2 + sub)); + } + }); +} + +static void repack_tiled_q3_K(void * data, const ggml_tensor * t, size_t offset, size_t size) { + unpack_kquant_from_tiled_grid<block_q3_K>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q3_K, + [](block_q3_K * b, const uint8_t * sb_tiles, size_t tile_size, int row) { + ggml_half sub_scales[16]; + for (int kt_local = 0; kt_local < 8; kt_local++) { + const uint8_t * tile = sb_tiles + kt_local * tile_size; + const uint8_t * lo_pl = tile; + const uint8_t * neg_pl = tile + 256; + const ggml_half * sc_pl = (const ggml_half *) (tile + 384); + + for (int lk = 0; lk < 32; lk++) { + const int e = kt_local * 32 + lk; + const int g = lk >> 2; + const int pos = row * 4 + (lk & 3); + const uint8_t lo = (lo_pl[(g >> 2) * 128 + pos] >> ((g & 3) * 2)) & 3; + + const int c = e / 128; + const int j = (e % 128) / 32; + const int l = e % 32; + b->qs[c * 32 + l] |= (uint8_t) (lo << (2 * j)); + if (!((neg_pl[pos] >> g) & 1)) { + b->hmask[l] |= (uint8_t) (1 << (c * 4 + j)); + } + } + + for (int sub = 0; sub < 2; sub++) { + sub_scales[kt_local * 2 + sub] = sc_pl[sub * 32 + row]; + } + } + + int ls[16]; + if (!hexagon_recover_k_scales(sub_scales, 16, -32, 31, &b->d, ls)) { + float max_scale = 0.0f; + for (int s = 0; s < 16; s++) { + if (fabsf(GGML_FP16_TO_FP32(sub_scales[s])) > fabsf(max_scale)) { + max_scale = GGML_FP16_TO_FP32(sub_scales[s]); + } + } + b->d = GGML_FP32_TO_FP16(-max_scale / 32.0f); + const float d_actual = GGML_FP16_TO_FP32(b->d); + const float inv_d = (d_actual != 0.0f) ? (1.0f / d_actual) : 0.0f; + for (int s = 0; s < 16; s++) { + ls[s] = (std::max)(-32, (std::min)(31, (int) roundf(GGML_FP16_TO_FP32(sub_scales[s]) * inv_d))); + } + } + + for (int s = 0; s < 16; s++) { + const int l = ls[s] + 32; + if (s < 8) { + b->scales[s] = l & 0xF; + } else { + b->scales[s - 8] |= (uint8_t) ((l & 0xF) << 4); + } + b->scales[s % 4 + 8] |= (uint8_t) ((l >> 4) << (2 * (s / 4))); + } + }); +} + +static void repack_q2_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { + repack_kquant_to_tiled_grid<block_q2_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q2_K, + [](uint8_t * tile, const block_q2_K * b, int row, int kt_local) { + uint8_t * lo_pl = tile; + uint8_t * sc_pl = tile + 256; + ggml_half * dd_pl = (ggml_half *) (tile + 320); + + for (int lk = 0; lk < 32; lk++) { + const int g = lk >> 2; + const int pos = row * 4 + (lk & 3); + lo_pl[(g >> 2) * 128 + pos] |= (uint8_t) (q2_3_K_get_low2(b->qs, kt_local * 32 + lk) << ((g & 3) * 2)); + } + for (int sub = 0; sub < 2; sub++) { + sc_pl[sub * 32 + row] = b->scales[kt_local * 2 + sub]; + } + if (row / 8 == kt_local % 4) { + dd_pl[row % 8] = kt_local < 4 ? b->d : (ggml_half) (b->dmin ^ 0x8000); + } + }); +} + +static void repack_tiled_q2_K(void * data, const ggml_tensor * t, size_t offset, size_t size) { + unpack_kquant_from_tiled_grid<block_q2_K>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q2_K, + [](block_q2_K * b, const uint8_t * sb_tiles, size_t tile_size, int row) { + for (int kt_local = 0; kt_local < 8; kt_local++) { + const uint8_t * tile = sb_tiles + kt_local * tile_size; + const uint8_t * lo_pl = tile; + const uint8_t * sc_pl = tile + 256; + const ggml_half * dd_pl = (const ggml_half *) (tile + 320); + + for (int lk = 0; lk < 32; lk++) { + const int e = kt_local * 32 + lk; + const int g = lk >> 2; + const int pos = row * 4 + (lk & 3); + const uint8_t lo = (lo_pl[(g >> 2) * 128 + pos] >> ((g & 3) * 2)) & 3; + b->qs[(e / 128) * 32 + e % 32] |= (uint8_t) (lo << (2 * ((e % 128) / 32))); + } + + for (int sub = 0; sub < 2; sub++) { + b->scales[kt_local * 2 + sub] = sc_pl[sub * 32 + row]; + } + if (row / 8 == kt_local % 4) { + if (kt_local < 4) { + b->d = dd_pl[row % 8]; + } else { + b->dmin = (ggml_half) (dd_pl[row % 8] ^ 0x8000); + } + } + } + }); +} + +static inline void get_scale_min_k4(int j, const uint8_t * q, uint8_t * d, uint8_t * m) { + if (j < 4) { + *d = q[j] & 63; + *m = q[j + 4] & 63; + } else { + *d = (q[j + 4] & 0xF) | ((q[j - 4] >> 6) << 4); + *m = (q[j + 4] >> 4) | ((q[j - 0] >> 6) << 4); + } +} + +static inline void pack_k4_scales_mins(uint8_t * scales, ggml_half * d, ggml_half * dmin, + const float * sub_scales, const float * sub_mins) { + float max_scale = 0.0f; + float max_min = 0.0f; + for (int j = 0; j < 8; j++) { + if (sub_scales[j] > max_scale) max_scale = sub_scales[j]; + if (sub_mins[j] > max_min) max_min = sub_mins[j]; + } + + float inv_scale = 0.0f; + if (max_scale > 0.0f) { + *d = GGML_FP32_TO_FP16(max_scale / 63.0f); + const float d_actual = GGML_FP16_TO_FP32(*d); + inv_scale = (d_actual > 0.0f) ? (1.0f / d_actual) : 0.0f; + } else { + *d = GGML_FP32_TO_FP16(0.0f); + } + + float inv_min = 0.0f; + if (max_min > 0.0f) { + *dmin = GGML_FP32_TO_FP16(max_min / 63.0f); + const float dmin_actual = GGML_FP16_TO_FP32(*dmin); + inv_min = (dmin_actual > 0.0f) ? (1.0f / dmin_actual) : 0.0f; + } else { + *dmin = GGML_FP32_TO_FP16(0.0f); + } + + for (int j = 0; j < 8; j++) { + uint8_t ls = (std::min)((uint8_t) 63, (uint8_t) roundf(inv_scale * sub_scales[j])); + uint8_t lm = (std::min)((uint8_t) 63, (uint8_t) roundf(inv_min * sub_mins[j])); + if (j < 4) { + scales[j] = ls; + scales[j + 4] = lm; + } else { + scales[j + 4] = (ls & 0xF) | ((lm & 0xF) << 4); + scales[j - 4] |= ((ls >> 4) << 6); + scales[j - 0] |= ((lm >> 4) << 6); + } + } +} + +static void repack_q4_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { + repack_kquant_to_tiled_grid<block_q4_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q4_K, + [](uint8_t * tile_dst, const block_q4_K * b, int row, int kt_local) { + const float d = GGML_FP16_TO_FP32(b->d); + const float dmin = GGML_FP16_TO_FP32(b->dmin); + + uint8_t sc, m; + get_scale_min_k4(kt_local, b->scales, &sc, &m); + + const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32; + const int shift = (kt_local & 1) ? 4 : 0; + + for (int cp = 0; cp < 16; cp++) { + const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F; + const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F; + tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0); + } + + ggml_half * scale_dst = (ggml_half *) (tile_dst + 512); + scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(d * (float) sc); + scale_dst[2 * row + 1] = GGML_FP32_TO_FP16(-dmin * (float) m); + }); +} + +static void repack_tiled_q4_K(void * data, const ggml_tensor * t, size_t offset, size_t size) { + unpack_kquant_from_tiled_grid<block_q4_K>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q4_K, + [](block_q4_K * b, const uint8_t * sb_tiles, size_t tile_size, int row) { + float sub_scales[8]; + float sub_mins[8]; + + for (int kt_local = 0; kt_local < 8; kt_local++) { + const uint8_t * tile_src = sb_tiles + kt_local * tile_size; + const ggml_half * scale_src = (const ggml_half *) (tile_src + 512); + + uint8_t * qs_sub = b->qs + (kt_local / 2) * 32; + const int shift = (kt_local & 1) ? 4 : 0; + + for (int cp = 0; cp < 16; cp++) { + const uint8_t val = tile_src[cp * 32 + row]; + qs_sub[2 * cp + 0] |= (uint8_t) ((val & 0x0F) << shift); + qs_sub[2 * cp + 1] |= (uint8_t) ((val >> 4) << shift); + } + + const float D = GGML_FP16_TO_FP32(scale_src[2 * row + 0]); + const float M = GGML_FP16_TO_FP32(scale_src[2 * row + 1]); + sub_scales[kt_local] = (D > 0.0f) ? D : 0.0f; + sub_mins[kt_local] = (-M > 0.0f) ? -M : 0.0f; + } + + pack_k4_scales_mins(b->scales, &b->d, &b->dmin, sub_scales, sub_mins); + }); +} + +static void repack_q5_K_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { + repack_kquant_to_tiled_grid<block_q5_K>(t, data, offset, size, HTP_TENSOR_TILE_SIZE_Q5_K, + [](uint8_t * tile_dst, const block_q5_K * b, int row, int kt_local) { + const float d = GGML_FP16_TO_FP32(b->d); + const float dmin = GGML_FP16_TO_FP32(b->dmin); + uint8_t * plane = tile_dst + 640; + + uint8_t sc, m; + get_scale_min_k4(kt_local, b->scales, &sc, &m); + + const uint8_t * qs_sub = b->qs + (kt_local / 2) * 32; + const int shift = (kt_local & 1) ? 4 : 0; + const uint8_t hbit = (uint8_t) (1 << kt_local); + + for (int cp = 0; cp < 16; cp++) { + const uint8_t q0 = (qs_sub[2 * cp + 0] >> shift) & 0x0F; + const uint8_t q1 = (qs_sub[2 * cp + 1] >> shift) & 0x0F; + tile_dst[cp * 32 + row] = (uint8_t) ((q1 << 4) | q0); + + const int i = cp / 4; + const int lane = (cp % 4) * 32 + row; + if (b->qh[2 * cp + 0] & hbit) { + plane[lane] |= (uint8_t) (1 << (2 * i)); + } + if (b->qh[2 * cp + 1] & hbit) { + plane[lane] |= (uint8_t) (1 << (2 * i + 1)); + } + } + + ggml_half * scale_dst = (ggml_half *) (tile_dst + 512); + scale_dst[2 * row + 0] = GGML_FP32_TO_FP16(d * (float) sc); + scale_dst[2 * row + 1] = GGML_FP32_TO_FP16(-dmin * (float) m); + }); +} + +static void repack_tiled_q5_K(void * data, const ggml_tensor * t, size_t offset, size_t size) { + unpack_kquant_from_tiled_grid<block_q5_K>(data, t, offset, size, HTP_TENSOR_TILE_SIZE_Q5_K, + [](block_q5_K * b, const uint8_t * sb_tiles, size_t tile_size, int row) { + float sub_scales[8]; + float sub_mins[8]; + + for (int kt_local = 0; kt_local < 8; kt_local++) { + const uint8_t * tile_src = sb_tiles + kt_local * tile_size; + const uint8_t * plane = tile_src + 640; + const ggml_half * scale_src = (const ggml_half *) (tile_src + 512); + + uint8_t * qs_sub = b->qs + (kt_local / 2) * 32; + const int shift = (kt_local & 1) ? 4 : 0; + const uint8_t hbit = (uint8_t) (1 << kt_local); + + for (int cp = 0; cp < 16; cp++) { + const uint8_t val = tile_src[cp * 32 + row]; + qs_sub[2 * cp + 0] |= (uint8_t) ((val & 0x0F) << shift); + qs_sub[2 * cp + 1] |= (uint8_t) ((val >> 4) << shift); + + const int i = cp / 4; + const int lane = (cp % 4) * 32 + row; + if (plane[lane] & (1 << (2 * i))) { + b->qh[2 * cp + 0] |= hbit; + } + if (plane[lane] & (1 << (2 * i + 1))) { + b->qh[2 * cp + 1] |= hbit; + } + } + + const float D = GGML_FP16_TO_FP32(scale_src[2 * row + 0]); + const float M = GGML_FP16_TO_FP32(scale_src[2 * row + 1]); + sub_scales[kt_local] = (D > 0.0f) ? D : 0.0f; + sub_mins[kt_local] = (-M > 0.0f) ? -M : 0.0f; + } + + pack_k4_scales_mins(b->scales, &b->d, &b->dmin, sub_scales, sub_mins); + }); +} + +// --- Top-level Repack / Unpack Dispatchers --- + +void htp_repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t size) { + switch (tensor->type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_IQ4_NL: + repack_q4_0_tiled(tensor, data, 0, size); + break; + case GGML_TYPE_Q4_1: + repack_q4_1_tiled(tensor, data, 0, size); + break; + case GGML_TYPE_Q8_0: + repack_q8_0_tiled(tensor, data, 0, size); + break; + case GGML_TYPE_MXFP4: + repack_mxfp4_tiled(tensor, data, 0, size); + break; + case GGML_TYPE_Q2_K: + repack_q2_K_tiled(tensor, data, 0, size); + break; + case GGML_TYPE_Q3_K: + repack_q3_K_tiled(tensor, data, 0, size); + break; + case GGML_TYPE_Q4_K: + repack_q4_K_tiled(tensor, data, 0, size); + break; + case GGML_TYPE_Q5_K: + repack_q5_K_tiled(tensor, data, 0, size); + break; + case GGML_TYPE_Q6_K: + repack_q6_K_tiled(tensor, data, 0, size); + break; + case GGML_TYPE_F16: + repack_float_to_f16_tiled<uint16_t>(tensor, data, 0, size, + [](uint16_t v) { return v; }); + break; + case GGML_TYPE_F32: + repack_float_to_f16_tiled<float>(tensor, data, 0, size, + [](float v) { return (uint16_t) ggml_fp32_to_fp16(v); }); + break; + case GGML_TYPE_BF16: + repack_float_to_f16_tiled<ggml_bf16_t>(tensor, data, 0, size, + [](ggml_bf16_t v) { return (uint16_t) ggml_fp32_to_fp16(ggml_bf16_to_fp32(v)); }); + break; + default: + break; + } +} + +void htp_unpack_tensor_tiled(void * data, const ggml_tensor * tensor, size_t offset, size_t size) { + GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); + + switch (tensor->type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_IQ4_NL: + repack_tiled_q4_0(data, tensor, offset, size); + break; + case GGML_TYPE_Q4_1: + repack_tiled_q4_1(data, tensor, offset, size); + break; + case GGML_TYPE_Q8_0: + repack_tiled_q8_0(data, tensor, offset, size); + break; + case GGML_TYPE_MXFP4: + repack_tiled_mxfp4(data, tensor, offset, size); + break; + case GGML_TYPE_Q2_K: + repack_tiled_q2_K(data, tensor, offset, size); + break; + case GGML_TYPE_Q3_K: + repack_tiled_q3_K(data, tensor, offset, size); + break; + case GGML_TYPE_Q4_K: + repack_tiled_q4_K(data, tensor, offset, size); + break; + case GGML_TYPE_Q5_K: + repack_tiled_q5_K(data, tensor, offset, size); + break; + case GGML_TYPE_Q6_K: + repack_tiled_q6_K(data, tensor, offset, size); + break; + case GGML_TYPE_F16: + unpack_f16_tiled_to_float<uint16_t>(data, tensor, offset, size, + [](uint16_t v) { return v; }); + break; + case GGML_TYPE_F32: + unpack_f16_tiled_to_float<float>(data, tensor, offset, size, + [](uint16_t v) { return ggml_fp16_to_fp32((ggml_fp16_t) v); }); + break; + case GGML_TYPE_BF16: + unpack_f16_tiled_to_float<ggml_bf16_t>(data, tensor, offset, size, + [](uint16_t v) { return ggml_fp32_to_bf16(ggml_fp16_to_fp32((ggml_fp16_t) v)); }); + break; + default: + memcpy(data, (const char *) tensor->data + offset, size); + break; + } +} diff --git a/ggml/src/ggml-hexagon/htp-repack.h b/ggml/src/ggml-hexagon/htp-repack.h new file mode 100644 index 000000000000..ae66faf61fed --- /dev/null +++ b/ggml/src/ggml-hexagon/htp-repack.h @@ -0,0 +1,64 @@ +#ifndef HTP_REPACK_H +#define HTP_REPACK_H + +#include <stddef.h> +#include <stdint.h> + +#include "ggml.h" +#include "htp/htp-tensor.h" + +static inline bool htp_is_tiled_weight_type(enum ggml_type type) { + return type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q4_1 || + type == GGML_TYPE_Q8_0 || type == GGML_TYPE_IQ4_NL || + type == GGML_TYPE_MXFP4 || type == GGML_TYPE_Q6_K || + type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K || + type == GGML_TYPE_Q3_K || type == GGML_TYPE_Q2_K || + type == GGML_TYPE_F16 || type == GGML_TYPE_F32 || + type == GGML_TYPE_BF16; +} + +static inline bool htp_is_tiled_weight_tensor(const struct ggml_tensor * t) { + if (!ggml_is_contiguous(t) || ggml_is_view(t)) { + return false; + } + if (!htp_is_tiled_weight_type(t->type)) { + return false; + } + if (t->ne[0] < 32) { + return false; + } + if (!ggml_is_quantized(t->type)) { + return (t->ne[0] % 32 == 0) && (t->ne[1] % 32 == 0) && (t->ne[2] == 1) && (t->ne[3] == 1); + } + return t->ne[3] == 1; +} + +// Size of one tiled row in the DSP tiled layout. The Q3_K..Q6_K tiles store uncompressed scales/mins, +// so they are larger than the ggml blocks. For the other tiled types (Q2_K included) the tile has the same size. +static inline size_t htp_tiled_row_size(enum ggml_type type, int64_t ne0) { + if (type == GGML_TYPE_F16 || type == GGML_TYPE_F32 || type == GGML_TYPE_BF16) { + return (size_t) ne0 * sizeof(ggml_fp16_t); + } + if (type == GGML_TYPE_Q6_K) { + return (size_t) (ne0 / 32) * (HTP_TENSOR_TILE_SIZE_Q6_K / 32); + } + if (type == GGML_TYPE_Q4_K) { + return (size_t) (ne0 / 32) * (HTP_TENSOR_TILE_SIZE_Q4_K / 32); + } + if (type == GGML_TYPE_Q5_K) { + return (size_t) (ne0 / 32) * (HTP_TENSOR_TILE_SIZE_Q5_K / 32); + } + if (type == GGML_TYPE_Q3_K) { + return (size_t) (ne0 / 32) * (HTP_TENSOR_TILE_SIZE_Q3_K / 32); + } + return ggml_row_size(type, ne0); +} + +static inline bool htp_is_hmx_weight_type(enum ggml_type type) { + return htp_is_tiled_weight_type(type); +} + +void htp_repack_tensor_tiled(struct ggml_tensor * t, const void * data, size_t size); +void htp_unpack_tensor_tiled(void * data, const struct ggml_tensor * t, size_t offset, size_t size); + +#endif // HTP_REPACK_H diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index bc1474ae8320..b1131317f5e0 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -20,7 +20,7 @@ add_library(${HTP_LIB} SHARED work-queue.c dma-queue.c hmx-queue.c - htp-tensor.c + htp-dcache.c matmul-ops.c flash-attn-ops.c gated-delta-net-ops.c diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 2ee55feb9aba..4b273469521e 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -361,7 +361,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir); for (uint32_t r = 0; r < current_block_size; ++r) { __fp16 * row_k = (__fp16 *)(k_base + r * factx->size_k_row_padded); - hvx_dequantize_row_q8_0_f16(row_k, row_k, DK); + hvx_dequantize_row_flat_q8_0_f16(row_k, row_k, DK); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir); } @@ -369,7 +369,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir); for (uint32_t r = 0; r < current_block_size; ++r) { __fp16 * row_v = (__fp16 *)(v_base + r * factx->size_v_row_padded); - hvx_dequantize_row_q8_0_f16(row_v, row_v, DV); + hvx_dequantize_row_flat_q8_0_f16(row_v, row_v, DV); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir); } @@ -625,7 +625,7 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) if (factx->k->type == HTP_TYPE_Q8_0) { for (uint32_t r = start; r < end; ++r) { __fp16 * row_k = (__fp16 *)((char *)args->curr_k + r * args->src_stride * sizeof(__fp16)); - hvx_dequantize_row_q8_0_f16(row_k, row_k, factx->DK); + hvx_dequantize_row_flat_q8_0_f16(row_k, row_k, factx->DK); } } hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK_pad, @@ -679,7 +679,7 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data) if (factx->v->type == HTP_TYPE_Q8_0) { for (uint32_t r = start; r < end; ++r) { __fp16 * row_v = (__fp16 *)((char *)args->v_src + r * args->src_stride * sizeof(__fp16)); - hvx_dequantize_row_q8_0_f16(row_v, row_v, factx->DV); + hvx_dequantize_row_flat_q8_0_f16(row_v, row_v, factx->DV); } } hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV_pad, diff --git a/ggml/src/ggml-hexagon/htp/get-rows-ops.c b/ggml/src/ggml-hexagon/htp/get-rows-ops.c index 90f38fa56b51..27948c37acff 100644 --- a/ggml/src/ggml-hexagon/htp/get-rows-ops.c +++ b/ggml/src/ggml-hexagon/htp/get-rows-ops.c @@ -17,21 +17,32 @@ #include "htp-tensor.h" #include "hvx-utils.h" #include "hvx-quant.h" -#include "matmul-ops.h" #include "get-rows-ops.h" #include "work-queue.h" +struct get_rows_context; + +typedef uint32_t (*get_rows_read_index_fn_t)(const void * src1, uint32_t idx, int idx_type); +typedef void (*get_rows_flat_compute_fn_t)(float * dst_spad, const void * src_spad, uint32_t cur_elems); +typedef void (*get_rows_tiled_compute_fn_t)(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row); + struct get_rows_context { struct htp_ops_context * octx; const struct htp_get_rows_kernel_params * kparams; struct htp_get_rows_vtcm_layout vtcm_layout; - uint8_t * vtcm_base; + uint8_t * vtcm_base; + const void * src1_data; + get_rows_read_index_fn_t read_index; + get_rows_flat_compute_fn_t compute_flat; + get_rows_tiled_compute_fn_t compute_tiled; uint32_t task_start; uint32_t tasks; uint32_t tasks_per_thread; uint32_t tile_size; uint32_t tile_stride; - bool index_i32; + uint32_t src0_row_size; + uint32_t dst_row_size; + int idx_type; }; #define get_rows_preamble \ @@ -64,229 +75,169 @@ struct get_rows_context { \ const uint32_t nr = ne10 * ne11 * ne12; -#define GET_ROWS_THREAD_ST_FN(IDX_TYPE) \ -static void get_rows_thread_st_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ - struct get_rows_context * grctx = (struct get_rows_context *)data; \ - struct htp_ops_context * octx = grctx->octx; \ - const struct htp_get_rows_kernel_params * kparams = grctx->kparams; \ - get_rows_preamble; \ - const uint32_t dr = grctx->tasks_per_thread; \ - const uint32_t ir0 = grctx->task_start + dr * ith; \ - if (ir0 >= grctx->task_start + grctx->tasks) { \ - return; \ - } \ - const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks); \ - const uint32_t row_size_bytes = htp_tensor_get_row_size(octx->src[0]->type, ne00); \ - dma_queue * dma_q = octx->ctx->dma[ith]; \ - for (uint32_t i = ir0; i < ir1; ++i) { \ - const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11); \ - const uint32_t rem = i - i12 * ne11 * ne10; \ - const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ - const uint32_t i10 = rem - i11 * ne10; \ - const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \ - const uint32_t i01 = (uint32_t)*src1_ptr; \ - assert(i01 < ne01); \ - const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); \ - const uint32_t i02 = i11 - q02 * ne02; \ - const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); \ - const uint32_t i03 = i12 - q03 * ne03; \ - const dma_addr_t src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03; \ - const dma_addr_t dst_data = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3; \ - while (!dma_queue_push(dma_q, dma_make_data(dst_data, src0_data), nb1, nb01, \ - row_size_bytes, 1)) { \ - dma_queue_pop(dma_q); \ - } \ - } \ - dma_queue_flush(dma_q); \ +static uint32_t get_rows_read_index_ddr( + const void * src1, + uint32_t idx, + int idx_type +) { + const uint32_t off = idx * ((idx_type == HTP_TYPE_I64) ? sizeof(int64_t) : sizeof(int32_t)); + return *(const uint32_t *) ((const uint8_t *) src1 + off); } -GET_ROWS_THREAD_ST_FN(int32_t) -GET_ROWS_THREAD_ST_FN(int64_t) - -#define GET_ROWS_THREAD_DT_FN(TYPE_NAME, SRC0_SIZE_EXPR, IDX_TYPE, COMPUTE_EXPR) \ -static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ - struct get_rows_context * grctx = (struct get_rows_context *)data; \ - struct htp_ops_context * octx = grctx->octx; \ - const struct htp_get_rows_kernel_params * kparams = grctx->kparams; \ - get_rows_preamble; \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - const uint32_t dr = grctx->tasks_per_thread; \ - const uint32_t ir0 = grctx->task_start + dr * ith; \ - if (ir0 >= grctx->task_start + grctx->tasks) { \ - return; \ - } \ - const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks); \ - const uint32_t chunks_per_row = kparams->chunks_per_row; \ - const uint32_t chunk_size = kparams->chunk_size; \ - dma_queue * dma_q = octx->ctx->dma[ith]; \ - const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout; \ - uint8_t * vtcm_src0 = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread; \ - uint8_t * vtcm_dst = grctx->vtcm_base + vtcm_layout->off_dst + ith * vtcm_layout->dst_bytes_per_thread; \ - for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, spad_idx++) { \ - const uint32_t i = ir0 + step; \ - const uint32_t row_idx = fastdiv(i, &kparams->div_chunks_per_row); \ - const uint32_t chunk_idx = i - row_idx * chunks_per_row; \ - const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11); \ - const uint32_t rem = row_idx - i12 * ne11 * ne10; \ - const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ - const uint32_t i10 = rem - i11 * ne10; \ - const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \ - const uint32_t i01 = (uint32_t)*src1_ptr; \ - assert(i01 < ne01); \ - const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); \ - const uint32_t i02 = i11 - q02 * ne02; \ - const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); \ - const uint32_t i03 = i12 - q03 * ne03; \ - const uint32_t offset = chunk_idx * chunk_size; \ - const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0; \ - const uint32_t cur_src0_bytes = SRC0_SIZE_EXPR(cur_elems); \ - const uint32_t cur_dst_bytes = cur_elems * sizeof(float); \ - const dma_addr_t src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03 + SRC0_SIZE_EXPR(offset); \ - dma_queue_push(dma_q, \ - dma_make_data(octx->dst->data, \ - vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size), \ - cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 0); \ - dma_queue_push(dma_q, \ - dma_make_data(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size, src0_data), \ - vtcm_layout->src0_spad_half_size, cur_src0_bytes, cur_src0_bytes, 1); \ - } \ - for (uint32_t step = 0; step < ir1 - ir0; ++step) { \ - const uint32_t i = ir0 + step; \ - void * dst_spad = (void *) dma_queue_pop(dma_q).src; \ - void * src_spad = (void *) dma_queue_pop(dma_q).dst; \ - const uint32_t row_idx = fastdiv(i, &kparams->div_chunks_per_row); \ - const uint32_t chunk_idx = i - row_idx * chunks_per_row; \ - const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11); \ - const uint32_t rem = row_idx - i12 * ne11 * ne10; \ - const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ - const uint32_t i10 = rem - i11 * ne10; \ - const uint32_t offset = chunk_idx * chunk_size; \ - const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0; \ - const uint32_t cur_dst_bytes = cur_elems * sizeof(float); \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, i); \ - COMPUTE_EXPR; \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, i); \ - const dma_addr_t dst_data = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3 + offset * sizeof(float); \ - dma_queue_push(dma_q, \ - dma_make_data(dst_data, dst_spad), \ - cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 1); \ - const uint32_t next_step = step + 2; \ - if (next_step < ir1 - ir0) { \ - const uint32_t pi = ir0 + next_step; \ - const uint32_t prow_idx = fastdiv(pi, &kparams->div_chunks_per_row); \ - const uint32_t pchunk_idx = pi - prow_idx * chunks_per_row; \ - const uint32_t pi12 = fastdiv(prow_idx, &kparams->div_ne10_ne11); \ - const uint32_t prem = prow_idx - pi12 * ne11 * ne10; \ - const uint32_t pi11 = fastdiv(prem, &kparams->div_ne10); \ - const uint32_t pi10 = prem - pi11 * ne10; \ - const IDX_TYPE * psrc1_ptr = (const IDX_TYPE *)(uintptr_t)(octx->src[1]->data + pi10*nb10 + pi11*nb11 + pi12*nb12); \ - const uint32_t pi01 = (uint32_t)*psrc1_ptr; \ - assert(pi01 < ne01); \ - const uint32_t pq02 = fastdiv(pi11, &kparams->div_ne02); \ - const uint32_t pi02 = pi11 - pq02 * ne02; \ - const uint32_t pq03 = fastdiv(pi12, &kparams->div_ne03); \ - const uint32_t pi03 = pi12 - pq03 * ne03; \ - const uint32_t poffset = pchunk_idx * chunk_size; \ - const uint32_t pcur_elems = (poffset < ne00) ? MIN(chunk_size, ne00 - poffset) : 0; \ - const uint32_t pcur_src0_bytes = SRC0_SIZE_EXPR(pcur_elems); \ - const dma_addr_t psrc0_data = \ - octx->src[0]->data + pi01*nb01 + pi02*nb02 + pi03*nb03 + SRC0_SIZE_EXPR(poffset); \ - dma_queue_push(dma_q, \ - dma_make_data(src_spad, psrc0_data), \ - vtcm_layout->src0_spad_half_size, pcur_src0_bytes, pcur_src0_bytes, 1); \ - } \ - } \ - dma_queue_flush(dma_q); \ +static __attribute__((noinline)) uint32_t get_rows_read_index_vtcm( + const void * src1, + uint32_t idx, + int idx_type +) { + const uint32_t off = idx * ((idx_type == HTP_TYPE_I64) ? sizeof(int64_t) : sizeof(int32_t)); + const HVX_Vector v = hvx_vmem((const uint8_t *) src1 + off); // off & ~127u in HW + // vextract is normally slow, but here it avoids an HVX->L2 store and L1D miss + return (uint32_t) Q6_R_vextract_VR(v, off); } -#define F16_BYTES(n) ((n) * sizeof(__fp16)) -#define Q8_0_BYTES(n) (((n) / 32) * sizeof(block_q8_0)) +static void get_rows_thread_sametype(unsigned int nth, unsigned int ith, void * data) { + struct get_rows_context * grctx = (struct get_rows_context *) data; + struct htp_ops_context * octx = grctx->octx; + const struct htp_get_rows_kernel_params * kparams = grctx->kparams; + get_rows_preamble; -static __attribute__((noinline)) void compute_get_rows_f16(float * dst_spad, const void * src_spad, uint32_t cur_elems) { - hvx_dequantize_row_f16_f32(dst_spad, src_spad, cur_elems); -} + const uint32_t dr = grctx->tasks_per_thread; + const uint32_t ir0 = grctx->task_start + dr * ith; + if (ir0 >= grctx->task_start + grctx->tasks) { + return; + } + const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks); + const uint32_t row_size_bytes = grctx->src0_row_size; + dma_queue * dma_q = octx->ctx->dma[ith]; -static __attribute__((noinline)) void compute_get_rows_q8_0(float * dst_spad, const void * src_spad, uint32_t cur_elems) { - hvx_dequantize_row_q8_0_f32(dst_spad, src_spad, cur_elems); -} + for (uint32_t i = ir0; i < ir1; ++i) { + const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11); + const uint32_t rem = i - i12 * ne11 * ne10; + const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); + const uint32_t i10 = rem - i11 * ne10; -GET_ROWS_THREAD_DT_FN(f16, F16_BYTES, int32_t, { compute_get_rows_f16((float *)dst_spad, src_spad, cur_elems); }) -GET_ROWS_THREAD_DT_FN(f16, F16_BYTES, int64_t, { compute_get_rows_f16((float *)dst_spad, src_spad, cur_elems); }) - -GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int32_t, { compute_get_rows_q8_0((float *)dst_spad, src_spad, cur_elems); }) -GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int64_t, { compute_get_rows_q8_0((float *)dst_spad, src_spad, cur_elems); }) - - -static __attribute__((noinline)) void compute_get_rows_tiled(float * dst, const uint8_t * tile, uint32_t row, bool q4, bool q4_k) { - const HVX_VectorPred first2 = Q6_Q_vsetq_R(2); - const HVX_VectorPred first4 = Q6_Q_vsetq_R(4); - HVX_Vector vq = Q6_V_vzero(); - if (q4) { - const HVX_VectorPred first1 = Q6_Q_vsetq_R(1); - const HVX_VectorPred first3 = Q6_Q_vsetq_R(3); - for (int group = 3; group >= 0; --group) { - const HVX_Vector v = Q6_V_vror_VR(hvx_vmem(tile + group * VLEN), row); - // Four planes contribute bytes at 0, 32, 64 and 96 after rotation. - HVX_Vector packed = Q6_V_vmux_QVV(first1, v, Q6_V_vror_VR(v, 31)); - packed = Q6_V_vmux_QVV(first2, packed, Q6_V_vror_VR(v, 62)); - packed = Q6_V_vmux_QVV(first3, packed, Q6_V_vror_VR(v, 93)); - vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4)); - } - const HVX_Vector lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F)); - const HVX_Vector hi = Q6_Vub_vlsr_VubR(vq, 4); - vq = Q6_V_lo_W(Q6_W_vshuff_VVR(hi, lo, -1)); - if (!q4_k) { - vq = Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(8)); - } - } else { - for (int group = 7; group >= 0; --group) { - const HVX_Vector v = Q6_V_vror_VR(hvx_vmem(tile + group * VLEN), 2 * row); - // Two planes contribute halfwords at 0 and 64 after rotation. - const HVX_Vector packed = Q6_V_vmux_QVV(first2, v, Q6_V_vror_VR(v, 62)); - vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4)); + const uint32_t i01 = grctx->read_index(grctx->src1_data, i, grctx->idx_type); + assert(i01 < ne01); + + const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); + const uint32_t i02 = i11 - q02 * ne02; + const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); + const uint32_t i03 = i12 - q03 * ne03; + + const dma_addr_t src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03; + const dma_addr_t dst_data = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3; + + while (!dma_queue_push(dma_q, dma_make_data(dst_data, src0_data), nb1, nb01, row_size_bytes, 1)) { + dma_queue_pop(dma_q); } } - const HVX_Vector scales = hvx_vmem(tile + (q4 ? 512 : 1024)); - const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, (q4_k ? 4 : 2) * row)); - const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf)); - const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq); - const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16)); - HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale); - if (q4_k) { - const HVX_Vector offset_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 4 * row + 2)); - const HVX_Vector offset = Q6_V_lo_W(hvx_vec_f16_to_f32(offset_hf)); - values = hvx_vec_add_f32_f32(values, offset); + dma_queue_flush(dma_q); +} + + + +struct get_rows_flat_task { + dma_addr_t src0_data; + dma_addr_t dst_data; +}; + +static inline struct get_rows_flat_task get_rows_flat_calc_task( + const struct htp_ops_context * octx, + const struct get_rows_context * grctx, + uint32_t i, + bool need_src0 +) { + const struct htp_get_rows_kernel_params * kparams = grctx->kparams; + get_rows_preamble; + + const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11); + const uint32_t rem = i - i12 * ne11 * ne10; + const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); + const uint32_t i10 = rem - i11 * ne10; + + struct get_rows_flat_task task; + task.src0_data = 0; + task.dst_data = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3; + + if (need_src0) { + const uint32_t i01 = grctx->read_index(grctx->src1_data, i, grctx->idx_type); + assert(i01 < ne01); + + const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); + const uint32_t i02 = i11 - q02 * ne02; + const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); + const uint32_t i03 = i12 - q03 * ne03; + + task.src0_data = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03; } - *(HVX_Vector *) dst = values; + + return task; } -static __attribute__((noinline)) void compute_get_rows_q6_k(float * dst, const uint8_t * tile, uint32_t row) { - const HVX_VectorPred first4 = Q6_Q_vsetq_R(4); - const HVX_VectorPred first16 = Q6_Q_vsetq_R(16 * sizeof(float)); - const HVX_Vector mask_0f = Q6_Vb_vsplat_R(0x0F); - const HVX_Vector mask_03 = Q6_Vb_vsplat_R(0x03); - HVX_Vector vq = Q6_V_vzero(); - - for (int group = 7; group >= 0; --group) { - const HVX_Vector lo_plane = Q6_V_vror_VR(hvx_vmem(tile + (group >> 1) * VLEN), 4 * row); - const HVX_Vector hi_plane = Q6_V_vror_VR(hvx_vmem(tile + 512 + (group >> 2) * VLEN), 4 * row); - const HVX_Vector lo = (group & 1) ? Q6_Vub_vlsr_VubR(lo_plane, 4) : Q6_V_vand_VV(lo_plane, mask_0f); - const HVX_Vector hi = Q6_Vub_vlsr_VubR(hi_plane, 2 * (group & 3)); - const HVX_Vector packed = Q6_V_vor_VV(lo, Q6_Vw_vasl_VwR(Q6_V_vand_VV(hi, mask_03), 4)); - vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4)); +static void get_rows_thread_flat(unsigned int nth, unsigned int ith, void * data) { + struct get_rows_context * grctx = (struct get_rows_context *) data; + struct htp_ops_context * octx = grctx->octx; + get_rows_preamble; + + const uint32_t dr = grctx->tasks_per_thread; + const uint32_t ir0 = grctx->task_start + dr * ith; + if (ir0 >= grctx->task_start + grctx->tasks) { + return; + } + const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks); + + dma_queue * dma_q = octx->ctx->dma[ith]; + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout; + uint8_t * vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_src0) + ith * vtcm_layout->src0_bytes_per_thread; + uint8_t * vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_dst) + ith * vtcm_layout->dst_bytes_per_thread; + const uint32_t src0_bytes = grctx->src0_row_size; + const uint32_t dst_bytes = grctx->dst_row_size; + get_rows_flat_compute_fn_t compute = grctx->compute_flat; + + for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, spad_idx++) { + const uint32_t i = ir0 + step; + struct get_rows_flat_task task = get_rows_flat_calc_task(octx, grctx, i, true); + + dma_queue_push(dma_q, + dma_make_data(octx->dst->data, + vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size), + dst_bytes, vtcm_layout->dst_spad_half_size, dst_bytes, 0); + dma_queue_push(dma_q, + dma_make_data(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size, task.src0_data), + vtcm_layout->src0_spad_half_size, src0_bytes, src0_bytes, 1); + } + + for (uint32_t step = 0; step < ir1 - ir0; ++step) { + const uint32_t i = ir0 + step; + void * dst_spad = (void *) dma_queue_pop(dma_q).src; + void * src_spad = (void *) dma_queue_pop(dma_q).dst; + + struct get_rows_flat_task task = get_rows_flat_calc_task(octx, grctx, i, false); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); + compute((float *) dst_spad, src_spad, ne00); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); + + dma_queue_push(dma_q, + dma_make_data(task.dst_data, dst_spad), + dst_bytes, vtcm_layout->dst_spad_half_size, dst_bytes, 1); + + const uint32_t next_step = step + 2; + if (next_step < ir1 - ir0) { + const uint32_t pi = ir0 + next_step; + struct get_rows_flat_task next_task = get_rows_flat_calc_task(octx, grctx, pi, true); + dma_queue_push(dma_q, + dma_make_data(src_spad, next_task.src0_data), + vtcm_layout->src0_spad_half_size, src0_bytes, src0_bytes, 1); + } } - const HVX_Vector scales = hvx_vmem(tile + 768); - const HVX_Vector scale_lo_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 2 * row)); - const HVX_Vector scale_hi_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 64 + 2 * row)); - const HVX_Vector scale_lo = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_lo_hf)); - const HVX_Vector scale_hi = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hi_hf)); - const HVX_Vector scale = Q6_V_vmux_QVV(first16, scale_lo, scale_hi); - const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(32))); - const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16)); - *(HVX_Vector *) dst = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale); + dma_queue_flush(dma_q); } + + struct get_rows_tiled_task { dma_addr_t tile_src_base; dma_addr_t dst_data; @@ -307,16 +258,15 @@ static inline struct get_rows_tiled_task get_rows_tiled_calc_task( const uint32_t rem = i - i12 * ne11 * ne10; const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); const uint32_t i10 = rem - i11 * ne10; - const dma_addr_t src1_data = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12; - const uint32_t i01 = grctx->index_i32 ? *(const int32_t *)(uintptr_t) src1_data : (uint32_t) *(const int64_t *)(uintptr_t) src1_data; + const uint32_t i01 = grctx->read_index(grctx->src1_data, i, grctx->idx_type); assert(i01 < ne01); const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); const uint32_t i02 = i11 - q02 * ne02; const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); const uint32_t i03 = i12 - q03 * ne03; - const uint32_t column_tile = i01 / HTP_MM_HMX_TILE_N_ROWS; - const uint32_t row = i01 % HTP_MM_HMX_TILE_N_ROWS; + const uint32_t column_tile = i01 / HTP_TENSOR_TILE_N_ROWS; + const uint32_t row = i01 % HTP_TENSOR_TILE_N_ROWS; const dma_addr_t matrix = octx->src[0]->data + i02*nb02 + i03*nb03; struct get_rows_tiled_task task; @@ -329,7 +279,6 @@ static inline struct get_rows_tiled_task get_rows_tiled_calc_task( static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * data) { struct get_rows_context * grctx = (struct get_rows_context *) data; struct htp_ops_context * octx = grctx->octx; - const struct htp_get_rows_kernel_params * kparams = grctx->kparams; get_rows_preamble; const uint32_t dr = grctx->tasks_per_thread; @@ -339,19 +288,17 @@ static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * dat } const uint32_t ir1 = MIN(ir0 + dr, grctx->task_start + grctx->tasks); - const uint32_t n_k_tiles = ne00 / HTP_MM_HMX_TILE_N_COLS; + const uint32_t n_k_tiles = ne00 / HTP_TENSOR_TILE_N_COLS; const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout; - uint8_t * src_spad_base = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread; - uint8_t * dst_spad_base = grctx->vtcm_base + vtcm_layout->off_dst + ith * vtcm_layout->dst_bytes_per_thread; + uint8_t * src_spad_base = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_src0) + ith * vtcm_layout->src0_bytes_per_thread; + uint8_t * dst_spad_base = VTCM_LAYOUT_PTR(uint8_t, grctx->vtcm_base, vtcm_layout->off_dst) + ith * vtcm_layout->dst_bytes_per_thread; dma_queue * dma_q = octx->ctx->dma[ith]; struct htp_thread_trace * tr = &octx->ctx->trace[ith]; const uint32_t tile_size = grctx->tile_size; const uint32_t tile_stride = grctx->tile_stride; - const uint32_t dst_bytes = ne00 * sizeof(float); - const bool is_q4 = octx->src[0]->type == HTP_TYPE_Q4_0 || octx->src[0]->type == HTP_TYPE_Q4_K; - const bool is_q4_k = octx->src[0]->type == HTP_TYPE_Q4_K; - const bool is_q6_k = octx->src[0]->type == HTP_TYPE_Q6_K; + const uint32_t dst_bytes = grctx->dst_row_size; + get_rows_tiled_compute_fn_t compute_tiled = grctx->compute_tiled; for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, ++spad_idx) { const uint32_t i = ir0 + step; @@ -376,15 +323,7 @@ static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * dat struct get_rows_tiled_task task = get_rows_tiled_calc_task(octx, grctx, i, n_k_tiles, tile_size); htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); - for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) { - const uint8_t * tile = src_spad + k_tile * tile_stride; - float * dst_block = dst_spad + k_tile * HTP_MM_HMX_TILE_N_COLS; - if (is_q6_k) { - compute_get_rows_q6_k(dst_block, tile, task.row); - } else { - compute_get_rows_tiled(dst_block, tile, task.row, is_q4, is_q4_k); - } - } + compute_tiled(dst_spad, src_spad, n_k_tiles, tile_stride, task.row); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i); // Real writeback of dst_spad @@ -405,16 +344,40 @@ static void get_rows_thread_tiled(unsigned int nth, unsigned int ith, void * dat dma_queue_flush(dma_q); } +static __attribute__((noinline)) void get_rows_dma_src1( + struct htp_ops_context * octx, + uint8_t * vtcm_src1, + uint32_t elem_size +) { + const uint32_t ne10 = octx->src[1]->ne[0]; + const uint32_t ne11 = octx->src[1]->ne[1]; + const uint32_t ne12 = octx->src[1]->ne[2]; + const uint32_t ne13 = octx->src[1]->ne[3]; + const uint32_t nb11 = octx->src[1]->nb[1]; + const uint32_t nb12 = octx->src[1]->nb[2]; + const uint32_t nb13 = octx->src[1]->nb[3]; + + dma_queue * dma_q = octx->ctx->dma[0]; + uint8_t * cur_dst = vtcm_src1; + for (uint32_t i13 = 0; i13 < ne13; ++i13) { + for (uint32_t i12 = 0; i12 < ne12; ++i12) { + const dma_addr_t src_plane = octx->src[1]->data + i12 * nb12 + i13 * nb13; + while (!dma_queue_push(dma_q, dma_make_data(cur_dst, src_plane), + ne10 * elem_size, nb11, ne10 * elem_size, ne11)) { + dma_queue_pop(dma_q); + } + cur_dst += ne10 * ne11 * elem_size; + } + } + dma_queue_flush(dma_q); +} + int op_get_rows(struct htp_ops_context * octx) { const struct htp_get_rows_kernel_params * kparams = (const struct htp_get_rows_kernel_params *) octx->kernel_params; - if (octx->src[0]->type != HTP_TYPE_F32 && - octx->src[0]->type != HTP_TYPE_F16 && - octx->src[0]->type != HTP_TYPE_Q4_0 && - octx->src[0]->type != HTP_TYPE_Q4_K && - octx->src[0]->type != HTP_TYPE_Q6_K && - octx->src[0]->type != HTP_TYPE_Q8_0 && - octx->src[0]->type != HTP_TYPE_I32) { + const uint32_t ne00 = octx->src[0]->ne[0]; + const uint32_t src0_row_size = htp_tensor_get_row_size(octx->src[0]->type, ne00); + if (src0_row_size == 0) { return HTP_STATUS_NO_SUPPORT; } @@ -432,20 +395,16 @@ int op_get_rows(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - if (htp_tensor_is_extended(octx->src[1])) { - return HTP_STATUS_NO_SUPPORT; - } - const struct htp_tensor * dst = octx->dst; const uint32_t total_tasks = kparams->total_tasks; - const size_t dst_row_size = htp_tensor_get_row_size(dst->type, dst->ne[0]); + const uint32_t dst_row_size = (kparams->kernel_type == HTP_GET_ROWS_KERNEL_SAMETYPE) ? src0_row_size : ne00 * sizeof(float); uint32_t task_start = 0; uint32_t tasks = total_tasks; if (octx->ctx->mdev.count > 1) { uint32_t tasks_per_chunk = 1; - htp_tensor_mdev_rows_per_chunk(dst, dst_row_size / dst->ne[0], (uint32_t) dst_row_size, &tasks_per_chunk); + htp_tensor_mdev_rows_per_chunk(dst, dst_row_size / dst->ne[0], dst_row_size, &tasks_per_chunk); const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(total_tasks, tasks_per_chunk, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); task_start = range.start; tasks = range.count; @@ -465,15 +424,18 @@ int op_get_rows(struct htp_ops_context * octx) { grctx.octx = octx; grctx.kparams = kparams; grctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base; + grctx.compute_flat = NULL; + grctx.compute_tiled = NULL; grctx.task_start = task_start; grctx.tasks = tasks; grctx.tasks_per_thread = octx->ctx->mdev.count == 1 ? kparams->tasks_per_thread : fastdiv(tasks + n_threads - 1, &octx->n_threads_div); - grctx.tile_size = htp_mm_get_weight_tile_size(octx->src[0]->type); - grctx.tile_stride = (grctx.tile_size + 127) & ~127; - grctx.index_i32 = octx->src[1]->type == HTP_TYPE_I32; + grctx.tile_size = 0; + grctx.tile_stride = 0; + grctx.src0_row_size = src0_row_size; + grctx.dst_row_size = dst_row_size; + grctx.idx_type = octx->src[1]->type; - const uint32_t ne00 = octx->src[0]->ne[0]; - htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, kparams->kernel_type, octx->src[0]->type, ne00, n_threads); + htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, kparams->kernel_type, octx->src[0]->type, ne00, total_tasks, grctx.idx_type, n_threads); if (grctx.vtcm_layout.total_bytes > octx->ctx->vtcm_size) { FARF(ERROR, "get-rows: VTCM reservation %zu is too small, needed %zu\n", @@ -481,27 +443,56 @@ int op_get_rows(struct htp_ops_context * octx) { return HTP_STATUS_INVAL_PARAMS; } - const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); - work_queue_func_t q_func = NULL; switch (kparams->kernel_type) { case HTP_GET_ROWS_KERNEL_SAMETYPE: - q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_st_int32_t : get_rows_thread_st_int64_t); + q_func = get_rows_thread_sametype; break; case HTP_GET_ROWS_KERNEL_TILED: + switch (octx->src[0]->type) { + case HTP_TYPE_F16: grctx.compute_tiled = hvx_dequantize_row_tiled_f16_f32; break; + case HTP_TYPE_Q8_0: grctx.compute_tiled = hvx_dequantize_row_tiled_q8_0_f32; break; + case HTP_TYPE_Q4_0: grctx.compute_tiled = hvx_dequantize_row_tiled_q4_0_f32; break; + case HTP_TYPE_Q4_K: grctx.compute_tiled = hvx_dequantize_row_tiled_q4_k_f32; break; + case HTP_TYPE_Q5_K: grctx.compute_tiled = hvx_dequantize_row_tiled_q5_k_f32; break; + case HTP_TYPE_Q6_K: grctx.compute_tiled = hvx_dequantize_row_tiled_q6_k_f32; break; + default: return HTP_STATUS_NO_SUPPORT; + } + grctx.tile_size = htp_tensor_get_tile_size(octx->src[0]->type); + grctx.tile_stride = htp_tensor_get_tile_padded_size(octx->src[0]->type); q_func = get_rows_thread_tiled; break; case HTP_GET_ROWS_KERNEL_FLAT: switch (octx->src[0]->type) { - case HTP_TYPE_F16: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f16_int32_t : get_rows_thread_f16_int64_t); break; - case HTP_TYPE_Q8_0: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_q8_0_int32_t : get_rows_thread_q8_0_int64_t); break; + case HTP_TYPE_F16: grctx.compute_flat = hvx_dequantize_row_flat_f16_f32; break; + case HTP_TYPE_BF16: grctx.compute_flat = hvx_dequantize_row_flat_bf16_f32; break; + case HTP_TYPE_Q8_0: grctx.compute_flat = hvx_dequantize_row_flat_q8_0_f32; break; + case HTP_TYPE_Q4_0: grctx.compute_flat = hvx_dequantize_row_flat_q4_0_f32; break; + case HTP_TYPE_Q4_K: grctx.compute_flat = hvx_dequantize_row_flat_q4_k_f32; break; + case HTP_TYPE_Q5_K: grctx.compute_flat = hvx_dequantize_row_flat_q5_k_f32; break; + case HTP_TYPE_Q6_K: grctx.compute_flat = hvx_dequantize_row_flat_q6_k_f32; break; default: return HTP_STATUS_NO_SUPPORT; } + q_func = get_rows_thread_flat; break; default: return HTP_STATUS_NO_SUPPORT; } + const uint32_t elem_size = (grctx.idx_type == HTP_TYPE_I64) ? sizeof(int64_t) : sizeof(int32_t); + const size_t src1_bytes = (size_t) total_tasks * elem_size; + + if (!htp_tensor_is_extended(octx->src[1]) && src1_bytes <= 128 && + htp_tensor_is_contiguous(octx->src[1], elem_size)) { + grctx.src1_data = (const void *) (uintptr_t) octx->src[1]->data; + grctx.read_index = get_rows_read_index_ddr; + } else { + uint8_t * vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, grctx.vtcm_base, grctx.vtcm_layout.off_src1); + get_rows_dma_src1(octx, vtcm_src1, elem_size); + grctx.src1_data = vtcm_src1; + grctx.read_index = get_rows_read_index_vtcm; + } + FARF(HIGH, "get-rows: (%ux%ux%ux%u) x (%ux%ux%ux%u) -> (%ux%ux%ux%u) : src0-vtcm-size %zu dst-vtcm-size %zu kernel-type %d n-threads %d\n", octx->src[0]->ne[0], octx->src[0]->ne[1], octx->src[0]->ne[2], octx->src[0]->ne[3], octx->src[1]->ne[0], octx->src[1]->ne[1], octx->src[1]->ne[2], octx->src[1]->ne[3], diff --git a/ggml/src/ggml-hexagon/htp/get-rows-ops.h b/ggml/src/ggml-hexagon/htp/get-rows-ops.h index 69fa544649c3..f44e58d164f5 100644 --- a/ggml/src/ggml-hexagon/htp/get-rows-ops.h +++ b/ggml/src/ggml-hexagon/htp/get-rows-ops.h @@ -5,7 +5,8 @@ #include <string.h> #include "hex-fastdiv.h" -#include "matmul-ops.h" +#include "htp-tensor.h" +#include "htp-vtcm.h" enum htp_get_rows_kernel_type { HTP_GET_ROWS_KERNEL_SAMETYPE = 0, @@ -16,8 +17,6 @@ enum htp_get_rows_kernel_type { struct htp_get_rows_kernel_params { int32_t n_threads; int32_t kernel_type; - int32_t chunks_per_row; - int32_t chunk_size; int32_t total_tasks; int32_t tasks_per_thread; int32_t vtcm_size; @@ -25,13 +24,14 @@ struct htp_get_rows_kernel_params { // Fastdiv helpers struct fastdiv_values div_ne10; struct fastdiv_values div_ne10_ne11; - struct fastdiv_values div_chunks_per_row; struct fastdiv_values div_ne02; struct fastdiv_values div_ne03; }; struct htp_get_rows_vtcm_layout { size_t total_bytes; + size_t off_src1; + size_t src1_bytes; size_t off_src0; size_t off_dst; @@ -43,57 +43,44 @@ struct htp_get_rows_vtcm_layout { }; static inline void htp_get_rows_vtcm_layout_build( - struct htp_get_rows_vtcm_layout * vtcm_layout, + struct htp_get_rows_vtcm_layout * L, int kernel_type, int type, uint32_t ne00, + uint32_t total_tasks, + int idx_type, uint32_t n_threads) { - if (kernel_type == HTP_GET_ROWS_KERNEL_SAMETYPE) { - memset(vtcm_layout, 0, sizeof(*vtcm_layout)); - return; + memset(L, 0, sizeof(*L)); + + const size_t elem_size = (idx_type == HTP_TYPE_I64) ? sizeof(int64_t) : sizeof(int32_t); + const size_t src1_raw_bytes = total_tasks * elem_size; + const size_t src1_bytes = ((src1_raw_bytes + 127) & ~127) + 128; + + L->src1_bytes = src1_bytes; + const size_t src1_alloc_bytes = (src1_bytes + 255) & ~255; + + if (kernel_type != HTP_GET_ROWS_KERNEL_SAMETYPE) { + size_t src0_row_size = 0; + if (kernel_type == HTP_GET_ROWS_KERNEL_TILED) { + const size_t tile_stride = htp_tensor_get_tile_padded_size(type); + const uint32_t n_k_tiles = ne00 / HTP_TENSOR_TILE_N_COLS; + src0_row_size = n_k_tiles > 0 ? (n_k_tiles * tile_stride) : tile_stride; + } else { + src0_row_size = htp_tensor_get_row_size(type, ne00); + } + + L->src0_spad_half_size = (src0_row_size + 255) & ~255; + L->dst_spad_half_size = (ne00 * sizeof(float) + 255) & ~255; + L->src0_bytes_per_thread = 2 * L->src0_spad_half_size; + L->dst_bytes_per_thread = 2 * L->dst_spad_half_size; } - if (kernel_type == HTP_GET_ROWS_KERNEL_TILED) { - const size_t tile_size = htp_mm_get_weight_tile_size(type); - const size_t tile_stride = (tile_size + 127) & ~127; - const uint32_t n_k_tiles = ne00 / HTP_MM_HMX_TILE_N_COLS; - const size_t row_tiles_size = n_k_tiles > 0 ? (n_k_tiles * tile_stride) : tile_stride; - vtcm_layout->src0_spad_half_size = (row_tiles_size + 255) & ~255; - vtcm_layout->dst_spad_half_size = (ne00 * sizeof(float) + 255) & ~255; - vtcm_layout->src0_bytes_per_thread = 2 * vtcm_layout->src0_spad_half_size; - vtcm_layout->dst_bytes_per_thread = 2 * vtcm_layout->dst_spad_half_size; - vtcm_layout->off_src0 = 0; - vtcm_layout->off_dst = vtcm_layout->src0_bytes_per_thread * n_threads; - vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads; - return; - } - - uint32_t src0_row_size = 0; - switch (type) { - case HTP_TYPE_F16: - src0_row_size = ne00 * 2; - break; - case HTP_TYPE_Q8_0: - src0_row_size = (ne00 / 32) * 34; - break; - default: - src0_row_size = 0; - break; - } - - size_t src0_row_size_aligned = (src0_row_size + 255) & ~255; - size_t dst_row_size_aligned = (ne00 * sizeof(float) + 255) & ~255; - - vtcm_layout->src0_spad_half_size = src0_row_size_aligned; - vtcm_layout->dst_spad_half_size = dst_row_size_aligned; - - vtcm_layout->src0_bytes_per_thread = src0_row_size_aligned * 2; - vtcm_layout->dst_bytes_per_thread = dst_row_size_aligned * 2; - - vtcm_layout->off_src0 = 0; - vtcm_layout->off_dst = vtcm_layout->off_src0 + vtcm_layout->src0_bytes_per_thread * n_threads; - vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads; + size_t off = 0; + VTCM_LAYOUT_ALLOC(off, off_src1, src1_alloc_bytes); + VTCM_LAYOUT_ALLOC(off, off_src0, L->src0_bytes_per_thread * n_threads); + VTCM_LAYOUT_ALLOC(off, off_dst, L->dst_bytes_per_thread * n_threads); + L->total_bytes = off; } #if defined(__cplusplus) diff --git a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h index f0a0d9de19bc..309ec487779f 100644 --- a/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hmx-mm-kernels-tiled.h @@ -14,11 +14,15 @@ static const __fp16 iq4_nl_to_fp16_lut[64] __attribute__((aligned(VLEN))) = { // --- tiled format dequantizers --- -typedef struct { +typedef struct tiled_dequantize_state tiled_dequantize_state_t; +typedef void (*tiled_dequant_task_fn_t)(const tiled_dequantize_state_t * state, uint32_t start_tile, uint32_t end_tile); + +struct tiled_dequantize_state { struct htp_context * ctx; struct htp_thread_trace * traces; __fp16 * dst; const uint8_t * src; + tiled_dequant_task_fn_t task_fn; struct fastdiv_values n_k_tiles_div; uint32_t n_k_tiles; @@ -31,7 +35,7 @@ typedef struct { uint32_t k_block; size_t row_stride; uint32_t weight_type; -} tiled_dequantize_state_t; +}; // Dequantize a single tile from tiled weight data (already in VTCM) to tile-major FP16. static void dequantize_tiled_weight_to_fp16_task_q4_0( @@ -611,7 +615,7 @@ static void dequantize_tiled_weight_to_fp16_task_q5_k( } } -// Q6_K stores 6-bit weights and one fp16 scale per 16 k, see HTP_MM_WEIGHT_TILE_SIZE_Q6_K. +// Q6_K stores 6-bit weights and one fp16 scale per 16 k, see HTP_TENSOR_TILE_SIZE_Q6_K. // A k-group holds 4 k per row, the HMX tile holds 2, so each group is dealt into two tiles. static void dequantize_tiled_weight_to_fp16_task_q6_k( const tiled_dequantize_state_t *state, @@ -669,7 +673,7 @@ static void dequantize_tiled_weight_to_fp16_task_q6_k( } } -// Q3_K stores 3-bit weights and one fp16 scale per 16 k, see HTP_MM_WEIGHT_TILE_SIZE_Q3_K. +// Q3_K stores 3-bit weights and one fp16 scale per 16 k, see HTP_TENSOR_TILE_SIZE_Q3_K. static void dequantize_tiled_weight_to_fp16_task_q3_k( const tiled_dequantize_state_t *state, uint32_t start_tile, uint32_t end_tile) { @@ -701,7 +705,7 @@ static void dequantize_tiled_weight_to_fp16_task_q3_k( } } -// Q2_K stores 2-bit weights, the ggml scale bytes and d / -dmin per super-block, see HTP_MM_WEIGHT_TILE_SIZE_Q2_K. +// Q2_K stores 2-bit weights, the ggml scale bytes and d / -dmin per super-block, see HTP_TENSOR_TILE_SIZE_Q2_K. static void dequantize_tiled_weight_to_fp16_task_q2_k( const tiled_dequantize_state_t *state, uint32_t start_tile, uint32_t end_tile) { @@ -837,6 +841,56 @@ void quantize_f32_weight_to_fp16_tiles_task( } } +static __attribute__((noinline)) +void convert_bf16_weight_to_fp16_tiles_task( + const tiled_dequantize_state_t *state, + uint32_t start_tile, uint32_t end_tile) { + + const uint32_t n_k_tiles = state->n_k_tiles; + const struct fastdiv_values n_k_tiles_div = state->n_k_tiles_div; + + const HVX_Vector v_scat_base = hvx_vmem(hmx_transpose_scatter_offsets); + const HVX_Vector v_scat_step = Q6_V_vsplat_R(4); + const HVX_VectorPred q_mask64 = Q6_Q_vsetq_R(64); + + unsigned ct = fastdiv((unsigned)start_tile, &n_k_tiles_div); + unsigned kt = fastmodulo((unsigned)start_tile, n_k_tiles, &n_k_tiles_div); + + for (unsigned t = start_tile; t < (unsigned)end_tile; ) { + if (kt >= (unsigned)n_k_tiles) { kt = 0; ct++; } + + __fp16 *tile_base = state->dst + t * HTP_MM_HMX_TILE_N_ELMS; + { + uint32_t byte_off = kt * 32 * sizeof(uint16_t); + + HVX_Vector v_off = v_scat_base; + for (uint32_t r = 0; r < HTP_MM_HMX_TILE_N_ROWS; r += 2) { + uint32_t row0 = ct * HTP_MM_HMX_TILE_N_COLS + r; + uint32_t row1 = row0 + 1; + + const uint8_t *r0 = state->src + row0 * state->row_stride; + const uint8_t *r1 = state->src + row1 * state->row_stride; + + HVX_Vector v0_raw = (row0 < state->n_cols) ? hvx_vmemu((const void *)(r0 + byte_off)) : Q6_V_vzero(); + HVX_Vector v1_raw = (row1 < state->n_cols) ? hvx_vmemu((const void *)(r1 + byte_off)) : Q6_V_vzero(); + + HVX_Vector v0_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(Q6_Ww_vunpack_Vh(v0_raw)), 16); + HVX_Vector v1_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(Q6_Ww_vunpack_Vh(v1_raw)), 16); + + HVX_Vector v_out = hvx_vec_f32_to_f16(v0_f32, v1_f32); + + Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v_out); + v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step); + + HVX_Vector v_out_hi = Q6_V_vror_VR(v_out, 64); + Q6_vscatter_QRMVwV(q_mask64, (size_t)tile_base, HTP_MM_HMX_TILE_SIZE - 1, v_off, v_out_hi); + v_off = Q6_Vw_vadd_VwVw(v_off, v_scat_step); + } + } + ++t; ++kt; + } +} + // --- End tiled dequantizers --- // dot-chunk functions require external HMX lock diff --git a/ggml/src/ggml-hexagon/htp/htp-ctx.h b/ggml/src/ggml-hexagon/htp/htp-ctx.h index b36322ae81ed..22a05b9179fc 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ctx.h +++ b/ggml/src/ggml-hexagon/htp/htp-ctx.h @@ -3,6 +3,7 @@ #include "dma-queue.h" #include "hmx-queue.h" +#include "htp-dcache.h" #include "htp-ops.h" #include "hex-profile.h" #include "work-queue.h" @@ -18,8 +19,7 @@ #define HTP_MAX_NTHREADS 10 #endif -#define HTP_MAX_MMAPS 64 -#define HTP_MAX_DIRTY_RANGES 64 +#define HTP_MAX_MMAPS 64 // Memory mapping struct htp_mmap { @@ -29,11 +29,6 @@ struct htp_mmap { uint32_t flags; }; -struct htp_dirty_range { - uint32_t start; - uint32_t end; -}; - // Scratchpad state struct htp_spad { const struct htp_tensor * src; // original src of the data (for reuse) diff --git a/ggml/src/ggml-hexagon/htp/htp-tensor.c b/ggml/src/ggml-hexagon/htp/htp-dcache.c similarity index 97% rename from ggml/src/ggml-hexagon/htp/htp-tensor.c rename to ggml/src/ggml-hexagon/htp/htp-dcache.c index 03b0070d78a7..840b1edb7a7c 100644 --- a/ggml/src/ggml-hexagon/htp/htp-tensor.c +++ b/ggml/src/ggml-hexagon/htp/htp-dcache.c @@ -1,4 +1,4 @@ -#include "htp-tensor.h" +#include "htp-dcache.h" #include <qurt.h> #include <qurt_memory.h> @@ -94,7 +94,7 @@ static void merge_dirty_ranges(struct htp_context * ctx) { } } -void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) { +void htp_dcache_mark_dirty_tensors(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) { const struct htp_tensor * pending[HTP_OP_MAX_OUTPUTS]; uint32_t n_pending = 0; @@ -301,7 +301,7 @@ static void flush_dirty_ranges(struct htp_context * ctx, const struct htp_dirty_ } } -void htp_flush_dirty_ranges(struct htp_context * ctx) { +void htp_dcache_flush_dirty_ranges(struct htp_context * ctx) { struct htp_dirty_range ranges[HTP_MAX_DIRTY_RANGES]; uint32_t n_ranges = 0; uint64_t total_dirty = 0; @@ -328,7 +328,7 @@ void htp_flush_dirty_ranges(struct htp_context * ctx) { memset(ctx->dirty_ranges, 0, sizeof(ctx->dirty_ranges)); } -void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) { +void htp_dcache_flush_dirty_tensors(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n) { const struct htp_tensor * dirty_tensors[HTP_OP_MAX_INPUTS]; struct htp_dirty_range ranges[HTP_OP_MAX_INPUTS]; uint32_t n_dirty = 0; diff --git a/ggml/src/ggml-hexagon/htp/htp-dcache.h b/ggml/src/ggml-hexagon/htp/htp-dcache.h new file mode 100644 index 000000000000..1996ab258d3e --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/htp-dcache.h @@ -0,0 +1,20 @@ +#ifndef HTP_DCACHE_H +#define HTP_DCACHE_H + +#include <stdint.h> +#include "htp-ops.h" + +#define HTP_MAX_DIRTY_RANGES 64 + +struct htp_dirty_range { + uint32_t start; + uint32_t end; +}; + +struct htp_context; + +void htp_dcache_flush_dirty_ranges(struct htp_context * ctx); +void htp_dcache_flush_dirty_tensors(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n); +void htp_dcache_mark_dirty_tensors(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n); + +#endif // HTP_DCACHE_H diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h index 70d56e181a55..bb358d09ad18 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -30,21 +30,14 @@ enum htp_data_type { HTP_TYPE_IQ4_NL = 20, HTP_TYPE_I32 = 26, HTP_TYPE_I64 = 27, + HTP_TYPE_BF16 = 30, HTP_TYPE_MXFP4 = 39, - // types used internally for repack, dyn.quant, etc - HTP_TYPE_Q4_0_TILED = 200, - HTP_TYPE_Q4_1_TILED, - HTP_TYPE_Q8_0_TILED, - HTP_TYPE_MXFP4_TILED, - HTP_TYPE_INVALID }; -// Constats for internal types -#define QK_Q4_0_TILED 256 // 32x32 Q4_0 tiled layout +// Constants for internal types #define QK_Q8_0_TILED 128 // 32x32 Q8_0 tiled layout -#define QK_MXFP4_TILED 256 // 32x32 MXFP4 tiled layout @@ -141,7 +134,7 @@ enum htp_op_code { enum htp_tensor_flags { HTP_TENSOR_WEIGHT = (1U << 0), // Tensor buffer model weight data (not compute) - HTP_TENSOR_REPACK = (1U << 1), // Tensor is in repacked tiled format + HTP_TENSOR_TILED = (1U << 1), // Tensor is in tiled format HTP_TENSOR_FENCE = (1U << 2) // Tensor is synchronization fence (explicitly managed) }; diff --git a/ggml/src/ggml-hexagon/htp/htp-tensor.h b/ggml/src/ggml-hexagon/htp/htp-tensor.h index f7a96683615a..0bce68dd06a9 100644 --- a/ggml/src/ggml-hexagon/htp/htp-tensor.h +++ b/ggml/src/ggml-hexagon/htp/htp-tensor.h @@ -8,27 +8,59 @@ #include "hex-common.h" #include "hex-fastdiv.h" -enum { - HTP_TENSOR_MDEV_LINE_SIZE = 128, -}; +#define HTP_TENSOR_MDEV_LINE_SIZE 128 + +// Currently HMX requires 32x32 tile formats (do not change) +#define HTP_TENSOR_TILE_N_COLS 32 +#define HTP_TENSOR_TILE_N_ROWS 32 + +#define HTP_TENSOR_TILE_SIZE_F16 2048 +#define HTP_TENSOR_TILE_SIZE_Q4_0 576 +#define HTP_TENSOR_TILE_SIZE_Q4_1 640 +#define HTP_TENSOR_TILE_SIZE_Q8_0 1088 +#define HTP_TENSOR_TILE_SIZE_IQ4_NL 576 +#define HTP_TENSOR_TILE_SIZE_MXFP4 544 +#define HTP_TENSOR_TILE_SIZE_Q4_K 640 +#define HTP_TENSOR_TILE_SIZE_Q5_K 768 +#define HTP_TENSOR_TILE_SIZE_Q6_K 896 +#define HTP_TENSOR_TILE_SIZE_Q3_K 512 +#define HTP_TENSOR_TILE_SIZE_Q2_K 336 + +// Q2_K, Q3_K, Q6_K use vrmpy byte order within each 128B HVX vector: +// k = 0..31 is split into 8 groups g = k/4 of 4 k each; byte 4*r + (k%4) holds row r +// +// Q2_K (336B, lossless tiling with same size as block_q2_K): +// bytes 0..255 (vec 0..1): unsigned 2-bit quants, vec m holds groups 4m..4m+3 at bit offsets 0,2,4,6 +// bytes 256..319: raw scales[] bytes (min << 4 | sc): [0..31] for k 0..15, [32..63] for k 16..31 +// bytes 320..335: 8 fp16 in tile j (0..7) of super-block: d (j < 4) or -dmin (j >= 4) for rows 8*(j%4)..8*(j%4)+7 +// +// Q3_K (512B): +// bytes 0..255 (vec 0..1): low 2 bits, vec m holds groups 4m..4m+3 at bit offsets 0,2,4,6 +// bytes 256..383 (vec 2): sign mask, bit g set when hmask is 0 (quant = low2 - 4) +// bytes 384..511 (vec 3): fp16 scales d * (scales[] - 32): fp16[0..31] for k 0..15, fp16[32..63] for k 16..31 +// +// Q4_K (640B, same layout as Q4_1 tile; each 256-k super-block splits into 8 tiles of 32 k): +// bytes 0..511 (vec 0..3): 4-bit quants, byte 32*cp + r holds k = 2*cp (low nibble) and 2*cp + 1 (high nibble) +// bytes 512..639 (vec 4): 32x2 fp16 per row: (d * sc, -dmin * m) +// +// Q5_K (768B): +// bytes 0..639 (vec 0..4): Q4_K tile (low 4 bits + fp16 scale/min pairs) +// bytes 640..767 (vec 5): 5th-bit plane; byte l holds the 8 high bits for lane l across vec i = 0..3 (bit 2i = lo, 2i+1 = hi) +// +// Q6_K (896B): +// bytes 0..511 (vec 0..3): low 4 bits, vec i holds group 2i (low nibble) and 2i+1 (high nibble) +// bytes 512..767 (vec 4..5): high 2 bits, vec 4+m holds groups 4m..4m+3 at bit offsets 0,2,4,6 +// bytes 768..895 (vec 6): fp16 scales d * scales[]: fp16[0..31] for k 0..15, fp16[32..63] for k 16..31 struct htp_tensor_mdev_range { uint32_t start; uint32_t count; }; -static inline void * htp_tensor_data(const struct htp_tensor * t) { - return (void *) (uintptr_t) t->data; -} - static inline bool htp_tensor_is_extended(const struct htp_tensor * t) { return t && (t->data >> 32) != 0; } -static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) { - return (uint32_t *) &t->flags; -} - static inline bool htp_tensor_is_contiguous(const struct htp_tensor * t, uint32_t type_size) { uint32_t next_nb = type_size; if (t->ne[0] != 1 && t->nb[0] != next_nb) { @@ -125,19 +157,50 @@ static inline struct htp_tensor_mdev_range htp_tensor_mdev_partition(uint32_t to return range; } -static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) { +static inline uint32_t htp_tensor_get_tile_size(int type) { switch (type) { - case HTP_TYPE_F32: return ne00 * 4; - case HTP_TYPE_F16: return ne00 * 2; - case HTP_TYPE_Q8_0: return (ne00 / 32) * 34; - case HTP_TYPE_I32: return ne00 * 4; - default: return 0; + case HTP_TYPE_Q4_0: return HTP_TENSOR_TILE_SIZE_Q4_0; + case HTP_TYPE_IQ4_NL: return HTP_TENSOR_TILE_SIZE_IQ4_NL; + case HTP_TYPE_Q4_1: return HTP_TENSOR_TILE_SIZE_Q4_1; + case HTP_TYPE_Q4_K: return HTP_TENSOR_TILE_SIZE_Q4_K; + case HTP_TYPE_Q8_0: return HTP_TENSOR_TILE_SIZE_Q8_0; + case HTP_TYPE_Q5_K: return HTP_TENSOR_TILE_SIZE_Q5_K; + case HTP_TYPE_Q6_K: return HTP_TENSOR_TILE_SIZE_Q6_K; + case HTP_TYPE_Q3_K: return HTP_TENSOR_TILE_SIZE_Q3_K; + case HTP_TYPE_Q2_K: return HTP_TENSOR_TILE_SIZE_Q2_K; + case HTP_TYPE_MXFP4: return HTP_TENSOR_TILE_SIZE_MXFP4; + case HTP_TYPE_F16: return HTP_TENSOR_TILE_SIZE_F16; + default: return 0; } } -struct htp_context; -void htp_flush_dirty_ranges(struct htp_context * ctx); -void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n); -void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n); +static inline uint32_t htp_tensor_get_tile_padded_size(int type) { + return (htp_tensor_get_tile_size(type) + 127u) & ~127u; +} + +static inline bool htp_tensor_is_tiled(const struct htp_tensor * t) { + return t && (t->flags & HTP_TENSOR_TILED) != 0; +} + +static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) { + switch (type) { + case HTP_TYPE_F32: return ne00 * 4; + case HTP_TYPE_F16: + case HTP_TYPE_BF16: return ne00 * 2; + case HTP_TYPE_Q8_0: return (ne00 / 32) * 34; + case HTP_TYPE_Q4_0: + case HTP_TYPE_IQ4_NL: return (ne00 / 32) * 18; + case HTP_TYPE_Q4_1: return (ne00 / 32) * 20; + case HTP_TYPE_MXFP4: return (ne00 / 32) * 17; + case HTP_TYPE_Q2_K: return (ne00 / 256) * 84; + case HTP_TYPE_Q3_K: return (ne00 / 256) * 110; + case HTP_TYPE_Q4_K: return (ne00 / 256) * 144; + case HTP_TYPE_Q5_K: return (ne00 / 256) * 176; + case HTP_TYPE_Q6_K: return (ne00 / 256) * 210; + case HTP_TYPE_I32: return ne00 * 4; + case HTP_TYPE_I64: return ne00 * 8; + default: return 0; + } +} #endif // HTP_TENSOR_H diff --git a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-float.h b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-float.h index 605892aa777c..d7488976bb19 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-float.h +++ b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-float.h @@ -76,7 +76,7 @@ static inline void quantize_f16_f16_kernel( #define HVX_OP_MUL_F32(a, b) Q6_Vsf_vmpy_VsfVsf(a, b) #endif -static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { +static inline void flat_vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { const HVX_Vector * restrict x = (const HVX_Vector *) vx; const HVX_Vector * restrict y = (const HVX_Vector *) vy; @@ -104,7 +104,7 @@ static inline void vec_dot_f32_f32_aa_1x1(const uint32_t n, float * restrict s, *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum)); } -static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0, +static inline void flat_vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0, const void * restrict vx0, const void * restrict vx1, const void * restrict vy0) { const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; @@ -143,7 +143,7 @@ static inline void vec_dot_f32_f32_aa_2x1(const uint32_t n, float * restrict s0, hvx_vec_store_u(s0, 8, rsum); } -static inline void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1, +static inline void flat_vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx0, const void * restrict vx1, const void * restrict vy0, const void * restrict vy1) { const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; @@ -196,10 +196,203 @@ static inline void vec_dot_f32_f32_aa_2x2(const uint32_t n, float * restrict s0, hvx_vec_store_u(s1, 8, r0_r1_c1_sum); } +static inline void flat_vec_dot_bf16_f32_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { + const HVX_Vector * restrict x = (const HVX_Vector *) vx; + const HVX_Vector * restrict y = (const HVX_Vector *) vy; + + uint32_t nvec = n / 64; + uint32_t rem = n % 64; + + HVX_Vector rsum0 = Q6_V_vzero(); + HVX_Vector rsum1 = Q6_V_vzero(); + + uint32_t i = 0; + #pragma unroll(2) + for (i = 0; i < nvec; i++) { + HVX_Vector v_bf16 = x[i]; + HVX_VectorPair p = Q6_Ww_vunpack_Vh(v_bf16); + HVX_Vector w0_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p), 16); + HVX_Vector w1_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p), 16); + + rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w0_f32, y[2 * i + 0])); + rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w1_f32, y[2 * i + 1])); + } + + if (rem) { + HVX_VectorPred bmask = Q6_Q_vsetq2_R(rem * 2); + HVX_Vector v_bf16 = Q6_V_vand_QV(bmask, x[i]); + HVX_VectorPair p = Q6_Ww_vunpack_Vh(v_bf16); + HVX_Vector w0_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p), 16); + HVX_Vector w1_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p), 16); + + if (rem <= 32) { + HVX_VectorPred a_mask = Q6_Q_vsetq2_R(rem * 4); + HVX_Vector y0 = Q6_V_vand_QV(a_mask, y[2 * i + 0]); + rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w0_f32, y0)); + } else { + HVX_Vector y0 = y[2 * i + 0]; + rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w0_f32, y0)); + + HVX_VectorPred a_mask = Q6_Q_vsetq2_R((rem - 32) * 4); + HVX_Vector y1 = Q6_V_vand_QV(a_mask, y[2 * i + 1]); + rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w1_f32, y1)); + } + } + + HVX_Vector rsum = HVX_OP_ADD_F32(rsum0, rsum1); + *s = hvx_vec_get_f32(hvx_vec_reduce_sum_f32(rsum)); +} + +static inline void flat_vec_dot_bf16_f32_aa_2x1(const uint32_t n, float * restrict s0, + const void * restrict vx0, const void * restrict vx1, + const void * restrict vy0) { + const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; + const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; + const HVX_Vector * restrict y = (const HVX_Vector *) vy0; + + uint32_t nvec = n / 64; + uint32_t rem = n % 64; + + HVX_Vector rsum0 = Q6_V_vzero(); + HVX_Vector rsum1 = Q6_V_vzero(); + + uint32_t i = 0; + #pragma unroll(2) + for (i = 0; i < nvec; i++) { + HVX_VectorPair p0 = Q6_Ww_vunpack_Vh(x0[i]); + HVX_VectorPair p1 = Q6_Ww_vunpack_Vh(x1[i]); + HVX_Vector w00_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p0), 16); + HVX_Vector w01_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p0), 16); + HVX_Vector w10_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p1), 16); + HVX_Vector w11_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p1), 16); + + HVX_Vector y0 = y[2 * i + 0]; + HVX_Vector y1 = y[2 * i + 1]; + + rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w00_f32, y0)); + rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w10_f32, y0)); + rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w01_f32, y1)); + rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w11_f32, y1)); + } + + if (rem) { + HVX_VectorPred bmask = Q6_Q_vsetq2_R(rem * 2); + HVX_VectorPair p0 = Q6_Ww_vunpack_Vh(Q6_V_vand_QV(bmask, x0[i])); + HVX_VectorPair p1 = Q6_Ww_vunpack_Vh(Q6_V_vand_QV(bmask, x1[i])); + HVX_Vector w00_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p0), 16); + HVX_Vector w01_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p0), 16); + HVX_Vector w10_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p1), 16); + HVX_Vector w11_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p1), 16); + + if (rem <= 32) { + HVX_VectorPred a_mask = Q6_Q_vsetq2_R(rem * 4); + HVX_Vector y0 = Q6_V_vand_QV(a_mask, y[2 * i + 0]); + rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w00_f32, y0)); + rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w10_f32, y0)); + } else { + HVX_Vector y0 = y[2 * i + 0]; + rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w00_f32, y0)); + rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w10_f32, y0)); + + HVX_VectorPred a_mask = Q6_Q_vsetq2_R((rem - 32) * 4); + HVX_Vector y1 = Q6_V_vand_QV(a_mask, y[2 * i + 1]); + rsum0 = HVX_OP_ADD_F32(rsum0, HVX_OP_MUL_F32(w01_f32, y1)); + rsum1 = HVX_OP_ADD_F32(rsum1, HVX_OP_MUL_F32(w11_f32, y1)); + } + } + + HVX_Vector rsum = hvx_vec_reduce_sum_f32x2(rsum0, rsum1); + hvx_vec_store_u(s0, 8, rsum); +} + +static inline void flat_vec_dot_bf16_f32_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1, + const void * restrict vx0, const void * restrict vx1, + const void * restrict vy0, const void * restrict vy1) { + const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; + const HVX_Vector * restrict x1 = (const HVX_Vector *) vx1; + const HVX_Vector * restrict y0 = (const HVX_Vector *) vy0; + const HVX_Vector * restrict y1 = (const HVX_Vector *) vy1; + + uint32_t nvec = n / 64; + uint32_t rem = n % 64; + + HVX_Vector r0_c0_sum = Q6_V_vzero(); + HVX_Vector r0_c1_sum = Q6_V_vzero(); + HVX_Vector r1_c0_sum = Q6_V_vzero(); + HVX_Vector r1_c1_sum = Q6_V_vzero(); + + uint32_t i = 0; + #pragma unroll(2) + for (i = 0; i < nvec; i++) { + HVX_VectorPair p0 = Q6_Ww_vunpack_Vh(x0[i]); + HVX_VectorPair p1 = Q6_Ww_vunpack_Vh(x1[i]); + HVX_Vector w00_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p0), 16); + HVX_Vector w01_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p0), 16); + HVX_Vector w10_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p1), 16); + HVX_Vector w11_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p1), 16); + + HVX_Vector c0_0 = y0[2 * i + 0]; + HVX_Vector c0_1 = y0[2 * i + 1]; + HVX_Vector c1_0 = y1[2 * i + 0]; + HVX_Vector c1_1 = y1[2 * i + 1]; + + r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(w00_f32, c0_0)); + r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(w00_f32, c1_0)); + r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(w10_f32, c0_0)); + r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(w10_f32, c1_0)); + + r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(w01_f32, c0_1)); + r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(w01_f32, c1_1)); + r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(w11_f32, c0_1)); + r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(w11_f32, c1_1)); + } + + if (rem) { + HVX_VectorPred bmask = Q6_Q_vsetq2_R(rem * 2); + HVX_VectorPair p0 = Q6_Ww_vunpack_Vh(Q6_V_vand_QV(bmask, x0[i])); + HVX_VectorPair p1 = Q6_Ww_vunpack_Vh(Q6_V_vand_QV(bmask, x1[i])); + HVX_Vector w00_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p0), 16); + HVX_Vector w01_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p0), 16); + HVX_Vector w10_f32 = Q6_Vw_vasl_VwR(Q6_V_lo_W(p1), 16); + HVX_Vector w11_f32 = Q6_Vw_vasl_VwR(Q6_V_hi_W(p1), 16); + + if (rem <= 32) { + HVX_VectorPred a_mask = Q6_Q_vsetq2_R(rem * 4); + HVX_Vector c0_0 = Q6_V_vand_QV(a_mask, y0[2 * i + 0]); + HVX_Vector c1_0 = Q6_V_vand_QV(a_mask, y1[2 * i + 0]); + r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(w00_f32, c0_0)); + r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(w00_f32, c1_0)); + r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(w10_f32, c0_0)); + r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(w10_f32, c1_0)); + } else { + HVX_Vector c0_0 = y0[2 * i + 0]; + HVX_Vector c1_0 = y1[2 * i + 0]; + r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(w00_f32, c0_0)); + r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(w00_f32, c1_0)); + r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(w10_f32, c0_0)); + r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(w10_f32, c1_0)); + + HVX_VectorPred a_mask = Q6_Q_vsetq2_R((rem - 32) * 4); + HVX_Vector c0_1 = Q6_V_vand_QV(a_mask, y0[2 * i + 1]); + HVX_Vector c1_1 = Q6_V_vand_QV(a_mask, y1[2 * i + 1]); + r0_c0_sum = HVX_OP_ADD_F32(r0_c0_sum, HVX_OP_MUL_F32(w01_f32, c0_1)); + r0_c1_sum = HVX_OP_ADD_F32(r0_c1_sum, HVX_OP_MUL_F32(w01_f32, c1_1)); + r1_c0_sum = HVX_OP_ADD_F32(r1_c0_sum, HVX_OP_MUL_F32(w11_f32, c0_1)); + r1_c1_sum = HVX_OP_ADD_F32(r1_c1_sum, HVX_OP_MUL_F32(w11_f32, c1_1)); + } + } + + HVX_Vector r0_r1_c0_sum = hvx_vec_reduce_sum_f32x2(r0_c0_sum, r1_c0_sum); + HVX_Vector r0_r1_c1_sum = hvx_vec_reduce_sum_f32x2(r0_c1_sum, r1_c1_sum); + + hvx_vec_store_u(s0, 8, r0_r1_c0_sum); + hvx_vec_store_u(s1, 8, r0_r1_c1_sum); +} + #undef HVX_OP_ADD_F32 #undef HVX_OP_MUL_F32 -static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { +static inline void flat_vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy) { const HVX_Vector * restrict x = (const HVX_Vector *) vx; const HVX_Vector * restrict y = (const HVX_Vector *) vy; @@ -226,7 +419,7 @@ static inline void vec_dot_f16_f16_aa_1x1(const uint32_t n, float * restrict s, hvx_vec_store_u(s, 4, hvx_vec_reduce_sum_f32(rsum)); } -static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0, +static inline void flat_vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0, const void * restrict vx0, const void * restrict vx1, const void * restrict vy0) { const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; @@ -263,7 +456,7 @@ static inline void vec_dot_f16_f16_aa_2x1(const uint32_t n, float * restrict s0, hvx_vec_store_u(s0, 8, rsum); } -static inline void vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1, +static inline void flat_vec_dot_f16_f16_aa_2x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx0, const void * restrict vx1, const void * restrict vy0, const void * restrict vy1) { const HVX_Vector * restrict x0 = (const HVX_Vector *) vx0; diff --git a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h index 0c86e407efe7..91e5117255f4 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h +++ b/ggml/src/ggml-hexagon/htp/hvx-mm-kernels-tiled.h @@ -274,6 +274,72 @@ static void quantize_row_f32_q8_1_s16_tiled(float * restrict x, uint8_t * restri } } +static inline void quantize_half_f16_tiled_splat(HVX_Vector v_src, HVX_Vector * restrict dst) { + static const uint8_t __attribute__((aligned(128))) repl[128] = { + 0x00, 0x00, 0x00, 0x00, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, + 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, + 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, + 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, + 0x40, 0x40, 0x40, 0x40, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, + 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, + 0x20, 0x20, 0x20, 0x20, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, + 0x10, 0x10, 0x10, 0x10, 0x04, 0x04, 0x04, 0x04, 0x08, 0x08, 0x08, 0x08, 0x04, 0x04, 0x04, 0x04, + }; + HVX_Vector v_repl_ctrl = *(const HVX_Vector *) repl; + + #pragma unroll(4) + for (int cp = 0; cp < 32; cp++) { + dst[cp] = Q6_V_vdelta_VV(cp == 0 ? v_src : Q6_V_vror_VR(v_src, cp * 4), v_repl_ctrl); + } +} + +static inline void quantize_block_f32_f16_tiled(const float * restrict x, uint8_t * restrict y_block) { + assert((unsigned long) x % 128 == 0); + assert((unsigned long) y_block % 128 == 0); + + const HVX_Vector * vx = (const HVX_Vector *) x; + HVX_Vector * restrict dst = (HVX_Vector *) y_block; + + HVX_Vector vx01_hf = hvx_vec_f32_to_f16(vx[0], vx[1]); + quantize_half_f16_tiled_splat(vx01_hf, dst); + + HVX_Vector vx23_hf = hvx_vec_f32_to_f16(vx[2], vx[3]); + quantize_half_f16_tiled_splat(vx23_hf, dst + 32); +} + +static inline void quantize_block_f16_f16_tiled(const __fp16 * restrict x, uint8_t * restrict y_block) { + assert((unsigned long) x % 128 == 0); + assert((unsigned long) y_block % 128 == 0); + + const HVX_Vector * vx = (const HVX_Vector *) x; + HVX_Vector * restrict dst = (HVX_Vector *) y_block; + + quantize_half_f16_tiled_splat(vx[0], dst); + quantize_half_f16_tiled_splat(vx[1], dst + 32); +} + +static void quantize_row_f32_f16_tiled(const float * restrict x, uint8_t * restrict y, uint32_t k) { + assert(k % 32 == 0); + const uint32_t qk = QK_Q8_0_TILED; + const uint32_t nb = (k + qk - 1) / qk; + + for (uint32_t i = 0; i < nb; i++) { + uint8_t * restrict y_block = y + i * 4 * 2048; + quantize_block_f32_f16_tiled(x + i * qk, y_block); + } +} + +static void quantize_row_f16_f16_tiled(const __fp16 * restrict x, uint8_t * restrict y, uint32_t k) { + assert(k % 32 == 0); + const uint32_t qk = QK_Q8_0_TILED; + const uint32_t nb = (k + qk - 1) / qk; + + for (uint32_t i = 0; i < nb; i++) { + uint8_t * restrict y_block = y + i * 4 * 2048; + quantize_block_f16_f16_tiled(x + i * qk, y_block); + } +} + // Dot kernels & helpers that consume tiled activations static inline HVX_Vector hvx_vec_mul_f16_f16_to_f32_lower32(HVX_Vector v1, HVX_Vector v2) { @@ -454,7 +520,7 @@ static inline HVX_VectorPair accum_q8_0_32x2( return Q6_W_vcombine_VV(v_sum1, v_sum0); } -// Q5_K: OR 0x10 into every lane of v whose flag j is set in the plane (see HTP_MM_WEIGHT_TILE_SIZE_Q5_K) +// Q5_K: OR 0x10 into every lane of v whose flag j is set in the plane (see HTP_TENSOR_TILE_SIZE_Q5_K) static inline HVX_Vector hvx_q5k_or_hibit(HVX_Vector v, HVX_Vector v_plane, int j) { HVX_VectorPred q = Q6_Q_vand_VR(v_plane, 0x01010101u << j); return Q6_V_vandor_VQR(v, q, 0x10101010); @@ -526,7 +592,7 @@ static inline HVX_VectorPair accum_5bit_32x2( return Q6_W_vcombine_VV(v_sum1, v_sum0); } -// Q6_K weights are stored unsigned (0..63), see HTP_MM_WEIGHT_TILE_SIZE_Q6_K. Unpack k-group g of a tile to signed bytes (q - 32) +// Q6_K weights are stored unsigned (0..63), see HTP_TENSOR_TILE_SIZE_Q6_K. Unpack k-group g of a tile to signed bytes (q - 32) static inline HVX_Vector unpack_q6_k_group(const HVX_Vector * restrict vptr, int g, HVX_Vector mask_0f, HVX_Vector mask_03, HVX_Vector i32) { HVX_Vector v_lo = (g & 1) ? Q6_Vub_vlsr_VubR(vptr[g >> 1], 4) : Q6_V_vand_VV(vptr[g >> 1], mask_0f); HVX_Vector v_hi = (g & 3) ? Q6_Vub_vlsr_VubR(vptr[4 + (g >> 2)], 2 * (g & 3)) : vptr[4 + (g >> 2)]; @@ -593,7 +659,7 @@ static inline HVX_Vector scale_q6_k_32x1(HVX_VectorPair v_sums, HVX_Vector v_sca return hvx_vec_add_f32_f32(v_lo, v_hi); } -// Q3_K / Q2_K: low 2 bits of k-group g, see HTP_MM_WEIGHT_TILE_SIZE_Q3_K +// Q3_K / Q2_K: low 2 bits of k-group g, see HTP_TENSOR_TILE_SIZE_Q3_K static inline HVX_Vector unpack_q3_k_low2(const HVX_Vector * restrict vptr, int g, HVX_Vector mask_03) { HVX_Vector v = vptr[g >> 2]; if ((g & 3) == 3) { @@ -1482,6 +1548,59 @@ static void tiled_vec_dot_mxfp4_32x2(const uint32_t n, float * restrict s0, floa } } +static void tiled_vec_dot_f16_32x1(const uint32_t n, float * restrict s, const void * restrict vx, const void * restrict vy, uint32_t valid_rows, const float * restrict sz) { + const HVX_Vector * restrict vptr = (const HVX_Vector *) vx; + const HVX_Vector * restrict v_act = (const HVX_Vector *) vy; + const uint32_t n_vecs = ((n + 31) / 32) * 16; + + HVX_VectorPair acc = Q6_W_vzero(); + + #pragma unroll(4) + for (uint32_t i = 0; i < n_vecs; i++) { + acc = hvx_vec_mpyacc_f32_f16(acc, vptr[i], v_act[i]); + } + + HVX_Vector v_sum = hvx_vec_add_f32_f32(Q6_V_lo_W(acc), Q6_V_hi_W(acc)); + + if (sz) { + hvx_vec_store_u(s, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum, hvx_vmemu(sz))); + } else { + hvx_vec_store_u(s, valid_rows * sizeof(float), v_sum); + } +} + +static void tiled_vec_dot_f16_32x2(const uint32_t n, float * restrict s0, float * restrict s1, const void * restrict vx, const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, const float * restrict sz0, const float * restrict sz1) { + const HVX_Vector * restrict vptr = (const HVX_Vector *) vx; + const HVX_Vector * restrict v_act0 = (const HVX_Vector *) vy0; + const HVX_Vector * restrict v_act1 = (const HVX_Vector *) vy1; + const uint32_t n_vecs = ((n + 31) / 32) * 16; + + HVX_VectorPair acc0 = Q6_W_vzero(); + HVX_VectorPair acc1 = Q6_W_vzero(); + + #pragma unroll(4) + for (uint32_t i = 0; i < n_vecs; i++) { + HVX_Vector v_w = vptr[i]; + acc0 = hvx_vec_mpyacc_f32_f16(acc0, v_w, v_act0[i]); + acc1 = hvx_vec_mpyacc_f32_f16(acc1, v_w, v_act1[i]); + } + + HVX_Vector v_sum0 = hvx_vec_add_f32_f32(Q6_V_lo_W(acc0), Q6_V_hi_W(acc0)); + HVX_Vector v_sum1 = hvx_vec_add_f32_f32(Q6_V_lo_W(acc1), Q6_V_hi_W(acc1)); + + if (sz0) { + hvx_vec_store_u(s0, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum0, hvx_vmemu(sz0))); + } else { + hvx_vec_store_u(s0, valid_rows * sizeof(float), v_sum0); + } + if (sz1) { + hvx_vec_store_u(s1, valid_rows * sizeof(float), hvx_vec_add_f32_f32(v_sum1, hvx_vmemu(sz1))); + } else { + hvx_vec_store_u(s1, valid_rows * sizeof(float), v_sum1); + } +} + + static inline void quantize_f32_q8_0_tiled_kernel( const uint8_t * restrict src_data, uint8_t * restrict dst_data, @@ -1533,6 +1652,40 @@ static inline void quantize_f32_q8_1_s16_tiled_kernel( } } +static inline void quantize_f32_f16_tiled_kernel( + const uint8_t * restrict src_data, + uint8_t * restrict dst_data, + uint8_t * restrict tmp_data, + uint32_t ne0, + uint32_t nrows, + size_t src_row_size, + size_t dst_row_size +) { + (void) tmp_data; + for (uint32_t i = 0; i < nrows; ++i) { + quantize_row_f32_f16_tiled((const float *) src_data, dst_data, ne0); + dst_data += dst_row_size; + src_data += src_row_size; + } +} + +static inline void quantize_f16_f16_tiled_kernel( + const uint8_t * restrict src_data, + uint8_t * restrict dst_data, + uint8_t * restrict tmp_data, + uint32_t ne0, + uint32_t nrows, + size_t src_row_size, + size_t dst_row_size +) { + (void) tmp_data; + for (uint32_t i = 0; i < nrows; ++i) { + quantize_row_f16_f16_tiled((const __fp16 *) src_data, dst_data, ne0); + dst_data += dst_row_size; + src_data += src_row_size; + } +} + static inline void quantize_f32_q8_0_tiled_block_kernel( const float * restrict src, uint8_t * restrict dst, @@ -1622,3 +1775,63 @@ static inline void quantize_f32_q8_1_s16_tiled_block_kernel( } } } + +static inline void quantize_f32_f16_tiled_block_kernel( + const float * restrict src, + uint8_t * restrict dst, + uint8_t * restrict tmp_data, + uint32_t ne0, + uint32_t ib_first, + uint32_t ib_last, + size_t src_row_size, + size_t dst_row_size, + uint32_t r, + uint32_t c +) { + (void) tmp_data; + const uint32_t qk = QK_Q8_0_TILED; + const uint32_t nb = (ne0 + qk - 1) / qk; + + for (uint32_t ib = ib_first; ib < ib_last; ++ib) { + const float * restrict src_ptr = (const float *) ((const uint8_t *) src + r * src_row_size + c * qk * sizeof(float)); + uint8_t * restrict dst_ptr = dst + r * dst_row_size + c * 4 * 2048; + + quantize_block_f32_f16_tiled(src_ptr, dst_ptr); + + c++; + if (c == nb) { + c = 0; + r++; + } + } +} + +static inline void quantize_f16_f16_tiled_block_kernel( + const __fp16 * restrict src, + uint8_t * restrict dst, + uint8_t * restrict tmp_data, + uint32_t ne0, + uint32_t ib_first, + uint32_t ib_last, + size_t src_row_size, + size_t dst_row_size, + uint32_t r, + uint32_t c +) { + (void) tmp_data; + const uint32_t qk = QK_Q8_0_TILED; + const uint32_t nb = (ne0 + qk - 1) / qk; + + for (uint32_t ib = ib_first; ib < ib_last; ++ib) { + const __fp16 * restrict src_ptr = (const __fp16 *) ((const uint8_t *) src + r * src_row_size + c * qk * sizeof(__fp16)); + uint8_t * restrict dst_ptr = dst + r * dst_row_size + c * 4 * 2048; + + quantize_block_f16_f16_tiled(src_ptr, dst_ptr); + + c++; + if (c == nb) { + c = 0; + r++; + } + } +} diff --git a/ggml/src/ggml-hexagon/htp/hvx-quant.h b/ggml/src/ggml-hexagon/htp/hvx-quant.h index 6b172cd63c3a..b7837967a828 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-quant.h +++ b/ggml/src/ggml-hexagon/htp/hvx-quant.h @@ -5,6 +5,7 @@ #include <stdint.h> #include <string.h> +#include "htp-tensor.h" #include "hvx-arith.h" #include "hvx-base.h" #include "hvx-reduce.h" @@ -96,7 +97,7 @@ static inline void hvx_quantize_row_q8_0_f32(void * restrict dst_ptr, const floa } } -static inline void hvx_dequantize_row_q8_0_f32(float * restrict dst_ptr, const void * restrict src_ptr, int n) { +static inline void hvx_dequantize_row_flat_q8_0_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) { const int nb = n / QK8_0; const block_q8_0 * src = (const block_q8_0 *) src_ptr; @@ -120,7 +121,7 @@ static inline void hvx_dequantize_row_q8_0_f32(float * restrict dst_ptr, const v } } -static inline void hvx_dequantize_row_q8_0_f16(__fp16 * restrict dst_ptr, const void * restrict src_ptr, int n) { +static inline void hvx_dequantize_row_flat_q8_0_f16(__fp16 * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) { const int nb = n / QK8_0; const block_q8_0 * src = (const block_q8_0 *) src_ptr; @@ -146,20 +147,460 @@ static inline void hvx_dequantize_row_q8_0_f16(__fp16 * restrict dst_ptr, const } } -static inline void hvx_dequantize_row_f16_f32(float * restrict dst_ptr, const void * restrict src_ptr, int n) { - const int nb = n / 32; +static inline void hvx_dequantize_row_flat_f16_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) { + const uint32_t nvec = n / VLEN_FP32; + const uint32_t nloe = n % VLEN_FP32; const _Float16 * src = (const _Float16 *) src_ptr; - for (int i = 0; i < nb; i++) { - HVX_Vector v_f16 = *(const HVX_UVector *) (src + i * 32); + for (uint32_t i = 0; i < nvec; i++) { + HVX_Vector v_f16 = *(const HVX_UVector *) (src + i * VLEN_FP32); HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(v_f16); HVX_Vector res = Q6_V_lo_W(vp_f32); - float * block_dst = dst_ptr + i * 32; - hvx_vmem(block_dst) = res; + hvx_vmem(dst_ptr + i * VLEN_FP32) = res; + } + + if (nloe) { + HVX_Vector v_f16 = *(const HVX_UVector *) (src + nvec * VLEN_FP32); + HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(v_f16); + HVX_Vector res = Q6_V_lo_W(vp_f32); + + hvx_vec_store_a(dst_ptr + nvec * VLEN_FP32, nloe * sizeof(float), res); + } +} + +static inline void hvx_dequantize_row_flat_bf16_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) { + const uint32_t nvec = n / VLEN_FP32; + const uint32_t nloe = n % VLEN_FP32; + const uint16_t * src = (const uint16_t *) src_ptr; + + for (uint32_t i = 0; i < nvec; i++) { + HVX_Vector v_bf16 = *(const HVX_UVector *) (src + i * VLEN_FP32); + HVX_VectorPair p = Q6_Ww_vunpack_Vh(v_bf16); + HVX_Vector res = Q6_Vw_vasl_VwR(Q6_V_lo_W(p), 16); + + hvx_vmem(dst_ptr + i * VLEN_FP32) = res; + } + + if (nloe) { + HVX_Vector v_bf16 = *(const HVX_UVector *) (src + nvec * VLEN_FP32); + HVX_VectorPair p = Q6_Ww_vunpack_Vh(v_bf16); + HVX_Vector res = Q6_Vw_vasl_VwR(Q6_V_lo_W(p), 16); + + hvx_vec_store_a(dst_ptr + nvec * VLEN_FP32, nloe * sizeof(float), res); + } +} + +static inline void hvx_dequantize_row_flat_q4_0_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) { + const int nb = n / QK4_0; + const block_q4_0 * src = (const block_q4_0 *) src_ptr; + + const HVX_VectorPred q16 = Q6_Q_vsetq_R(16); + const HVX_Vector v_eight = Q6_Vb_vsplat_R(8); + const HVX_Vector v_zero = Q6_V_vzero(); + + for (int i = 0; i < nb; i++) { + HVX_Vector vd_f16 = Q6_Vh_vsplat_R(*(const int16_t *) &src[i].d); + HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(vd_f16); + HVX_Vector vd = Q6_V_lo_W(vp_f32); + + HVX_Vector vq = *(const HVX_UVector *) src[i].qs; + HVX_Vector v_lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F)); + HVX_Vector v_hi = Q6_Vub_vlsr_VubR(vq, 4); + + HVX_Vector v_lo_s8 = Q6_Vb_vsub_VbVb(v_lo, v_eight); + HVX_Vector v_hi_s8 = Q6_Vb_vsub_VbVb(v_hi, v_eight); + + HVX_Vector v32 = Q6_V_vmux_QVV(q16, v_lo_s8, Q6_V_vlalign_VVR(v_hi_s8, v_zero, 16)); + + HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(v32); + HVX_Vector v_i16 = Q6_V_lo_W(p16); + HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(v_i16); + HVX_Vector v_i32 = Q6_V_lo_W(p32); + + HVX_Vector v_f32 = Q6_Vsf_equals_Vw(v_i32); + HVX_Vector res = hvx_vec_mul_f32_f32(v_f32, vd); + + hvx_vmem(dst_ptr + i * QK4_0) = res; + } +} + +static inline void hvx_dequantize_row_flat_q4_k_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) { + const int nb = n / QK_K; + const block_q4_K * src = (const block_q4_K *) src_ptr; + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(src[i].d); + const float min = GGML_FP16_TO_FP32(src[i].dmin); + const uint8_t * sc = src[i].scales; + + float * y = dst_ptr + i * QK_K; + + for (int step = 0; step < 4; ++step) { + uint8_t sc1, m1_u, sc2, m2_u; + if (step < 2) { + sc1 = sc[2 * step + 0] & 63; + m1_u = sc[2 * step + 4] & 63; + sc2 = sc[2 * step + 1] & 63; + m2_u = sc[2 * step + 5] & 63; + } else { + const int s = step - 2; + sc1 = (sc[8 + 2 * s + 0] & 0xF) | ((sc[2 * s + 0] >> 6) << 4); + m1_u = (sc[8 + 2 * s + 0] >> 4) | ((sc[4 + 2 * s + 0] >> 6) << 4); + sc2 = (sc[8 + 2 * s + 1] & 0xF) | ((sc[2 * s + 1] >> 6) << 4); + m2_u = (sc[8 + 2 * s + 1] >> 4) | ((sc[4 + 2 * s + 1] >> 6) << 4); + } + + HVX_Vector vq = *(const HVX_UVector *) (src[i].qs + step * 32); + HVX_Vector v_lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F)); + HVX_Vector v_hi = Q6_Vub_vlsr_VubR(vq, 4); + + HVX_VectorPair p16_lo = Q6_Wuh_vunpack_Vub(v_lo); + HVX_VectorPair p32_lo = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_lo)); + HVX_Vector v_i32_lo = Q6_V_lo_W(p32_lo); + HVX_Vector v_f32_lo = Q6_Vsf_equals_Vw(v_i32_lo); + + HVX_Vector vd1 = hvx_vec_splat_f32(d * sc1); + HVX_Vector vm1 = hvx_vec_splat_f32(min * m1_u); + HVX_Vector v_out1 = hvx_vec_sub_f32_f32(hvx_vec_mul_f32_f32(v_f32_lo, vd1), vm1); + hvx_vmem(y + 0) = v_out1; + + HVX_VectorPair p16_hi = Q6_Wuh_vunpack_Vub(v_hi); + HVX_VectorPair p32_hi = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_hi)); + HVX_Vector v_i32_hi = Q6_V_lo_W(p32_hi); + HVX_Vector v_f32_hi = Q6_Vsf_equals_Vw(v_i32_hi); + + HVX_Vector vd2 = hvx_vec_splat_f32(d * sc2); + HVX_Vector vm2 = hvx_vec_splat_f32(min * m2_u); + HVX_Vector v_out2 = hvx_vec_sub_f32_f32(hvx_vec_mul_f32_f32(v_f32_hi, vd2), vm2); + hvx_vmem(y + 32) = v_out2; + + y += 64; + } } } +static inline void hvx_dequantize_row_flat_q5_k_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) { + const int nb = n / QK_K; + const block_q5_K * src = (const block_q5_K *) src_ptr; + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(src[i].d); + const float min = GGML_FP16_TO_FP32(src[i].dmin); + const uint8_t * sc = src[i].scales; + + HVX_Vector v_qh = *(const HVX_UVector *) src[i].qh; + float * y = dst_ptr + i * QK_K; + + for (int step = 0; step < 4; ++step) { + uint8_t sc1, m1_u, sc2, m2_u; + if (step < 2) { + sc1 = sc[2 * step + 0] & 63; + m1_u = sc[2 * step + 4] & 63; + sc2 = sc[2 * step + 1] & 63; + m2_u = sc[2 * step + 5] & 63; + } else { + const int s = step - 2; + sc1 = (sc[8 + 2 * s + 0] & 0xF) | ((sc[2 * s + 0] >> 6) << 4); + m1_u = (sc[8 + 2 * s + 0] >> 4) | ((sc[4 + 2 * s + 0] >> 6) << 4); + sc2 = (sc[8 + 2 * s + 1] & 0xF) | ((sc[2 * s + 1] >> 6) << 4); + m2_u = (sc[8 + 2 * s + 1] >> 4) | ((sc[4 + 2 * s + 1] >> 6) << 4); + } + + HVX_Vector v_ql = *(const HVX_UVector *) (src[i].qs + step * 32); + HVX_Vector v_lo = Q6_V_vand_VV(v_ql, Q6_Vb_vsplat_R(0x0F)); + HVX_Vector v_hi = Q6_Vub_vlsr_VubR(v_ql, 4); + + HVX_Vector v_h1 = Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(v_qh, 2 * step), Q6_Vb_vsplat_R(1)), 4); + HVX_Vector v_val1 = Q6_V_vor_VV(v_lo, v_h1); + + HVX_VectorPair p16_1 = Q6_Wuh_vunpack_Vub(v_val1); + HVX_VectorPair p32_1 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_1)); + HVX_Vector v_i32_1 = Q6_V_lo_W(p32_1); + HVX_Vector v_f32_1 = Q6_Vsf_equals_Vw(v_i32_1); + + HVX_Vector vd1 = hvx_vec_splat_f32(d * sc1); + HVX_Vector vm1 = hvx_vec_splat_f32(min * m1_u); + HVX_Vector v_out1 = hvx_vec_sub_f32_f32(hvx_vec_mul_f32_f32(v_f32_1, vd1), vm1); + hvx_vmem(y + 0) = v_out1; + + HVX_Vector v_h2 = Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(v_qh, 2 * step + 1), Q6_Vb_vsplat_R(1)), 4); + HVX_Vector v_val2 = Q6_V_vor_VV(v_hi, v_h2); + + HVX_VectorPair p16_2 = Q6_Wuh_vunpack_Vub(v_val2); + HVX_VectorPair p32_2 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_2)); + HVX_Vector v_i32_2 = Q6_V_lo_W(p32_2); + HVX_Vector v_f32_2 = Q6_Vsf_equals_Vw(v_i32_2); + + HVX_Vector vd2 = hvx_vec_splat_f32(d * sc2); + HVX_Vector vm2 = hvx_vec_splat_f32(min * m2_u); + HVX_Vector v_out2 = hvx_vec_sub_f32_f32(hvx_vec_mul_f32_f32(v_f32_2, vd2), vm2); + hvx_vmem(y + 32) = v_out2; + + y += 64; + } + } +} + +static inline void hvx_dequantize_row_flat_q6_k_f32(float * restrict dst_ptr, const void * restrict src_ptr, uint32_t n) { + const int nb = n / QK_K; + const block_q6_K * src = (const block_q6_K *) src_ptr; + + const HVX_VectorPred q64 = Q6_Q_vsetq_R(64); + const HVX_Vector v32_b = Q6_Vb_vsplat_R(32); + const HVX_Vector mask_0f = Q6_Vb_vsplat_R(0x0F); + const HVX_Vector mask_03 = Q6_Vb_vsplat_R(0x03); + + for (int i = 0; i < nb; i++) { + const float d = GGML_FP16_TO_FP32(src[i].d); + + const uint8_t * restrict ql = src[i].ql; + const uint8_t * restrict qh = src[i].qh; + const int8_t * restrict sc = src[i].scales; + + float * y = dst_ptr + i * QK_K; + + for (int half = 0; half < 2; ++half) { + float s[8]; + for (int j = 0; j < 8; ++j) { + s[j] = d * sc[j]; + } + + HVX_Vector vscale1 = Q6_V_vmux_QVV(q64, hvx_vec_splat_f32(s[0]), hvx_vec_splat_f32(s[1])); + HVX_Vector vscale2 = Q6_V_vmux_QVV(q64, hvx_vec_splat_f32(s[2]), hvx_vec_splat_f32(s[3])); + HVX_Vector vscale3 = Q6_V_vmux_QVV(q64, hvx_vec_splat_f32(s[4]), hvx_vec_splat_f32(s[5])); + HVX_Vector vscale4 = Q6_V_vmux_QVV(q64, hvx_vec_splat_f32(s[6]), hvx_vec_splat_f32(s[7])); + + HVX_Vector v_ql0 = *(const HVX_UVector *) (ql); + HVX_Vector v_ql1 = *(const HVX_UVector *) (ql + 32); + HVX_Vector v_qh = *(const HVX_UVector *) (qh); + + HVX_Vector ql0_lo = Q6_V_vand_VV(v_ql0, mask_0f); + HVX_Vector qh_0 = Q6_Vw_vasl_VwR(Q6_V_vand_VV(v_qh, mask_03), 4); + HVX_Vector v_q1 = Q6_Vb_vsub_VbVb(Q6_V_vor_VV(ql0_lo, qh_0), v32_b); + + HVX_Vector ql1_lo = Q6_V_vand_VV(v_ql1, mask_0f); + HVX_Vector qh_1 = Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(v_qh, 2), mask_03), 4); + HVX_Vector v_q2 = Q6_Vb_vsub_VbVb(Q6_V_vor_VV(ql1_lo, qh_1), v32_b); + + HVX_Vector ql0_hi = Q6_Vub_vlsr_VubR(v_ql0, 4); + HVX_Vector qh_2 = Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(v_qh, 4), mask_03), 4); + HVX_Vector v_q3 = Q6_Vb_vsub_VbVb(Q6_V_vor_VV(ql0_hi, qh_2), v32_b); + + HVX_Vector ql1_hi = Q6_Vub_vlsr_VubR(v_ql1, 4); + HVX_Vector qh_3 = Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(v_qh, 6), mask_03), 4); + HVX_Vector v_q4 = Q6_Vb_vsub_VbVb(Q6_V_vor_VV(ql1_hi, qh_3), v32_b); + + HVX_VectorPair p16_1 = Q6_Wh_vunpack_Vb(v_q1); + HVX_VectorPair p32_1 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_1)); + HVX_Vector v_f32_1 = Q6_Vsf_equals_Vw(Q6_V_lo_W(p32_1)); + hvx_vmem(y + 0) = hvx_vec_mul_f32_f32(v_f32_1, vscale1); + + HVX_VectorPair p16_2 = Q6_Wh_vunpack_Vb(v_q2); + HVX_VectorPair p32_2 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_2)); + HVX_Vector v_f32_2 = Q6_Vsf_equals_Vw(Q6_V_lo_W(p32_2)); + hvx_vmem(y + 32) = hvx_vec_mul_f32_f32(v_f32_2, vscale2); + + HVX_VectorPair p16_3 = Q6_Wh_vunpack_Vb(v_q3); + HVX_VectorPair p32_3 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_3)); + HVX_Vector v_f32_3 = Q6_Vsf_equals_Vw(Q6_V_lo_W(p32_3)); + hvx_vmem(y + 64) = hvx_vec_mul_f32_f32(v_f32_3, vscale3); + + HVX_VectorPair p16_4 = Q6_Wh_vunpack_Vb(v_q4); + HVX_VectorPair p32_4 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16_4)); + HVX_Vector v_f32_4 = Q6_Vsf_equals_Vw(Q6_V_lo_W(p32_4)); + hvx_vmem(y + 96) = hvx_vec_mul_f32_f32(v_f32_4, vscale4); + + y += 128; + ql += 64; + qh += 32; + sc += 8; + } + } +} + +// bytes row, row + 32, row + 64 and row + 96 of v, to bytes 0..3 +static inline HVX_Vector hvx_dequantize_gather4(HVX_Vector v, uint32_t row) { + v = Q6_V_vror_VR(v, row); + HVX_Vector packed = Q6_V_vmux_QVV(Q6_Q_vsetq_R(1), v, Q6_V_vror_VR(v, 31)); + packed = Q6_V_vmux_QVV(Q6_Q_vsetq_R(2), packed, Q6_V_vror_VR(v, 62)); + return Q6_V_vmux_QVV(Q6_Q_vsetq_R(3), packed, Q6_V_vror_VR(v, 93)); +} + +static inline void hvx_dequantize_row_tiled_q8_0_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) { + const HVX_VectorPred first2 = Q6_Q_vsetq_R(2); + const HVX_VectorPred first4 = Q6_Q_vsetq_R(4); + + for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) { + const uint8_t * tile = src_spad + k_tile * tile_stride; + float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS; + + HVX_Vector vq = Q6_V_vzero(); + for (int group = 7; group >= 0; --group) { + const HVX_Vector v = Q6_V_vror_VR(hvx_vmem(tile + group * VLEN), 2 * row); + const HVX_Vector packed = Q6_V_vmux_QVV(first2, v, Q6_V_vror_VR(v, 62)); + vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4)); + } + + const HVX_Vector scales = hvx_vmem(tile + 1024); + const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 2 * row)); + const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf)); + const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq); + const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16)); + HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale); + *(HVX_Vector *) dst = values; + } +} + +static inline void hvx_dequantize_row_tiled_q4_0_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) { + const HVX_VectorPred first4 = Q6_Q_vsetq_R(4); + + for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) { + const uint8_t * tile = src_spad + k_tile * tile_stride; + float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS; + + HVX_Vector vq = Q6_V_vzero(); + for (int group = 3; group >= 0; --group) { + vq = Q6_V_vmux_QVV(first4, hvx_dequantize_gather4(hvx_vmem(tile + group * VLEN), row), Q6_V_vror_VR(vq, VLEN - 4)); + } + + const HVX_Vector lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F)); + const HVX_Vector hi = Q6_Vub_vlsr_VubR(vq, 4); + vq = Q6_V_lo_W(Q6_W_vshuff_VVR(hi, lo, -1)); + vq = Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(8)); + + const HVX_Vector scales = hvx_vmem(tile + 512); + const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 2 * row)); + const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf)); + const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq); + const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16)); + HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale); + *(HVX_Vector *) dst = values; + } +} + +static inline void hvx_dequantize_row_tiled_q4_k_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) { + const HVX_VectorPred first4 = Q6_Q_vsetq_R(4); + + for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) { + const uint8_t * tile = src_spad + k_tile * tile_stride; + float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS; + + HVX_Vector vq = Q6_V_vzero(); + for (int group = 3; group >= 0; --group) { + vq = Q6_V_vmux_QVV(first4, hvx_dequantize_gather4(hvx_vmem(tile + group * VLEN), row), Q6_V_vror_VR(vq, VLEN - 4)); + } + + const HVX_Vector lo = Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F)); + const HVX_Vector hi = Q6_Vub_vlsr_VubR(vq, 4); + vq = Q6_V_lo_W(Q6_W_vshuff_VVR(hi, lo, -1)); + + const HVX_Vector scales = hvx_vmem(tile + 512); + const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 4 * row)); + const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf)); + const HVX_Vector offset_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 4 * row + 2)); + const HVX_Vector offset = Q6_V_lo_W(hvx_vec_f16_to_f32(offset_hf)); + + const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq); + const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16)); + HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale); + *(HVX_Vector *) dst = hvx_vec_add_f32_f32(values, offset); + } +} + +static inline void hvx_dequantize_row_tiled_q5_k_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) { + const HVX_VectorPred first4 = Q6_Q_vsetq_R(4); + const HVX_Vector one = Q6_Vb_vsplat_R(1); + + for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) { + const uint8_t * tile = src_spad + k_tile * tile_stride; + float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS; + + const HVX_Vector p5 = hvx_dequantize_gather4(hvx_vmem(tile + 640), row); + HVX_Vector vq = Q6_V_vzero(); + HVX_Vector vh = Q6_V_vzero(); + for (int group = 3; group >= 0; --group) { + vq = Q6_V_vmux_QVV(first4, hvx_dequantize_gather4(hvx_vmem(tile + group * VLEN), row), Q6_V_vror_VR(vq, VLEN - 4)); + vh = Q6_V_vmux_QVV(first4, Q6_Vub_vlsr_VubR(p5, 2 * group), Q6_V_vror_VR(vh, VLEN - 4)); + } + + const HVX_Vector lo = Q6_V_vor_VV(Q6_V_vand_VV(vq, Q6_Vb_vsplat_R(0x0F)), Q6_Vw_vasl_VwR(Q6_V_vand_VV(vh, one), 4)); + const HVX_Vector hi = Q6_V_vor_VV(Q6_Vub_vlsr_VubR(vq, 4), Q6_Vw_vasl_VwR(Q6_V_vand_VV(Q6_Vub_vlsr_VubR(vh, 1), one), 4)); + vq = Q6_V_lo_W(Q6_W_vshuff_VVR(hi, lo, -1)); + + const HVX_Vector scales = hvx_vmem(tile + 512); + const HVX_Vector scale_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 4 * row)); + const HVX_Vector scale = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hf)); + const HVX_Vector offset_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 4 * row + 2)); + const HVX_Vector offset = Q6_V_lo_W(hvx_vec_f16_to_f32(offset_hf)); + + const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq); + const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16)); + HVX_Vector values = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale); + *(HVX_Vector *) dst = hvx_vec_add_f32_f32(values, offset); + } +} + +static inline void hvx_dequantize_row_tiled_q6_k_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) { + const HVX_VectorPred first4 = Q6_Q_vsetq_R(4); + const HVX_VectorPred first16 = Q6_Q_vsetq_R(16 * sizeof(float)); + const HVX_Vector mask_0f = Q6_Vb_vsplat_R(0x0F); + const HVX_Vector mask_03 = Q6_Vb_vsplat_R(0x03); + + for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) { + const uint8_t * tile = src_spad + k_tile * tile_stride; + float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS; + + HVX_Vector vq = Q6_V_vzero(); + for (int group = 7; group >= 0; --group) { + const HVX_Vector lo_plane = Q6_V_vror_VR(hvx_vmem(tile + (group >> 1) * VLEN), 4 * row); + const HVX_Vector hi_plane = Q6_V_vror_VR(hvx_vmem(tile + 512 + (group >> 2) * VLEN), 4 * row); + const HVX_Vector lo = (group & 1) ? Q6_Vub_vlsr_VubR(lo_plane, 4) : Q6_V_vand_VV(lo_plane, mask_0f); + const HVX_Vector hi = Q6_Vub_vlsr_VubR(hi_plane, 2 * (group & 3)); + const HVX_Vector packed = Q6_V_vor_VV(lo, Q6_Vw_vasl_VwR(Q6_V_vand_VV(hi, mask_03), 4)); + vq = Q6_V_vmux_QVV(first4, packed, Q6_V_vror_VR(vq, VLEN - 4)); + } + + const HVX_Vector scales = hvx_vmem(tile + 768); + const HVX_Vector scale_lo_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 2 * row)); + const HVX_Vector scale_hi_hf = hvx_vec_repl_f16(Q6_V_vror_VR(scales, 64 + 2 * row)); + const HVX_Vector scale_lo = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_lo_hf)); + const HVX_Vector scale_hi = Q6_V_lo_W(hvx_vec_f16_to_f32(scale_hi_hf)); + const HVX_Vector scale = Q6_V_vmux_QVV(first16, scale_lo, scale_hi); + const HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(Q6_Vb_vsub_VbVb(vq, Q6_Vb_vsplat_R(32))); + const HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(Q6_V_lo_W(p16)); + *(HVX_Vector *) dst = hvx_vec_mul_f32_f32(Q6_Vsf_equals_Vw(Q6_V_lo_W(p32)), scale); + } +} + +static inline HVX_Vector hvx_dequantize_tiled_f16_gather4(const uint8_t * p, uint32_t r_byte) { + HVX_Vector w0 = Q6_V_vror_VR(hvx_vmem(p + 0 * VLEN), r_byte); + HVX_Vector w1 = Q6_V_vror_VR(hvx_vmem(p + 1 * VLEN), r_byte); + HVX_Vector w01 = Q6_V_lo_W(Q6_W_vshuff_VVR(w1, w0, 4)); + HVX_Vector w2 = Q6_V_vror_VR(hvx_vmem(p + 2 * VLEN), r_byte); + HVX_Vector w3 = Q6_V_vror_VR(hvx_vmem(p + 3 * VLEN), r_byte); + HVX_Vector w23 = Q6_V_lo_W(Q6_W_vshuff_VVR(w3, w2, 4)); + return Q6_V_lo_W(Q6_W_vshuff_VVR(w23, w01, 8)); +} + +static inline void hvx_dequantize_row_tiled_f16_f32(float * dst_spad, const uint8_t * src_spad, uint32_t n_k_tiles, uint32_t tile_stride, uint32_t row) { + const uint32_t r_byte = 4 * row; + + for (uint32_t k_tile = 0; k_tile < n_k_tiles; ++k_tile) { + const uint8_t * tile = src_spad + k_tile * tile_stride; + float * dst = dst_spad + k_tile * HTP_TENSOR_TILE_N_COLS; + + HVX_Vector g0 = hvx_dequantize_tiled_f16_gather4(tile + 0 * VLEN, r_byte); + HVX_Vector g1 = hvx_dequantize_tiled_f16_gather4(tile + 4 * VLEN, r_byte); + HVX_Vector g01 = Q6_V_lo_W(Q6_W_vshuff_VVR(g1, g0, 16)); + + HVX_Vector g2 = hvx_dequantize_tiled_f16_gather4(tile + 8 * VLEN, r_byte); + HVX_Vector g3 = hvx_dequantize_tiled_f16_gather4(tile + 12 * VLEN, r_byte); + HVX_Vector g23 = Q6_V_lo_W(Q6_W_vshuff_VVR(g3, g2, 16)); + + HVX_Vector g = Q6_V_lo_W(Q6_W_vshuff_VVR(g23, g01, 32)); + *(HVX_Vector *) dst = Q6_V_lo_W(hvx_vec_f16_to_f32(g)); + } +} #endif // HVX_QUANT_H diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index df7c4a5bd1ff..2374add3f531 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -28,8 +28,8 @@ #include "ggml-common.h" #include "hex-bitmap.h" #include "htp-ctx.h" +#include "htp-dcache.h" #include "htp-ops.h" -#include "htp-tensor.h" #include "htp_iface.h" #include "work-queue.h" #include "hex-profile.h" @@ -835,7 +835,7 @@ static int op_fence(struct htp_ops_context * octx) { atomic_uint * sync_fence = (atomic_uint *) (uintptr_t) sync->data; if (mode == 1) { - htp_flush_dirty_ranges(ctx); + htp_dcache_flush_dirty_ranges(ctx); htp_mdev_group_barrier(octx); @@ -1211,7 +1211,7 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs src->ne[0], src->ne[1], src->ne[2], src->ne[3]); } - htp_tensor_flush_all(octx->ctx, octx->src, HTP_OP_MAX_INPUTS); + htp_dcache_flush_dirty_tensors(octx->ctx, octx->src, HTP_OP_MAX_INPUTS); // Prep output tensors for (uint32_t i = 0; i < HTP_OP_MAX_OUTPUTS; i++) { @@ -1227,7 +1227,7 @@ static int proc_op_req(struct htp_ops_context * octx, struct htp_buf_desc * bufs dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3]); } - htp_tensor_dirty_all(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS); + htp_dcache_mark_dirty_tensors(octx->ctx, octx->dsts, HTP_OP_MAX_OUTPUTS); htp_mdev_group_barrier(octx); diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 28aed4680df4..3837c290d62c 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -38,8 +38,12 @@ typedef struct { int act_stride; uint32_t act_elem_size; // 4=F32 src1, 2=F16 src1 int weight_stride; + int weight_type; int dst_stride; + int dst_cols; uint32_t bias_stride; + int tile_size; + int aligned_tile_size; int ne02; int ne03; int ne12; @@ -56,7 +60,7 @@ typedef struct { int r3; struct fastdiv_values div_r2; struct fastdiv_values div_r3; -} hmx_mm_f16_f32_batched_params_t; +} hmx_mm_batched_params_t; static bool htp_matmul_has_extended_weight(const struct htp_ops_context * octx, uint32_t n_weights) { for (uint32_t i = 0; i < n_weights; ++i) { @@ -89,6 +93,11 @@ struct htp_mm_context { const void * restrict vy, uint32_t valid_rows, const float * restrict sz); + void (*vec_dot_32x2)(const uint32_t n, float * restrict s0, float * restrict s1, + const void * restrict vx, + const void * restrict vy0, const void * restrict vy1, uint32_t valid_rows, + const float * restrict sz0, const float * restrict sz1); + // Precomputed values uint32_t src0_nrows_per_thread; uint32_t src0_row_start; @@ -238,314 +247,301 @@ static const uint8_t __attribute__((aligned(VLEN))) kvalues_mxfp4_lut[] = { #include "hvx-mm-kernels-float.h" #include "hmx-mm-kernels-tiled.h" -// Specialized repacked matmul macros -#define MATMUL_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ -static void hvx_mm_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ - htp_matmul_preamble; \ - \ - const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; \ - const uint32_t act_nrows = mmctx->cur_m_rows ? mmctx->cur_m_rows : (ne11 * ne12 * ne13); \ - const uint32_t cur_m_start = mmctx->cur_m_start; \ - \ - const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; \ - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); \ - \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - \ - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ - const uint32_t n_prefetch = kparams->n_prefetch; \ - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ - \ - const size_t dst_row_size = nb1; \ - const size_t act_row_size = nb11; \ - const size_t act_stride = mmctx->vtcm_act_stride; \ - const size_t src2_stride = src2 ? ((src2->ne[1] == 1) ? 0 : src2->nb[1]) : 0; \ - \ - uint8_t * restrict vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; \ - uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ - uint8_t * restrict act_data = mmctx->vtcm_act; \ - \ - const dma_addr_t src0_row = src0->data; \ - \ - const uint32_t tile_size = TILE_SIZE; \ - const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ - \ - uint32_t n_k_tiles_w = ne00 / 32; \ - uint32_t n_k_tiles_a = ne10 / 32; \ - uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ - uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ - \ - uint32_t ct_start = src0_start_row / 32; \ - uint32_t ct_end = (src0_end_row + 31) / 32; \ - \ - uint32_t push_ct = ct_start; \ - if (src0_start_row < src0_end_row) { \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ - dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ - src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - } \ - } \ - \ - if (src0_start_row >= src0_end_row) { \ - return; \ - } \ - \ - for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ - const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; \ - \ - int valid_rows = (int)ne0 - (int)(ct * 32); \ - valid_rows = MIN(32, MAX(0, valid_rows)); \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - uint32_t ir1 = 0; \ - for (; ir1 + 1 < act_nrows; ir1 += 2) { \ - const uint8_t * restrict act_col0 = (const uint8_t *) (act_data + (ir1+0) * act_stride); \ - const uint8_t * restrict act_col1 = (const uint8_t *) (act_data + (ir1+1) * act_stride); \ - float * restrict dst_row0 = (float *) (dst->data + ((cur_m_start + ir1+0) * dst_row_size)); \ - float * restrict dst_row1 = (float *) (dst->data + ((cur_m_start + ir1+1) * dst_row_size)); \ - \ - float * dst_ptr0 = &dst_row0[ct * 32]; \ - float * dst_ptr1 = &dst_row1[ct * 32]; \ - \ - const float * src2_ptr0 = NULL; \ - const float * src2_ptr1 = NULL; \ - if (src2) { \ - const float * restrict src2_row0 = (const float *) ((const uint8_t *) src2->data + ((cur_m_start + ir1+0) * src2_stride)); \ - const float * restrict src2_row1 = (const float *) ((const uint8_t *) src2->data + ((cur_m_start + ir1+1) * src2_stride)); \ - src2_ptr0 = &src2_row0[ct * 32]; \ - src2_ptr1 = &src2_row1[ct * 32]; \ - } \ - DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, act_col0, act_col1, valid_rows, src2_ptr0, src2_ptr1); \ - } \ - \ - for (; ir1 < act_nrows; ++ir1) { \ - const uint8_t * restrict act_col = (const uint8_t *) (act_data + ir1 * act_stride); \ - float * restrict dst_row = (float *) (dst->data + ((cur_m_start + ir1) * dst_row_size)); \ - float * dst_ptr = &dst_row[ct * 32]; \ - \ - const float * src2_ptr = NULL; \ - if (src2) { \ - const float * restrict src2_row = (const float *) ((const uint8_t *) src2->data + ((cur_m_start + ir1) * src2_stride)); \ - src2_ptr = &src2_row[ct * 32]; \ - } \ - DOT_2X1(ne10, dst_ptr, w_tile, act_col, valid_rows, src2_ptr); \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - \ - if (push_ct < ct_end) { \ - dma_queue_push(dma_q, dma_make_data(w_tile, src0_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - push_ct++; \ - } \ - } \ -} +static void hvx_mm_2d_tiled(unsigned int nth, unsigned int ith, void * data) { + htp_matmul_preamble; + + const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; + const uint32_t act_nrows = mmctx->cur_m_rows ? mmctx->cur_m_rows : (ne11 * ne12 * ne13); + const uint32_t cur_m_start = mmctx->cur_m_start; + + const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); + + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_prefetch = kparams->n_prefetch; + assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); + + const size_t dst_row_size = nb1; + const size_t act_row_size = nb11; + const size_t act_stride = mmctx->vtcm_act_stride; + const size_t src2_stride = src2 ? ((src2->ne[1] == 1) ? 0 : src2->nb[1]) : 0; + + uint8_t * restrict vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; + uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; + uint8_t * restrict act_data = mmctx->vtcm_act; + + const dma_addr_t src0_row = src0->data; + + const uint32_t tile_size = htp_tensor_get_tile_size(src0->type); + const uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(src0->type); + + uint32_t n_k_tiles_w = ne00 / 32; + uint32_t n_k_tiles_a = ne10 / 32; + uint32_t tile_row_stride = n_k_tiles_w * tile_size; + uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; + + uint32_t ct_start = src0_start_row / 32; + uint32_t ct_end = (src0_end_row + 31) / 32; + + uint32_t push_ct = ct_start; + if (src0_start_row < src0_end_row) { + for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, + src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + } + } + + if (src0_start_row >= src0_end_row) { + return; + } + + for (uint32_t ct = ct_start; ct < ct_end; ct++) { + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; + + int valid_rows = (int)ne0 - (int)(ct * 32); + valid_rows = MIN(32, MAX(0, valid_rows)); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); + uint32_t ir1 = 0; + for (; ir1 + 1 < act_nrows; ir1 += 2) { + const uint8_t * restrict act_col0 = (const uint8_t *) (act_data + (ir1+0) * act_stride); + const uint8_t * restrict act_col1 = (const uint8_t *) (act_data + (ir1+1) * act_stride); + float * restrict dst_row0 = (float *) (dst->data + ((cur_m_start + ir1+0) * dst_row_size)); + float * restrict dst_row1 = (float *) (dst->data + ((cur_m_start + ir1+1) * dst_row_size)); + + float * dst_ptr0 = &dst_row0[ct * 32]; + float * dst_ptr1 = &dst_row1[ct * 32]; -#define MATVEC_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X1) \ -static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ - htp_matmul_preamble; \ - \ - const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; \ - \ - const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; \ - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); \ - \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - \ - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ - const uint32_t n_prefetch = kparams->n_prefetch; \ - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ - \ - const size_t dst_row_size = nb1; \ - const size_t act_row_size = nb11; \ - const size_t act_stride = mmctx->vtcm_act_stride; \ - \ - uint8_t * vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; \ - uint8_t * vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ - uint8_t * act_data = mmctx->vtcm_act; \ - \ - float * tmp = (float *) vtcm_dst_ptr; \ - \ - const dma_addr_t src0_row = src0->data; \ - \ - const uint8_t * restrict act_col = (const uint8_t *) act_data; \ - float * restrict dst_col = (float *) dst->data; \ - \ - const uint32_t tile_size = TILE_SIZE; \ - const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ - \ - uint32_t n_k_tiles_w = ne00 / 32; \ - uint32_t n_k_tiles_a = ne10 / 32; \ - uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ - uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ - \ - uint32_t ct_start = src0_start_row / 32; \ - uint32_t ct_end = (src0_end_row + 31) / 32; \ - \ - uint32_t push_ct = ct_start; \ - if (src0_start_row < src0_end_row) { \ - if (src2) { \ - float * vtcm_bias_ptr = (float *) mmctx->vtcm_bias + src0_start_row; \ - const dma_addr_t src2_addr = src2->data + src0_start_row * sizeof(float); \ - int slice_size = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \ - if (slice_size > 0) { \ - dma_queue_push(dma_q, dma_make_data(vtcm_bias_ptr, src2_addr), \ - slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1); \ - dma_queue_pop_nowait(dma_q); \ - } \ - } \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ - dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ - src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - } \ - } \ - \ - if (src0_start_row >= src0_end_row) { \ - return; \ - } \ - \ - for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ - const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; \ - \ - float * dst_ptr = &tmp[ct * 32 - src0_start_row]; \ - int valid_rows = (int)ne0 - (int)(ct * 32); \ - valid_rows = MIN(32, MAX(0, valid_rows)); \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - DOT_2X1(ne10, dst_ptr, w_tile, act_col, valid_rows, NULL); \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - \ - if (push_ct < ct_end) { \ - dma_queue_push(dma_q, dma_make_data(w_tile, src0_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - push_ct++; \ - } \ - } \ - \ - int copy_cnt = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; \ - if (copy_cnt > 0) { \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \ - if (src2) { \ - hvx_add_f32_uuu((uint8_t *) &dst_col[src0_start_row], \ - (const uint8_t *) tmp, \ - (const uint8_t *) ((const float *) mmctx->vtcm_bias + src0_start_row), \ - copy_cnt); \ - } else { \ - hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); \ - } \ + const float * src2_ptr0 = NULL; + const float * src2_ptr1 = NULL; + if (src2) { + const float * restrict src2_row0 = (const float *) ((const uint8_t *) src2->data + ((cur_m_start + ir1+0) * src2_stride)); + const float * restrict src2_row1 = (const float *) ((const uint8_t *) src2->data + ((cur_m_start + ir1+1) * src2_stride)); + src2_ptr0 = &src2_row0[ct * 32]; + src2_ptr1 = &src2_row1[ct * 32]; + } + mmctx->vec_dot_32x2(ne10, dst_ptr0, dst_ptr1, w_tile, act_col0, act_col1, valid_rows, src2_ptr0, src2_ptr1); + } + + for (; ir1 < act_nrows; ++ir1) { + const uint8_t * restrict act_col = (const uint8_t *) (act_data + ir1 * act_stride); + float * restrict dst_row = (float *) (dst->data + ((cur_m_start + ir1) * dst_row_size)); + float * dst_ptr = &dst_row[ct * 32]; + + const float * src2_ptr = NULL; + if (src2) { + const float * restrict src2_row = (const float *) ((const uint8_t *) src2->data + ((cur_m_start + ir1) * src2_stride)); + src2_ptr = &src2_row[ct * 32]; + } + mmctx->vec_dot_32x1(ne10, dst_ptr, w_tile, act_col, valid_rows, src2_ptr); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); + + if (push_ct < ct_end) { + dma_queue_push(dma_q, dma_make_data(w_tile, src0_row + push_ct * tile_row_stride), + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + push_ct++; + } + } } -#define MATMUL_NX_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ -static void hvx_mm_nx_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_mm_context * mmctx = data; \ - struct htp_ops_context * octx = mmctx->octx; \ - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ - const uint32_t n_weights = kparams->n_weights; \ - \ - const struct htp_tensor * restrict act = octx->src[n_weights]; /* x */ \ - const uint32_t ne10 = act->ne[0]; \ - const uint32_t act_nrows = act->ne[1] * act->ne[2] * act->ne[3]; \ - const size_t act_stride = mmctx->vtcm_act_stride; \ - \ - uint8_t * restrict vtcm_weight_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ - uint8_t * restrict act_data = mmctx->vtcm_act; \ - \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - const uint32_t n_prefetch = kparams->n_prefetch; \ - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ - \ - const uint32_t tile_size = TILE_SIZE; \ - const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ - uint32_t n_k_tiles_a = ne10 / 32; \ - uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ - \ - for (uint32_t widx = 0; widx < n_weights; widx++) { \ - const struct htp_tensor * restrict src_w = octx->src[widx]; \ - const struct htp_tensor * restrict dst = octx->dsts[widx]; \ - if (!src_w || !dst) continue; \ - dma_queue * dma_q = octx->ctx->dma[ith]; \ - \ - const uint32_t ne00 = src_w->ne[0]; \ - const uint32_t ne01 = src_w->ne[1]; \ - const size_t dst_row_size = dst->nb[1]; \ - const dma_addr_t src_w_row = src_w->data; \ - \ - uint32_t n_k_tiles_w = ne00 / 32; \ - uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ - \ - uint32_t src0_start_row = 0; \ - uint32_t src0_end_row = ne01; \ - if (octx->ctx->mdev.count > 1) { \ - const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float)); \ - const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, \ - octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); \ - src0_start_row = range.start; \ - src0_end_row = range.start + range.count; \ - } \ - \ - const uint32_t nrows = src0_end_row - src0_start_row; \ - uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div); \ - src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32); \ - \ - const uint32_t start_row = src0_start_row + src0_nrows_per_thread * ith; \ - const uint32_t end_row = MIN(start_row + src0_nrows_per_thread, src0_end_row); \ - if (start_row >= end_row) continue; \ - \ - uint32_t ct_start = start_row / 32; \ - uint32_t ct_end = (end_row + 31) / 32; \ - \ - uint32_t push_ct = ct_start; \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ - dma_queue_push(dma_q, dma_make_data(vtcm_weight_ptr + d * tile_row_transfer_size_aligned, \ - src_w_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - } \ - \ - for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ - const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; \ - int valid_rows = (int)ne01 - (int)(ct * 32); \ - valid_rows = MIN(32, MAX(0, valid_rows)); \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - uint32_t ir1 = 0; \ - for (; ir1 + 1 < act_nrows; ir1 += 2) { \ - const uint8_t * restrict act_col0 = (const uint8_t *) (act_data + (ir1+0) * act_stride); \ - const uint8_t * restrict act_col1 = (const uint8_t *) (act_data + (ir1+1) * act_stride); \ - \ - float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); \ - float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); \ - float * dst_ptr0 = &dst_row0[ct * 32]; \ - float * dst_ptr1 = &dst_row1[ct * 32]; \ - \ - DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, act_col0, act_col1, valid_rows, NULL, NULL); \ - } \ - \ - for (; ir1 < act_nrows; ++ir1) { \ - const uint8_t * restrict act_col = (const uint8_t *) (act_data + ir1 * act_stride); \ - float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); \ - float * dst_ptr = &dst_row[ct * 32]; \ - DOT_2X1(ne10, dst_ptr, w_tile, act_col, valid_rows, NULL); \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - \ - if (push_ct < ct_end) { \ - dma_queue_push(dma_q, dma_make_data(w_tile, src_w_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - push_ct++; \ - } \ - } \ - } \ +static void hvx_mv_2d_tiled(unsigned int nth, unsigned int ith, void * data) { + htp_matmul_preamble; + + const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; + + const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); + + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_prefetch = kparams->n_prefetch; + assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); + + const size_t dst_row_size = nb1; + const size_t act_row_size = nb11; + const size_t act_stride = mmctx->vtcm_act_stride; + + uint8_t * vtcm_dst_ptr = mmctx->vtcm_dst + mmctx->vtcm_dst_size_per_thread * ith; + uint8_t * vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; + uint8_t * act_data = mmctx->vtcm_act; + + float * tmp = (float *) vtcm_dst_ptr; + + const dma_addr_t src0_row = src0->data; + + const uint8_t * restrict act_col = (const uint8_t *) act_data; + float * restrict dst_col = (float *) dst->data; + + const uint32_t tile_size = htp_tensor_get_tile_size(src0->type); + const uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(src0->type); + + uint32_t n_k_tiles_w = ne00 / 32; + uint32_t n_k_tiles_a = ne10 / 32; + uint32_t tile_row_stride = n_k_tiles_w * tile_size; + uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; + + uint32_t ct_start = src0_start_row / 32; + uint32_t ct_end = (src0_end_row + 31) / 32; + + uint32_t push_ct = ct_start; + if (src0_start_row < src0_end_row) { + if (src2) { + float * vtcm_bias_ptr = (float *) mmctx->vtcm_bias + src0_start_row; + const dma_addr_t src2_addr = src2->data + src0_start_row * sizeof(float); + int slice_size = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; + if (slice_size > 0) { + dma_queue_push(dma_q, dma_make_data(vtcm_bias_ptr, src2_addr), + slice_size * sizeof(float), slice_size * sizeof(float), slice_size * sizeof(float), 1); + dma_queue_pop_nowait(dma_q); + } + } + for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, + src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + } + } + + if (src0_start_row >= src0_end_row) { + return; + } + + for (uint32_t ct = ct_start; ct < ct_end; ct++) { + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; + + float * dst_ptr = &tmp[ct * 32 - src0_start_row]; + int valid_rows = (int)ne0 - (int)(ct * 32); + valid_rows = MIN(32, MAX(0, valid_rows)); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); + mmctx->vec_dot_32x1(ne10, dst_ptr, w_tile, act_col, valid_rows, NULL); + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); + + if (push_ct < ct_end) { + dma_queue_push(dma_q, dma_make_data(w_tile, src0_row + push_ct * tile_row_stride), + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + push_ct++; + } + } + + int copy_cnt = (int)MIN(src0_end_row, ne0) - (int)src0_start_row; + if (copy_cnt > 0) { + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); + if (src2) { + hvx_add_f32_uuu((uint8_t *) &dst_col[src0_start_row], + (const uint8_t *) tmp, + (const uint8_t *) ((const float *) mmctx->vtcm_bias + src0_start_row), + copy_cnt); + } else { + hvx_copy_f32_ua((uint8_t *) &dst_col[src0_start_row], (uint8_t *) tmp, copy_cnt); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct_end); + } } -MATMUL_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) -MATMUL_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) -MATMUL_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) -MATMUL_2D_REPACKED_IMPL(q6_k, 896, tiled_vec_dot_q6_k_32x2, tiled_vec_dot_q6_k_32x1) -MATMUL_2D_REPACKED_IMPL(q5_k, 768, tiled_vec_dot_q5_k_32x2, tiled_vec_dot_q5_k_32x1) -MATMUL_2D_REPACKED_IMPL(q3_k, 512, tiled_vec_dot_q3_k_32x2, tiled_vec_dot_q3_k_32x1) -MATMUL_2D_REPACKED_IMPL(q2_k, 336, tiled_vec_dot_q2_k_32x2, tiled_vec_dot_q2_k_32x1) -MATMUL_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) -MATMUL_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) +static void hvx_mm_nx_2d_tiled(unsigned int nth, unsigned int ith, void * data) { + struct htp_mm_context * mmctx = data; + struct htp_ops_context * octx = mmctx->octx; + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_weights = kparams->n_weights; + + const struct htp_tensor * restrict act = octx->src[n_weights]; /* x */ + const uint32_t ne10 = act->ne[0]; + const uint32_t act_nrows = act->ne[1] * act->ne[2] * act->ne[3]; + const size_t act_stride = mmctx->vtcm_act_stride; + + uint8_t * restrict vtcm_weight_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; + uint8_t * restrict act_data = mmctx->vtcm_act; + + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + const uint32_t n_prefetch = kparams->n_prefetch; + assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); + + uint32_t n_k_tiles_a = ne10 / 32; + + for (uint32_t widx = 0; widx < n_weights; widx++) { + const struct htp_tensor * restrict src_w = octx->src[widx]; + const struct htp_tensor * restrict dst = octx->dsts[widx]; + if (!src_w || !dst) continue; + dma_queue * dma_q = octx->ctx->dma[ith]; + + const uint32_t tile_size = htp_tensor_get_tile_size(src_w->type); + const uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(src_w->type); + const uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; + + const uint32_t ne00 = src_w->ne[0]; + const uint32_t ne01 = src_w->ne[1]; + const size_t dst_row_size = dst->nb[1]; + const dma_addr_t src_w_row = src_w->data; + + uint32_t n_k_tiles_w = ne00 / 32; + uint32_t tile_row_stride = n_k_tiles_w * tile_size; + + uint32_t src0_start_row = 0; + uint32_t src0_end_row = ne01; + if (octx->ctx->mdev.count > 1) { + const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float)); + const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, + octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); + src0_start_row = range.start; + src0_end_row = range.start + range.count; + } + + const uint32_t nrows = src0_end_row - src0_start_row; + uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div); + src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32); + + const uint32_t start_row = src0_start_row + src0_nrows_per_thread * ith; + const uint32_t end_row = MIN(start_row + src0_nrows_per_thread, src0_end_row); + if (start_row >= end_row) continue; + + uint32_t ct_start = start_row / 32; + uint32_t ct_end = (end_row + 31) / 32; + + uint32_t push_ct = ct_start; + for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { + dma_queue_push(dma_q, dma_make_data(vtcm_weight_ptr + d * tile_row_transfer_size_aligned, + src_w_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + } + + for (uint32_t ct = ct_start; ct < ct_end; ct++) { + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; + int valid_rows = (int)dst->ne[0] - (int)(ct * 32); + valid_rows = MIN(32, MAX(0, valid_rows)); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); + uint32_t ir1 = 0; + for (; ir1 + 1 < act_nrows; ir1 += 2) { + const uint8_t * restrict act_col0 = (const uint8_t *) (act_data + (ir1+0) * act_stride); + const uint8_t * restrict act_col1 = (const uint8_t *) (act_data + (ir1+1) * act_stride); + + float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); + float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); + float * dst_ptr0 = &dst_row0[ct * 32]; + float * dst_ptr1 = &dst_row1[ct * 32]; + + mmctx->vec_dot_32x2(ne10, dst_ptr0, dst_ptr1, w_tile, act_col0, act_col1, valid_rows, NULL, NULL); + } + + for (; ir1 < act_nrows; ++ir1) { + const uint8_t * restrict act_col = (const uint8_t *) (act_data + ir1 * act_stride); + float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); + float * dst_ptr = &dst_row[ct * 32]; + mmctx->vec_dot_32x1(ne10, dst_ptr, w_tile, act_col, valid_rows, NULL); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); + + if (push_ct < ct_end) { + dma_queue_push(dma_q, dma_make_data(w_tile, src_w_row + push_ct * tile_row_stride), + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + push_ct++; + } + } + } +} static void hvx_mm_transfer_act_dma( struct htp_ops_context * octx, @@ -649,12 +645,14 @@ static void name(unsigned int nth, unsigned int ith, void * data) { htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, ir_first); \ } -QUANTIZE_IMPL(quantize_f32_q8_0_tiled, "quantize-f32-q8_0_tiled", quantize_f32_q8_0_tiled_kernel, htp_mm_q8_0_tiled_row_size(ne0)) -QUANTIZE_IMPL(quantize_f32_q8_1_tiled, "quantize-f32-q8_1_tiled", quantize_f32_q8_1_tiled_kernel, htp_mm_q8_1_tiled_row_size(ne0)) +QUANTIZE_IMPL(quantize_f32_q8_0_tiled, "quantize-f32-q8_0_tiled", quantize_f32_q8_0_tiled_kernel, htp_mm_q8_0_tiled_row_size(ne0)) +QUANTIZE_IMPL(quantize_f32_q8_1_tiled, "quantize-f32-q8_1_tiled", quantize_f32_q8_1_tiled_kernel, htp_mm_q8_1_tiled_row_size(ne0)) QUANTIZE_IMPL(quantize_f32_q8_1_s16_tiled, "quantize-f32-q8_1_s16_tiled", quantize_f32_q8_1_s16_tiled_kernel, htp_mm_q8_1_tiled_row_size(ne0)) -QUANTIZE_IMPL(quantize_f32_f32, "quantize-f32-f32", quantize_f32_f32_kernel, mmctx->vtcm_act_stride) -QUANTIZE_IMPL(quantize_f32_f16, "quantize-f32-f16", quantize_f32_f16_kernel, mmctx->vtcm_act_stride) -QUANTIZE_IMPL(quantize_f16_f16, "quantize-f16-f16", quantize_f16_f16_kernel, mmctx->vtcm_act_stride) +QUANTIZE_IMPL(quantize_f32_f16_tiled, "quantize-f32-f16_tiled", quantize_f32_f16_tiled_kernel, htp_mm_f16_tiled_row_size(ne0)) +QUANTIZE_IMPL(quantize_f16_f16_tiled, "quantize-f16-f16_tiled", quantize_f16_f16_tiled_kernel, htp_mm_f16_tiled_row_size(ne0)) +QUANTIZE_IMPL(quantize_f32_f32, "quantize-f32-f32", quantize_f32_f32_kernel, mmctx->vtcm_act_stride) +QUANTIZE_IMPL(quantize_f32_f16, "quantize-f32-f16", quantize_f32_f16_kernel, mmctx->vtcm_act_stride) +QUANTIZE_IMPL(quantize_f16_f16, "quantize-f16-f16", quantize_f16_f16_kernel, mmctx->vtcm_act_stride) static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, void * data) { (void) nth; @@ -740,178 +738,209 @@ static void quantize_f32_q8_1_s16_tiled_block(unsigned int nth, unsigned int ith htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); } -// q8_1 for weight types with offsets (q8_1_s16 for Q2_K), otherwise q8_0 -static inline worker_callback_t htp_mm_act_quant_row_func(int weight_type) { +static void quantize_f32_f16_tiled_block(unsigned int nth, unsigned int ith, void * data) { + (void) nth; + struct htp_mm_context * mmctx = data; + if (mmctx->quant_ib_first[ith] >= mmctx->quant_ib_last[ith]) { + return; + } + struct htp_ops_context * octx = mmctx->octx; + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); + + const struct htp_tensor * src = mmctx->act; + + quantize_f32_f16_tiled_block_kernel( + (const float *) mmctx->vtcm_act_raw, + mmctx->vtcm_act, + NULL, + src->ne[0], + mmctx->quant_ib_first[ith], + mmctx->quant_ib_last[ith], + mmctx->vtcm_act_raw_stride, + htp_mm_f16_tiled_row_size(src->ne[0]), + mmctx->quant_r[ith], + mmctx->quant_c[ith] + ); + + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); +} + +static void quantize_f16_f16_tiled_block(unsigned int nth, unsigned int ith, void * data) { + (void) nth; + struct htp_mm_context * mmctx = data; + if (mmctx->quant_ib_first[ith] >= mmctx->quant_ib_last[ith]) { + return; + } + struct htp_ops_context * octx = mmctx->octx; + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); + + const struct htp_tensor * src = mmctx->act; + + quantize_f16_f16_tiled_block_kernel( + (const __fp16 *) mmctx->vtcm_act_raw, + mmctx->vtcm_act, + NULL, + src->ne[0], + mmctx->quant_ib_first[ith], + mmctx->quant_ib_last[ith], + mmctx->vtcm_act_raw_stride, + htp_mm_f16_tiled_row_size(src->ne[0]), + mmctx->quant_r[ith], + mmctx->quant_c[ith] + ); + + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); +} + +// q8_1 for weight types with offsets (q8_1_s16 for Q2_K), f16_tiled for F16, otherwise q8_0 +static inline worker_callback_t htp_mm_act_quant_row_func(int weight_type, int act_type) { + if (weight_type == HTP_TYPE_F16) { + return (act_type == HTP_TYPE_F16) ? quantize_f16_f16_tiled : quantize_f32_f16_tiled; + } if (weight_type == HTP_TYPE_Q2_K) { return quantize_f32_q8_1_s16_tiled; } return htp_mm_weight_has_offset(weight_type) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; } -static inline worker_callback_t htp_mm_act_quant_block_func(int weight_type) { +static inline worker_callback_t htp_mm_act_quant_block_func(int weight_type, int act_type) { + if (weight_type == HTP_TYPE_F16) { + return (act_type == HTP_TYPE_F16) ? quantize_f16_f16_tiled_block : quantize_f32_f16_tiled_block; + } if (weight_type == HTP_TYPE_Q2_K) { return quantize_f32_q8_1_s16_tiled_block; } return htp_mm_weight_has_offset(weight_type) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; } -MATVEC_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x1) -MATVEC_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x1) -MATVEC_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x1) -MATVEC_2D_REPACKED_IMPL(q5_k, 768, tiled_vec_dot_q5_k_32x1) -MATVEC_2D_REPACKED_IMPL(q6_k, 896, tiled_vec_dot_q6_k_32x1) -MATVEC_2D_REPACKED_IMPL(q3_k, 512, tiled_vec_dot_q3_k_32x1) -MATVEC_2D_REPACKED_IMPL(q2_k, 336, tiled_vec_dot_q2_k_32x1) -MATVEC_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x1) -MATVEC_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x1) - -MATMUL_NX_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) -MATMUL_NX_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) -MATMUL_NX_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) -MATMUL_NX_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) -MATMUL_NX_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) -MATMUL_NX_2D_REPACKED_IMPL(q5_k, 768, tiled_vec_dot_q5_k_32x2, tiled_vec_dot_q5_k_32x1) -MATMUL_NX_2D_REPACKED_IMPL(q3_k, 512, tiled_vec_dot_q3_k_32x2, tiled_vec_dot_q3_k_32x1) -MATMUL_NX_2D_REPACKED_IMPL(q2_k, 336, tiled_vec_dot_q2_k_32x2, tiled_vec_dot_q2_k_32x1) - -#define MATMUL_4D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ -static void hvx_mm_4d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ - htp_matmul_preamble; \ - \ - const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; \ - const uint32_t cur_m_rows = mmctx->cur_m_rows ? mmctx->cur_m_rows : (ne11 * ne12 * ne13); \ - const uint32_t cur_m_start = mmctx->cur_m_start; \ - \ - const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; \ - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); \ - \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - \ - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ - const uint32_t n_prefetch = kparams->n_prefetch; \ - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ - \ - const size_t dst_row_size = nb1; \ - const size_t act_stride = mmctx->vtcm_act_stride; \ - \ - uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ - uint8_t * restrict act_data = mmctx->vtcm_act; \ - \ - const uint32_t tile_size = TILE_SIZE; \ - const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ - \ - const uint32_t n_k_tiles_w = ne00 / 32; \ - const uint32_t n_k_tiles_a = ne10 / 32; \ - const uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ - const uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ - const uint32_t src0_slice_stride = ((ne01 + 31) / 32) * tile_row_stride; \ - \ - const uint32_t ct_start = src0_start_row / 32; \ - const uint32_t ct_end = (src0_end_row + 31) / 32; \ - \ - \ - if (src0_start_row >= src0_end_row || cur_m_rows == 0) { \ - return; \ - } \ - \ - const uint32_t total_batches = ne12 * ne13; \ - const uint32_t b_start = fastdiv(cur_m_start, &kparams->div_ne1); \ - uint32_t b_end = fastdiv(cur_m_start + cur_m_rows + ne11 - 1, &kparams->div_ne1); \ - b_end = MIN(b_end, total_batches); \ - \ - uint32_t b_grp_start = b_start; \ - while (b_grp_start < b_end) { \ - const uint32_t b3 = fastdiv(b_grp_start, &kparams->div_ne12); \ - const uint32_t b2 = b_grp_start - b3 * ne12; \ - const uint32_t i02 = fastdiv(b2, &kparams->div_r2); \ - const uint32_t i03 = fastdiv(b3, &kparams->div_r3); \ - \ - uint32_t b_grp_end = b_grp_start + 1; \ - while (b_grp_end < b_end) { \ - const uint32_t cur_b3 = fastdiv(b_grp_end, &kparams->div_ne12); \ - const uint32_t cur_b2 = b_grp_end - cur_b3 * ne12; \ - const uint32_t cur_i02 = fastdiv(cur_b2, &kparams->div_r2); \ - const uint32_t cur_i03 = fastdiv(cur_b3, &kparams->div_r3); \ - if (cur_i02 != i02 || cur_i03 != i03) { \ - break; \ - } \ - b_grp_end++; \ - } \ - \ - const uint32_t slice_idx = i03 * ne02 + i02; \ - const dma_addr_t src0_slice = src0->data + (size_t) slice_idx * src0_slice_stride; \ - \ - uint32_t push_ct = ct_start; \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ - dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ - src0_slice + (size_t) push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - } \ - \ - for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ - const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; \ - \ - int valid_rows = (int)ne0 - (int)(ct * 32); \ - valid_rows = MIN(32, MAX(0, valid_rows)); \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - for (uint32_t b = b_grp_start; b < b_grp_end; b++) { \ - const uint32_t b_m_start = b * ne11; \ - const uint32_t m_first = MAX(cur_m_start, b_m_start); \ - const uint32_t m_last = MIN(cur_m_start + cur_m_rows, b_m_start + ne11); \ - if (m_first >= m_last) continue; \ - \ - const uint32_t cur_b3 = fastdiv(b, &kparams->div_ne12); \ - const uint32_t cur_b2 = b - cur_b3 * ne12; \ - uint8_t * dst_batch_base = (uint8_t *) dst->data + (size_t) cur_b2 * nb2 + (size_t) cur_b3 * nb3; \ - \ - const uint32_t chunk_m_offset = m_first - cur_m_start; \ - const uint32_t dst_m_offset = m_first - b_m_start; \ - const uint32_t batch_nrows = m_last - m_first; \ - \ - uint32_t ir1 = 0; \ - for (; ir1 + 1 < batch_nrows; ir1 += 2) { \ - const uint8_t * restrict act_col0 = (const uint8_t *) (act_data + (chunk_m_offset + ir1 + 0) * act_stride); \ - const uint8_t * restrict act_col1 = (const uint8_t *) (act_data + (chunk_m_offset + ir1 + 1) * act_stride); \ - float * restrict dst_row0 = (float *) (dst_batch_base + (dst_m_offset + ir1 + 0) * dst_row_size); \ - float * restrict dst_row1 = (float *) (dst_batch_base + (dst_m_offset + ir1 + 1) * dst_row_size); \ - float * dst_ptr0 = &dst_row0[ct * 32]; \ - float * dst_ptr1 = &dst_row1[ct * 32]; \ - DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, act_col0, act_col1, valid_rows, NULL, NULL); \ - } \ - for (; ir1 < batch_nrows; ++ir1) { \ - const uint8_t * restrict act_col = (const uint8_t *) (act_data + (chunk_m_offset + ir1) * act_stride); \ - float * restrict dst_row = (float *) (dst_batch_base + (dst_m_offset + ir1) * dst_row_size); \ - float * dst_ptr = &dst_row[ct * 32]; \ - DOT_2X1(ne10, dst_ptr, w_tile, act_col, valid_rows, NULL); \ - } \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - \ - if (push_ct < ct_end) { \ - dma_queue_push(dma_q, dma_make_data(w_tile, src0_slice + (size_t) push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - push_ct++; \ - } \ - } \ - b_grp_start = b_grp_end; \ - } \ - if (src2) { \ - hvx_tensor_add_f32_grid(dst, src2, cur_m_start, cur_m_start + cur_m_rows, src0_start_row, src0_end_row, &kparams->div_ne12_ne1, &kparams->div_ne1); \ - } \ +static void hvx_mm_4d_tiled(unsigned int nth, unsigned int ith, void * data) { + htp_matmul_preamble; + + const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; + const uint32_t cur_m_rows = mmctx->cur_m_rows ? mmctx->cur_m_rows : (ne11 * ne12 * ne13); + const uint32_t cur_m_start = mmctx->cur_m_start; + + const uint32_t src0_start_row = mmctx->src0_row_start + src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, mmctx->src0_row_end); + + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; + + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_prefetch = kparams->n_prefetch; + assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); + + const size_t dst_row_size = nb1; + const size_t act_stride = mmctx->vtcm_act_stride; + + uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; + uint8_t * restrict act_data = mmctx->vtcm_act; + + const uint32_t tile_size = htp_tensor_get_tile_size(src0->type); + const uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(src0->type); + + const uint32_t n_k_tiles_w = ne00 / 32; + const uint32_t n_k_tiles_a = ne10 / 32; + const uint32_t tile_row_stride = n_k_tiles_w * tile_size; + const uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; + const uint32_t src0_slice_stride = ((ne01 + 31) / 32) * tile_row_stride; + + const uint32_t ct_start = src0_start_row / 32; + const uint32_t ct_end = (src0_end_row + 31) / 32; + + if (src0_start_row >= src0_end_row || cur_m_rows == 0) { + return; + } + + const uint32_t total_batches = ne12 * ne13; + const uint32_t b_start = fastdiv(cur_m_start, &kparams->div_ne1); + uint32_t b_end = fastdiv(cur_m_start + cur_m_rows + ne11 - 1, &kparams->div_ne1); + b_end = MIN(b_end, total_batches); + + uint32_t b_grp_start = b_start; + while (b_grp_start < b_end) { + const uint32_t b3 = fastdiv(b_grp_start, &kparams->div_ne12); + const uint32_t b2 = b_grp_start - b3 * ne12; + const uint32_t i02 = fastdiv(b2, &kparams->div_r2); + const uint32_t i03 = fastdiv(b3, &kparams->div_r3); + + uint32_t b_grp_end = b_grp_start + 1; + while (b_grp_end < b_end) { + const uint32_t cur_b3 = fastdiv(b_grp_end, &kparams->div_ne12); + const uint32_t cur_b2 = b_grp_end - cur_b3 * ne12; + const uint32_t cur_i02 = fastdiv(cur_b2, &kparams->div_r2); + const uint32_t cur_i03 = fastdiv(cur_b3, &kparams->div_r3); + if (cur_i02 != i02 || cur_i03 != i03) { + break; + } + b_grp_end++; + } + + const uint32_t slice_idx = i03 * ne02 + i02; + const dma_addr_t src0_slice = src0->data + (size_t) slice_idx * src0_slice_stride; + + uint32_t push_ct = ct_start; + for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { + dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, + src0_slice + (size_t) push_ct * tile_row_stride), + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + } + + for (uint32_t ct = ct_start; ct < ct_end; ct++) { + const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; + + int valid_rows = (int)ne0 - (int)(ct * 32); + valid_rows = MIN(32, MAX(0, valid_rows)); + + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); + for (uint32_t b = b_grp_start; b < b_grp_end; b++) { + const uint32_t b_m_start = b * ne11; + const uint32_t m_first = MAX(cur_m_start, b_m_start); + const uint32_t m_last = MIN(cur_m_start + cur_m_rows, b_m_start + ne11); + if (m_first >= m_last) continue; + + const uint32_t cur_b3 = fastdiv(b, &kparams->div_ne12); + const uint32_t cur_b2 = b - cur_b3 * ne12; + uint8_t * dst_batch_base = (uint8_t *) dst->data + (size_t) cur_b2 * nb2 + (size_t) cur_b3 * nb3; + + const uint32_t chunk_m_offset = m_first - cur_m_start; + const uint32_t dst_m_offset = m_first - b_m_start; + const uint32_t batch_nrows = m_last - m_first; + + uint32_t ir1 = 0; + for (; ir1 + 1 < batch_nrows; ir1 += 2) { + const uint8_t * restrict act_col0 = (const uint8_t *) (act_data + (chunk_m_offset + ir1 + 0) * act_stride); + const uint8_t * restrict act_col1 = (const uint8_t *) (act_data + (chunk_m_offset + ir1 + 1) * act_stride); + float * restrict dst_row0 = (float *) (dst_batch_base + (dst_m_offset + ir1 + 0) * dst_row_size); + float * restrict dst_row1 = (float *) (dst_batch_base + (dst_m_offset + ir1 + 1) * dst_row_size); + float * dst_ptr0 = &dst_row0[ct * 32]; + float * dst_ptr1 = &dst_row1[ct * 32]; + mmctx->vec_dot_32x2(ne10, dst_ptr0, dst_ptr1, w_tile, act_col0, act_col1, valid_rows, NULL, NULL); + } + for (; ir1 < batch_nrows; ++ir1) { + const uint8_t * restrict act_col = (const uint8_t *) (act_data + (chunk_m_offset + ir1) * act_stride); + float * restrict dst_row = (float *) (dst_batch_base + (dst_m_offset + ir1) * dst_row_size); + float * dst_ptr = &dst_row[ct * 32]; + mmctx->vec_dot_32x1(ne10, dst_ptr, w_tile, act_col, valid_rows, NULL); + } + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); + + if (push_ct < ct_end) { + dma_queue_push(dma_q, dma_make_data(w_tile, src0_slice + (size_t) push_ct * tile_row_stride), + aligned_tile_size, tile_size, tile_size, n_k_tiles_a); + push_ct++; + } + } + b_grp_start = b_grp_end; + } + if (src2) { + hvx_tensor_add_f32_grid(dst, src2, cur_m_start, cur_m_start + cur_m_rows, src0_start_row, src0_end_row, &kparams->div_ne12_ne1, &kparams->div_ne1); + } } -MATMUL_4D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) -MATMUL_4D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) -MATMUL_4D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) -MATMUL_4D_REPACKED_IMPL(q6_k, 896, tiled_vec_dot_q6_k_32x2, tiled_vec_dot_q6_k_32x1) -MATMUL_4D_REPACKED_IMPL(q5_k, 768, tiled_vec_dot_q5_k_32x2, tiled_vec_dot_q5_k_32x1) -MATMUL_4D_REPACKED_IMPL(q3_k, 512, tiled_vec_dot_q3_k_32x2, tiled_vec_dot_q3_k_32x1) -MATMUL_4D_REPACKED_IMPL(q2_k, 336, tiled_vec_dot_q2_k_32x2, tiled_vec_dot_q2_k_32x1) -MATMUL_4D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) -MATMUL_4D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) - -static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { +static void hvx_mm_2d_flat(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; @@ -1014,7 +1043,7 @@ static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { } } -static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { +static void hvx_mv_2d_flat(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; const uint32_t src0_nrows = mmctx->src0_row_end - mmctx->src0_row_start; @@ -1119,7 +1148,7 @@ static void hvx_mv_2d(unsigned int nth, unsigned int ith, void * data) { } } -static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) { +static void hvx_mm_4d_flat(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; @@ -1268,7 +1297,7 @@ static void hvx_mm_4d(unsigned int nth, unsigned int ith, void * data) { #define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id) * mmctx->mapping_stride + (i1)] -static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { +static void hvx_mm_id_tiled(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; const struct htp_tensor * restrict ids = octx->src[2]; @@ -1312,8 +1341,8 @@ static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { const dma_addr_t src0_row = src0->data + cur_a * nb02; - const uint32_t tile_size = htp_mm_get_weight_tile_size(src0->type); - const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(src0->type); + const uint32_t tile_size = htp_tensor_get_tile_size(src0->type); + const uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(src0->type); const uint32_t n_k_tiles_w = ne00 / 32; const uint32_t n_k_tiles_a = ne10 / 32; const uint32_t tile_row_stride = n_k_tiles_w * tile_size; @@ -1331,7 +1360,7 @@ static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { for (uint32_t ct = ct_start; ct < ct_end; ct++) { const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; - int valid_rows = (int)ne01 - (int)(ct * 32); + int valid_rows = (int)ne0 - (int)(ct * 32); valid_rows = MIN(32, MAX(0, valid_rows)); htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); @@ -1357,7 +1386,7 @@ static void hvx_mm_id(unsigned int nth, unsigned int ith, void * data) { } } -static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) { +static void hvx_mv_id_tiled(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; const struct htp_tensor * restrict ids = octx->src[2]; @@ -1400,8 +1429,8 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) { const uint8_t * restrict act_col = (const uint8_t *) act_data; float * restrict dst_row = (float *) (dst->data + ie1 * nb1); - const uint32_t tile_size = htp_mm_get_weight_tile_size(src0->type); - const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(src0->type); + const uint32_t tile_size = htp_tensor_get_tile_size(src0->type); + const uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(src0->type); const uint32_t n_k_tiles_w = ne00 / 32; const uint32_t n_k_tiles_a = ne10 / 32; const uint32_t tile_row_stride = n_k_tiles_w * tile_size; @@ -1419,7 +1448,7 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) { for (uint32_t ct = ct_start; ct < ct_end; ct++) { const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; - int valid_rows = (int)ne01 - (int)(ct * 32); + int valid_rows = (int)ne0 - (int)(ct * 32); valid_rows = MIN(32, MAX(0, valid_rows)); htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); @@ -1435,7 +1464,7 @@ static void hvx_mv_id(unsigned int nth, unsigned int ith, void * data) { } } -static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { +static void hvx_mv_id_nx_tiled(unsigned int nth, unsigned int ith, void * data) { struct htp_mm_context * mmctx = (struct htp_mm_context *) data; struct htp_ops_context * octx = mmctx->octx; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; @@ -1489,8 +1518,8 @@ static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { const uint8_t * restrict act_col = (const uint8_t *) act_data; float * restrict dst_row = (float *) (dst->data + ie1 * dst->nb[1]); - const uint32_t tile_size = htp_mm_get_weight_tile_size(src_w->type); - const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(src_w->type); + const uint32_t tile_size = htp_tensor_get_tile_size(src_w->type); + const uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(src_w->type); const uint32_t n_k_tiles_w = src_w->ne[0] / 32; const uint32_t n_k_tiles_a = act->ne[0] / 32; const uint32_t tile_row_stride = n_k_tiles_w * tile_size; @@ -1508,7 +1537,7 @@ static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { for (uint32_t ct = ct_start; ct < ct_end; ct++) { const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; - int valid_rows = (int)src_w->ne[1] - (int)(ct * 32); + int valid_rows = (int)dst->ne[0] - (int)(ct * 32); valid_rows = MIN(32, MAX(0, valid_rows)); htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); @@ -1525,7 +1554,7 @@ static void hvx_mv_id_nx(unsigned int nth, unsigned int ith, void * data) { } } -static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) { +static void hvx_mm_id_nx_tiled(unsigned int nth, unsigned int ith, void * data) { struct htp_mm_context * mmctx = (struct htp_mm_context *) data; struct htp_ops_context * octx = mmctx->octx; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; @@ -1580,8 +1609,8 @@ static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) { const dma_addr_t src0_row = src_w->data + cur_a * src_w->nb[2]; - const uint32_t tile_size = htp_mm_get_weight_tile_size(src_w->type); - const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(src_w->type); + const uint32_t tile_size = htp_tensor_get_tile_size(src_w->type); + const uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(src_w->type); const uint32_t n_k_tiles_w = src_w->ne[0] / 32; const uint32_t n_k_tiles_a = act->ne[0] / 32; const uint32_t tile_row_stride = n_k_tiles_w * tile_size; @@ -1599,7 +1628,7 @@ static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) { for (uint32_t ct = ct_start; ct < ct_end; ct++) { const uint8_t * w_tile = (void *) dma_queue_pop(dma_q).dst; - int valid_rows = (int)src_w->ne[1] - (int)(ct * 32); + int valid_rows = (int)dst->ne[0] - (int)(ct * 32); valid_rows = MIN(32, MAX(0, valid_rows)); htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); @@ -1626,44 +1655,58 @@ static void hvx_mm_id_nx(unsigned int nth, unsigned int ith, void * data) { } } -static int hvx_mm_init_vec_dot(struct htp_mm_context * mmctx, enum htp_data_type type) { +static int hvx_mm_init_tiled_vec_dot(struct htp_mm_context * mmctx, enum htp_data_type type) { switch (type) { case HTP_TYPE_Q4_0: mmctx->type = "q4_0_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q4_0_32x1; + mmctx->vec_dot_32x2 = tiled_vec_dot_q4_0_32x2; return 0; case HTP_TYPE_Q4_1: case HTP_TYPE_Q4_K: mmctx->type = "q4_1_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q4_1_32x1; + mmctx->vec_dot_32x2 = tiled_vec_dot_q4_1_32x2; return 0; case HTP_TYPE_Q8_0: mmctx->type = "q8_0_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q8_0_32x1; + mmctx->vec_dot_32x2 = tiled_vec_dot_q8_0_32x2; return 0; case HTP_TYPE_Q5_K: mmctx->type = "q5_k_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q5_k_32x1; + mmctx->vec_dot_32x2 = tiled_vec_dot_q5_k_32x2; return 0; case HTP_TYPE_Q6_K: mmctx->type = "q6_k_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q6_k_32x1; + mmctx->vec_dot_32x2 = tiled_vec_dot_q6_k_32x2; return 0; case HTP_TYPE_Q3_K: mmctx->type = "q3_k_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q3_k_32x1; + mmctx->vec_dot_32x2 = tiled_vec_dot_q3_k_32x2; return 0; case HTP_TYPE_Q2_K: mmctx->type = "q2_k_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_q2_k_32x1; + mmctx->vec_dot_32x2 = tiled_vec_dot_q2_k_32x2; return 0; case HTP_TYPE_IQ4_NL: mmctx->type = "iq4nl_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_iq4nl_32x1; + mmctx->vec_dot_32x2 = tiled_vec_dot_iq4nl_32x2; return 0; case HTP_TYPE_MXFP4: mmctx->type = "mxfp4_tiled-f32"; mmctx->vec_dot_32x1 = tiled_vec_dot_mxfp4_32x1; + mmctx->vec_dot_32x2 = tiled_vec_dot_mxfp4_32x2; + return 0; + case HTP_TYPE_F16: + mmctx->type = "f16_tiled-f16"; + mmctx->vec_dot_32x1 = tiled_vec_dot_f16_32x1; + mmctx->vec_dot_32x2 = tiled_vec_dot_f16_32x2; return 0; default: return -1; @@ -1705,15 +1748,11 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { mmctx->src0_row_start = src0_row_start; mmctx->src0_row_end = src0_row_end; - bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 || - src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL || - src0->type == HTP_TYPE_MXFP4 || src0->type == HTP_TYPE_Q6_K || - src0->type == HTP_TYPE_Q4_K || src0->type == HTP_TYPE_Q5_K || - src0->type == HTP_TYPE_Q3_K || src0->type == HTP_TYPE_Q2_K); + const bool is_tiled = htp_tensor_is_tiled(src0); // Compute src0_nrows_per_thread mmctx->src0_nrows_per_thread = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div); - if (is_repacked) { + if (is_tiled) { mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); } else { mmctx->src0_nrows_per_thread += (mmctx->src0_nrows_per_thread & 1); // round up to even @@ -1730,59 +1769,11 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { uint32_t n_quant_tasks = 1; const bool is_batched = (ne12 > 1 || ne13 > 1 || ne02 > 1 || ne03 > 1); if (is_batched) { - if (is_repacked) { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_4d_repacked_q4_0; break; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: matmul_job_func = hvx_mm_4d_repacked_q4_1; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_4d_repacked_q8_0; break; - case HTP_TYPE_Q6_K: matmul_job_func = hvx_mm_4d_repacked_q6_k; break; - case HTP_TYPE_Q5_K: matmul_job_func = hvx_mm_4d_repacked_q5_k; break; - case HTP_TYPE_Q3_K: matmul_job_func = hvx_mm_4d_repacked_q3_k; break; - case HTP_TYPE_Q2_K: matmul_job_func = hvx_mm_4d_repacked_q2_k; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_4d_repacked_iq4nl; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_4d_repacked_mxfp4; break; - default: return HTP_STATUS_NO_SUPPORT; - } - } else { - matmul_job_func = hvx_mm_4d; - } + matmul_job_func = is_tiled ? hvx_mm_4d_tiled : hvx_mm_4d_flat; } else if (act_nrows > 1) { - if (is_repacked) { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_2d_repacked_q4_0; break; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: matmul_job_func = hvx_mm_2d_repacked_q4_1; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_2d_repacked_q8_0; break; - case HTP_TYPE_Q6_K: matmul_job_func = hvx_mm_2d_repacked_q6_k; break; - case HTP_TYPE_Q5_K: matmul_job_func = hvx_mm_2d_repacked_q5_k; break; - case HTP_TYPE_Q3_K: matmul_job_func = hvx_mm_2d_repacked_q3_k; break; - case HTP_TYPE_Q2_K: matmul_job_func = hvx_mm_2d_repacked_q2_k; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_2d_repacked_iq4nl; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_2d_repacked_mxfp4; break; - default: return HTP_STATUS_NO_SUPPORT; - } - } else { - matmul_job_func = hvx_mm_2d; - } + matmul_job_func = is_tiled ? hvx_mm_2d_tiled : hvx_mm_2d_flat; } else { - if (is_repacked) { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mv_2d_repacked_q4_0; break; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: matmul_job_func = hvx_mv_2d_repacked_q4_1; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mv_2d_repacked_q8_0; break; - case HTP_TYPE_Q5_K: matmul_job_func = hvx_mv_2d_repacked_q5_k; break; - case HTP_TYPE_Q6_K: matmul_job_func = hvx_mv_2d_repacked_q6_k; break; - case HTP_TYPE_Q3_K: matmul_job_func = hvx_mv_2d_repacked_q3_k; break; - case HTP_TYPE_Q2_K: matmul_job_func = hvx_mv_2d_repacked_q2_k; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mv_2d_repacked_iq4nl; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mv_2d_repacked_mxfp4; break; - default: return HTP_STATUS_NO_SUPPORT; - } - } else { - matmul_job_func = hvx_mv_2d; - } + matmul_job_func = is_tiled ? hvx_mv_2d_tiled : hvx_mv_2d_flat; } bool need_quant = true; @@ -1792,9 +1783,9 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { quant_task_func = (act->type == HTP_TYPE_F32) ? quantize_f32_f16 : quantize_f16_f16; need_quant = (act->type == HTP_TYPE_F32); mmctx->type = (act->type == HTP_TYPE_F32) ? "f32-f16" : "f16-f16"; - mmctx->vec_dot_1x1 = vec_dot_f16_f16_aa_1x1; - mmctx->vec_dot_2x1 = vec_dot_f16_f16_aa_2x1; - mmctx->vec_dot_2x2 = vec_dot_f16_f16_aa_2x2; + mmctx->vec_dot_1x1 = flat_vec_dot_f16_f16_aa_1x1; + mmctx->vec_dot_2x1 = flat_vec_dot_f16_f16_aa_2x1; + mmctx->vec_dot_2x2 = flat_vec_dot_f16_f16_aa_2x2; act_row_size = kparams->act_row_size; break; @@ -1802,16 +1793,26 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { quant_task_func = NULL; need_quant = false; mmctx->type = "f32-f32"; - mmctx->vec_dot_1x1 = vec_dot_f32_f32_aa_1x1; - mmctx->vec_dot_2x1 = vec_dot_f32_f32_aa_2x1; - mmctx->vec_dot_2x2 = vec_dot_f32_f32_aa_2x2; + mmctx->vec_dot_1x1 = flat_vec_dot_f32_f32_aa_1x1; + mmctx->vec_dot_2x1 = flat_vec_dot_f32_f32_aa_2x1; + mmctx->vec_dot_2x2 = flat_vec_dot_f32_f32_aa_2x2; + act_row_size = kparams->act_row_size; + break; + + case HTP_MM_KERNEL_HVX_BF16_F32_VTCM: + quant_task_func = NULL; + need_quant = false; + mmctx->type = "bf16-f32"; + mmctx->vec_dot_1x1 = flat_vec_dot_bf16_f32_aa_1x1; + mmctx->vec_dot_2x1 = flat_vec_dot_bf16_f32_aa_2x1; + mmctx->vec_dot_2x2 = flat_vec_dot_bf16_f32_aa_2x2; act_row_size = kparams->act_row_size; break; case HTP_MM_KERNEL_HVX_QUANT_BLOCK: case HTP_MM_KERNEL_HVX_QUANT_ROW: - default: - if (hvx_mm_init_vec_dot(mmctx, src0->type) != 0) { + default: { + if (hvx_mm_init_tiled_vec_dot(mmctx, src0->type) != 0) { return HTP_STATUS_NO_SUPPORT; } @@ -1821,7 +1822,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { if (act_nrows < octx->n_threads && !is_batched) { n_quant_tasks = MIN(total_nb, octx->n_threads); - quant_task_func = htp_mm_act_quant_block_func(src0->type); + quant_task_func = htp_mm_act_quant_block_func(src0->type, act->type); for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { uint32_t ib_first = (total_nb * ith) / n_quant_tasks; uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks; @@ -1832,10 +1833,11 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { } } else { n_quant_tasks = MIN(act_nrows, octx->n_threads); - quant_task_func = htp_mm_act_quant_row_func(src0->type); + quant_task_func = htp_mm_act_quant_row_func(src0->type, act->type); } act_row_size = kparams->act_row_size; break; + } } const uint32_t m_chunk = (kparams->m_chunk > 0 && (uint32_t) kparams->m_chunk < act_nrows) @@ -1848,6 +1850,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM || + kparams->kernel_type == HTP_MM_KERNEL_HVX_BF16_F32_VTCM || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { mmctx->vtcm_act_size_per_thread = L.act_bytes; @@ -1908,7 +1911,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { work_queue_func_t q_func; if (cur_m_rows < octx->n_threads && (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK || kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW)) { quant_tasks = MIN(total_nb, octx->n_threads); - q_func = htp_mm_act_quant_block_func(src0->type); + q_func = htp_mm_act_quant_block_func(src0->type, act->type); for (uint32_t ith = 0; ith < quant_tasks; ++ith) { uint32_t ib_first = (total_nb * ith) / quant_tasks; uint32_t ib_last = (total_nb * (ith + 1)) / quant_tasks; @@ -1949,158 +1952,18 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { return HTP_STATUS_OK; } -static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) { - struct htp_mm_context * mmctx = data; - struct htp_ops_context * octx = mmctx->octx; - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; - const uint32_t n_weights = kparams->n_weights; - - const struct htp_tensor * restrict act = octx->src[n_weights]; - const uint32_t act_nrows = act->ne[1] * act->ne[2] * act->ne[3]; - const size_t act_stride = mmctx->vtcm_act_stride; - - uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; - uint8_t * restrict act_data = mmctx->vtcm_act; - - const uint32_t n_prefetch = kparams->n_prefetch; - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); - const uint32_t prefetch_mask = n_prefetch - 1; - - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; - - - for (uint32_t widx = 0; widx < n_weights; widx++) { - const struct htp_tensor * restrict src_w = octx->src[widx]; - const struct htp_tensor * restrict dst = octx->dsts[widx]; - if (!src_w || !dst) continue; - dma_queue * dma_q = octx->ctx->dma[ith]; - - const uint32_t ne00 = src_w->ne[0]; - const uint32_t ne01 = src_w->ne[1]; - uint32_t start_row = 0; - uint32_t end_row = ne01; - if (octx->ctx->mdev.count > 1) { - const bool can_split = htp_tensor_can_row_partition(dst, sizeof(float)); - const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(ne01, can_split ? 32 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div); - start_row = range.start; - end_row = range.start + range.count; - } - - const uint32_t nrows = end_row - start_row; - uint32_t src0_nrows_per_thread = fastdiv(nrows + nth - 1, &octx->n_threads_div); - src0_nrows_per_thread += (src0_nrows_per_thread & 1); - - const uint32_t src0_start_row = start_row + src0_nrows_per_thread * ith; - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, end_row); - const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U); - if (src0_start_row >= src0_end_row) continue; - - const size_t dst_row_size = dst->nb[1]; - const size_t src0_row_size = src_w->nb[1]; - const size_t src0_stride = hex_round_up(src0_row_size, 128); - - const dma_addr_t src0_row = src_w->data; - - for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { - const int is0 = (ir0 - src0_start_row); - if (is0 >= (int)n_prefetch) break; - dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 2); - } - - for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { - const uint8_t * ss0 = (void *) dma_queue_pop(dma_q).dst; - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); - uint32_t ir1 = 0; - for (; ir1 + 1 < act_nrows; ir1 += 2) { - const uint8_t * restrict act_col0 = (const uint8_t *) (act_data + (ir1+0) * act_stride); - const uint8_t * restrict act_col1 = (const uint8_t *) (act_data + (ir1+1) * act_stride); - float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); - float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); - mmctx->vec_dot_2x2(ne00, &dst_row0[ir0], &dst_row1[ir0], ss0, ss0 + src0_stride, act_col0, act_col1); - } - for (; ir1 < act_nrows; ++ir1) { - const uint8_t * restrict act_col = (const uint8_t *) (act_data + ir1 * act_stride); - float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); - mmctx->vec_dot_2x1(ne00, &dst_row[ir0], ss0, ss0 + src0_stride, act_col); - } - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); - - const int pr0 = (ir0 + n_prefetch); - const int is0 = (pr0 - src0_start_row) & prefetch_mask; - if (pr0 < src0_end_row_x2) { - dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 2); - } - } - - if (src0_end_row != src0_end_row_x2) { - uint32_t ir0 = src0_end_row_x2; - const int is0 = (ir0 - src0_start_row) & prefetch_mask; - dma_queue_push(dma_q, dma_make_data(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 1); - const uint8_t * ss0 = (void *) dma_queue_pop(dma_q).dst; - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); - for (uint32_t ir1 = 0; ir1 < act_nrows; ++ir1) { - const uint8_t * restrict act_col = (const uint8_t *) (act_data + ir1 * act_stride); - float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); - mmctx->vec_dot_1x1(ne00, &dst_row[ir0], ss0, act_col); - } - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); - } - } -} - -#define DEQUANTIZE_WORKER_LOOP_IMPL(SUFFIX) \ -static void dequantize_tiled_worker_loop_##SUFFIX(unsigned int n, unsigned int i, void *data) { \ - tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data; \ - struct htp_thread_trace * tr = &state->traces[i]; \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i); \ - for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) { \ - int start = task_id * state->n_tiles_per_task; \ - int end = hex_smin(start + state->n_tiles_per_task, state->n_tot_tiles); \ - dequantize_tiled_weight_to_fp16_task_##SUFFIX(state, start, end); \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i); \ -} - -DEQUANTIZE_WORKER_LOOP_IMPL(q4_0) -DEQUANTIZE_WORKER_LOOP_IMPL(q4_1) -DEQUANTIZE_WORKER_LOOP_IMPL(iq4_nl) -DEQUANTIZE_WORKER_LOOP_IMPL(mxfp4) -DEQUANTIZE_WORKER_LOOP_IMPL(q8_0) -DEQUANTIZE_WORKER_LOOP_IMPL(q6_k) -DEQUANTIZE_WORKER_LOOP_IMPL(q5_k) -DEQUANTIZE_WORKER_LOOP_IMPL(q3_k) -DEQUANTIZE_WORKER_LOOP_IMPL(q2_k) - -static void convert_f16_worker_loop(unsigned int n, unsigned int i, void *data) { - tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data; +static void dequantize_tiled_worker_loop(unsigned int n, unsigned int i, void * data) { + tiled_dequantize_state_t * state = (tiled_dequantize_state_t *) data; struct htp_thread_trace * tr = &state->traces[i]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i); - for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) { + for (unsigned int task_id = i; task_id < (unsigned int) state->n_tasks; task_id += n) { int start = task_id * state->n_tiles_per_task; int end = hex_smin(start + state->n_tiles_per_task, state->n_tot_tiles); - convert_f16_weight_to_fp16_tiles_task(state, start, end); + state->task_fn(state, start, end); } htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_W_DEQUANT, i); } -static void quantize_f32_worker_loop(unsigned int n, unsigned int i, void *data) { - tiled_dequantize_state_t *state = (tiled_dequantize_state_t *)data; - - struct htp_thread_trace * tr = &state->traces[i]; - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, i); - - for (unsigned int task_id = i; task_id < (unsigned int)state->n_tasks; task_id += n) { - int start = task_id * state->n_tiles_per_task; - int end = hex_smin(start + state->n_tiles_per_task, state->n_tot_tiles); - quantize_f32_weight_to_fp16_tiles_task(state, start, end); - } - - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_A_QUANT, i); -} - static void transfer_output_chunk_worker_fn(unsigned int n, unsigned int i, void *data) { output_transfer_task_state_t *st = (output_transfer_task_state_t *) data; @@ -2450,13 +2313,36 @@ static void transfer_output_chunk_scattered_worker_fn(unsigned int n, unsigned i // --- HMX Dispatchers & Entry Points --- +static inline tiled_dequant_task_fn_t hmx_get_dequant_worker_fn(int weight_type, bool is_tiled) { + switch (weight_type) { + case HTP_TYPE_Q4_0: return dequantize_tiled_weight_to_fp16_task_q4_0; + case HTP_TYPE_IQ4_NL: return dequantize_tiled_weight_to_fp16_task_iq4_nl; + case HTP_TYPE_Q4_1: + case HTP_TYPE_Q4_K: return dequantize_tiled_weight_to_fp16_task_q4_1; + case HTP_TYPE_MXFP4: return dequantize_tiled_weight_to_fp16_task_mxfp4; + case HTP_TYPE_Q8_0: return dequantize_tiled_weight_to_fp16_task_q8_0; + case HTP_TYPE_Q5_K: return dequantize_tiled_weight_to_fp16_task_q5_k; + case HTP_TYPE_Q6_K: return dequantize_tiled_weight_to_fp16_task_q6_k; + case HTP_TYPE_Q3_K: return dequantize_tiled_weight_to_fp16_task_q3_k; + case HTP_TYPE_Q2_K: return dequantize_tiled_weight_to_fp16_task_q2_k; + case HTP_TYPE_F16: return is_tiled ? NULL : convert_f16_weight_to_fp16_tiles_task; + case HTP_TYPE_BF16: return convert_bf16_weight_to_fp16_tiles_task; + case HTP_TYPE_F32: return quantize_f32_weight_to_fp16_tiles_task; + default: return NULL; + } +} + static void dequantize_tiled_weight_chunk_to_fp16_tiles( struct htp_context *ctx, __fp16 *vtcm_dst, const void *weight_src_ddr, int n_cols, int k_block, size_t row_stride, int weight_type, int n_k_tiles, struct fastdiv_values n_k_tiles_div, - worker_callback_t dequant_worker_fn, int n_threads) { + tiled_dequant_task_fn_t dequant_worker_fn, int n_threads) { + + if (!dequant_worker_fn) { + return; + } assert(k_block % HTP_MM_HMX_TILE_N_COLS == 0); @@ -2471,6 +2357,7 @@ static void dequantize_tiled_weight_chunk_to_fp16_tiles( state.n_tiles_per_task = n_tiles_per_task; state.dst = vtcm_dst; state.src = (const uint8_t *)weight_src_ddr; + state.task_fn = dequant_worker_fn; state.n_cols = n_cols; state.k_block = k_block; state.row_stride = row_stride; @@ -2480,14 +2367,14 @@ static void dequantize_tiled_weight_chunk_to_fp16_tiles( state.traces = ctx->trace; state.ctx = ctx; - state.tile_size = htp_mm_get_weight_tile_size(weight_type); - state.aligned_tile_size = htp_mm_get_weight_aligned_tile_size(weight_type); + state.tile_size = htp_tensor_get_tile_size(weight_type); + state.aligned_tile_size = htp_tensor_get_tile_padded_size(weight_type); if (state.n_tasks == 1 || n_threads == 1) { - dequant_worker_fn(1, 0, &state); + dequantize_tiled_worker_loop(1, 0, &state); } else { int n_tasks = hex_smin((int) state.n_tasks, n_threads); - worker_pool_run_func(ctx->worker_pool, dequant_worker_fn, &state, n_tasks); + worker_pool_run_func(ctx->worker_pool, dequantize_tiled_worker_loop, &state, n_tasks); } } @@ -2764,53 +2651,35 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, struct htp_thread_trace * tr = &ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); - // Quantized weights are repacked and padded to 32, so we it has to be 32-aligned. - // Only F16/F32 weights can be non-32-aligned, they will be padded in the following kernel. - const bool wtype_is_quant = (weight_type != HTP_TYPE_F16 && weight_type != HTP_TYPE_F32); - if (k % 32 != 0 || (wtype_is_quant && n % 32 != 0)) { + const bool is_tiled = (tile_size > 0); + + if (k % 32 != 0) { return -1; } - if (!hex_is_aligned(dst, VLEN) || (act_dma_addr & (VLEN - 1)) != 0) { return -1; } + if (!hex_is_aligned(dst, sizeof(float))) { return -1; } - size_t row_stride = htp_mm_get_tiled_row_stride(weight_type, k); + size_t row_stride = htp_tensor_get_row_size(weight_type, k); if (row_stride == 0) { return -1; } - worker_callback_t dequant_worker_fn = NULL; - switch (weight_type) { - case HTP_TYPE_Q4_0: dequant_worker_fn = dequantize_tiled_worker_loop_q4_0; break; - case HTP_TYPE_IQ4_NL: dequant_worker_fn = dequantize_tiled_worker_loop_iq4_nl; break; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: dequant_worker_fn = dequantize_tiled_worker_loop_q4_1; break; - case HTP_TYPE_MXFP4: dequant_worker_fn = dequantize_tiled_worker_loop_mxfp4; break; - case HTP_TYPE_Q8_0: dequant_worker_fn = dequantize_tiled_worker_loop_q8_0; break; - case HTP_TYPE_Q5_K: dequant_worker_fn = dequantize_tiled_worker_loop_q5_k; break; - case HTP_TYPE_Q6_K: dequant_worker_fn = dequantize_tiled_worker_loop_q6_k; break; - case HTP_TYPE_Q3_K: dequant_worker_fn = dequantize_tiled_worker_loop_q3_k; break; - case HTP_TYPE_Q2_K: dequant_worker_fn = dequantize_tiled_worker_loop_q2_k; break; - case HTP_TYPE_F16: dequant_worker_fn = convert_f16_worker_loop; break; - case HTP_TYPE_F32: dequant_worker_fn = quantize_f32_worker_loop; break; - default: - return -1; - } + tiled_dequant_task_fn_t dequant_worker_fn = hmx_get_dequant_worker_fn(weight_type, is_tiled); const int n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS; const struct fastdiv_values n_k_tiles_div = init_fastdiv_values(n_k_tiles); - const bool is_quant = (weight_type != HTP_TYPE_F16 && weight_type != HTP_TYPE_F32); const size_t vec_dot_size = k * sizeof(__fp16); const size_t vtcm_budget = ctx->vtcm_size; - const uint32_t dma_dst_stride = is_quant ? aligned_tile_size : row_stride; - const uint32_t dma_src_stride = is_quant ? tile_size : weight_stride; - const uint32_t dma_width_bytes = is_quant ? tile_size : row_stride; + const uint32_t dma_dst_stride = is_tiled ? aligned_tile_size : row_stride; + const uint32_t dma_src_stride = is_tiled ? tile_size : weight_stride; + const uint32_t dma_width_bytes = is_tiled ? tile_size : row_stride; size_t m_chunk_n_rows = m_chunk; size_t n_chunk_n_cols = n_chunk; size_t vtcm_used = vtcm_size; - const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0; + const size_t qweight_row_stride = is_tiled ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0; struct htp_mm_hmx_vtcm_layout L; htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, pipeline, act_threads, aligned_tile_size, src2_bytes); @@ -2827,7 +2696,6 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, VTCM_LAYOUT_PTR(__fp16, base, L.off_weight[0]), VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_weight[1], pipeline) }; - __fp16 *vtcm_f16_act = VTCM_LAYOUT_PTR(__fp16, base, L.off_act); float *vtcm_f32_act = VTCM_LAYOUT_PTR(float, base, L.off_act_f32); __fp16 *vtcm_output = VTCM_LAYOUT_PTR(__fp16, base, L.off_dst[0]); @@ -2855,13 +2723,23 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, if (pipeline) { // --- Asynchronous Pipelined Loop --- hmx_matmul_job_t job_slots[2]; // persistent double-buffered job descriptors + void *vtcm_raw_bufs[4] = { + vtcm_weight_raw[0], + vtcm_weight_raw[1], + dequant_worker_fn ? vtcm_weight_raw[0] : vtcm_scratch0, + dequant_worker_fn ? vtcm_weight_raw[1] : vtcm_scratch1, + }; + void *vtcm_weight_bufs[4] = { + dequant_worker_fn ? vtcm_scratch0 : vtcm_weight_raw[0], + dequant_worker_fn ? vtcm_scratch1 : vtcm_weight_raw[1], + vtcm_scratch0, + vtcm_scratch1, + }; + void *vtcm_output_bufs[2] = { vtcm_output, vtcm_scratch2 }; for (size_t mr = 0; mr < m; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows); - void *vtcm_weight_bufs[2] = { vtcm_scratch0, vtcm_scratch1 }; - void *vtcm_output_bufs[2] = { vtcm_output, vtcm_scratch2 }; - struct activation_transfer_params act_params = { .ctx = ctx, .dst = vtcm_f16_act, @@ -2881,14 +2759,14 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // Prologue: push A0 and optionally A1 (if n_chunk_cnt > 1) const size_t n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols); - const uint32_t height_A0 = is_quant ? hmx_ceil_div(n_cols_A0, 32) * n_k_tiles : n_cols_A0; - dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[0], weight), + const uint32_t height_A0 = is_tiled ? hmx_ceil_div(n_cols_A0, 32) * n_k_tiles : n_cols_A0; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height_A0); if (1 < n_chunk_cnt) { const size_t n_cols_A1 = hex_smin(n - 1 * n_chunk_n_cols, n_chunk_n_cols); - const uint32_t height_A1 = is_quant ? hmx_ceil_div(n_cols_A1, 32) * n_k_tiles : n_cols_A1; - dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride), + const uint32_t height_A1 = is_tiled ? hmx_ceil_div(n_cols_A1, 32) * n_k_tiles : n_cols_A1; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[1], weight + n_chunk_n_cols * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_A1); } @@ -2905,20 +2783,20 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // 2. dequantize A_i dequantize_tiled_weight_chunk_to_fp16_tiles( - ctx, vtcm_weight_bufs[i % 2], curr_raw, + ctx, vtcm_weight_bufs[i % 4], curr_raw, n_cols, k, row_stride, weight_type, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); // 3. push A_{i+2} (if i+2 < n_chunk_cnt) if (i + 2 < n_chunk_cnt) { - const uint32_t height_p2 = is_quant ? hmx_ceil_div(n_cols_p2, 32) * n_k_tiles : n_cols_p2; - dma_queue_push(weight_dma, dma_make_data(curr_raw, weight + nc_p2 * weight_stride), + const uint32_t height_p2 = is_tiled ? hmx_ceil_div(n_cols_p2, 32) * n_k_tiles : n_cols_p2; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[(i + 2) % 4], weight + nc_p2 * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_p2); } // 4. submit C_i hmx_matmul_job_init(&job_slots[i % 2], (__fp16 *) vtcm_output_bufs[i % 2], - (__fp16 *) vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[i % 2], + (__fp16 *) vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[i % 4], vtcm_scales, hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS), hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS); hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[i % 2])); @@ -2951,6 +2829,9 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, } else { // --- Synchronous loop (m <= 32 or fallback) --- hmx_matmul_job_t job; + void *vtcm_raw_bufs[2] = { vtcm_weight_raw[0], dequant_worker_fn ? vtcm_weight_raw[0] : vtcm_scratch0 }; + void *vtcm_weight_bufs[2] = { dequant_worker_fn ? vtcm_scratch0 : vtcm_weight_raw[0], vtcm_scratch0 }; + for (size_t mr = 0; mr < m; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin(m - mr, m_chunk_n_rows); @@ -2974,11 +2855,12 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // A0: Pre-fetch the first weight chunk (nc = 0) if (n > 0) { const size_t n_cols = hex_smin(n, n_chunk_n_cols); - const uint32_t height = is_quant ? hmx_ceil_div(n_cols, 32) * n_k_tiles : n_cols; - dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); + const uint32_t height = is_tiled ? hmx_ceil_div(n_cols, 32) * n_k_tiles : n_cols; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); } - for (size_t nc = 0; nc < n; nc += n_chunk_n_cols) { + for (int i = 0; i < n_chunk_cnt; ++i) { + const size_t nc = (size_t) i * n_chunk_n_cols; const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols); const size_t n_row_tiles = hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS); const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS); @@ -2988,7 +2870,7 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, // B: Weight Dequantize (Threaded) dequantize_tiled_weight_chunk_to_fp16_tiles( - ctx, vtcm_scratch0, curr_raw, + ctx, vtcm_weight_bufs[i % 2], curr_raw, n_cols, k, row_stride, weight_type, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); @@ -2996,12 +2878,12 @@ static int hmx_mm_2d_f32(struct htp_context *ctx, const size_t nc_next = nc + n_chunk_n_cols; if (nc_next < n) { const size_t n_cols_next = hex_smin(n - nc_next, n_chunk_n_cols); - const uint32_t height_next = is_quant ? hmx_ceil_div(n_cols_next, 32) * n_k_tiles : n_cols_next; - dma_queue_push(weight_dma, dma_make_data(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); + const uint32_t height_next = is_tiled ? hmx_ceil_div(n_cols_next, 32) * n_k_tiles : n_cols_next; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[(i + 1) % 2], weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); } // C: HMX Compute (Queue-based) - hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); + hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[i % 2], vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job)); hmx_queue_pop(ctx->hmx_queue); @@ -3041,35 +2923,20 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const dma_addr_t act_dma_addr = act->data; if (k % 32 != 0) { return HTP_STATUS_NO_SUPPORT; } - if ((act_dma_addr & (VLEN - 1)) != 0) { return HTP_STATUS_NO_SUPPORT; } - size_t row_stride = htp_mm_get_tiled_row_stride(weight_type, k); + size_t row_stride = htp_tensor_get_row_size(weight_type, k); if (row_stride == 0) { return HTP_STATUS_NO_SUPPORT; } - worker_callback_t dequant_worker_fn = NULL; - switch (weight_type) { - case HTP_TYPE_Q4_0: dequant_worker_fn = dequantize_tiled_worker_loop_q4_0; break; - case HTP_TYPE_IQ4_NL: dequant_worker_fn = dequantize_tiled_worker_loop_iq4_nl; break; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: dequant_worker_fn = dequantize_tiled_worker_loop_q4_1; break; - case HTP_TYPE_MXFP4: dequant_worker_fn = dequantize_tiled_worker_loop_mxfp4; break; - case HTP_TYPE_Q8_0: dequant_worker_fn = dequantize_tiled_worker_loop_q8_0; break; - case HTP_TYPE_Q5_K: dequant_worker_fn = dequantize_tiled_worker_loop_q5_k; break; - case HTP_TYPE_Q6_K: dequant_worker_fn = dequantize_tiled_worker_loop_q6_k; break; - case HTP_TYPE_Q3_K: dequant_worker_fn = dequantize_tiled_worker_loop_q3_k; break; - case HTP_TYPE_Q2_K: dequant_worker_fn = dequantize_tiled_worker_loop_q2_k; break; - case HTP_TYPE_F16: dequant_worker_fn = convert_f16_worker_loop; break; - case HTP_TYPE_F32: dequant_worker_fn = quantize_f32_worker_loop; break; - default: - return HTP_STATUS_NO_SUPPORT; - } - const int n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS; const struct fastdiv_values n_k_tiles_div = init_fastdiv_values(n_k_tiles); - const bool is_quant = (weight_type != HTP_TYPE_F16 && weight_type != HTP_TYPE_F32); + const int tile_size = kparams->tile_size; + const int aligned_tile_size = kparams->aligned_tile_size; + + const bool is_tiled = (tile_size > 0); + tiled_dequant_task_fn_t dequant_worker_fn = hmx_get_dequant_worker_fn(weight_type, is_tiled); const size_t vtcm_budget = ctx->vtcm_size; const int m_chunk_n_rows = kparams->m_chunk; @@ -3079,11 +2946,9 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const int act_threads = kparams->n_act_threads; const struct fastdiv_values * act_threads_div = &kparams->div_n_act_threads; const struct fastdiv_values * k_div = &kparams->div_ne00_padded; - const int tile_size = kparams->tile_size; - const int aligned_tile_size = kparams->aligned_tile_size; - const uint32_t dma_dst_stride = is_quant ? aligned_tile_size : row_stride; - const uint32_t dma_width_bytes = is_quant ? tile_size : row_stride; + const uint32_t dma_dst_stride = is_tiled ? aligned_tile_size : row_stride; + const uint32_t dma_width_bytes = is_tiled ? tile_size : row_stride; struct htp_mm_hmx_vtcm_layout L; htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_2D, weight_type, k, m_chunk_n_rows, n_chunk_n_cols, 1, pipeline, act_threads, aligned_tile_size, 0); @@ -3099,7 +2964,6 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k VTCM_LAYOUT_PTR(__fp16, base, L.off_weight[0]), VTCM_LAYOUT_PTR_OPTIONAL(__fp16, base, L.off_weight[1], pipeline) }; - __fp16 *vtcm_f16_act = VTCM_LAYOUT_PTR(__fp16, base, L.off_act); float *vtcm_f32_act = VTCM_LAYOUT_PTR(float, base, L.off_act_f32); __fp16 *vtcm_output = VTCM_LAYOUT_PTR(__fp16, base, L.off_dst[0]); @@ -3141,13 +3005,23 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k if (pipeline) { hmx_matmul_job_t job_slots[2]; + void *vtcm_raw_bufs[4] = { + vtcm_weight_raw[0], + vtcm_weight_raw[1], + dequant_worker_fn ? vtcm_weight_raw[0] : vtcm_scratch0, + dequant_worker_fn ? vtcm_weight_raw[1] : vtcm_scratch1, + }; + void *vtcm_weight_bufs[4] = { + dequant_worker_fn ? vtcm_scratch0 : vtcm_weight_raw[0], + dequant_worker_fn ? vtcm_scratch1 : vtcm_weight_raw[1], + vtcm_scratch0, + vtcm_scratch1, + }; + void *vtcm_output_bufs[2] = { vtcm_output, vtcm_scratch2 }; for (size_t mr = (size_t) m_start; mr < mr_end; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin(mr_end - mr, m_chunk_n_rows); - void *vtcm_weight_bufs[2] = { vtcm_scratch0, vtcm_scratch1 }; - void *vtcm_output_bufs[2] = { vtcm_output, vtcm_scratch2 }; - struct activation_transfer_params act_params = { .ctx = ctx, .dst = vtcm_f16_act, @@ -3180,17 +3054,17 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const int dst_cols = (int) dst->ne[0]; const int n_chunk_cnt = hmx_ceil_div(n, n_chunk_n_cols); - const uint32_t dma_src_stride = is_quant ? tile_size : weight_stride; + const uint32_t dma_src_stride = is_tiled ? tile_size : weight_stride; const size_t n_cols_A0 = hex_smin(n - 0 * n_chunk_n_cols, n_chunk_n_cols); - const uint32_t height_A0 = is_quant ? hmx_ceil_div(n_cols_A0, 32) * n_k_tiles : n_cols_A0; - dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[0], weight), + const uint32_t height_A0 = is_tiled ? hmx_ceil_div(n_cols_A0, 32) * n_k_tiles : n_cols_A0; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height_A0); if (1 < n_chunk_cnt) { const size_t n_cols_A1 = hex_smin(n - 1 * n_chunk_n_cols, n_chunk_n_cols); - const uint32_t height_A1 = is_quant ? hmx_ceil_div(n_cols_A1, 32) * n_k_tiles : n_cols_A1; - dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[1], weight + n_chunk_n_cols * weight_stride), + const uint32_t height_A1 = is_tiled ? hmx_ceil_div(n_cols_A1, 32) * n_k_tiles : n_cols_A1; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[1], weight + n_chunk_n_cols * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_A1); } @@ -3204,18 +3078,18 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k void * curr_raw = (void *) dma_queue_pop(weight_dma).dst; dequantize_tiled_weight_chunk_to_fp16_tiles( - ctx, vtcm_weight_bufs[i % 2], curr_raw, + ctx, vtcm_weight_bufs[i % 4], curr_raw, n_cols, k, row_stride, weight_type, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); if (i + 2 < n_chunk_cnt) { - const uint32_t height_p2 = is_quant ? hmx_ceil_div(n_cols_p2, 32) * n_k_tiles : n_cols_p2; - dma_queue_push(weight_dma, dma_make_data(curr_raw, weight + nc_p2 * weight_stride), + const uint32_t height_p2 = is_tiled ? hmx_ceil_div(n_cols_p2, 32) * n_k_tiles : n_cols_p2; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[(i + 2) % 4], weight + nc_p2 * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_p2); } hmx_matmul_job_init(&job_slots[i % 2], (__fp16 *) vtcm_output_bufs[i % 2], - (__fp16 *) vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[i % 2], + (__fp16 *) vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[i % 4], vtcm_scales, hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS), hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS), k / HTP_MM_HMX_TILE_N_ROWS); hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job_slots[i % 2])); @@ -3244,6 +3118,9 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k } } else { hmx_matmul_job_t job; + void *vtcm_raw_bufs[2] = { vtcm_weight_raw[0], dequant_worker_fn ? vtcm_weight_raw[0] : vtcm_scratch0 }; + void *vtcm_weight_bufs[2] = { dequant_worker_fn ? vtcm_scratch0 : vtcm_weight_raw[0], vtcm_scratch0 }; + for (size_t mr = (size_t) m_start; mr < mr_end; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin(mr_end - mr, m_chunk_n_rows); @@ -3277,16 +3154,18 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k const size_t weight_stride = src_w->nb[1]; const size_t dst_stride = dst->nb[1] / sizeof(float); const int dst_cols = (int) dst->ne[0]; + const int n_chunk_cnt = hmx_ceil_div(n, n_chunk_n_cols); - const uint32_t dma_src_stride = is_quant ? tile_size : weight_stride; + const uint32_t dma_src_stride = is_tiled ? tile_size : weight_stride; if (n > 0) { const size_t n_cols = hex_smin(n, n_chunk_n_cols); - const uint32_t height = is_quant ? hmx_ceil_div(n_cols, 32) * n_k_tiles : n_cols; - dma_queue_push(weight_dma, dma_make_data(vtcm_weight_raw[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); + const uint32_t height = is_tiled ? hmx_ceil_div(n_cols, 32) * n_k_tiles : n_cols; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); } - for (size_t nc = 0; nc < n; nc += n_chunk_n_cols) { + for (int i = 0; i < n_chunk_cnt; ++i) { + const size_t nc = (size_t) i * n_chunk_n_cols; const size_t n_cols = hex_smin(n - nc, n_chunk_n_cols); const size_t n_row_tiles = hmx_ceil_div(n_rows, HTP_MM_HMX_TILE_N_ROWS); const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS); @@ -3294,18 +3173,18 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k void * curr_raw = (void *) dma_queue_pop(weight_dma).dst; dequantize_tiled_weight_chunk_to_fp16_tiles( - ctx, vtcm_scratch0, curr_raw, + ctx, vtcm_weight_bufs[i % 2], curr_raw, n_cols, k, row_stride, weight_type, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads); const size_t nc_next = nc + n_chunk_n_cols; if (nc_next < n) { const size_t n_cols_next = hex_smin(n - nc_next, n_chunk_n_cols); - const uint32_t height_next = is_quant ? hmx_ceil_div(n_cols_next, 32) * n_k_tiles : n_cols_next; - dma_queue_push(weight_dma, dma_make_data(curr_raw, weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); + const uint32_t height_next = is_tiled ? hmx_ceil_div(n_cols_next, 32) * n_k_tiles : n_cols_next; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[(i + 1) % 2], weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); } - hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); + hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[i % 2], vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job)); hmx_queue_pop(ctx->hmx_queue); @@ -3322,30 +3201,30 @@ static int hmx_mm_nx_2d_f32(struct htp_ops_context * octx, const struct htp_mm_k return HTP_STATUS_OK; } -static inline dma_addr_t hmx_mm_weight_batch_data(const hmx_mm_f16_f32_batched_params_t *params, +static inline dma_addr_t hmx_mm_weight_batch_data(const hmx_mm_batched_params_t *params, int dst_b2, int dst_b3) { const size_t b2_idx = (params->r2 <= 1) ? (size_t) dst_b2 : (size_t) fastdiv((uint32_t) dst_b2, ¶ms->div_r2); const size_t b3_idx = (params->r3 <= 1) ? (size_t) dst_b3 : (size_t) fastdiv((uint32_t) dst_b3, ¶ms->div_r3); return params->weight + b2_idx * params->src0_nb2 + b3_idx * params->src0_nb3; } -static inline dma_addr_t hmx_mm_act_batch_addr(const hmx_mm_f16_f32_batched_params_t *params, +static inline dma_addr_t hmx_mm_act_batch_addr(const hmx_mm_batched_params_t *params, int dst_b2, int dst_b3) { return params->act_dma_addr + dst_b2 * params->act_nb2 + dst_b3 * params->act_nb3; } -static inline float *hmx_mm_dst_batch_ptr(const hmx_mm_f16_f32_batched_params_t *params, +static inline float *hmx_mm_dst_batch_ptr(const hmx_mm_batched_params_t *params, int dst_b2, int dst_b3) { return (float *) ((uint8_t *) params->dst + (size_t) dst_b2 * params->dst_nb2 + (size_t) dst_b3 * params->dst_nb3); } -static int hmx_mm_f16_f32_batched_simple(struct htp_context *ctx, - const hmx_mm_f16_f32_batched_params_t *params, - int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, int vtcm_size, - const struct fastdiv_values * act_threads_div, const struct fastdiv_values * k_div) { +static int hmx_mm_batched_simple(struct htp_context *ctx, + const hmx_mm_batched_params_t *params, + int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, int vtcm_size, + const struct fastdiv_values * act_threads_div, const struct fastdiv_values * k_div) { int ret = 0; for (int b3 = 0; b3 < params->ne13 && ret == 0; ++b3) { for (int b2 = 0; b2 < params->ne12 && ret == 0; ++b2) { @@ -3357,54 +3236,53 @@ static int hmx_mm_f16_f32_batched_simple(struct htp_context *ctx, hmx_mm_act_batch_addr(params, b2, b3), hmx_mm_weight_batch_data(params, b2, b3), params->m, params->k, params->n, - params->act_stride, params->act_elem_size, params->weight_stride * (int)sizeof(__fp16), - HTP_TYPE_F16, params->k, params->dst_stride, params->bias_stride, params->n, + params->act_stride, params->act_elem_size, params->weight_stride, + params->weight_type, params->k, params->dst_stride, params->bias_stride, params->dst_cols, m_chunk, n_chunk, pipeline, n_threads, act_threads, - act_threads_div, k_div, 0, 0, vtcm_size); + act_threads_div, k_div, params->tile_size, params->aligned_tile_size, vtcm_size); } } return ret; } -static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_batched_params_t *params, - int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, - const struct fastdiv_values * act_threads_div, - const struct fastdiv_values * k_div, - int vtcm_size) { - if (params->act_stride < params->k || params->weight_stride < params->k || params->dst_stride < params->n) { return -1; } +static int hmx_mm_batched(struct htp_context *ctx, const hmx_mm_batched_params_t *params, + int m_chunk, int n_chunk, int pipeline, int n_threads, int act_threads, + const struct fastdiv_values * act_threads_div, + const struct fastdiv_values * k_div, + int vtcm_size) { + const size_t row_stride = htp_tensor_get_row_size(params->weight_type, params->k); + if (row_stride == 0) { return -1; } + if (params->act_stride < params->k || (params->tile_size == 0 && (size_t) params->weight_stride < row_stride) || params->dst_stride < params->dst_cols) { return -1; } if (params->ne02 <= 0 || params->ne03 <= 0 || params->ne12 <= 0 || params->ne13 <= 0) { return -1; } if (params->ne12 % params->ne02 != 0 || params->ne13 % params->ne03 != 0) { return -1; } // N (the weight row count) does not have to be 32-aligned: if (params->k % 32 != 0) { return -1; } - if (!hex_is_aligned(params->dst, VLEN) || (params->act_dma_addr & (VLEN - 1)) != 0) { return -1; } + if (!hex_is_aligned(params->dst, sizeof(float))) { return -1; } const int group_size = params->r2; const size_t vtcm_budget = ctx->vtcm_size; // Check if the precomputed parameters are grouped or simple. // If simple, or if group_size <= 1, we use simple fallback loop. - // Grouped path is only valid if group_size > 1 and it fits within VTCM budget. - bool run_grouped = (group_size > 1 && (size_t)vtcm_size <= vtcm_budget); + // Grouped path is only valid for flat F16 if group_size > 1 and it fits within VTCM budget. + bool run_grouped = (params->weight_type == HTP_TYPE_F16 && params->tile_size == 0 && group_size > 1 && (size_t)vtcm_size <= vtcm_budget); if (!run_grouped) { - return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size, act_threads_div, k_div); + return hmx_mm_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size, act_threads_div, k_div); } struct htp_thread_trace * tr = &ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); - const size_t vec_dot_size = params->k * sizeof(__fp16); - size_t m_chunk_n_rows = m_chunk; size_t n_chunk_n_cols = n_chunk; - size_t vtcm_used = vtcm_size; struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, HTP_TYPE_F16, params->k, m_chunk_n_rows, n_chunk_n_cols, group_size, false, act_threads, 0, params->bias_bytes); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_BATCHED, HTP_TYPE_F16, params->k, m_chunk_n_rows, n_chunk_n_cols, group_size, false, act_threads, 0, params->bias_bytes); if (L.total_bytes > vtcm_budget) { FARF(HIGH, "%s: grouped layout overflowed VTCM, falling back to simple batched loop", __func__); htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); - return hmx_mm_f16_f32_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size, act_threads_div, k_div); + return hmx_mm_batched_simple(ctx, params, m_chunk, n_chunk, pipeline, n_threads, act_threads, vtcm_size, act_threads_div, k_div); } uint8_t * const base = (uint8_t *) ctx->vtcm_base; @@ -3431,7 +3309,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ L.total_bytes, vtcm_budget); const size_t fp16_row_bytes = (size_t) params->k * sizeof(__fp16); - const size_t weight_row_bytes = (size_t) params->weight_stride * sizeof(__fp16); + const size_t weight_row_bytes = (size_t) params->weight_stride; htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); @@ -3481,7 +3359,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ } if (n_chunk_n_cols < (size_t) params->n) { const size_t n_cols_second = hex_smin((size_t) params->n - n_chunk_n_cols, n_chunk_n_cols); - const dma_addr_t second_weight_chunk = weight_group + n_chunk_n_cols * params->weight_stride * sizeof(__fp16); + const dma_addr_t second_weight_chunk = weight_group + n_chunk_n_cols * weight_row_bytes; dma_queue_push(weight_dma, dma_make_data(vtcm_scratch1, second_weight_chunk), fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_second); } @@ -3500,7 +3378,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ const size_t nc_next = nc + n_chunk_n_cols * 2; if (nc_next < (size_t) params->n) { const size_t n_cols_next = hex_smin((size_t) params->n - nc_next, n_chunk_n_cols); - const dma_addr_t next_weight_chunk = weight_group + nc_next * params->weight_stride * sizeof(__fp16); + const dma_addr_t next_weight_chunk = weight_group + nc_next * weight_row_bytes; dma_queue_push(weight_dma, dma_make_data(curr_raw, next_weight_chunk), fp16_row_bytes, weight_row_bytes, fp16_row_bytes, n_cols_next); @@ -3519,7 +3397,7 @@ static int hmx_mm_f16_f32_batched(struct htp_context *ctx, const hmx_mm_f16_f32_ { float *output = hmx_mm_dst_batch_ptr(params, b2_base + g, b3) + mr * params->dst_stride + nc; const float *bias_chunk = has_bias ? (vtcm_bias + mr * params->bias_stride + nc) : NULL; - int chunk_dst_cols = params->n - (int)nc; + int chunk_dst_cols = params->dst_cols - (int)nc; if (chunk_dst_cols > 0) { transfer_output_chunk_threaded(ctx, output, bias_chunk, vtcm_output, (int) n_rows, (int) n_cols, params->dst_stride, params->bias_stride, chunk_dst_cols, n_threads); @@ -3642,6 +3520,8 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, size_t dst_nb1, size_t dst_nb2, int weight_stride, int weight_type, + int tile_size, + int aligned_tile_size, const struct mmid_row_mapping *matrix_rows, int cur_a, int mapping_stride, @@ -3658,47 +3538,27 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, if (k % 32 != 0 || n % 32 != 0) { return -1; } if (!hex_is_aligned(dst, VLEN) || !hex_is_aligned(activation, VLEN)) { return -1; } - size_t row_stride = htp_mm_get_tiled_row_stride(weight_type, k); + size_t row_stride = htp_tensor_get_row_size(weight_type, k); if (row_stride == 0) { return -1; } - worker_callback_t dequant_worker_fn = NULL; - switch (weight_type) { - case HTP_TYPE_Q4_0: dequant_worker_fn = dequantize_tiled_worker_loop_q4_0; break; - case HTP_TYPE_IQ4_NL: dequant_worker_fn = dequantize_tiled_worker_loop_iq4_nl; break; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: dequant_worker_fn = dequantize_tiled_worker_loop_q4_1; break; - case HTP_TYPE_MXFP4: dequant_worker_fn = dequantize_tiled_worker_loop_mxfp4; break; - case HTP_TYPE_Q8_0: dequant_worker_fn = dequantize_tiled_worker_loop_q8_0; break; - case HTP_TYPE_Q5_K: dequant_worker_fn = dequantize_tiled_worker_loop_q5_k; break; - case HTP_TYPE_Q6_K: dequant_worker_fn = dequantize_tiled_worker_loop_q6_k; break; - case HTP_TYPE_Q3_K: dequant_worker_fn = dequantize_tiled_worker_loop_q3_k; break; - case HTP_TYPE_Q2_K: dequant_worker_fn = dequantize_tiled_worker_loop_q2_k; break; - case HTP_TYPE_F16: dequant_worker_fn = convert_f16_worker_loop; break; - case HTP_TYPE_F32: dequant_worker_fn = quantize_f32_worker_loop; break; - default: - return -1; - } - const int n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS; const struct fastdiv_values n_k_tiles_div = init_fastdiv_values(n_k_tiles); - const bool is_quant = (weight_type != HTP_TYPE_F16 && weight_type != HTP_TYPE_F32); - const size_t vec_dot_size = k * sizeof(__fp16); const size_t vtcm_budget = ctx->vtcm_size; size_t vtcm_used = 0; - int tile_size = htp_mm_get_weight_tile_size(weight_type); - int aligned_tile_size = htp_mm_get_weight_aligned_tile_size(weight_type); + const bool is_tiled = (tile_size > 0); + tiled_dequant_task_fn_t dequant_worker_fn = hmx_get_dequant_worker_fn(weight_type, is_tiled); - const uint32_t dma_dst_stride = is_quant ? aligned_tile_size : row_stride; - const uint32_t dma_src_stride = is_quant ? tile_size : weight_stride; - const uint32_t dma_width_bytes = is_quant ? tile_size : row_stride; + const uint32_t dma_dst_stride = is_tiled ? aligned_tile_size : row_stride; + const uint32_t dma_src_stride = is_tiled ? tile_size : weight_stride; + const uint32_t dma_width_bytes = is_tiled ? tile_size : row_stride; - const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0; - const size_t weight_row_stride = is_quant ? qweight_row_stride : row_stride; + const size_t qweight_row_stride = is_tiled ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0; + const size_t weight_row_stride = is_tiled ? qweight_row_stride : row_stride; size_t size_per_n = 0, size_per_m = 0, size_per_mn = 0; htp_mm_hmx_get_2d_chunk_costs(weight_type, k, /*pipeline=*/false, aligned_tile_size, @@ -3738,6 +3598,9 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); hmx_matmul_job_t job; + void *vtcm_raw_bufs[2] = { vtcm_weight, dequant_worker_fn ? vtcm_weight : vtcm_scratch0 }; + void *vtcm_weight_bufs[2] = { dequant_worker_fn ? vtcm_scratch0 : vtcm_weight, vtcm_scratch0 }; + const int n_chunk_cnt = hmx_ceil_div(n, n_chunk_n_cols); for (size_t mr = (size_t) m_start; mr < (size_t) m_end; mr += m_chunk_n_rows) { const size_t n_rows = hex_smin((size_t) m_end - mr, m_chunk_n_rows); @@ -3750,12 +3613,13 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, // A0: Pre-fetch the first weight chunk (nc = 0) if (n > 0) { const size_t n_cols = hex_smin((size_t) n, n_chunk_n_cols); - const uint32_t height = is_quant ? hmx_ceil_div(n_cols, 32) * n_k_tiles : n_cols; - dma_queue_push(weight_dma, dma_make_data(vtcm_weight, weight), + const uint32_t height = is_tiled ? hmx_ceil_div(n_cols, 32) * n_k_tiles : n_cols; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[0], weight), dma_dst_stride, dma_src_stride, dma_width_bytes, height); } - for (size_t nc = 0; nc < (size_t) n; nc += n_chunk_n_cols) { + for (int i = 0; i < n_chunk_cnt; ++i) { + const size_t nc = (size_t) i * n_chunk_n_cols; const size_t n_cols = hex_smin((size_t) n - nc, n_chunk_n_cols); const size_t n_col_tiles = hmx_ceil_div(n_cols, HTP_MM_HMX_TILE_N_COLS); @@ -3764,7 +3628,7 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, // B: Weight Dequantize (Threaded) dequantize_tiled_weight_chunk_to_fp16_tiles( - ctx, vtcm_scratch0, curr_raw, + ctx, vtcm_weight_bufs[i % 2], curr_raw, n_cols, k, row_stride, weight_type, n_k_tiles, n_k_tiles_div, dequant_worker_fn, n_threads ); @@ -3773,13 +3637,13 @@ static int hmx_mm_id_2d_f32(struct htp_context *ctx, const size_t nc_next = nc + n_chunk_n_cols; if (nc_next < (size_t) n) { const size_t n_cols_next = hex_smin((size_t) n - nc_next, n_chunk_n_cols); - const uint32_t height_next = is_quant ? hmx_ceil_div(n_cols_next, 32) * n_k_tiles : n_cols_next; - dma_queue_push(weight_dma, dma_make_data(curr_raw, weight + nc_next * weight_stride), + const uint32_t height_next = is_tiled ? hmx_ceil_div(n_cols_next, 32) * n_k_tiles : n_cols_next; + dma_queue_push(weight_dma, dma_make_data(vtcm_raw_bufs[(i + 1) % 2], weight + nc_next * weight_stride), dma_dst_stride, dma_src_stride, dma_width_bytes, height_next); } // C: HMX Compute (Queue-based) - hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, vtcm_scratch0, vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); + hmx_matmul_job_init(&job, vtcm_output, vtcm_f16_act, (__fp16 *) vtcm_weight_bufs[i % 2], vtcm_scales, n_row_tiles, n_col_tiles, k / HTP_MM_HMX_TILE_N_ROWS); hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_matmul_worker_fn, &job)); hmx_queue_pop(ctx->hmx_queue); @@ -3803,7 +3667,6 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k const int m_total = (int) act->ne[1]; const uint32_t act_elem_size = (act->type == HTP_TYPE_F16) ? sizeof(__fp16) : sizeof(float); const int act_stride = (int) (act->nb[1] / act_elem_size); - const int wgt_stride = (int)(src0->nb[1] / sizeof(__fp16)); int m_start = 0; int m_rows = m_total; @@ -3842,46 +3705,50 @@ static int hmx_mm_op_matmul(struct htp_ops_context * octx, const struct htp_mm_k int ret = -1; const int n_threads = kparams->n_threads; - if (kparams->kernel_type == HTP_MM_KERNEL_HMX_F16_BATCHED) { - hmx_mm_f16_f32_batched_params_t batch_params = { - .dst = dst_ptr, - .bias_addr = src2_addr, - .bias_bytes = src2_bytes, - .act_dma_addr = act_addr, - .weight = src0->data, - .weight_dma = octx->ctx->dma[0], - .m = m_rows, - .k = k, - .n = n, - .act_stride = act_stride, - .act_elem_size = act_elem_size, - .weight_stride = wgt_stride, - .dst_stride = dst_stride, - .bias_stride = src2_stride, - .ne02 = ne02, - .ne03 = ne03, - .ne12 = ne12, - .ne13 = ne13, - .src0_nb2 = src0->nb[2], - .src0_nb3 = src0->nb[3], - .act_nb2 = act->nb[2], - .act_nb3 = act->nb[3], - .dst_nb2 = dst->nb[2], - .dst_nb3 = dst->nb[3], - .bias_nb2 = src2_nb2, - .bias_nb3 = src2_nb3, - .r2 = (ne02 > 0) ? (ne12 / ne02) : 1, - .r3 = (ne03 > 0) ? (ne13 / ne03) : 1, - .div_r2 = kparams->div_r2, - .div_r3 = kparams->div_r3, + if (kparams->kernel_type == HTP_MM_KERNEL_HMX_BATCHED) { + hmx_mm_batched_params_t batch_params = { + .dst = dst_ptr, + .bias_addr = src2_addr, + .bias_bytes = src2_bytes, + .act_dma_addr = act_addr, + .weight = src0->data, + .weight_dma = octx->ctx->dma[0], + .m = m_rows, + .k = k, + .n = n, + .act_stride = act_stride, + .act_elem_size = act_elem_size, + .weight_stride = (int) src0->nb[1], + .weight_type = (int) src0->type, + .dst_stride = dst_stride, + .dst_cols = (int) dst->ne[0], + .bias_stride = src2_stride, + .tile_size = kparams->tile_size, + .aligned_tile_size = kparams->aligned_tile_size, + .ne02 = ne02, + .ne03 = ne03, + .ne12 = ne12, + .ne13 = ne13, + .src0_nb2 = src0->nb[2], + .src0_nb3 = src0->nb[3], + .act_nb2 = act->nb[2], + .act_nb3 = act->nb[3], + .dst_nb2 = dst->nb[2], + .dst_nb3 = dst->nb[3], + .bias_nb2 = src2_nb2, + .bias_nb3 = src2_nb3, + .r2 = (ne02 > 0) ? (ne12 / ne02) : 1, + .r3 = (ne03 > 0) ? (ne13 / ne03) : 1, + .div_r2 = kparams->div_r2, + .div_r3 = kparams->div_r3, }; - ret = hmx_mm_f16_f32_batched(octx->ctx, &batch_params, - kparams->m_chunk, kparams->n_chunk, - kparams->pipeline, n_threads, - kparams->n_act_threads, - &kparams->div_n_act_threads, - &kparams->div_ne00_padded, - kparams->vtcm_size); + ret = hmx_mm_batched(octx->ctx, &batch_params, + kparams->m_chunk, kparams->n_chunk, + kparams->pipeline, n_threads, + kparams->n_act_threads, + &kparams->div_n_act_threads, + &kparams->div_ne00_padded, + kparams->vtcm_size); } else { ret = hmx_mm_2d_f32( octx->ctx, octx->ctx->dma[0], dst_ptr, src2_addr, src2_bytes, @@ -4014,6 +3881,7 @@ static int hmx_mm_op_matmul_id( nb11, nb12, nb1, nb2, (int) src0->nb[1], (int) src0->type, + kparams->tile_size, kparams->aligned_tile_size, matrix_rows, cur_a, mmctx->mapping_stride, m_start, m_end, (int) octx->n_threads); if (ret != 0) { @@ -4049,7 +3917,7 @@ static int hvx_mm_matmul_id( uint32_t n_quant_tasks = 1; if (act_nrows < octx->n_threads) { n_quant_tasks = MIN(total_nb, octx->n_threads); - quant_task_func = htp_mm_act_quant_block_func(src0->type); + quant_task_func = htp_mm_act_quant_block_func(src0->type, act->type); for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { uint32_t ib_first = (total_nb * ith) / n_quant_tasks; uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks; @@ -4060,9 +3928,9 @@ static int hvx_mm_matmul_id( } } else { n_quant_tasks = MIN(act_nrows, octx->n_threads); - quant_task_func = htp_mm_act_quant_row_func(src0->type); + quant_task_func = htp_mm_act_quant_row_func(src0->type, act->type); } - size_t act_row_size = htp_mm_weight_has_offset(src0->type) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + size_t act_row_size = htp_mm_tiled_act_row_size(src0->type, ne10); struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, act_nrows, octx->n_threads, @@ -4177,6 +4045,7 @@ static int hmx_mm_op_matmul_id_nx( act->nb[1], act->nb[2], dst->nb[1], dst->nb[2], (int) src_w->nb[1], (int) src_w->type, + kparams->tile_size, kparams->aligned_tile_size, matrix_rows, cur_a, mmctx->mapping_stride, m_start, m_end, (int) octx->n_threads); if (ret != 0) { @@ -4215,7 +4084,7 @@ static int hvx_mm_matmul_id_nx( uint32_t n_quant_tasks = 1; if (act_nrows < octx->n_threads) { n_quant_tasks = MIN(total_nb, octx->n_threads); - quant_task_func = htp_mm_act_quant_block_func(src0->type); + quant_task_func = htp_mm_act_quant_block_func(src0->type, act->type); for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { uint32_t ib_first = (total_nb * ith) / n_quant_tasks; uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks; @@ -4226,9 +4095,9 @@ static int hvx_mm_matmul_id_nx( } } else { n_quant_tasks = MIN(act_nrows, octx->n_threads); - quant_task_func = htp_mm_act_quant_row_func(src0->type); + quant_task_func = htp_mm_act_quant_row_func(src0->type, act->type); } - size_t act_row_size = htp_mm_weight_has_offset(src0->type) ? htp_mm_q8_1_tiled_row_size(act->ne[0]) : htp_mm_q8_0_tiled_row_size(act->ne[0]); + size_t act_row_size = htp_mm_tiled_act_row_size(src0->type, act->ne[0]); struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], act_nrows, octx->n_threads, @@ -4451,8 +4320,8 @@ int op_matmul_id(struct htp_ops_context * octx) { mmctx->src0_nrows_per_thread = fastdiv(nrows + octx->n_threads - 1, &octx->n_threads_div); mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); - if (hvx_mm_init_vec_dot(mmctx, src0->type) == 0) { - s = hvx_mm_matmul_id(octx, mmctx, act_nrows > 1 ? hvx_mm_id : hvx_mv_id); + if (hvx_mm_init_tiled_vec_dot(mmctx, src0->type) == 0) { + s = hvx_mm_matmul_id(octx, mmctx, act_nrows > 1 ? hvx_mm_id_tiled : hvx_mv_id_tiled); } else { s = HTP_STATUS_NO_SUPPORT; } @@ -4551,8 +4420,8 @@ int op_matmul_id_nx(struct htp_ops_context * octx) { if (kparams->n_hmx) { s = hmx_mm_op_matmul_id_nx(octx, mmctx); } else { - if (hvx_mm_init_vec_dot(mmctx, src0->type) == 0) { - s = hvx_mm_matmul_id_nx(octx, mmctx, act_nrows > 1 ? hvx_mm_id_nx : hvx_mv_id_nx); + if (hvx_mm_init_tiled_vec_dot(mmctx, src0->type) == 0) { + s = hvx_mm_matmul_id_nx(octx, mmctx, act_nrows > 1 ? hvx_mm_id_nx_tiled : hvx_mv_id_nx_tiled); } else { s = HTP_STATUS_NO_SUPPORT; } @@ -4584,12 +4453,6 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) { const struct htp_tensor * restrict src0 = octx->src[0]; // first weight const struct htp_tensor * restrict act = octx->src[n_weights]; // activation x - bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 || - src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL || - src0->type == HTP_TYPE_MXFP4 || src0->type == HTP_TYPE_Q4_K || - src0->type == HTP_TYPE_Q5_K || src0->type == HTP_TYPE_Q3_K || - src0->type == HTP_TYPE_Q2_K); - struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; mmctx->octx = octx; @@ -4601,9 +4464,8 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) { mmctx->cur_m_rows = act_nrows; const size_t src0_row_size = src0->nb[1]; - const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); - if (hvx_mm_init_vec_dot(mmctx, src0->type) != 0) { + if (hvx_mm_init_tiled_vec_dot(mmctx, src0->type) != 0) { return HTP_STATUS_NO_SUPPORT; } @@ -4615,7 +4477,7 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) { uint32_t n_quant_tasks = 1; if (act_nrows < octx->n_threads) { n_quant_tasks = MIN(total_nb, octx->n_threads); - quant_task_func = htp_mm_act_quant_block_func(src0->type); + quant_task_func = htp_mm_act_quant_block_func(src0->type, act->type); for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { uint32_t ib_first = (total_nb * ith) / n_quant_tasks; uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks; @@ -4626,12 +4488,10 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) { } } else { n_quant_tasks = MIN(act_nrows, octx->n_threads); - quant_task_func = htp_mm_act_quant_row_func(src0->type); + quant_task_func = htp_mm_act_quant_row_func(src0->type, act->type); } - const size_t act_row_size = htp_mm_weight_has_offset(src0->type) - ? htp_mm_q8_1_tiled_row_size(act->ne[0]) - : htp_mm_q8_0_tiled_row_size(act->ne[0]); + const size_t act_row_size = htp_mm_tiled_act_row_size(src0->type, act->ne[0]); struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], act_nrows, octx->n_threads, @@ -4651,7 +4511,7 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) { mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); mmctx->vtcm_act_raw = VTCM_LAYOUT_PTR(uint8_t, base, L.off_act_raw); - mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded; + mmctx->vtcm_src0_stride = 0; mmctx->vtcm_act_stride = act_row_size; mmctx->vtcm_act_raw_stride = hex_round_up(act->ne[0] * sizeof(float), QK_Q8_0_TILED * sizeof(float)); @@ -4659,28 +4519,6 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) { mmctx->vtcm_act_size_per_thread = L.act_bytes; mmctx->vtcm_dst_size_per_thread = fastdiv(L.dst_bytes, &octx->n_threads_div); - // Run fused matmul - const uint32_t n_matmul_jobs = octx->n_threads; - worker_callback_t matmul_job_func; - if (is_repacked) { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_nx_2d_repacked_q4_0; break; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: matmul_job_func = hvx_mm_nx_2d_repacked_q4_1; break; - case HTP_TYPE_Q5_K: matmul_job_func = hvx_mm_nx_2d_repacked_q5_k; break; - case HTP_TYPE_Q3_K: matmul_job_func = hvx_mm_nx_2d_repacked_q3_k; break; - case HTP_TYPE_Q2_K: matmul_job_func = hvx_mm_nx_2d_repacked_q2_k; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_nx_2d_repacked_q8_0; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_nx_2d_repacked_iq4nl; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_nx_2d_repacked_mxfp4; break; - default: - htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); - return HTP_STATUS_NO_SUPPORT; - } - } else { - matmul_job_func = hvx_mm_nx_2d; - } - htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); hvx_mm_transfer_act_dma(octx, kparams, act, mmctx->vtcm_act_raw, mmctx->vtcm_act_raw_stride, 0, act_nrows); @@ -4689,7 +4527,7 @@ static int op_matmul_nx_impl(struct htp_ops_context * octx) { mmctx->n_quant_tasks = n_quant_tasks; work_queue_run(octx->ctx->work_queue, quant_task_func, mmctx, n_quant_tasks); - work_queue_run(octx->ctx->work_queue, matmul_job_func, mmctx, n_matmul_jobs); + work_queue_run(octx->ctx->work_queue, hvx_mm_nx_2d_tiled, mmctx, octx->n_threads); return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.h b/ggml/src/ggml-hexagon/htp/matmul-ops.h index 522d8779feb6..339eea4a7b8b 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.h +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.h @@ -4,6 +4,7 @@ #include <stdint.h> #include <stddef.h> #include "htp-ops.h" +#include "htp-tensor.h" #include "hex-fastdiv.h" #include "hex-common.h" #include "htp-vtcm.h" @@ -19,45 +20,10 @@ extern "C" { #define HTP_MM_HMX_TILE_N_ELMS 1024 #define HTP_MM_HMX_MIN_NROWS 4 -// --- Weight Repacked Tile Sizes --- -#define HTP_MM_WEIGHT_TILE_SIZE_Q4_0 576 -#define HTP_MM_WEIGHT_TILE_SIZE_Q4_1 640 -#define HTP_MM_WEIGHT_TILE_SIZE_Q8_0 1088 -#define HTP_MM_WEIGHT_TILE_SIZE_IQ4_NL 576 -#define HTP_MM_WEIGHT_TILE_SIZE_MXFP4 544 -// Q5_K: the Q4_1 tile (640) followed by a 128-byte plane with the 5th bit of every quant, transposed so that -// plane byte l holds the eight flags of lane l: bit 2i = low nibble of nibble vector i, bit 2i+1 = high nibble -#define HTP_MM_WEIGHT_TILE_SIZE_Q5_K 768 -// Q6_K native 6-bit tile (32 rows x 32 k), vrmpy-ready: byte 4*row+b of a vector holds k = 4*group+b -// vectors 0..3: low nibbles, vector i holds group 2i (low nibble) and group 2i+1 (high nibble) -// vectors 4..5: high 2 bits, vector m holds groups 4m..4m+3 at bit offsets 0,2,4,6 -// vector 6: fp16 scales per row, d * scales[]: k 0..15 in lanes 0..31, k 16..31 in lanes 32..63 -#define HTP_MM_WEIGHT_TILE_SIZE_Q6_K 896 -// Q3_K native 3-bit tile, vrmpy-ready like Q6_K -// vectors 0..1: low 2 bits, vector m holds groups 4m..4m+3 at bit offsets 0,2,4,6 -// vector 2: bit g set where the hmask bit of group g is clear (quant = low 2 bits - 4) -// vector 3: fp16 scales per row, d * (scales[] - 32): k 0..15 in lanes 0..31, k 16..31 in lanes 32..63 -#define HTP_MM_WEIGHT_TILE_SIZE_Q3_K 512 -// Q2_K native 2-bit tile with the ggml scale bytes, vrmpy-ready like Q6_K (same bits per weight as block_q2_K) -// bytes 0..255: unsigned 2-bit quants, vector m holds groups 4m..4m+3 at bit offsets 0,2,4,6 -// bytes 256..319: scales[] bytes (min << 4 | scale), sub-block s of row r at 32 * s + r -// bytes 320..335: 8 fp16, tile j of a super-block: d (j < 4) or -dmin (j >= 4) of rows 8 * (j % 4) .. 8 * (j % 4) + 7 -#define HTP_MM_WEIGHT_TILE_SIZE_Q2_K 336 - -// --- Weight Repacked Aligned Tile Sizes --- -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_0 640 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_1 640 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q8_0 1152 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_IQ4_NL 640 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_MXFP4 640 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q5_K 768 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q6_K 896 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q3_K 512 -#define HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q2_K 384 - // --- Activation Tiled Block Sizes (including padding) --- #define HTP_MM_ACT_TILE_SIZE_Q8_0 1152 #define HTP_MM_ACT_TILE_SIZE_Q8_1 1280 +#define HTP_MM_ACT_TILE_SIZE_F16 2048 #define HTP_MM_MAX_PREFETCH 16 @@ -74,11 +40,12 @@ enum htp_mm_kernel_type { // HMX paths HTP_MM_KERNEL_HMX_2D, - HTP_MM_KERNEL_HMX_F16_BATCHED, + HTP_MM_KERNEL_HMX_BATCHED, // HVX floating-point paths HTP_MM_KERNEL_HVX_F16_F16_VTCM, HTP_MM_KERNEL_HVX_F32_F32_VTCM, + HTP_MM_KERNEL_HVX_BF16_F32_VTCM, // HVX quantized paths HTP_MM_KERNEL_HVX_QUANT_ROW, // standard row-wise parallel quantization @@ -212,57 +179,6 @@ static inline int htp_mm_hmx_compute_chunks(size_t vtcm_total, return 0; } -// --- Tile Size Helpers --- -static inline uint32_t htp_mm_get_weight_tile_size(int weight_type) { - switch (weight_type) { - case HTP_TYPE_Q4_0: - case HTP_TYPE_IQ4_NL: - return HTP_MM_WEIGHT_TILE_SIZE_Q4_0; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: - return HTP_MM_WEIGHT_TILE_SIZE_Q4_1; - case HTP_TYPE_Q8_0: - return HTP_MM_WEIGHT_TILE_SIZE_Q8_0; - case HTP_TYPE_Q5_K: - return HTP_MM_WEIGHT_TILE_SIZE_Q5_K; - case HTP_TYPE_Q6_K: - return HTP_MM_WEIGHT_TILE_SIZE_Q6_K; - case HTP_TYPE_Q3_K: - return HTP_MM_WEIGHT_TILE_SIZE_Q3_K; - case HTP_TYPE_Q2_K: - return HTP_MM_WEIGHT_TILE_SIZE_Q2_K; - case HTP_TYPE_MXFP4: - return HTP_MM_WEIGHT_TILE_SIZE_MXFP4; - default: - return 0; - } -} - -static inline uint32_t htp_mm_get_weight_aligned_tile_size(int weight_type) { - switch (weight_type) { - case HTP_TYPE_Q4_0: - case HTP_TYPE_IQ4_NL: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_0; - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q4_1; - case HTP_TYPE_Q8_0: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q8_0; - case HTP_TYPE_Q5_K: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q5_K; - case HTP_TYPE_Q6_K: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q6_K; - case HTP_TYPE_Q3_K: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q3_K; - case HTP_TYPE_Q2_K: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_Q2_K; - case HTP_TYPE_MXFP4: - return HTP_MM_WEIGHT_ALIGNED_TILE_SIZE_MXFP4; - default: - return 0; - } -} - // weight types whose tiles carry a per-block offset (x = d * q + m): the activations need block sums (q8_1) // (Q2_K: per-16 k sums, q8_1_s16) static inline bool htp_mm_weight_has_offset(int weight_type) { @@ -283,27 +199,17 @@ static inline size_t htp_mm_q8_1_tiled_row_size(uint32_t ne) { return nb_32 * HTP_MM_ACT_TILE_SIZE_Q8_1; } -static inline size_t htp_mm_get_tiled_row_stride(int weight_type, uint32_t k) { - uint32_t nb = (k + QK_Q4_0_TILED - 1) / QK_Q4_0_TILED; - switch (weight_type) { - case HTP_TYPE_Q4_0: - case HTP_TYPE_IQ4_NL: - case HTP_TYPE_Q4_1: - case HTP_TYPE_Q4_K: - case HTP_TYPE_Q8_0: - case HTP_TYPE_Q5_K: - case HTP_TYPE_Q6_K: - case HTP_TYPE_Q3_K: - case HTP_TYPE_Q2_K: - case HTP_TYPE_MXFP4: - return (size_t) nb * htp_mm_get_weight_tile_size(weight_type); - case HTP_TYPE_F16: - return (size_t) k * sizeof(__fp16); - case HTP_TYPE_F32: - return (size_t) k * sizeof(float); - default: - return 0; +static inline size_t htp_mm_f16_tiled_row_size(uint32_t ne) { + const uint32_t ne_padded = ((ne + 127) / 128) * 128; + const uint32_t nb_32 = ne_padded / 32; + return nb_32 * HTP_MM_ACT_TILE_SIZE_F16; +} + +static inline size_t htp_mm_tiled_act_row_size(int weight_type, uint32_t ne) { + if (weight_type == HTP_TYPE_F16) { + return htp_mm_f16_tiled_row_size(ne); } + return htp_mm_weight_has_offset(weight_type) ? htp_mm_q8_1_tiled_row_size(ne) : htp_mm_q8_0_tiled_row_size(ne); } static inline size_t htp_mm_round_up(size_t n, size_t m) { @@ -318,8 +224,8 @@ static inline void htp_mm_hmx_get_2d_chunk_costs( int wtype, uint32_t k, bool pipeline, uint32_t aligned_tile_size, size_t * size_per_n_out, size_t * size_per_m_out, size_t * size_per_mn_out ) { - const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32); - const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k); + const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32 && wtype != HTP_TYPE_BF16); + const size_t row_stride = htp_tensor_get_row_size(wtype, k); const size_t vec_dot_size = k * sizeof(uint16_t); const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS; const size_t qweight_row_stride = is_quant ? (size_t)(n_k_tiles * aligned_tile_size) / 32 : 0; @@ -404,7 +310,7 @@ static inline void htp_mm_hmx_vtcm_layout_build( ) { size_t off = 0; - if (kernel_type == HTP_MM_KERNEL_HMX_F16_BATCHED) { + if (kernel_type == HTP_MM_KERNEL_HMX_BATCHED) { const size_t vec_dot_size = k * sizeof(uint16_t); const size_t act_head_stride = mc * k; const size_t weight_area_size = hex_align_up(nc * vec_dot_size, HTP_MM_HMX_TILE_SIZE); @@ -450,8 +356,8 @@ static inline void htp_mm_hmx_vtcm_layout_build( off = off_group_a + hex_smax(group_b_size, group_c_size); } else { // HTP_MM_KERNEL_HMX_2D - const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32); - const size_t row_stride = htp_mm_get_tiled_row_stride(wtype, k); + const bool is_quant = (wtype != HTP_TYPE_F16 && wtype != HTP_TYPE_F32 && wtype != HTP_TYPE_BF16); + const size_t row_stride = htp_tensor_get_row_size(wtype, k); const size_t vec_dot_size = k * sizeof(uint16_t); const uint32_t n_k_tiles = k / HTP_MM_HMX_TILE_N_COLS; @@ -527,28 +433,13 @@ static inline void htp_mm_hvx_vtcm_layout_build( size_t dst_sz = 0; size_t act_raw_sz = 0; - const bool is_repack = (wtype == HTP_TYPE_Q4_0 || wtype == HTP_TYPE_Q4_1 || - wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL || - wtype == HTP_TYPE_MXFP4 || wtype == HTP_TYPE_Q6_K || - wtype == HTP_TYPE_Q4_K || wtype == HTP_TYPE_Q5_K || - wtype == HTP_TYPE_Q3_K || wtype == HTP_TYPE_Q2_K); - if (is_fused_nx) { - const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); + uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(wtype); + uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32; + uint32_t tile_row_size = n_k_tiles * aligned_tile_size; + size_t weight_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); - size_t weight_sz_per_thread = 0; - - if (is_repack) { - uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype); - uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32; - uint32_t tile_row_size = n_k_tiles * aligned_tile_size; - - weight_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); - } else { - weight_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); - } - - size_t tiled_act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + size_t tiled_act_row_size = htp_mm_tiled_act_row_size(wtype, ne10); size_t q_act_sz = hex_round_up(tiled_act_row_size * act_nrows, 128); size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); @@ -558,20 +449,13 @@ static inline void htp_mm_hvx_vtcm_layout_build( dst_sz = 0; act_raw_sz = hex_round_up(raw_row_size * act_nrows, 128); } else if (is_matmul_id) { - const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); - const size_t act_row_size_tiled = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) - : htp_mm_q8_0_tiled_row_size(ne10); - - size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256); - act_sz = htp_mm_round_up(act_row_size_tiled * act_nrows, 256); - - if (is_repack) { - const uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype); - const uint32_t n_k_tiles = ne10 / 32; - const uint32_t tile_row_size = n_k_tiles * aligned_tile_size; - size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256); - src0_sz_per_thread = repacked_vtcm_size; - } + const size_t act_row_size_tiled = htp_mm_tiled_act_row_size(wtype, ne10); + const uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(wtype); + const uint32_t n_k_tiles = ne10 / 32; + const uint32_t tile_row_size = n_k_tiles * aligned_tile_size; + size_t src0_sz_per_thread = htp_mm_round_up(n_prefetch * tile_row_size, 256); + + act_sz = htp_mm_round_up(act_row_size_tiled * act_nrows, 256); size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); @@ -592,7 +476,8 @@ static inline void htp_mm_hvx_vtcm_layout_build( act_raw_sz = hex_round_up(hex_round_up(ne10 * sizeof(float), 128) * act_nrows, 128); break; } - case HTP_MM_KERNEL_HVX_F32_F32_VTCM: { + case HTP_MM_KERNEL_HVX_F32_F32_VTCM: + case HTP_MM_KERNEL_HVX_BF16_F32_VTCM: { size_t f32_act_row_size = htp_mm_round_up(ne10 * 4, 128); act_sz = htp_mm_round_up(f32_act_row_size * act_nrows, 256); src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256) * n_threads; @@ -602,21 +487,15 @@ static inline void htp_mm_hvx_vtcm_layout_build( } case HTP_MM_KERNEL_HVX_QUANT_BLOCK: case HTP_MM_KERNEL_HVX_QUANT_ROW: { - size_t q_act_row_size = htp_mm_weight_has_offset(wtype) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + size_t q_act_row_size = htp_mm_tiled_act_row_size(wtype, ne10); + uint32_t aligned_tile_size = htp_tensor_get_tile_padded_size(wtype); + uint32_t n_k_tiles = (ne10 + 31) / 32; + uint32_t tile_row_size = n_k_tiles * aligned_tile_size; + size_t tiled_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256); - src0_sz = htp_mm_round_up(n_prefetch * src0_row_size_padded, 256); + src0_sz = tiled_vtcm_size * n_threads; act_sz = htp_mm_round_up(q_act_row_size * act_nrows, 256); - src0_sz = src0_sz * n_threads; - - if (is_repack) { - uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype); - uint32_t n_k_tiles = ne10 / 32; - uint32_t tile_row_size = n_k_tiles * aligned_tile_size; - size_t repacked_vtcm_size = htp_mm_round_up(n_prefetch * tile_row_size, 256); - src0_sz = repacked_vtcm_size * n_threads; - } - size_t dst_slice_per_thread = (dst_nrows > 0 && act_nrows == 1) ? htp_mm_round_up((dst_row_size + n_threads - 1) / n_threads, 128) : 0; dst_sz = dst_slice_per_thread * n_threads; size_t raw_row_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)); @@ -687,9 +566,7 @@ static inline bool htp_mm_hvx_solve_vtcm_params( const size_t avail_act = vtcm_budget - fixed_bytes; size_t row_size = 0; if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW || kernel_type == HTP_MM_KERNEL_HVX_QUANT_BLOCK) { - row_size = htp_mm_weight_has_offset(wtype) - ? htp_mm_q8_1_tiled_row_size(ne10) - : htp_mm_q8_0_tiled_row_size(ne10); + row_size = htp_mm_tiled_act_row_size(wtype, ne10); } else if (kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM) { row_size = hex_round_up(ne10 * 2, 128); } else { @@ -749,7 +626,7 @@ static inline size_t htp_mm_hmx_get_batched_vtcm_size( int wtype, uint32_t k, size_t mc, size_t nc, uint32_t group_size, bool pipeline, uint32_t act_threads, size_t bias_size) { (void)pipeline; struct htp_mm_hmx_vtcm_layout L; - htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_F16_BATCHED, wtype, k, mc, nc, group_size, false, act_threads, 0, bias_size); + htp_mm_hmx_vtcm_layout_build(&L, HTP_MM_KERNEL_HMX_BATCHED, wtype, k, mc, nc, group_size, false, act_threads, 0, bias_size); return L.total_bytes; } diff --git a/scripts/snapdragon/run.py b/scripts/snapdragon/run.py index dec4b996fc2d..7e11d90de842 100755 --- a/scripts/snapdragon/run.py +++ b/scripts/snapdragon/run.py @@ -27,6 +27,7 @@ "GGML_HEXAGON_OPPOLL", "GGML_HEXAGON_OPFILTER", "GGML_HEXAGON_OPFUSION", + "GGML_HEXAGON_RT_REPACK", "GGML_HEXAGON_VMEM", "GGML_HEXAGON_MBUF", "GGML_HEXAGON_MM_SELECT", @@ -167,6 +168,7 @@ def main(): parser.add_argument("--hex-oppoll", default="1", help="Enable (1) or Disable (0) polling for NPU opbatch completion (GGML_HEXAGON_OPPOLL) (default: 1)") parser.add_argument("--hex-opfilter", help="Regex pattern to filter/select which operators are offloaded to NPU (GGML_HEXAGON_OPFILTER)") parser.add_argument("--hex-opfusion", help="NPU graph node fusion optimization level (0: disabled, 1: enabled) (GGML_HEXAGON_OPFUSION)") + parser.add_argument("--hex-rt-repack", help="Runtime weight repacking mode (0: disabled, 1: quantized only, 2: always) (GGML_HEXAGON_RT_REPACK)") parser.add_argument("--hex-vmem", help="Maximum NPU VMEM size limit in MB to allocate (GGML_HEXAGON_VMEM)") parser.add_argument("--hex-mbuf", help="Host buffer size limits in MB (supports K/M/G suffix): <dyn>[,<static>[,<total>]] (default: 512,1024,0) (GGML_HEXAGON_MBUF)") parser.add_argument("--hex-mm-select", help="Select MUL_MAT and MUL_MAT_ID kernel (GGML_HEXAGON_MM_SELECT) 2:HMX,1:HVX,0:disable") @@ -311,6 +313,7 @@ def set_env(env_name, opt_val): set_env("GGML_HEXAGON_OPPOLL", args.hex_oppoll) set_env("GGML_HEXAGON_OPFILTER", args.hex_opfilter) set_env("GGML_HEXAGON_OPFUSION", args.hex_opfusion) + set_env("GGML_HEXAGON_RT_REPACK", args.hex_rt_repack) set_env("GGML_HEXAGON_VMEM", args.hex_vmem) set_env("GGML_HEXAGON_MBUF", args.hex_mbuf) set_env("GGML_HEXAGON_MM_SELECT", args.hex_mm_select) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index ba7f7eb4e56c..b50364660ed9 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10741,11 +10741,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() { // testing very large N test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 16, 2105344, 32, {1, 1}, {1, 1})); + +#if 0 test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 16, 8400000, 32, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 16, 8400000, 32, {1, 1}, {1, 1})); // same shapes with a repeat in dim 3 test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 16, 8400000, 32, {1, 1}, {2, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 16, 8400000, 32, {1, 1}, {2, 1})); +#endif test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 31, 509, 2051, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 32, 509, 2112, {1, 1}, {1, 1})); From 5e03a1ff450cae37a774684de55d14fb48b1f90b Mon Sep 17 00:00:00 2001 From: Mikhail Podvitskii <podvitskiymichael@gmail.com> Date: Fri, 16 Oct 2026 05:32:26 +0200 Subject: [PATCH 042/423] graph : add precision handling for MELLUM architecture (#30584) --- src/llama-graph.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 80be4bb57373..454c21404266 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2452,7 +2452,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, slots); // [n_embd, n_expert_used, n_tokens] - if (arch == LLM_ARCH_MISTRAL4) { + if (arch == LLM_ARCH_MISTRAL4 || arch == LLM_ARCH_MELLUM) { // src1 can exceed F16 range ggml_prec_set_src(experts, GGML_PREC_F32, 1); } From f12aaa344c971556ec35be054d7aac1d901fd510 Mon Sep 17 00:00:00 2001 From: Karthik Vijayan <mailkv23@gmail.com> Date: Fri, 16 Oct 2026 04:46:55 +0000 Subject: [PATCH 043/423] Fix failing RocM tests (#30611) Co-authored-by: Karthik Vijayan <kavijay@microsoft.com> --- tests/test-llama-archs.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 2bfa1cf75bc1..f2c1ca90358b 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -459,7 +459,6 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { if (arch == LLM_ARCH_MAION_CODER) { ms.add_kv(LLM_KV_MOE_LATENT_SIZE, n_embd / 2); - ms.add_kv(LLM_KV_EXPERT_COUNT, uint32_t(4)); ms.add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, std::vector<uint32_t>{n_ff, n_ff / 2, n_ff, n_ff}); } From a8d9a7e1e55f5b613f297504d9e159807ec3a99b Mon Sep 17 00:00:00 2001 From: Ruben Ortlam <rortlam@redhat.com> Date: Fri, 16 Oct 2026 08:29:23 +0200 Subject: [PATCH 044/423] vulkan: fix gpu-visible host memory issues (#30468) * vulkan: fix log spam if host pointer import fails * exclude MoltenVK from shared memory import * remove unreachable code, assert instead. Shorten warnings * disable host memory buffer sharing on AMD due to amdgpu perf issue --- ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp | 16 +++------ ggml/src/ggml-vulkan/ggml-vulkan-types.h | 2 ++ ggml/src/ggml-vulkan/ggml-vulkan.cpp | 36 ++++++++++++++------ 3 files changed, 32 insertions(+), 22 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp b/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp index eabd8974af77..b687b139082a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp @@ -1,6 +1,5 @@ #include "ggml-vulkan-common.h" -#include <atomic> #include <cstdint> #include <cstdlib> #include <mutex> @@ -363,21 +362,16 @@ void ggml_vk_host_get(const vk_device& device, const void * ptr, vk_buffer& buf, vk_host_allocation & entry = vk_instance.host_allocations[index]; if (entry.views[device->idx] == nullptr) { - if (!vk_instance.multi_device) { - return; - } - if (!device->external_memory_host) { - static std::atomic<uint32_t> warned_mask{0}; - const uint32_t bit = 1u << device->idx; - if (!(warned_mask.fetch_or(bit) & bit)) { - GGML_LOG_WARN("ggml_vulkan: %s does not support VK_EXT_external_memory_host; host memory " - "will be staged, not pinned, on this device\n", device->name.c_str()); - } + const uint32_t bit = 1u << device->idx; + if (!vk_instance.multi_device || !vk_instance.device_host_buffer_import[device->idx] || (vk_instance.host_import_failed & bit)) { return; } + GGML_ASSERT(device->external_memory_host); // const only propagates from the const context; the pointee is mutable entry.views[device->idx] = ggml_vk_buffer_from_host_ptr(const_cast<vk_device&>(device), entry.base, entry.size); if (entry.views[device->idx] == nullptr) { + vk_instance.host_import_failed |= bit; + GGML_LOG_WARN("ggml_vulkan: %s: host memory import failed, using staging\n", device->name.c_str()); return; } } diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h index a41ce596625a..00d20e843b78 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h @@ -1329,12 +1329,14 @@ struct vk_instance_t { std::vector<size_t> device_indices; std::vector<bool> device_supports_membudget; + std::vector<bool> device_host_buffer_import; vk_device devices[GGML_VK_MAX_DEVICES]; std::vector<vk_host_allocation> host_allocations; std::shared_mutex host_allocations_mutex; size_t host_import_alignment = 0; + uint32_t host_import_failed = 0; // bitmask by device->idx, guarded by host_allocations_mutex bool multi_device = false; }; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 436bfe3bb054..fa038bb5b61f 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3987,6 +3987,17 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } } +static bool ggml_vk_external_memory_host_usable(vk::DriverId driver_id) { + // Disable external_memory_host until https://github.com/KhronosGroup/MoltenVK/pull/2622 + // is available in the Vulkan SDK. + return driver_id != vk::DriverId::eMoltenvk; +} + +static bool ggml_vk_host_buffer_import_usable(uint32_t vendor_id) { + // amdgpu re-faults all pages of every imported host allocation on each submission + return vendor_id != VK_VENDOR_ID_AMD; +} + vk_device ggml_vk_get_device(size_t idx) { VK_LOG_DEBUG("ggml_vk_get_device(" << idx << ")"); @@ -4174,11 +4185,7 @@ vk_device ggml_vk_get_device(size_t idx) { device->vendor_id = device->properties.vendorID; device->driver_id = driver_props.driverID; - if (device->driver_id == vk::DriverId::eMoltenvk) { - // Disable external_memory_host until https://github.com/KhronosGroup/MoltenVK/pull/2622 - // is available in the Vulkan SDK. - device->external_memory_host = false; - } + device->external_memory_host = device->external_memory_host && ggml_vk_external_memory_host_usable(device->driver_id); // Implementing the async backend interfaces seems broken on older Intel HW, // see https://github.com/ggml-org/llama.cpp/issues/17302. @@ -5421,7 +5428,7 @@ void ggml_vk_instance_init() { } GGML_LOG_DEBUG("ggml_vulkan: Found %zu Vulkan devices:\n", vk_instance.device_indices.size()); - bool all_external_memory_host = vk_instance.device_indices.size() > 1; + size_t host_buffer_import_count = 0; for (size_t i = 0; i < vk_instance.device_indices.size(); i++) { vk::PhysicalDevice vkdev = devices[vk_instance.device_indices[i]]; std::vector<vk::ExtensionProperties> extensionprops = vkdev.enumerateDeviceExtensionProperties(); @@ -5438,20 +5445,27 @@ void ggml_vk_instance_init() { vk_instance.device_supports_membudget.push_back(membudget_supported); + bool host_buffer_import = false; if (external_memory_host_supported) { vk::PhysicalDeviceProperties2 props2; + vk::PhysicalDeviceDriverProperties driver_props; vk::PhysicalDeviceExternalMemoryHostPropertiesEXT ext_props; - props2.pNext = &ext_props; + props2.pNext = &driver_props; + driver_props.pNext = &ext_props; vkdev.getProperties2(&props2); - vk_instance.host_import_alignment = std::max<size_t>(vk_instance.host_import_alignment, ext_props.minImportedHostPointerAlignment); - } else { - all_external_memory_host = false; + host_buffer_import = ggml_vk_external_memory_host_usable(driver_props.driverID) && + ggml_vk_host_buffer_import_usable(props2.properties.vendorID); + if (host_buffer_import) { + vk_instance.host_import_alignment = std::max<size_t>(vk_instance.host_import_alignment, ext_props.minImportedHostPointerAlignment); + host_buffer_import_count++; + } } + vk_instance.device_host_buffer_import.push_back(host_buffer_import); ggml_vk_print_gpu_info(i); } - vk_instance.multi_device = all_external_memory_host; + vk_instance.multi_device = host_buffer_import_count > 1 && vk_instance.device_host_buffer_import[0]; } void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) { From 88aaf325a7026ac1127d0897394b0a3d6413dc2a Mon Sep 17 00:00:00 2001 From: Georgi Gerganov <ggerganov@gmail.com> Date: Fri, 16 Oct 2026 09:35:19 +0300 Subject: [PATCH 045/423] tests : catch vanished rows in the fusion baseline (#30618) * tests : update the Metal fusion baseline Re-record the MTL fusion baseline with the current build: 276 rows over 131 dummy models. Assisted-by: pi:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) * tests : fail test-fusion when a baseline row stops matching The report was built from the counters the run produced, and the baseline was only ever used as a lookup table, so a fusion that stopped matching produced no row and was never compared. qwen4exp and mellum lost their fusions this way while the Fusion / metal job stayed green. Report a baseline row that the run does not reproduce as a regression. Scope the pass to the arches that were actually run so --model keeps working. Assisted-by: pi:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) * pi : never write inline comments unless the user asks Assisted-by: pi:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) * tests : keep the mellum dummy model softmax-gated mellum reads expert_gating_func from the GGUF since #30578, while the fixture writes sigmoid for every MoE arch, so the Mellum dummy model stopped building the softmax MoE routing fusion. Write softmax for mellum, matching the arch default and the GGUFs that conversion/mellum.py produces. Re-record the MTL baseline: the routing row is back, 277 ok / 0 failed. Assisted-by: pi:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) * pi : use emoji in the AI stats disclosure Assisted-by: pi:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) * pi : list the disclosure model after the hardware Assisted-by: pi:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) * pi : use monochrome symbols in the AI stats disclosure Assisted-by: pi:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) * tests : add muse-spark to the Metal fusion baseline Assisted-by: pi:llama.cpp/ggml-org/Qwen3.8-Flash-Next-GGUF (MacBook, M5 Max, 128GB) --- .pi/gg/SYSTEM.md | 9 ++++++--- tests/fusion/MTL.csv | 7 ++++++- tests/fusion/README.md | 4 ++-- tests/test-fusion.cpp | 39 +++++++++++++++++++++++++++++++++----- tests/test-llama-archs.cpp | 15 ++++++++++++++- 5 files changed, 62 insertions(+), 12 deletions(-) diff --git a/.pi/gg/SYSTEM.md b/.pi/gg/SYSTEM.md index b086d72162bc..927224fbfc20 100644 --- a/.pi/gg/SYSTEM.md +++ b/.pi/gg/SYSTEM.md @@ -11,6 +11,7 @@ General: - Never read the `AGENTS.md` file Coding: +- Never write inline comments in the code unless the user tells you to - When in doubt, always refer to the CONTRIBUTING.md file of the project - In `test-backend-ops.cpp`, do not mention specific backends (e.g. Metal, CUDA) in comments - When referencing issues or PRs in comments, use the format: @@ -23,13 +24,15 @@ Pull requests (PRs): - Don't explicitly wrap lines in the PR description (each paragraph and bullet is a single line) - When creating a pull request, look for the repository's PR template and follow it - For the AI usage disclosure section, write "YES", then these 4 bullets: - - Model: `<model>`, linked to its Hugging Face page - the link text is the plain model name - Harness: `pi:llama.cpp` - Hardware: `<hardware>` + - Model: `<model>`, linked to its Hugging Face page - the link text is the plain model name - Stats: `<stats>` - Run `~/.pi/stats.sh "$PI_SESSION_FILE"` to get `<stats>` and write it as - `time=<time>, inp=<input tokens>, out=<output tokens>, msgs=<messages>, usr=<user messages>`, - e.g. `time=12.3m, inp=120k, out=45k, msgs=87, usr=6` + `◷ <time> · ↑ <input tokens> · ↓ <output tokens> · ≡ <messages> · ⌨ <user messages>`, + e.g. `◷ 12.3m · ↑ 120k · ↓ 45k · ≡ 87 · ⌨ 6` +- Write these symbols as literal unicode in the PR description, never `:shortcodes:`, and keep + commit trailers ASCII - If `PI_MODEL_NAME` or `PI_HARDWARE` is not set, ask the user and use their answer for `<model>` or `<hardware>` - Always create the pull requests in draft mode - Never reply to review comments or post comments on issues/PRs without explicit permission from the user diff --git a/tests/fusion/MTL.csv b/tests/fusion/MTL.csv index e839e20ce3f5..d31a9f23e5f0 100644 --- a/tests/fusion/MTL.csv +++ b/tests/fusion/MTL.csv @@ -156,6 +156,10 @@ llama ,0 ,any ,SOFT_MAX+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, llama4 ,0 ,any ,ADD+ADD , 2 llama4 ,0 ,any ,RMS_NORM+MUL , 9 maincoder ,0 ,any ,RMS_NORM+MUL , 9 +maion-coder ,0 ,any ,MUL+ADD , 4 +maion-coder ,0 ,any ,RMS_NORM+MUL , 17 +maion-coder ,0 ,any ,RMS_NORM+MUL+ADD , 8 +maion-coder ,0 ,any ,SOFT_MAX+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, 2 mamba ,0 ,any ,RMS_NORM+MUL , 3 mamba2 ,0 ,any ,RMS_NORM+MUL , 5 maple ,0 ,any ,MUL+ADD , 4 @@ -189,6 +193,8 @@ mistral4 ,0 ,any ,RMS_NORM+MUL , 9 mpt ,0 ,any ,NORM+MUL+ADD , 5 muse-glimmer ,0 ,any ,RMS_NORM+MUL , 10 muse-glimmer ,0 ,any ,RMS_NORM+MUL+ADD , 3 +muse-spark ,0 ,any ,MUL+ADD , 4 +muse-spark ,0 ,any ,RMS_NORM+MUL , 14 nanbeige ,0 ,any ,RMS_NORM+MUL , 5 nemotron ,0 ,any ,NORM+MUL+ADD , 5 nemotron_h ,0 ,any ,RMS_NORM+MUL , 5 @@ -251,7 +257,6 @@ qwen3vlmoe ,1 ,any ,MUL+ADD , 4 qwen3vlmoe ,1 ,any ,RMS_NORM+MUL , 9 qwen3vlmoe ,1 ,any ,SOFT_MAX+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV, 2 qwen4exp ,0 ,any ,ADD+ADD+ADD , 2 -qwen4exp ,0 ,any ,ADD+ADD+ADD+ADD+ADD+ADD+ADD , 9 qwen4exp ,0 ,any ,GATED_DELTA_NET+CPY , 1 qwen4exp ,0 ,any ,MUL+ADD , 4 qwen4exp ,0 ,any ,RMS_NORM+MUL , 13 diff --git a/tests/fusion/README.md b/tests/fusion/README.md index 3ac02b061e9a..9951304fbfa6 100644 --- a/tests/fusion/README.md +++ b/tests/fusion/README.md @@ -22,5 +22,5 @@ rm -rf build-ci-models && mkdir -p build-ci-models ``` Non-zero exit means a row differs from the baseline. Use `--model FILE` to run a single -architecture. Note `--check` only sees present rows — a fusion that stops matching is not -reported, so diff the recorded CSV to catch removed patterns. +architecture. `--check` also reports baseline rows that the run no longer produces, so a fusion +that stops matching fails the run instead of silently dropping out of the CSV. diff --git a/tests/test-fusion.cpp b/tests/test-fusion.cpp index 55e4f9df46c7..c268472ca72e 100644 --- a/tests/test-fusion.cpp +++ b/tests/test-fusion.cpp @@ -32,6 +32,7 @@ #include <iostream> #include <map> #include <random> +#include <set> #include <string> #include <vector> @@ -191,6 +192,14 @@ static void read_counts(fusion_stats_get_t api_stats_get, ggml_backend_fusion_t api_stats_get(finfo, labels.data(), counts.data(), n); } +struct baseline_row { + std::string arch; + bool moe; + std::string mode; + std::string label; + uint64_t count; +}; + // one row of the per-label report struct fusion_row { std::string arch; @@ -329,8 +338,8 @@ int main(int argc, char ** argv) { const bool has_counts = true; - // load the baseline (if any): key arch|moe|mode|label -> expected count - std::map<std::string, uint64_t> baseline; + // load the baseline (if any): key arch|moe|mode|label -> expected row + std::map<std::string, baseline_row> baseline; if (!check_path.empty()) { std::ifstream in(check_path); if (!in) { @@ -352,12 +361,16 @@ int main(int argc, char ** argv) { if (cols.size() != 5) { continue; } - baseline[cols[0] + "|" + cols[1] + "|" + cols[2] + "|" + cols[3]] = std::stoull(cols[4]); + const std::string key = cols[0] + "|" + cols[1] + "|" + cols[2] + "|" + cols[3]; + baseline[key] = { cols[0], cols[1] == "1", cols[2], cols[3], std::stoull(cols[4]) }; } } std::vector<fusion_row> rows; + std::set<std::string> run_arches; + std::set<std::string> produced; + LOG_INF("%s: running fusion test over %zu models on '%s'\n", __func__, models.size(), base_name.c_str()); const size_t seed = 1; @@ -378,6 +391,8 @@ int main(int argc, char ** argv) { continue; } + run_arches.insert(arch); + struct mode_cfg { std::string name; std::vector<float> (*decode)(llama_model *, llama_context *, const std::vector<llama_token> &); @@ -480,7 +495,8 @@ int main(int argc, char ** argv) { if (match) { // one "any" row; use the worst NMSE across the two modes const std::string any_key = arch + "|" + (moe ? "1" : "0") + "|any|" + label; - const uint64_t expected = baseline.count(any_key) ? baseline.at(any_key) : 0; + produced.insert(any_key); + const uint64_t expected = baseline.count(any_key) ? baseline.at(any_key).count : 0; const bool ok_count = check_path.empty() || d[0].count_fused == expected; const bool ok_nmse = d[0].ok_nmse && d[1].ok_nmse; const double nmse_fus = std::max(d[0].nmse_fus, d[1].nmse_fus); @@ -495,7 +511,8 @@ int main(int argc, char ** argv) { } const mode_data & a = d[mi]; const std::string mode_key = arch + "|" + (moe ? "1" : "0") + "|" + modes[mi].name + "|" + label; - const uint64_t expected = baseline.count(mode_key) ? baseline.at(mode_key) : 0; + produced.insert(mode_key); + const uint64_t expected = baseline.count(mode_key) ? baseline.at(mode_key).count : 0; const bool ok_count = check_path.empty() || a.count_fused == expected; rows.push_back({ arch, moe, modes[mi].name, label, a.count_fused, a.count_unfused, expected, a.nmse_fus, a.nmse_dev, ok_count, a.ok_nmse }); @@ -507,6 +524,18 @@ int main(int argc, char ** argv) { LOG_INF("%s: %-20s (%s) done\n", __func__, arch.c_str(), model_path.c_str()); } + if (!check_path.empty()) { + for (const auto & kv : baseline) { + if (produced.count(kv.first) || !run_arches.count(kv.second.arch)) { + continue; + } + const baseline_row & b = kv.second; + LOG_WRN("%s: %s: fusion %s x%llu is no longer matched\n", __func__, b.arch.c_str(), + b.label.c_str(), (unsigned long long) b.count); + rows.push_back({ b.arch, b.moe, b.mode, b.label, 0, 0, b.count, 0.0, 0.0, false, true }); + } + } + // print the report { std::ofstream out(record_path); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index f2c1ca90358b..eb72553b3e81 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -448,7 +448,20 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_EXPERT_COUNT, uint32_t(host_experts_test(arch) ? 64 : 2)); ms.add_kv(LLM_KV_EXPERT_USED_COUNT, uint32_t(2)); ms.add_kv(LLM_KV_EXPERT_SHARED_COUNT, uint32_t(1)); - ms.add_kv(LLM_KV_EXPERT_GATING_FUNC, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(4) : uint32_t(2)); // sqrtsoftplus : sigmoid + + uint32_t gating_func = 2; + switch (arch) { + case LLM_ARCH_MELLUM: + gating_func = 1; + break; + case LLM_ARCH_DEEPSEEK4: + gating_func = 4; + break; + default: + break; + } + + ms.add_kv(LLM_KV_EXPERT_GATING_FUNC, gating_func); // sigmoid : softmax : sqrtsoftplus ms.add_kv(LLM_KV_EXPERT_GROUP_SCALE, 1.0f); ms.add_kv(LLM_KV_EXPERTS_PER_GROUP, uint32_t(1)); if (arch == LLM_ARCH_K2_HORIZON) { From 13de11e78fe102f04862d2ccdf4595843499cf5c Mon Sep 17 00:00:00 2001 From: Titaniumtown <titaniumtown@proton.me> Date: Thu, 15 Oct 2026 23:43:07 -0700 Subject: [PATCH 046/423] sycl: wide stores for the q4_K/q5_K weight dequant (#29696) --- ggml/src/ggml-sycl/convert.cpp | 48 ++------- ggml/src/ggml-sycl/dequantize.hpp | 163 ++++++++++++++++-------------- 2 files changed, 95 insertions(+), 116 deletions(-) diff --git a/ggml/src/ggml-sycl/convert.cpp b/ggml/src/ggml-sycl/convert.cpp index 08eca7a36a12..67dfeb9f943e 100644 --- a/ggml/src/ggml-sycl/convert.cpp +++ b/ggml/src/ggml-sycl/convert.cpp @@ -212,20 +212,10 @@ template <typename dst_t> static void dequantize_row_q4_K_sycl(const void *vx, dst_t *y, const int64_t k, dpct::queue_ptr stream) { const int64_t nb = k / QK_K; - { - dpct::has_capability_or_fail(stream->get_device(), - {sycl::aspect::fp16}); - - stream->submit([&](sycl::handler &cgh) { - sycl::local_accessor<uint8_t, 1> scale_local_acc(sycl::range<1>(12), cgh); - cgh.parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, nb) * - sycl::range<3>(1, 1, 32), - sycl::range<3>(1, 1, 32)), - [=](sycl::nd_item<3> item_ct1) { - dequantize_block_q4_K(vx, y, get_pointer(scale_local_acc), item_ct1); - }); - }); - } + dpct::has_capability_or_fail(stream->get_device(), { sycl::aspect::fp16 }); + stream->parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t) nb * SYCL_DEQK_WI_PER_BLOCK), + sycl::range<1>((size_t) SYCL_DEQK_WI_PER_BLOCK)), + [=](sycl::nd_item<1> it) { dequantize_block_q4_K_wide(vx, y, nb, it); }); } template <typename dst_t> @@ -250,32 +240,10 @@ template <typename dst_t> static void dequantize_row_q5_K_sycl(const void *vx, dst_t *y, const int64_t k, dpct::queue_ptr stream) { const int64_t nb = k / QK_K; -#if QK_K == 256 - { - dpct::has_capability_or_fail(stream->get_device(), - {sycl::aspect::fp16}); - - stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, nb) * - sycl::range<3>(1, 1, 64), - sycl::range<3>(1, 1, 64)), - [=](sycl::nd_item<3> item_ct1) { - dequantize_block_q5_K(vx, y, item_ct1); - }); - } -#else - { - dpct::has_capability_or_fail(stream->get_device(), - {sycl::aspect::fp16}); - - stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, nb) * - sycl::range<3>(1, 1, 32), - sycl::range<3>(1, 1, 32)), - [=](sycl::nd_item<3> item_ct1) { - dequantize_block_q5_K(vx, y, item_ct1); - }); - } - -#endif + dpct::has_capability_or_fail(stream->get_device(), { sycl::aspect::fp16 }); + stream->parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t) nb * SYCL_DEQK_WI_PER_BLOCK), + sycl::range<1>((size_t) SYCL_DEQK_WI_PER_BLOCK)), + [=](sycl::nd_item<1> it) { dequantize_block_q5_K_wide(vx, y, nb, it); }); } template <typename dst_t> diff --git a/ggml/src/ggml-sycl/dequantize.hpp b/ggml/src/ggml-sycl/dequantize.hpp index bc86d2c2af0c..14529a26beb3 100644 --- a/ggml/src/ggml-sycl/dequantize.hpp +++ b/ggml/src/ggml-sycl/dequantize.hpp @@ -1108,6 +1108,53 @@ static inline void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8 } #endif +// Work-items that cover one 256-weight K-quant block: four 64-element chunks, 8 work-items +// per chunk, 4 qs bytes per work-item. +static constexpr int SYCL_DEQK_WI_PER_BLOCK = 32; + +template <typename dst_t> +static inline void dequantize_q4_K_wide_one(dst_t * __restrict__ y, const uint8_t * __restrict__ qs, + const uint8_t * __restrict__ scales, const float dall, + const float dmin, const int tid) { + const int il = tid >> 3; // 64-element chunk, 0..3 + const int ir = (tid & 7) * 4; // 4 qs bytes within the chunk + + uint8_t sc, m; + get_scale_min_k4(2 * il + 0, scales, sc, m); + const float d1 = dall * sc; + const float m1 = dmin * m; + get_scale_min_k4(2 * il + 1, scales, sc, m); + const float d2 = dall * sc; + const float m2 = dmin * m; + + const sycl::vec<uint8_t, 4> q = vec_aligned_load<uint8_t, 4>(qs + 32 * il + ir); + + sycl::vec<dst_t, 4> lo, hi; +#pragma unroll + for (int l = 0; l < 4; ++l) { + lo[l] = d1 * (q[l] & 0xF) - m1; + hi[l] = d2 * (q[l] >> 4) - m2; + } + dst_t * yb = y + 64 * il + ir; + *reinterpret_cast<sycl::vec<dst_t, 4> *>(yb + 0) = lo; + *reinterpret_cast<sycl::vec<dst_t, 4> *>(yb + 32) = hi; +} + +template <typename dst_t> +static void dequantize_block_q4_K_wide(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t nb, + const sycl::nd_item<1> & it) { + static_assert(QK_K == 256, "wide K-quant dequant assumes QK_K == 256"); + const int64_t g = it.get_global_id(0); + const int64_t b = g / SYCL_DEQK_WI_PER_BLOCK; + const int tid = (int) (g % SYCL_DEQK_WI_PER_BLOCK); + if (b >= nb) { + return; + } + const block_q4_K * x = (const block_q4_K *) vx + b; + dequantize_q4_K_wide_one(yy + b * QK_K, x->qs, x->scales, (float) x->dm[0], (float) x->dm[1], tid); +} + + template <typename dst_t> inline void dequantize_q4_K_common(dst_t * __restrict__ y, const uint8_t * __restrict__ qs_ptr, const float dall, const float dmin, uint8_t * __restrict__ scales_local, int il, int ir) { @@ -1130,41 +1177,6 @@ inline void dequantize_q4_K_common(dst_t * __restrict__ y, const uint8_t * __res } } -template<typename dst_t> -static void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy, - uint8_t* scales_local, const sycl::nd_item<3> &item_ct1) { - const block_q4_K * x = (const block_q4_K *) vx; - - const int64_t i = item_ct1.get_group(2); - -#if QK_K == 256 - const int64_t tid = item_ct1.get_local_id(2); - const int64_t il = tid / 8; - const int64_t ir = tid % 8; - - dst_t * y = yy + i * QK_K + 64 * il + 4 * ir; - - const sycl::half2 dm = x[i].dm; - const float dall = dm[0]; - const float dmin = dm[1]; - - if (tid < 12) { - scales_local[tid] = x[i].scales[tid]; - } - - item_ct1.barrier(sycl::access::fence_space::local_space); - dequantize_q4_K_common(y, x[i].qs, dall, dmin, scales_local, il, ir); -#else - const int64_t tid = item_ct1.get_local_id(2); - const uint8_t * q = x[i].qs; - dst_t * y = yy + i*QK_K; - const float d = (float)x[i].dm[0]; - const float m = (float)x[i].dm[1]; - y[tid+ 0] = d * (x[i].scales[0] & 0xF) * (q[tid] & 0xF) - m * (x[i].scales[0] >> 4); - y[tid+32] = d * (x[i].scales[1] & 0xF) * (q[tid] >> 4) - m * (x[i].scales[1] >> 4); -#endif -} - template <typename dst_t> static void dequantize_block_q4_K_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy, uint8_t * scales_local, const sycl::nd_item<1> & item_ct1, int64_t nb) { @@ -1195,52 +1207,51 @@ static void dequantize_block_q4_K_reorder(const void * __restrict__ vx, dst_t * dequantize_q4_K_common(y, qs_ptr, dall, dmin, scales_local, il, ir); } -template<typename dst_t> -static void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy, - const sycl::nd_item<3> &item_ct1) { - const block_q5_K * x = (const block_q5_K *) vx; - - const int64_t i = item_ct1.get_group(2); - -#if QK_K == 256 - // assume 64 threads - this is very slightly better than the one below - const int64_t tid = item_ct1.get_local_id(2); - const int64_t il = tid/16; // il is in 0...3 - const int64_t ir = tid%16; // ir is in 0...15 - const int64_t is = 2*il; // is is in 0...6 +template <typename dst_t> +static inline void dequantize_q5_K_wide_one(dst_t * __restrict__ y, const uint8_t * __restrict__ qs, + const uint8_t * __restrict__ qh_ptr, + const uint8_t * __restrict__ scales, const float dall, + const float dmin, const int tid) { + const int il = tid >> 3; // 64-element chunk, 0..3 + const int ir = (tid & 7) * 4; // 4 qs bytes within the chunk; qh holds the 5th bit - dst_t * y = yy + i*QK_K + 64*il + 2*ir; + uint8_t sc, m; + get_scale_min_k4(2 * il + 0, scales, sc, m); + const float d1 = dall * sc; + const float m1 = dmin * m; + get_scale_min_k4(2 * il + 1, scales, sc, m); + const float d2 = dall * sc; + const float m2 = dmin * m; - const float dall = x[i].dm[0]; - const float dmin = x[i].dm[1]; + const sycl::vec<uint8_t, 4> q = vec_aligned_load<uint8_t, 4>(qs + 32 * il + ir); + const sycl::vec<uint8_t, 4> h = vec_aligned_load<uint8_t, 4>(qh_ptr + ir); - const uint8_t * ql = x[i].qs + 32*il + 2*ir; - const uint8_t * qh = x[i].qh + 2*ir; + const uint8_t hm1 = 1 << (2 * il); + const uint8_t hm2 = hm1 << 1; - uint8_t sc, m; - get_scale_min_k4(is + 0, x[i].scales, sc, m); - const float d1 = dall * sc; const float m1 = dmin * m; - get_scale_min_k4(is + 1, x[i].scales, sc, m); - const float d2 = dall * sc; const float m2 = dmin * m; + sycl::vec<dst_t, 4> lo, hi; +#pragma unroll + for (int l = 0; l < 4; ++l) { + lo[l] = d1 * ((q[l] & 0xF) + ((h[l] & hm1) ? 16 : 0)) - m1; + hi[l] = d2 * ((q[l] >> 4) + ((h[l] & hm2) ? 16 : 0)) - m2; + } + dst_t * yb = y + 64 * il + ir; + *reinterpret_cast<sycl::vec<dst_t, 4> *>(yb + 0) = lo; + *reinterpret_cast<sycl::vec<dst_t, 4> *>(yb + 32) = hi; +} - uint8_t hm = 1 << (2*il); - y[ 0] = d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1; - y[ 1] = d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1; - hm <<= 1; - y[32] = d2 * ((ql[ 0] >> 4) + (qh[ 0] & hm ? 16 : 0)) - m2; - y[33] = d2 * ((ql[ 1] >> 4) + (qh[ 1] & hm ? 16 : 0)) - m2; -#else - const int64_t tid = item_ct1.get_local_id(2); - const uint8_t q = x[i].qs[tid]; - const int64_t im = tid/8; // 0...3 - const int64_t in = tid%8; // 0...7 - const int64_t is = tid/16; // 0 or 1 - const uint8_t h = x[i].qh[in] >> im; - const float d = x[i].d; - dst_t * y = yy + i*QK_K + tid; - y[ 0] = d * x[i].scales[is+0] * ((q & 0xF) - ((h >> 0) & 1 ? 0 : 16)); - y[32] = d * x[i].scales[is+2] * ((q >> 4) - ((h >> 4) & 1 ? 0 : 16)); -#endif +template <typename dst_t> +static void dequantize_block_q5_K_wide(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t nb, + const sycl::nd_item<1> & it) { + static_assert(QK_K == 256, "wide K-quant dequant assumes QK_K == 256"); + const int64_t g = it.get_global_id(0); + const int64_t b = g / SYCL_DEQK_WI_PER_BLOCK; + const int tid = (int) (g % SYCL_DEQK_WI_PER_BLOCK); + if (b >= nb) { + return; + } + const block_q5_K * x = (const block_q5_K *) vx + b; + dequantize_q5_K_wide_one(yy + b * QK_K, x->qs, x->qh, x->scales, (float) x->dm[0], (float) x->dm[1], tid); } template <typename dst_t> From c0083e3052caf8e07bdd48d5504cffa741dc70fe Mon Sep 17 00:00:00 2001 From: Neo Zhang <zhang.jianyu@outlook.com> Date: Fri, 16 Oct 2026 14:44:24 +0800 Subject: [PATCH 047/423] sycl : reduce the build threads to half (#30236) --- .devops/intel.Dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.devops/intel.Dockerfile b/.devops/intel.Dockerfile index e20971d128d9..09daa5fe5d1a 100644 --- a/.devops/intel.Dockerfile +++ b/.devops/intel.Dockerfile @@ -50,7 +50,7 @@ RUN if [ "${GGML_SYCL_F16}" = "ON" ]; then \ fi && \ echo "Building with dynamic libs" && \ cmake -B build -DGGML_NATIVE=OFF -DGGML_SYCL=ON -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON -DLLAMA_BUILD_TESTS=OFF ${OPT_SYCL_F16} && \ - cmake --build build --config Release -j$(nproc) + cmake --build build --config Release -j$((($(nproc)+1)/2)) RUN mkdir -p /app/lib && \ find build -name "*.so*" -exec cp -P {} /app/lib \; From c804754529cf51f86e66847f093e7cf639c7b4d3 Mon Sep 17 00:00:00 2001 From: Aaron Teo <aaron.teo1@ibm.com> Date: Fri, 16 Oct 2026 14:45:35 +0800 Subject: [PATCH 048/423] ggml-cpu: refactor s309x repack helper functions (#30588) Signed-off-by: Aaron Teo <aaron.teo1@ibm.com> --- ggml/src/ggml-cpu/arch/s390/quants.c | 72 +++++++++++++------------- ggml/src/ggml-cpu/arch/s390/repack.cpp | 63 ++++++++-------------- ggml/src/ggml-cpu/ggml-cpu-impl.h | 54 +++++++++---------- 3 files changed, 87 insertions(+), 102 deletions(-) diff --git a/ggml/src/ggml-cpu/arch/s390/quants.c b/ggml/src/ggml-cpu/arch/s390/quants.c index 52344828e39d..22b8b83109c9 100644 --- a/ggml/src/ggml-cpu/arch/s390/quants.c +++ b/ggml/src/ggml-cpu/arch/s390/quants.c @@ -310,7 +310,7 @@ void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const voi const int8x16_t v_yl = vec_xl(0 , y[ib].qs); const int8x16_t v_yh = vec_xl(QK8_1/2, y[ib].qs); - const int32x4_t v_xy_ = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xl, v_yl), v_xh, v_yh); + const int32x4_t v_xy_ = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xl, v_yl), v_xh, v_yh); const float32x4_t v_xy = vec_float(v_xy_); const float32x4_t v_d = vec_splats(GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d)); @@ -379,8 +379,8 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo const int8x16_t v_y1l = vec_xl(0, y1->qs); const int8x16_t v_y1h = vec_xl(QK8_0/2, y1->qs); - const int32x4_t v_xy0 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x0l, v_y0l), v_x0h, v_y0h); - const int32x4_t v_xy1 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x1l, v_y1l), v_x1h, v_y1h); + const int32x4_t v_xy0 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x0l, v_y0l), v_x0h, v_y0h); + const int32x4_t v_xy1 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x1l, v_y1l), v_x1h, v_y1h); const float32x4_t v_xy0f = vec_float(v_xy0); const float32x4_t v_xy1f = vec_float(v_xy1); @@ -407,7 +407,7 @@ void ggml_vec_dot_mxfp4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo const int8x16_t v_yl = vec_xl(0, y0->qs); const int8x16_t v_yh = vec_xl(QK8_0/2, y0->qs); - const int32x4_t v_xy = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xl, v_yl), v_xh, v_yh); + const int32x4_t v_xy = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xl, v_yl), v_xh, v_yh); const float32x4_t v_xyf = vec_float(v_xy); const float32x4_t v_d = vec_splats(GGML_E8M0_TO_FP32_HALF(x0->e) * GGML_CPU_FP16_TO_FP32(y0->d)); @@ -502,8 +502,8 @@ void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi const int8x16_t v_y1l = vec_xl(0, (const int8_t *)y1->qs); const int8x16_t v_y1h = vec_xl(QK8_0/2, (const int8_t *)y1->qs); - const int32x4_t v_xy0 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x0lf, v_y0l), v_x0hf, v_y0h); - const int32x4_t v_xy1 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x1lf, v_y1l), v_x1hf, v_y1h); + const int32x4_t v_xy0 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x0lf, v_y0l), v_x0hf, v_y0h); + const int32x4_t v_xy1 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x1lf, v_y1l), v_x1hf, v_y1h); const float32x4_t v_xy0f = vec_float(v_xy0); const float32x4_t v_xy1f = vec_float(v_xy1); @@ -548,7 +548,7 @@ void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi const int8x16_t v_yl = vec_xl(0, (const int8_t *)y0->qs); const int8x16_t v_yh = vec_xl(QK8_0/2, (const int8_t *)y0->qs); - const int32x4_t v_xy = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xlf, v_yl), v_xhf, v_yh); + const int32x4_t v_xy = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xlf, v_yl), v_xhf, v_yh); const float32x4_t v_xyf = vec_float(v_xy); const float32x4_t v_d = vec_splats(GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d)); @@ -653,8 +653,8 @@ void ggml_vec_dot_q5_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const voi const int8x16_t v_y1l = vec_xl(0 , y1->qs); const int8x16_t v_y1h = vec_xl(QK8_1/2, y1->qs); - const int32x4_t v_xy0 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x0lf, v_y0l), v_x0hf, v_y0h); - const int32x4_t v_xy1 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x1lf, v_y1l), v_x1hf, v_y1h); + const int32x4_t v_xy0 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x0lf, v_y0l), v_x0hf, v_y0h); + const int32x4_t v_xy1 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x1lf, v_y1l), v_x1hf, v_y1h); const float32x4_t v_xy0f = vec_float(v_xy0); const float32x4_t v_xy1f = vec_float(v_xy1); @@ -701,7 +701,7 @@ void ggml_vec_dot_q5_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const voi const int8x16_t v_yl = vec_xl(0 , y0->qs); const int8x16_t v_yh = vec_xl(QK8_1/2, y0->qs); - const int32x4_t v_xy = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xlf, v_yl), v_xhf, v_yh); + const int32x4_t v_xy = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xlf, v_yl), v_xhf, v_yh); const float32x4_t v_xyf = vec_float(v_xy); const float32x4_t v_d = vec_splats(GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d)); @@ -751,7 +751,7 @@ void ggml_vec_dot_q8_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi const int8x16_t v_yl = vec_xl(0 , y[ib].qs); const int8x16_t v_yh = vec_xl(QK8_0/2, y[ib].qs); - const int32x4_t v_xy_ = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xl, v_yl), v_xh, v_yh); + const int32x4_t v_xy_ = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xl, v_yl), v_xh, v_yh); const float32x4_t v_xy = vec_float(v_xy_); const float32x4_t v_d = vec_splats(GGML_CPU_FP16_TO_FP32(x[ib].d) * GGML_CPU_FP16_TO_FP32(y[ib].d)); @@ -855,10 +855,10 @@ void ggml_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi q3bytes[2] = vec_sub((int8x16_t)vec_and(vec_sr(q3b[0], 2), v_3m), (int8x16_t)q3h[2]); q3bytes[3] = vec_sub((int8x16_t)vec_and(vec_sr(q3b[1], 2), v_3m), (int8x16_t)q3h[3]); - isum0 = ggml_vec_dot(v_z, q3bytes[0], q8bytes[0]); - isum1 = ggml_vec_dot(v_z, q3bytes[1], q8bytes[1]); - isum2 = ggml_vec_dot(v_z, q3bytes[2], q8bytes[2]); - isum3 = ggml_vec_dot(v_z, q3bytes[3], q8bytes[3]); + isum0 = ggml_vec_dot_s32(v_z, q3bytes[0], q8bytes[0]); + isum1 = ggml_vec_dot_s32(v_z, q3bytes[1], q8bytes[1]); + isum2 = ggml_vec_dot_s32(v_z, q3bytes[2], q8bytes[2]); + isum3 = ggml_vec_dot_s32(v_z, q3bytes[3], q8bytes[3]); isum += (isum0[0] + isum0[1] + isum0[2] + isum0[3]) * scale[0]; isum += (isum1[0] + isum1[1] + isum1[2] + isum1[3]) * scale[1]; @@ -877,10 +877,10 @@ void ggml_vec_dot_q3_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi q3bytes[2] = vec_sub((int8x16_t)vec_and(vec_sr(q3b[0], 6), v_3m), (int8x16_t)q3h[2]); q3bytes[3] = vec_sub((int8x16_t)vec_and(vec_sr(q3b[1], 6), v_3m), (int8x16_t)q3h[3]); - isum0 = ggml_vec_dot(v_z, q3bytes[0], q8bytes[4]); - isum1 = ggml_vec_dot(v_z, q3bytes[1], q8bytes[5]); - isum2 = ggml_vec_dot(v_z, q3bytes[2], q8bytes[6]); - isum3 = ggml_vec_dot(v_z, q3bytes[3], q8bytes[7]); + isum0 = ggml_vec_dot_s32(v_z, q3bytes[0], q8bytes[4]); + isum1 = ggml_vec_dot_s32(v_z, q3bytes[1], q8bytes[5]); + isum2 = ggml_vec_dot_s32(v_z, q3bytes[2], q8bytes[6]); + isum3 = ggml_vec_dot_s32(v_z, q3bytes[3], q8bytes[7]); isum += vec_hsum_i32x4(isum0) * scale[0]; isum += vec_hsum_i32x4(isum1) * scale[1]; @@ -981,7 +981,7 @@ void ggml_vec_dot_q4_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi v_xl[0] = (int8x16_t)vec_and(v_x[0], v_lm); v_xl[1] = (int8x16_t)vec_and(v_x[1], v_lm); - const int32x4_t p1 = ggml_vec_dot(ggml_vec_dot(v_z, v_xl[0], v_y[0]), v_xl[1], v_y[1]); + const int32x4_t p1 = ggml_vec_dot_s32(ggml_vec_dot_s32(v_z, v_xl[0], v_y[0]), v_xl[1], v_y[1]); sumi1 += vec_hsum_i32x4(p1) * scales[2*j+0]; v_y[0] = vec_xl(0 , y0); @@ -991,7 +991,7 @@ void ggml_vec_dot_q4_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi v_xl[0] = (int8x16_t)vec_sr(v_x[0], 4); v_xl[1] = (int8x16_t)vec_sr(v_x[1], 4); - const int32x4_t p2 = ggml_vec_dot(ggml_vec_dot(v_z, v_xl[0], v_y[0]), v_xl[1], v_y[1]); + const int32x4_t p2 = ggml_vec_dot_s32(ggml_vec_dot_s32(v_z, v_xl[0], v_y[0]), v_xl[1], v_y[1]); sumi2 += vec_hsum_i32x4(p2) * scales[2*j+1]; } @@ -1107,8 +1107,8 @@ void ggml_vec_dot_q5_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi q5b[2] = (int8x16_t)vec_or(vec_sr(v_xl[0], 4), q5h[2]); q5b[3] = (int8x16_t)vec_or(vec_sr(v_xl[1], 4), q5h[3]); - int32x4_t sumi0 = ggml_vec_dot(ggml_vec_dot(v_z, q5b[0], v_y[0]), q5b[1], v_y[1]); - int32x4_t sumi1 = ggml_vec_dot(ggml_vec_dot(v_z, q5b[2], v_y[2]), q5b[3], v_y[3]); + int32x4_t sumi0 = ggml_vec_dot_s32(ggml_vec_dot_s32(v_z, q5b[0], v_y[0]), q5b[1], v_y[1]); + int32x4_t sumi1 = ggml_vec_dot_s32(ggml_vec_dot_s32(v_z, q5b[2], v_y[2]), q5b[3], v_y[3]); sumi += vec_hsum_i32x4(sumi0) * *scales++; sumi += vec_hsum_i32x4(sumi1) * *scales++; @@ -1217,10 +1217,10 @@ void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi q6b[2] = (int8x16_t)(vec_or(vec_and(v_xl[2], v_lm), q6h[2])); q6b[3] = (int8x16_t)(vec_or(vec_and(v_xl[3], v_lm), q6h[3])); - int32x4_t summs0 = ggml_vec_dot(v_z, q6b[0], v_y[0]); - int32x4_t summs1 = ggml_vec_dot(v_z, q6b[1], v_y[1]); - int32x4_t summs2 = ggml_vec_dot(v_z, q6b[2], v_y[2]); - int32x4_t summs3 = ggml_vec_dot(v_z, q6b[3], v_y[3]); + int32x4_t summs0 = ggml_vec_dot_s32(v_z, q6b[0], v_y[0]); + int32x4_t summs1 = ggml_vec_dot_s32(v_z, q6b[1], v_y[1]); + int32x4_t summs2 = ggml_vec_dot_s32(v_z, q6b[2], v_y[2]); + int32x4_t summs3 = ggml_vec_dot_s32(v_z, q6b[3], v_y[3]); isum += vec_hsum_i32x4(summs0) * scale[0] + vec_hsum_i32x4(summs1) * scale[1] + @@ -1251,10 +1251,10 @@ void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi q6b[2] = (int8x16_t)(vec_or(vec_sr(v_xl[2], 4), q6h[2])); q6b[3] = (int8x16_t)(vec_or(vec_sr(v_xl[3], 4), q6h[3])); - summs0 = ggml_vec_dot(v_z, q6b[0], v_y[0]); - summs1 = ggml_vec_dot(v_z, q6b[1], v_y[1]); - summs2 = ggml_vec_dot(v_z, q6b[2], v_y[2]); - summs3 = ggml_vec_dot(v_z, q6b[3], v_y[3]); + summs0 = ggml_vec_dot_s32(v_z, q6b[0], v_y[0]); + summs1 = ggml_vec_dot_s32(v_z, q6b[1], v_y[1]); + summs2 = ggml_vec_dot_s32(v_z, q6b[2], v_y[2]); + summs3 = ggml_vec_dot_s32(v_z, q6b[3], v_y[3]); isum += vec_hsum_i32x4(summs0) * scale[0] + vec_hsum_i32x4(summs1) * scale[1] + @@ -1367,8 +1367,8 @@ void ggml_vec_dot_q6_K_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const voi // q2u2 = vec_mul(q2u2, q2s2); // q2u3 = vec_mul(q2u3, q2s3); -// const int32x4_t p1 = ggml_vec_dot(ggml_vec_dot(vec_splat_s32(0), q2u0, q8b0), q2u1, q8b1); -// const int32x4_t p2 = ggml_vec_dot(ggml_vec_dot(vec_splat_s32(0), q2u2, q8b2), q2u3, q8b3); +// const int32x4_t p1 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splat_s32(0), q2u0, q8b0), q2u1, q8b1); +// const int32x4_t p2 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splat_s32(0), q2u2, q8b2), q2u3, q8b3); // sumf1 += (p1[0] + p1[1] + p1[2] + p1[3]) * (0.5f + (aux32[1] >> 28)); // sumf2 += (p2[0] + p2[1] + p2[2] + p2[3]) * (0.5f + (aux32[3] >> 28)); @@ -1445,7 +1445,7 @@ void ggml_vec_dot_iq4_nl_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const v const int8x16_t v_yl = vec_xl(0 , y0->qs); const int8x16_t v_yh = vec_xl(QK8_0/2, y0->qs); - const int32x4_t v_xy = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_xl, v_yl), v_xh, v_yh); + const int32x4_t v_xy = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_xl, v_yl), v_xh, v_yh); sumf += GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d) * vec_hsum_i32x4(v_xy); } @@ -1508,8 +1508,8 @@ void ggml_vec_dot_iq4_xs_q8_K(int n, float * GGML_RESTRICT s, size_t bs, const v const int8x16_t v_y3 = vec_xl(48, q8); q8 += 64; - int32x4_t vsumi0 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x0l, v_y0), v_x0h, v_y1); - int32x4_t vsumi1 = ggml_vec_dot(ggml_vec_dot(vec_splats(0), v_x1l, v_y2), v_x1h, v_y3); + int32x4_t vsumi0 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x0l, v_y0), v_x0h, v_y1); + int32x4_t vsumi1 = ggml_vec_dot_s32(ggml_vec_dot_s32(vec_splats(0), v_x1l, v_y2), v_x1h, v_y3); int ls1 = ((x[ibl].scales_l[ib] & 0xF) | ((h << 4) & 0x30)) - 32; int ls2 = ((x[ibl].scales_l[ib] >> 4) | ((h << 2) & 0x30)) - 32; diff --git a/ggml/src/ggml-cpu/arch/s390/repack.cpp b/ggml/src/ggml-cpu/arch/s390/repack.cpp index abf3433adda2..c12fe215e80e 100644 --- a/ggml/src/ggml-cpu/arch/s390/repack.cpp +++ b/ggml/src/ggml-cpu/arch/s390/repack.cpp @@ -70,23 +70,6 @@ void ggml_quantize_mat_q8_0_4x4(const float * GGML_RESTRICT x, void * GGML_RESTR #endif } -#if defined(__VXE__) || defined(__VXE2__) -static inline int16x8_t vxe_dot_acc(const int8x16_t v_x, const int8x16_t v_y, const int16x8_t v_acc) { - return vec_meadd(v_x, v_y, vec_moadd(v_x, v_y, v_acc)); -} - -static inline int8x16_t vxe_splat_granule(const int8_t * qs) { - uint32_t g; - memcpy(&g, qs, sizeof(g)); - return (int8x16_t)vec_splats(g); -} - -static inline int32x4_t vxe_fold(const int16x8_t v_sumi) { - const int16x8_t v_ones = vec_splats((int16_t)1); - return vec_add(vec_mule(v_sumi, v_ones), vec_mulo(v_sumi, v_ones)); -} -#endif - void ggml_gemv_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, const void * GGML_RESTRICT vy, int nr, int nc) { const int qk = QK8_0; const int nb = n / qk; @@ -131,21 +114,21 @@ void ggml_gemv_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo int16x8_t v_sumi = vec_splats((int16_t)0); - v_sumi = vxe_dot_acc(v_x0l, vxe_splat_granule(y_lo + 0), v_sumi); - v_sumi = vxe_dot_acc(v_x1l, vxe_splat_granule(y_lo + 4), v_sumi); - v_sumi = vxe_dot_acc(v_x2l, vxe_splat_granule(y_lo + 8), v_sumi); - v_sumi = vxe_dot_acc(v_x3l, vxe_splat_granule(y_lo + 12), v_sumi); + v_sumi = ggml_vec_dot_s16(v_sumi, v_x0l, ggml_vec_xl_s8x4_splat(y_lo + 0)); + v_sumi = ggml_vec_dot_s16(v_sumi, v_x1l, ggml_vec_xl_s8x4_splat(y_lo + 4)); + v_sumi = ggml_vec_dot_s16(v_sumi, v_x2l, ggml_vec_xl_s8x4_splat(y_lo + 8)); + v_sumi = ggml_vec_dot_s16(v_sumi, v_x3l, ggml_vec_xl_s8x4_splat(y_lo + 12)); - v_sumi = vxe_dot_acc(v_x0h, vxe_splat_granule(y_hi + 0), v_sumi); - v_sumi = vxe_dot_acc(v_x1h, vxe_splat_granule(y_hi + 4), v_sumi); - v_sumi = vxe_dot_acc(v_x2h, vxe_splat_granule(y_hi + 8), v_sumi); - v_sumi = vxe_dot_acc(v_x3h, vxe_splat_granule(y_hi + 12), v_sumi); + v_sumi = ggml_vec_dot_s16(v_sumi, v_x0h, ggml_vec_xl_s8x4_splat(y_hi + 0)); + v_sumi = ggml_vec_dot_s16(v_sumi, v_x1h, ggml_vec_xl_s8x4_splat(y_hi + 4)); + v_sumi = ggml_vec_dot_s16(v_sumi, v_x2h, ggml_vec_xl_s8x4_splat(y_hi + 8)); + v_sumi = ggml_vec_dot_s16(v_sumi, v_x3h, ggml_vec_xl_s8x4_splat(y_hi + 12)); const float32x4_t v_yd = vec_splats(GGML_CPU_FP16_TO_FP32(a_ptr[l].d)); - const float32x4_t v_xd = __lzs_f16cx4_load(b_ptr[l].d); + const float32x4_t v_xd = ggml_vec_xl_f16x4_finite(b_ptr[l].d); const float32x4_t v_d = vec_mul(v_yd, v_xd); - v_sumf = vec_madd(vec_float(vxe_fold(v_sumi)), v_d, v_sumf); + v_sumf = vec_madd(vec_float(vec_paddl_s16(v_sumi)), v_d, v_sumf); } vec_xst(v_sumf, 0, res_ptr + x * ncols_interleaved); @@ -192,24 +175,24 @@ void ggml_gemm_q4_0_4x4_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo const int8_t * y_lo = a_ptr[l].qs + 16 * k; const int8_t * y_hi = y_lo + qk / 2 * 4; - v_sumi0 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo + 0), v_sumi0); - v_sumi1 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo + 4), v_sumi1); - v_sumi2 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo + 8), v_sumi2); - v_sumi3 = vxe_dot_acc(v_xl, vxe_splat_granule(y_lo + 12), v_sumi3); + v_sumi0 = ggml_vec_dot_s16(v_sumi0, v_xl, ggml_vec_xl_s8x4_splat(y_lo + 0)); + v_sumi1 = ggml_vec_dot_s16(v_sumi1, v_xl, ggml_vec_xl_s8x4_splat(y_lo + 4)); + v_sumi2 = ggml_vec_dot_s16(v_sumi2, v_xl, ggml_vec_xl_s8x4_splat(y_lo + 8)); + v_sumi3 = ggml_vec_dot_s16(v_sumi3, v_xl, ggml_vec_xl_s8x4_splat(y_lo + 12)); - v_sumi0 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi + 0), v_sumi0); - v_sumi1 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi + 4), v_sumi1); - v_sumi2 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi + 8), v_sumi2); - v_sumi3 = vxe_dot_acc(v_xh, vxe_splat_granule(y_hi + 12), v_sumi3); + v_sumi0 = ggml_vec_dot_s16(v_sumi0, v_xh, ggml_vec_xl_s8x4_splat(y_hi + 0)); + v_sumi1 = ggml_vec_dot_s16(v_sumi1, v_xh, ggml_vec_xl_s8x4_splat(y_hi + 4)); + v_sumi2 = ggml_vec_dot_s16(v_sumi2, v_xh, ggml_vec_xl_s8x4_splat(y_hi + 8)); + v_sumi3 = ggml_vec_dot_s16(v_sumi3, v_xh, ggml_vec_xl_s8x4_splat(y_hi + 12)); } const float32x4_t v_yd = __lzs_f16cx4_load(a_ptr[l].d); - const float32x4_t v_xd = __lzs_f16cx4_load(b_ptr[l].d); + const float32x4_t v_xd = ggml_vec_xl_f16x4_finite(b_ptr[l].d); - v_sumf[0] = vec_madd(vec_float(vxe_fold(v_sumi0)), vec_mul(v_xd, vec_splat(v_yd, 0)), v_sumf[0]); - v_sumf[1] = vec_madd(vec_float(vxe_fold(v_sumi1)), vec_mul(v_xd, vec_splat(v_yd, 1)), v_sumf[1]); - v_sumf[2] = vec_madd(vec_float(vxe_fold(v_sumi2)), vec_mul(v_xd, vec_splat(v_yd, 2)), v_sumf[2]); - v_sumf[3] = vec_madd(vec_float(vxe_fold(v_sumi3)), vec_mul(v_xd, vec_splat(v_yd, 3)), v_sumf[3]); + v_sumf[0] = vec_madd(vec_float(vec_paddl_s16(v_sumi0)), vec_mul(v_xd, vec_splat(v_yd, 0)), v_sumf[0]); + v_sumf[1] = vec_madd(vec_float(vec_paddl_s16(v_sumi1)), vec_mul(v_xd, vec_splat(v_yd, 1)), v_sumf[1]); + v_sumf[2] = vec_madd(vec_float(vec_paddl_s16(v_sumi2)), vec_mul(v_xd, vec_splat(v_yd, 2)), v_sumf[2]); + v_sumf[3] = vec_madd(vec_float(vec_paddl_s16(v_sumi3)), vec_mul(v_xd, vec_splat(v_yd, 3)), v_sumf[3]); } for (int m = 0; m < 4; m++) { diff --git a/ggml/src/ggml-cpu/ggml-cpu-impl.h b/ggml/src/ggml-cpu/ggml-cpu-impl.h index 82345b2063bd..0d48c9a22e3d 100644 --- a/ggml/src/ggml-cpu/ggml-cpu-impl.h +++ b/ggml/src/ggml-cpu/ggml-cpu-impl.h @@ -443,6 +443,13 @@ inline static ggml_int8x16x4_t ggml_vec_xl_s8x4(const int8_t * ptr) { return res; } +// load 4 bytes and replicate them into all 4 word lanes +inline static int8x16_t ggml_vec_xl_s8x4_splat(const int8_t * ptr) { + uint32_t tmp; + memcpy(&tmp, ptr, sizeof(tmp)); + return (int8x16_t)vec_splats(tmp); +} + typedef struct ggml_int16x8x2_t { int16x8_t val[2]; } ggml_int16x8x2_t; @@ -456,31 +463,17 @@ inline static ggml_int16x8x2_t ggml_vec_xl_s16x2(const int16_t * ptr) { return res; } -/* - ! WARNING: Very slow. Use vec_perm if possible. Refer to iq4_xs - ! or iq4_nl for example implementation. -*/ -inline static int8x16_t ggml_vec_tbl(int8x16_t a, uint8x16_t b) { - int8x16_t res; - - res[ 0] = a[b[ 0]]; - res[ 1] = a[b[ 1]]; - res[ 2] = a[b[ 2]]; - res[ 3] = a[b[ 3]]; - res[ 4] = a[b[ 4]]; - res[ 5] = a[b[ 5]]; - res[ 6] = a[b[ 6]]; - res[ 7] = a[b[ 7]]; - res[ 8] = a[b[ 8]]; - res[ 9] = a[b[ 9]]; - res[10] = a[b[10]]; - res[11] = a[b[11]]; - res[12] = a[b[12]]; - res[13] = a[b[13]]; - res[14] = a[b[14]]; - res[15] = a[b[15]]; - - return res; +// faster fp16 -> fp32 conversion using vector intrinsics +// WARNING: fp16 inf converts into 65536.0f. use this only for finite values. +inline static float32x4_t ggml_vec_xl_f16x4_finite(const ggml_fp16_t * x) { + uint64_t tmp; + memcpy(&tmp, x, sizeof(tmp)); + + const uint64x2_t v_x = { tmp, tmp }; + const uint32x4_t v_xh = vec_unpackh((uint16x8_t)v_x); + const uint32x4_t v_bits = vec_sl(vec_and(v_xh, vec_splats((uint32_t)0x7FFF)), 13); + const float32x4_t v_abs = vec_mul((float32x4_t)v_bits, vec_splats(0x1.0p+112f)); + return (float32x4_t)vec_or((uint32x4_t)v_abs, vec_sl(vec_and(v_xh, vec_splats((uint32_t)0x8000)), 16)); } inline static int16x8_t vec_padd_s16(int16x8_t a, int16x8_t b) { @@ -492,6 +485,11 @@ inline static int16x8_t vec_padd_s16(int16x8_t a, int16x8_t b) { return v_abo + v_abe; } +inline static int32x4_t vec_paddl_s16(int16x8_t a) { + const int16x8_t v_ones = vec_splats((int16_t)1); + return vec_mule(a, v_ones) + vec_mulo(a, v_ones); +} + /** * @see https://github.com/ggml-org/llama.cpp/pull/14037 */ @@ -505,7 +503,11 @@ inline static int32_t vec_hsum_i32x4(int32x4_t v) { return v_temp[0] + v_temp[1]; } -inline static int32x4_t ggml_vec_dot(int32x4_t acc, int8x16_t a, int8x16_t b) { +inline static int16x8_t ggml_vec_dot_s16(int16x8_t acc, int8x16_t a, int8x16_t b) { + return vec_meadd(a, b, vec_moadd(a, b, acc)); +} + +inline static int32x4_t ggml_vec_dot_s32(int32x4_t acc, int8x16_t a, int8x16_t b) { const int16x8_t p = vec_mule(a, b) + vec_mulo(a, b); return acc + (vec_unpackh(p) + vec_unpackl(p)); } From 3dbcb7fca56ab4769bd4af684a432e2331bc9ace Mon Sep 17 00:00:00 2001 From: Trivalik <3148279+trivalik@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:48:23 +0200 Subject: [PATCH 049/423] cpu : reduce AVX2 tiled accumulators together (#30296) Assisted-by: Codex --- ggml/src/ggml-cpu/tiled/tiled-kernel.cpp | 31 ++++++++++-------------- 1 file changed, 13 insertions(+), 18 deletions(-) diff --git a/ggml/src/ggml-cpu/tiled/tiled-kernel.cpp b/ggml/src/ggml-cpu/tiled/tiled-kernel.cpp index bffa737c7cd4..5819ac07ca31 100644 --- a/ggml/src/ggml-cpu/tiled/tiled-kernel.cpp +++ b/ggml/src/ggml-cpu/tiled/tiled-kernel.cpp @@ -244,16 +244,6 @@ static void tiled_run_microtile_avx2(const tiled_tile_src0 & src0, const tiled_t const int bs_stride = (nkr == 1) ? TILED_TILE_ROWS : TILED_MICRO; const int bs_off = seff * TILED_MICRO; - // Horizontal sum helper: converts 8x32-bit int lane inside a YMM register to scalar int32 - auto hsum256_epi32 = [](const __m256i v) -> int32_t { - __m128i low = _mm256_castsi256_si128(v); - __m128i high = _mm256_extracti128_si256(v, 1); - __m128i sum128 = _mm_add_epi32(low, high); - sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_SHUFFLE(2, 3, 0, 1))); - sum128 = _mm_add_epi32(sum128, _mm_shuffle_epi32(sum128, _MM_SHUFFLE(1, 0, 3, 2))); - return _mm_cvtsi128_si32(sum128); - }; - for (int i = 0; i < TILED_MICRO; i++) { const int ar = i0 + i; const int d_off = seff * TILED_MICRO + ar; @@ -362,15 +352,20 @@ static void tiled_run_microtile_avx2(const tiled_tile_src0 & src0, const tiled_t } } - // OPTIMIZATION 2: Vectorized Epilogue with FMA3 gating - // Reduce accumulators into a 256-bit vector across 8 columns - alignas(32) int32_t s1_vals[GROUP]; - #pragma GCC unroll 8 - for (int t = 0; t < GROUP; t++) { - s1_vals[t] = hsum256_epi32(acc[t]); - } - __m256i s1_vec = _mm256_load_si256((const __m256i *) s1_vals); + // Reduce the eight accumulators to column sums in acc[0]..acc[7] order, with additions modulo 2^32. + const __m256i ab = _mm256_hadd_epi32(acc[0], acc[1]); + const __m256i cd = _mm256_hadd_epi32(acc[2], acc[3]); + const __m256i ef = _mm256_hadd_epi32(acc[4], acc[5]); + const __m256i gh = _mm256_hadd_epi32(acc[6], acc[7]); + + // Each 128-bit half now holds four column sums over input lanes 0..3 or 4..7. + const __m256i abcd = _mm256_hadd_epi32(ab, cd); + const __m256i efgh = _mm256_hadd_epi32(ef, gh); + const __m256i lo = _mm256_permute2x128_si256(abcd, efgh, 0x20); + const __m256i hi = _mm256_permute2x128_si256(abcd, efgh, 0x31); + __m256i s1_vec = _mm256_add_epi32(lo, hi); + // OPTIMIZATION 2: Vectorized Epilogue with FMA3 gating if constexpr (BIAS != 0 && ACTBIAS) { // act-bias correction: raw = desired + 128*sum_s scales[s]*w_bsum[s]; corr precomputed in repack_src0 s1_vec = _mm256_sub_epi32(s1_vec, _mm256_set1_epi32(mins_row[0])); From 0f4376268f2ef33d007f331f25b32ffafa297cd5 Mon Sep 17 00:00:00 2001 From: FINAL-Bench <kkms1116@koreacu.ac.kr> Date: Fri, 16 Oct 2026 15:56:00 +0900 Subject: [PATCH 050/423] convert : support text-only checkpoints in pplx-decider (#30575) * pplx-decider-text-only convert : support text-only checkpoints in pplx-decider * Apply suggestions from code review * convert : remove unused import --- conversion/pplx_decider.py | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/conversion/pplx_decider.py b/conversion/pplx_decider.py index f9cfaf9d3de5..b71bd8a3ad23 100644 --- a/conversion/pplx_decider.py +++ b/conversion/pplx_decider.py @@ -3,7 +3,7 @@ import json from pathlib import Path -from typing import Any, Callable, Iterable, TYPE_CHECKING +from typing import Any, Iterable, TYPE_CHECKING import torch @@ -35,6 +35,7 @@ def _load_pplx_decider_hparams(dir_model: Path) -> dict[str, Any]: @ModelBase.register("PplxDeciderModel") @ModelBase.example("perplexity-ai/pplx-decider-v1-27b") +@ModelBase.example("FINAL-Bench/Darwin-27B-ZTC-v5") class PplxDeciderModel(Qwen3_5TextModel): model_arch = gguf.MODEL_ARCH.QWEN35 no_mtp = True # the checkpoint has no MTP head @@ -77,14 +78,6 @@ def set_gguf_parameters(self): for name in ("choice", "score", "noul"): self.gguf_writer.add_decision_temperature(name, self.hparams["decision"]["temperature"]) - @classmethod - def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: - name, gen = item - # the checkpoint is the bare backbone, its text tensors have no "model." prefix - if name.startswith("language_model."): - name = "model." + name - return super().filter_tensors((name, gen)) - def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: yield from super().generate_extra_tensors() from safetensors.torch import load_file @@ -92,7 +85,7 @@ def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: # the readout has one row per option label, store it as an LM head that is zero for the other tokens readout = load_file(self.dir_model / "readout.safetensors")["weight"] token_ids = self.hparams["decision"]["token_ids"] - n_vocab = self.hparams["text_config"]["vocab_size"] + n_vocab = self.hparams["vocab_size"] assert readout.shape[0] == len(token_ids) == len(set(token_ids)) lm_head = torch.zeros(n_vocab, readout.shape[1], dtype=readout.dtype) lm_head[token_ids] = readout From 13b80489204c4711ebb67e75e99aa75900c3ded8 Mon Sep 17 00:00:00 2001 From: hayder <tprod4242@gmail.com> Date: Fri, 16 Oct 2026 02:06:32 -0500 Subject: [PATCH 051/423] common : read ngram cache parts by reference instead of copying (#30061) --- common/ngram-cache.cpp | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/common/ngram-cache.cpp b/common/ngram-cache.cpp index dce54b364749..419a0ebc9833 100644 --- a/common/ngram-cache.cpp +++ b/common/ngram-cache.cpp @@ -68,7 +68,7 @@ static llama_token try_draft(common_ngram_cache & nc_static, const common_ngram if (part_static_it == nc_static.end()) { return LLAMA_TOKEN_NULL; } - const common_ngram_cache_part part_static = part_static_it->second; + const common_ngram_cache_part & part_static = part_static_it->second; int max_count_static = 0; int sum_count_static = 0; @@ -96,7 +96,7 @@ static llama_token try_draft(common_ngram_cache & nc_static, const common_ngram // Try to draft a token from primary cache (context/dynamic), validate with static cache: static llama_token try_draft( - common_ngram_cache & nc_primary, const std::vector<common_ngram> & ngrams_primary, common_ngram_cache_part & part_static, + common_ngram_cache & nc_primary, const std::vector<common_ngram> & ngrams_primary, const common_ngram_cache_part & part_static, const int * min_sample_size, const int * min_percent) { llama_token drafted_token = LLAMA_TOKEN_NULL; @@ -108,7 +108,7 @@ static llama_token try_draft( if (part_primary_it == nc_primary.end()) { continue; } - const common_ngram_cache_part part_primary = part_primary_it->second; + const common_ngram_cache_part & part_primary = part_primary_it->second; int max_count_primary = 0; int max_count_static = 0; @@ -118,7 +118,7 @@ static llama_token try_draft( for (std::pair<llama_token, int> token_count_primary : part_primary) { const llama_token token = token_count_primary.first; - common_ngram_cache_part::iterator token_count_static_it = part_static.find(token); + common_ngram_cache_part::const_iterator token_count_static_it = part_static.find(token); const int32_t count_primary = token_count_primary.second; const int32_t count_static = token_count_static_it != part_static.end() ? 100*token_count_static_it->second : 1; @@ -162,11 +162,9 @@ void common_ngram_cache_draft( for (int j = ngram_start_static; j < ngram_start_static + LLAMA_NGRAM_STATIC; ++j) { ngram_static.tokens[j-ngram_start_static] = get_token(inp, draft, j); } - common_ngram_cache::iterator part_static_it = nc_static.find(ngram_static); - common_ngram_cache_part part_static; - if (part_static_it != nc_static.end()) { - part_static = part_static_it->second; - } + static const common_ngram_cache_part part_static_empty; + common_ngram_cache::const_iterator part_static_it = nc_static.find(ngram_static); + const common_ngram_cache_part & part_static = part_static_it != nc_static.end() ? part_static_it->second : part_static_empty; // cd = context + dynamic std::vector<common_ngram> ngrams_cd; From 623c16393c7d53a37df8bb69c6ea8ed6e5171f62 Mon Sep 17 00:00:00 2001 From: uvos <carl@uvos.xyz> Date: Fri, 16 Oct 2026 09:07:34 +0200 Subject: [PATCH 052/423] CUDA: HIP: rms_norm: add variants that load float4 at a time and store x in registers (#29720) * CUDA: HIP: rms_norm: load float4 at a time and store x in registers * Add template paramter comments * use ggml_cuda_is_aligned for src * nits * also use ggml_cuda_is_aligned for src in fused variants * fix nits --- ggml/src/ggml-cuda/norm.cu | 279 +++++++++++++++++++++++-------------- 1 file changed, 173 insertions(+), 106 deletions(-) diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index d804b933bbb5..3f0ea6f2a474 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -81,7 +81,7 @@ static __global__ void group_norm_f32(const float * x, float * dst, const int gr } } -template <int block_size, bool do_multiply = false, bool do_add = false, bool do_scale = false> +template <int block_size, int vpt, bool do_vec, bool do_multiply = false, bool do_add = false, bool do_scale = false> static __global__ void rms_norm_f32(const float * x, float * dst, const int ncols, @@ -142,31 +142,82 @@ static __global__ void rms_norm_f32(const float * x, addc = add + add_sample * add_stride_sample + add_channel * add_stride_channel + add_row * add_stride_row; } - float tmp = 0.0f; // partial sum for thread in warp - - for (int col = tid; col < ncols; col += block_size) { - const float xi = xc[col]; - tmp += xi * xi; - } - - // sum up partial sums - tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum); + if constexpr (do_vec) { + const int ncols4 = ncols / 4; + const float4 * x4 = (const float4 *) xc; + const float4 * mul4 = (const float4 *) mulc; + const float4 * add4 = (const float4 *) addc; + float4 * dst4 = (float4 *) dstc; + + float4 xs[vpt]; + float tmp = 0.0f; + +#pragma unroll + for (int k = 0; k < vpt; k++) { + const int i = tid + k*block_size; + if (i < ncols4) { + xs[k] = x4[i]; + tmp += xs[k].x*xs[k].x + xs[k].y*xs[k].y + xs[k].z*xs[k].z + xs[k].w*xs[k].w; + } + } + tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum); + const float scale = rsqrtf(tmp / ncols + eps); + +#pragma unroll + for (int k = 0; k < vpt; k++) { + const int i = tid + k*block_size; + if (i < ncols4) { + float4 r = make_float4(scale*xs[k].x, scale*xs[k].y, scale*xs[k].z, scale*xs[k].w); + if constexpr (do_multiply) { + const float4 m = mul4[i]; + r.x *= m.x; + r.y *= m.y; + r.z *= m.z; + r.w *= m.w; + } + if constexpr (do_add) { + const float4 a = add4[i]; + r.x += a.x; + r.y += a.y; + r.z += a.z; + r.w += a.w; + } + if constexpr (do_scale) { + r.x *= scale_out; + r.y *= scale_out; + r.z *= scale_out; + r.w *= scale_out; + } + dst4[i] = r; + } + } + } else { + float tmp = 0.0f; // partial sum for thread in warp - const float mean = tmp / ncols; - const float scale = rsqrtf(mean + eps); + for (int col = tid; col < ncols; col += block_size) { + const float xi = xc[col]; + tmp += xi * xi; + } - for (int col = tid; col < ncols; col += block_size) { - if constexpr (do_multiply && do_add) { - const int mul_col = fastmodulo(col, mul_ncols_packed); - const int add_col = fastmodulo(col, add_ncols_packed); - dstc[col] = scale * xc[col] * mulc[mul_col] + addc[add_col]; - } else if constexpr (do_multiply) { - const int mul_col = fastmodulo(col, mul_ncols_packed); - dstc[col] = scale * xc[col] * mulc[mul_col]; - } else if constexpr (do_scale) { - dstc[col] = scale_out * (scale * xc[col]); - } else { - dstc[col] = scale * xc[col]; + // sum up partial sums + tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum); + + const float mean = tmp / ncols; + const float scale = rsqrtf(mean + eps); + + for (int col = tid; col < ncols; col += block_size) { + if constexpr (do_multiply && do_add) { + const int mul_col = fastmodulo(col, mul_ncols_packed); + const int add_col = fastmodulo(col, add_ncols_packed); + dstc[col] = scale * xc[col] * mulc[mul_col] + addc[add_col]; + } else if constexpr (do_multiply) { + const int mul_col = fastmodulo(col, mul_ncols_packed); + dstc[col] = scale * xc[col] * mulc[mul_col]; + } else if constexpr (do_scale) { + dstc[col] = scale_out * (scale * xc[col]); + } else { + dstc[col] = scale * xc[col]; + } } } @@ -332,34 +383,89 @@ static void group_norm_f32_cuda( } } +#define RMS_NORM_LAUNCH(BLOCK_SIZE, VPT, VEC) \ + do { \ + const dim3 block_dims(BLOCK_SIZE, 1, 1); \ + const ggml_cuda_kernel_launch_params launch_params{ blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float) : 0, stream }; \ + ggml_cuda_kernel_launch(rms_norm_f32<BLOCK_SIZE, VPT, VEC, do_multiply, do_add, do_scale>, launch_params, \ + x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, \ + mul, mul_stride_row, mul_stride_channel, mul_stride_sample, \ + mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, \ + add, add_stride_row, add_stride_channel, add_stride_sample, \ + add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, scale_out); \ + } while (0) + +template <bool do_multiply, bool do_add, bool do_scale> +static void rms_norm_f32_dispatch( + const dim3 blocks_num, + const ggml_tensor * x_t, float * dst, const int ncols, const int nchannels, const int nsamples, + const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, + const ggml_tensor * mul_t, const int64_t mul_stride_row, const int64_t mul_stride_channel, const int64_t mul_stride_sample, + const uint3 mul_ncols_packed, const uint3 mul_nrows_packed, const uint3 mul_nchannels_packed, const uint3 mul_nsamples_packed, + const ggml_tensor * add_t, const int64_t add_stride_row, const int64_t add_stride_channel, const int64_t add_stride_sample, + const uint3 add_ncols_packed, const uint3 add_nrows_packed, const uint3 add_nchannels_packed, const uint3 add_nsamples_packed, + cudaStream_t stream, const float scale_out = 1.0f) { + static_assert(!do_add || do_multiply, "fusing add is not supported without multiplying"); + static_assert(!do_scale || !do_multiply, "fusing scale is not supported with multiplying"); + + const float * x = (const float *) x_t->data; + const float * mul = mul_t ? (const float *) mul_t->data : nullptr; + const float * add = add_t ? (const float *) add_t->data : nullptr; + + // dst is written at row * ncols, i.e. always at a multiple of ncols, so only its base needs checking here + bool vec = ncols % 4 == 0 && ncols > 0 && ((uintptr_t) dst & 15) == 0; + vec = vec && ggml_cuda_is_aligned(x_t, 16); + if constexpr (do_multiply) { + vec = vec && mul_ncols_packed.z == (uint32_t) ncols + && ggml_cuda_is_aligned(mul_t, 16); + } + if constexpr (do_add) { + vec = vec && add_ncols_packed.z == (uint32_t) ncols + && ggml_cuda_is_aligned(add_t, 16); + } + + const int n4 = ncols / 4; + if (vec) { + if (n4 <= 32) { + RMS_NORM_LAUNCH(32, 1, true); + } else if (n4 <= 128) { + RMS_NORM_LAUNCH(128, 1, true); + } else if (n4 <= 256) { + RMS_NORM_LAUNCH(256, 1, true); + } else if (n4 <= 512) { + RMS_NORM_LAUNCH(512, 1, true); + } else if (n4 <= 1024) { + RMS_NORM_LAUNCH(1024, 1, true); + } else if (n4 <= 4096) { + RMS_NORM_LAUNCH(1024, 4, true); + } else { + RMS_NORM_LAUNCH(1024, 1, false); + } + } else if (ncols < 1024) { + RMS_NORM_LAUNCH(256, 1, false); + } else { + RMS_NORM_LAUNCH(1024, 1, false); + } +} +#undef RMS_NORM_LAUNCH + template <bool do_scale = false> static void rms_norm_f32_cuda( - const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, + const ggml_tensor * x_t, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream, const float scale_out = 1.0f) { const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX)); - if (ncols < 1024) { - const dim3 block_dims(256, 1, 1); - const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; - ggml_cuda_kernel_launch(rms_norm_f32<256, false, false, do_scale>, launch_params, - x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, - // underlying cudaLaunchKernelEx does not support default params + rms_norm_f32_dispatch</*do_multiply*/false, /*do_add*/false, do_scale>( + blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), - nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), scale_out); - } else { - const dim3 block_dims(1024, 1, 1); - const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; - ggml_cuda_kernel_launch(rms_norm_f32<1024, false, false, do_scale>, launch_params, x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, - // underlying cudaLaunchKernelEx does not support default params nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), - nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), scale_out); - } + stream, scale_out); } -static void rms_norm_mul_f32_cuda(const float * x, - const float * mul, - const float * add, - float * dst, +static void rms_norm_mul_f32_cuda(const ggml_tensor * x_t, + const ggml_tensor * mul_t, + const ggml_tensor * add_t, + float * dst, const int ncols, const int nrows, const int nchannels, @@ -384,59 +490,33 @@ static void rms_norm_mul_f32_cuda(const float * x, const float eps, cudaStream_t stream) { const dim3 blocks_num(nrows, MIN(nchannels, UINT16_MAX), MIN(nsamples, UINT16_MAX)); - if (mul == nullptr) { - rms_norm_f32_cuda(x, dst, ncols, nrows, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, stream); + if (mul_t == nullptr) { + rms_norm_f32_cuda(x_t, dst, ncols, nrows, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, stream); return; } - if (add == nullptr) { - const uint3 mul_ncols_packed = init_fastdiv_values(mul_ncols); - const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows); - const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels); - const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples); - if (ncols < 1024) { - const dim3 block_dims(256, 1, 1); - const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; - ggml_cuda_kernel_launch(rms_norm_f32<256, true>, launch_params, - x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, - mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, - // underlying cudaLaunchKernelEx does not support default params - nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), 1.0f); - } else { - const dim3 block_dims(1024, 1, 1); - const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; - ggml_cuda_kernel_launch(rms_norm_f32<1024, true>, launch_params, - x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, - mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, - // underlying cudaLaunchKernelEx does not support default params - nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), 1.0f); - } + const uint3 mul_ncols_packed = init_fastdiv_values(mul_ncols); + const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows); + const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels); + const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples); + const uint3 zero_packed = make_uint3(0, 0, 0); + + if (add_t == nullptr) { + rms_norm_f32_dispatch</*do_multiply*/true, /*do_add*/false, /*do_scale*/false>( + blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, + mul_t, mul_stride_row, mul_stride_channel, mul_stride_sample, + mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, + nullptr, 0, 0, 0, zero_packed, zero_packed, zero_packed, zero_packed, stream); } else { - const uint3 mul_ncols_packed = init_fastdiv_values(mul_ncols); - const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows); - const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels); - const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples); - const uint3 add_ncols_packed = init_fastdiv_values(add_ncols); const uint3 add_nrows_packed = init_fastdiv_values(add_nrows); const uint3 add_nchannels_packed = init_fastdiv_values(add_nchannels); const uint3 add_nsamples_packed = init_fastdiv_values(add_nsamples); - if (ncols < 1024) { - const dim3 block_dims(256, 1, 1); - const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims,block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; - ggml_cuda_kernel_launch(rms_norm_f32<256, true, true>, launch_params, - x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, - mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, - add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, - add_nchannels_packed, add_nsamples_packed, 1.0f); - } else { - const dim3 block_dims(1024, 1, 1); - const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; - ggml_cuda_kernel_launch(rms_norm_f32<1024, true, true>, launch_params, - x, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, - mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, add, - add_stride_row, add_stride_channel, add_stride_sample, add_ncols_packed, add_nrows_packed, - add_nchannels_packed, add_nsamples_packed, 1.0f); - } + rms_norm_f32_dispatch</*do_multiply*/true, /*do_add*/true, /*do_scale*/false>( + blocks_num, x_t, dst, ncols, nchannels, nsamples, stride_row, stride_channel, stride_sample, eps, + mul_t, mul_stride_row, mul_stride_channel, mul_stride_sample, + mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, + add_t, add_stride_row, add_stride_channel, add_stride_sample, + add_ncols_packed, add_nrows_packed, add_nchannels_packed, add_nsamples_packed, stream); } } @@ -510,7 +590,6 @@ void ggml_cuda_op_group_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const float * src0_d = (const float *) src0->data; float * dst_d = (float *) dst->data; cudaStream_t stream = ctx.stream(); @@ -529,12 +608,11 @@ void ggml_cuda_op_rms_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int64_t s02 = nb02 / ts0; const int64_t s03 = nb03 / ts0; - rms_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); + rms_norm_f32_cuda(src0, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); } void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * scale_tensor) { const ggml_tensor * src0 = dst->src[0]; - const float * src0_d = (const float *) src0->data; float * dst_d = (float *) scale_tensor->data; cudaStream_t stream = ctx.stream(); @@ -556,7 +634,7 @@ void ggml_cuda_op_rms_norm_scale_fused(ggml_backend_cuda_context & ctx, ggml_ten const int64_t s02 = nb02 / ts0; const int64_t s03 = nb03 / ts0; - rms_norm_f32_cuda<true>(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream, scale); + rms_norm_f32_cuda<true>(src0, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream, scale); } void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor) { @@ -565,15 +643,11 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * memcpy(&eps, dst->op_params, sizeof(float)); - const float * src0_d = (const float *) rms_norm_src->data; - const float * mul_d = nullptr; const ggml_tensor * mul_src = nullptr; if (mul_tensor->src[0] == dst) { - mul_d = (float *) mul_tensor->src[1]->data; mul_src = mul_tensor->src[1]; } else if(mul_tensor->src[1] == dst) { - mul_d = (float *) mul_tensor->src[0]->data; mul_src = mul_tensor->src[0]; } else { GGML_ASSERT(false); @@ -609,7 +683,7 @@ void ggml_cuda_op_rms_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * const int mul_nchannels = mul_src->ne[2]; const int mul_nsamples = mul_src->ne[3]; - rms_norm_mul_f32_cuda(src0_d, mul_d, nullptr, dst_d, + rms_norm_mul_f32_cuda(rms_norm_src, mul_src, nullptr, dst_d, ne00, ne01, ne02, ne03, /*s00*/ s01, s02, s03, /*mul_s00*/ mul_s01, mul_s02, mul_s03, @@ -628,28 +702,21 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, memcpy(&eps, dst->op_params, sizeof(float)); - const float * src0_d = (const float *) rms_norm_src->data; - const float * mul_d = nullptr; const ggml_tensor * mul_src = nullptr; if (mul_tensor->src[0] == dst) { - mul_d = (float *) mul_tensor->src[1]->data; mul_src = mul_tensor->src[1]; } else if (mul_tensor->src[1] == dst) { - mul_d = (float *) mul_tensor->src[0]->data; mul_src = mul_tensor->src[0]; } else { GGML_ASSERT(false); } - const float * add_d = nullptr; const ggml_tensor * add_src = nullptr; if (add_tensor->src[0] == mul_tensor) { - add_d = (float *) add_tensor->src[1]->data; add_src = add_tensor->src[1]; } else if (add_tensor->src[1] == mul_tensor) { - add_d = (float *) add_tensor->src[0]->data; add_src = add_tensor->src[0]; } else { GGML_ASSERT(false); @@ -697,7 +764,7 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, const int add_nchannels = add_src->ne[2]; const int add_nsamples = add_src->ne[3]; - rms_norm_mul_f32_cuda(src0_d, mul_d,add_d,dst_d, + rms_norm_mul_f32_cuda(rms_norm_src, mul_src,add_src,dst_d, ne00,ne01, ne02, ne03, /*s00*/ s01, s02, s03, /*mul_s00*/ mul_s01, mul_s02, mul_s03, From 495c5805ce92a3044b960830361c22aa75f30d38 Mon Sep 17 00:00:00 2001 From: Sh1m13 <93532735+SomeNobody21112@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:56:14 +0530 Subject: [PATCH 053/423] gguf-py : add numpy Q1_0 and Q2_0 de/quantization support (#30630) * gguf-py : add numpy Q1_0 and Q2_0 de/quantization support * cont: added q1_0 / q2_0 to the test list removing the small test Assisted-by: Claude Opus 5.5 <noreply@anthropic.com> --------- Assisted-by: Claude Opus 5.5 <noreply@anthropic.com> --- gguf-py/gguf/quants.py | 53 ++++++++++++++++++++++++++++++++++++ gguf-py/tests/test_quants.py | 1 + 2 files changed, 54 insertions(+) diff --git a/gguf-py/gguf/quants.py b/gguf-py/gguf/quants.py index 80966b6ef151..cc0772913bec 100644 --- a/gguf-py/gguf/quants.py +++ b/gguf-py/gguf/quants.py @@ -654,6 +654,59 @@ def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: return (d * qs.astype(np.float32)) +class Q1_0(__Quant, qtype=GGMLQuantizationType.Q1_0): + @classmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + # sequential sum like the C reference, np.sum uses pairwise summation and can round differently + d = np.cumsum(abs(blocks), axis=-1, dtype=np.float32)[..., -1:] / np.float32(cls.block_size) + qs = np.packbits(blocks >= 0, axis=-1, bitorder="little") + + d = d.astype(np.float16).view(np.uint8) + + return np.concatenate([d, qs], axis=-1) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + d, qs = np.hsplit(blocks, [2]) + + d = d.view(np.float16).astype(np.float32) + + qs = np.unpackbits(qs, axis=-1, bitorder="little") + + return np.where(qs != 0, d, -d) + + +class Q2_0(__Quant, qtype=GGMLQuantizationType.Q2_0): + @classmethod + def quantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d = abs(blocks).max(axis=-1, keepdims=True) + with np.errstate(divide="ignore"): + id = np.where(d == 0, 0, 1 / d) + qs = (np_roundf(blocks * id) + 1).clip(0, 3).astype(np.uint8) + + qs = qs.reshape((n_blocks, -1, 4)) << np.array([0, 2, 4, 6], dtype=np.uint8) + qs = qs[..., 0] | qs[..., 1] | qs[..., 2] | qs[..., 3] + + d = d.astype(np.float16).view(np.uint8) + + return np.concatenate([d, qs], axis=-1) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + d, qs = np.hsplit(blocks, [2]) + + d = d.view(np.float16).astype(np.float32) + + qs = qs.reshape((n_blocks, -1, 1)) >> np.array([0, 2, 4, 6], dtype=np.uint8) + qs = (qs & 0x03).reshape((n_blocks, -1)).astype(np.int8) - np.int8(1) + + return (d * qs.astype(np.float32)) + + class MXFP4(__Quant, qtype=GGMLQuantizationType.MXFP4): # e2m1 values (doubled) # ref: https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf diff --git a/gguf-py/tests/test_quants.py b/gguf-py/tests/test_quants.py index 9aa7c4ae2a5e..0d7729f95cc8 100755 --- a/gguf-py/tests/test_quants.py +++ b/gguf-py/tests/test_quants.py @@ -69,6 +69,7 @@ def __init__(self, libggml: Path): "tq1_0", "tq2_0", "mxfp4", "nvfp4", + "q1_0", "q2_0", "iq2_xxs", "iq2_xs", "iq2_s", "iq3_xxs", "iq3_s", "iq1_s", "iq1_m", "iq4_nl", "iq4_xs", ): From a30fb4924dcfec81c1b2c20634bd47fa4403061a Mon Sep 17 00:00:00 2001 From: Michael Taylor <162068037+mctylr-gh@users.noreply.github.com> Date: Fri, 16 Oct 2026 04:28:50 -0300 Subject: [PATCH 054/423] [SYCL]: fix memory errors in mul_mat, split buffer, host pool (#29889) --- ggml/src/ggml-sycl/common.cpp | 3 +- ggml/src/ggml-sycl/cpy.cpp | 2 + ggml/src/ggml-sycl/ggml-sycl.cpp | 118 ++++++++++++++++++++++--------- ggml/src/ggml-sycl/roll.cpp | 2 + 4 files changed, 90 insertions(+), 35 deletions(-) diff --git a/ggml/src/ggml-sycl/common.cpp b/ggml/src/ggml-sycl/common.cpp index 894006949d23..f01d0ff955e5 100644 --- a/ggml/src/ggml-sycl/common.cpp +++ b/ggml/src/ggml-sycl/common.cpp @@ -115,7 +115,8 @@ void * ggml_sycl_malloc_device(size_t size, sycl::queue &q, ggml_sycl_mem_type t #else ze_device_mem_alloc_desc_t alloc_desc = {ZE_STRUCTURE_TYPE_DEVICE_MEM_ALLOC_DESC, nullptr, 0, 0}; #endif - ze_result_t r = zeMemAllocDevice(ze_ctx, &alloc_desc, size, 64, ze_dev, &ptr); + // request SYCL_BUFFER_ALIGNMENT so tensor offsets aligned to it stay aligned in absolute terms + ze_result_t r = zeMemAllocDevice(ze_ctx, &alloc_desc, size, SYCL_BUFFER_ALIGNMENT, ze_dev, &ptr); if (r == ZE_RESULT_SUCCESS && ptr) { ggml_sycl_memtrace_add(type, ptr, size); return ptr; diff --git a/ggml/src/ggml-sycl/cpy.cpp b/ggml/src/ggml-sycl/cpy.cpp index ef7413abd880..8e1f1d50a934 100644 --- a/ggml/src/ggml-sycl/cpy.cpp +++ b/ggml/src/ggml-sycl/cpy.cpp @@ -1255,6 +1255,8 @@ void ggml_sycl_cpy(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, co scope_op_debug_print scope_dbg_print(__func__, src1, /*num_src=*/0, debug_get_tensor_str("\tsrc0", src0)); const int64_t ne = ggml_nelements(src0); GGML_ASSERT(ne == ggml_nelements(src1)); + // the cpy kernels index in 32-bit + GGML_ASSERT(ne <= INT32_MAX); GGML_TENSOR_BINARY_OP_LOCALS01; diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 1cfc84126412..b380cbedfdb5 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -899,6 +899,10 @@ static bool ggml_sycl_is_l0_discrete_gpu(int device) { static void memcpy_host_forward(sycl::queue &q_dst, sycl::queue &q_src, void *ptr_dst, const void *ptr_src, size_t size) { char *host_buf = (char *)malloc(size); + if (host_buf == nullptr) { + GGML_LOG_ERROR("%s: can't allocate %zu Bytes of memory for host staging\n", __func__, size); + GGML_ABORT("fatal error"); + } q_src.memcpy(host_buf, (const char *)ptr_src, size).wait(); q_dst.memcpy((char *)ptr_dst, host_buf, size).wait(); free(host_buf); @@ -1177,12 +1181,13 @@ bool is_bmg_g31_arch(int device) { } static size_t ggml_backend_sycl_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { - size_t max_alloc_size = dpct::get_current_device().get_max_mem_alloc_size(); + ggml_backend_sycl_buffer_type_context * ctx = (ggml_backend_sycl_buffer_type_context *)buft->context; + int device = ctx->device; + // use the buft device, the current device may differ on multi-GPU systems + size_t max_alloc_size = dpct::get_device(device).get_max_mem_alloc_size(); if (g_ggml_sycl_host_pinned_mem_2g) { return std::min(max_alloc_size, (size_t) 2LL*1024*1024*1024); } else { - ggml_backend_sycl_buffer_type_context * ctx = (ggml_backend_sycl_buffer_type_context *)buft->context; - int device = ctx->device; if(is_bmg_g31_arch(device)) { //Todo, it's workaround for BMG-G31, which has a known issue with large allocations. //The max alloc size is reduced to 60% of the reported max alloc size. @@ -1398,7 +1403,12 @@ ggml_backend_sycl_split_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor_extra_gpu * extra = new ggml_tensor_extra_gpu{}; ctx->tensor_extras.push_back(extra); - ctx->streams.push_back(&(dpct::get_current_device().default_queue())); + // consumers index the queues by device id, so size the vector once with one queue per device + if (ctx->streams.empty()) { + for (int i = 0; i < ggml_sycl_info().device_count; ++i) { + ctx->streams.push_back(&(dpct::get_device(i).default_queue())); + } + } for (int i = 0; i < ggml_sycl_info().device_count; ++i) { int64_t row_low, row_high; @@ -2089,8 +2099,6 @@ struct ggml_sycl_pool_host : public ggml_sycl_pool { queue_ptr qptr; int device; - inline static int counter{ 0 }; - struct ggml_sycl_buffer { void * ptr = nullptr; size_t size = 0; @@ -2100,6 +2108,7 @@ struct ggml_sycl_pool_host : public ggml_sycl_pool { static constexpr int MAX_POOL_SIZE{ 64 }; std::vector<ggml_sycl_buffer> buffer_pool = std::vector<ggml_sycl_buffer>(MAX_POOL_SIZE); size_t pool_size = 0; + int counter = 0; explicit ggml_sycl_pool_host(queue_ptr qptr_, int device_) : qptr(qptr_), device(device_) {} @@ -2113,41 +2122,55 @@ struct ggml_sycl_pool_host : public ggml_sycl_pool { b.size = 0; } } - counter = 0; } void * alloc(size_t size, size_t * actual_size) override { if (counter == MAX_POOL_SIZE) { - ggml_sycl_buffer b = buffer_pool[0]; - void * ptr = b.ptr; - *actual_size = b.size; - counter = 1; - return ptr; - } - ggml_sycl_buffer & b = buffer_pool[counter]; - - if (b.ptr == nullptr) { - void * ptr; - + // the pool wrapped around, no slot is reserved for this allocation + void * ptr = nullptr; SYCL_CHECK(CHECK_TRY_ERROR(ptr = (void *) sycl::malloc_host(size, *qptr))); if (!ptr) { GGML_LOG_ERROR("%s: can't allocate %zu Bytes of memory on host\n", __func__, size); return nullptr; } - pool_size += size; + pool_size += size; *actual_size = size; - counter = counter + 1; + counter = 0; return ptr; - } else { - ++counter; - b.size = size; - return b.ptr; } + ggml_sycl_buffer & b = buffer_pool[counter]; + ++counter; + + if (b.ptr != nullptr) { + if (b.size >= size) { + // hand out the cached buffer, free() puts it back into the pool + void * ptr = b.ptr; + *actual_size = b.size; + b.ptr = nullptr; + b.size = 0; + return ptr; + } + // the cached buffer is too small, replace it + SYCL_CHECK(CHECK_TRY_ERROR(sycl::free(b.ptr, *qptr))); + pool_size -= b.size; + b.ptr = nullptr; + b.size = 0; + } + + void * ptr; + SYCL_CHECK(CHECK_TRY_ERROR(ptr = (void *) sycl::malloc_host(size, *qptr))); + if (!ptr) { + GGML_LOG_ERROR("%s: can't allocate %zu Bytes of memory on host\n", __func__, size); + return nullptr; + } + pool_size += size; + *actual_size = size; + return ptr; } void free(void * ptr, size_t size) override { // if the pool is not completed add the pointer to it in place of the first nullptr found. - // Otherwise do nothing, pointers will be freed once the pool is deallocated. + // Otherwise free the buffer, all slots are cached. for (int i = 0; i < MAX_POOL_SIZE; ++i) { ggml_sycl_buffer & b = buffer_pool[i]; if (b.ptr == nullptr) { @@ -2156,6 +2179,8 @@ struct ggml_sycl_pool_host : public ggml_sycl_pool { return; } } + SYCL_CHECK(CHECK_TRY_ERROR(sycl::free(ptr, *qptr))); + pool_size -= size; } }; @@ -3154,8 +3179,8 @@ inline void ggml_sycl_op_mul_mat_sycl( to_fp16_sycl(src1_ddf_i, src1_as_f16.get(), ne, stream); } const sycl::half *src1_ptr = src1->type == GGML_TYPE_F16 - ? (const sycl::half *)src1->data + src1_padded_row_size - : src1_as_f16.get(); + ? (const sycl::half *)src1_ddf_i + : src1_as_f16.get(); ggml_sycl_pool_alloc<sycl::half> src0_as_f16(ctx.pool()); if (src0->type != GGML_TYPE_F16) { @@ -3252,6 +3277,8 @@ inline void ggml_sycl_op_sum(ggml_backend_sycl_context & ctx, ggml_tensor *dst) float * dst_dd = static_cast<float *>(dst->data); const int64_t ne = ggml_nelements(dst->src[0]); + // the kernel indexes in 32-bit + GGML_ASSERT(ne <= INT32_MAX); sum_rows_f32_sycl(src0_dd, dst_dd, ne, 1, main_stream); } @@ -3266,6 +3293,8 @@ inline void ggml_sycl_op_sum_rows(ggml_backend_sycl_context & ctx, ggml_tensor * const int64_t ncols = dst->src[0]->ne[0]; const int64_t nrows = ggml_nrows(dst->src[0]); + // the kernel indexes in 32-bit + GGML_ASSERT(ncols * nrows <= INT32_MAX); sum_rows_f32_sycl(src0_dd, dst_dd, ncols, nrows, main_stream); } @@ -3282,6 +3311,8 @@ inline void ggml_sycl_op_mean(ggml_backend_sycl_context & ctx, ggml_tensor * dst const int64_t ncols = dst->src[0]->ne[0]; const int64_t nrows = ggml_nrows(dst->src[0]); + // the kernel indexes in 32-bit + GGML_ASSERT(ncols * nrows <= INT32_MAX); sum_rows_f32_sycl(src0_dd, dst_dd, ncols, nrows, main_stream); @@ -3351,6 +3382,8 @@ inline void ggml_sycl_op_argmax(ggml_backend_sycl_context & ctx, ggml_tensor * d const int64_t ncols = dst->src[0]->ne[0]; const int64_t nrows = ggml_nrows(dst->src[0]); + // the kernel indexes in 32-bit + GGML_ASSERT(ncols * nrows <= INT32_MAX); argmax_f32_i32_sycl(src0_dd, dst_dd, ncols, nrows, main_stream); } @@ -3366,6 +3399,8 @@ inline void ggml_sycl_op_diag_mask_inf(ggml_backend_sycl_context & ctx, ggml_ten const int64_t ne00 = dst->src[0]->ne[0]; const int64_t ne01 = dst->src[0]->ne[1]; const int nrows0 = ggml_nrows(dst->src[0]); + // the kernel indexes in 32-bit + GGML_ASSERT(ne00 * nrows0 <= INT32_MAX); const int n_past = ((int32_t *) dst->op_params)[0]; @@ -3444,6 +3479,9 @@ inline void ggml_sycl_op_scale(ggml_backend_sycl_context & ctx, ggml_tensor * ds memcpy(&scale, (float *) dst->op_params + 0, sizeof(float)); memcpy(&bias, (float *) dst->op_params + 1, sizeof(float)); + // the kernel indexes in 32-bit + GGML_ASSERT(ggml_nelements(dst->src[0]) <= INT32_MAX); + scale_f32_sycl(src0_dd, dst_dd, scale, bias, ggml_nelements(dst->src[0]), main_stream); /* DPCT1010:87: SYCL uses exceptions to report errors and does not use the @@ -3524,7 +3562,6 @@ static void ggml_sycl_op_mul_mat(ggml_backend_sycl_context & ctx, const ggml_ten const size_t q8_1_bs = QK8_1; ggml_tensor_extra_gpu * src0_extra = (ggml_tensor_extra_gpu *) src0->extra; - ggml_tensor_extra_gpu * src1_extra = (ggml_tensor_extra_gpu *) src1->extra; const bool src0_is_contiguous = ggml_is_contiguous(src0); const bool src1_is_contiguous = ggml_is_contiguous(src1); @@ -3606,7 +3643,8 @@ static void ggml_sycl_op_mul_mat(ggml_backend_sycl_context & ctx, const ggml_ten queue_ptr stream = ctx.stream(i, 0); if (src0_is_contiguous) { - dev[i].src0_dd = (char *) src0->data; + // a split tensor keeps its rows in the per-device slices of the tensor extra + dev[i].src0_dd = split ? (char *) src0_extra->data_device[i] : (char *) src0->data; } else { dev[i].src0_dd = dev[i].src0_dd_alloc.alloc(ctx.pool(i), ggml_nbytes(src0)); } @@ -3702,7 +3740,8 @@ static void ggml_sycl_op_mul_mat(ggml_backend_sycl_context & ctx, const ggml_ten src1_ncols * src1_padded_col_size * q8_1_ts / q8_1_bs) .wait())); } else { - const char * src1_ddf_i_source = (const char *) src1_extra->data_device[ctx.device] + + // src1 is not split, on the main device it lives at dev[ctx.device].src1_ddf + const char * src1_ddf_i_source = (const char *) dev[ctx.device].src1_ddf + (i0 * ne11 + src1_col_0) * ne10 * ggml_type_size(src1->type); SYCL_CHECK( @@ -5269,6 +5308,7 @@ static bool ggml_sycl_mul_mat_glu_mmvq_plain(ggml_backend_sycl_context & ctx, gg (int) ne11, /*stride_col_y=*/src1_padded_cols / QK8_1, /*stride_col_dst=*/(int) glu->ne[0], stream); + GGML_UNUSED(gate); } // Fused dense-FFN mat-vec for the {mul_mat(gate), mul_mat(up), GLU} subgraph at node_idx. @@ -6860,7 +6900,9 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_GLU_OP_GEGLU_ERF: case GGML_GLU_OP_GEGLU_QUICK: case GGML_GLU_OP_SWIGLU_CLAMP: - return ggml_is_contiguous_1(op->src[0]); + // the fused GLU kernels only have f32/f16 instantiations + return ggml_is_contiguous_1(op->src[0]) && + (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16); default: return false; } @@ -7123,7 +7165,8 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_RMS_NORM_BACK: return ggml_is_contiguous(op->src[0]); case GGML_OP_SCALE: - return true; + // the scale kernel is f32 only + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32; case GGML_OP_CONT: return true; case GGML_OP_TRI: @@ -7156,12 +7199,19 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons #endif ) && op->src[0]->type == op->type; - case GGML_OP_CONV_3D: + case GGML_OP_CONV_3D: { + const int32_t * opts = (const int32_t *) op->op_params; + // zero kernel volume, channels or output makes the im2col GEMM degenerate (k == 0 or m == 0) + // and MKL BLAS rejects it; let those fall back to the CPU backend + const int64_t knl_n_total = op->src[0]->ne[0] * op->src[0]->ne[1] * op->src[0]->ne[2] * opts[9]; + const int64_t patch_total = op->ne[0] * op->ne[1] * op->ne[2] * opts[10]; return op->type == GGML_TYPE_F32 && (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) && op->src[1]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0]) && - ggml_is_contiguous(op->src[1]); + ggml_is_contiguous(op->src[1]) && + knl_n_total > 0 && patch_total > 0; + } case GGML_OP_SUM: case GGML_OP_SUM_ROWS: case GGML_OP_MEAN: diff --git a/ggml/src/ggml-sycl/roll.cpp b/ggml/src/ggml-sycl/roll.cpp index 1e05181789c2..a95cee9a974a 100644 --- a/ggml/src/ggml-sycl/roll.cpp +++ b/ggml/src/ggml-sycl/roll.cpp @@ -73,6 +73,8 @@ void ggml_sycl_roll(ggml_backend_sycl_context & ctx, ggml_tensor *dst) { const ggml_tensor *src = dst->src[0]; GGML_ASSERT(src && src->type == GGML_TYPE_F32); + // the kernel indexes in 32-bit + GGML_ASSERT(dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3] <= INT32_MAX); const int ne0 = (int) dst->ne[0]; const int ne1 = (int) dst->ne[1]; const int ne2 = (int) dst->ne[2]; From b84ea5cf5e56dd934e0d52c6ffdfffbf06f58d2f Mon Sep 17 00:00:00 2001 From: Aaron Teo <aaron.teo1@ibm.com> Date: Fri, 16 Oct 2026 16:18:29 +0800 Subject: [PATCH 055/423] convert : speed up model endian conversion and add more datatype support (#30647) * convert: speed up model endian conversion Signed-off-by: Aaron Teo <aaron.teo1@ibm.com> * convert: explicitly state noop Signed-off-by: Aaron Teo <aaron.teo1@ibm.com> * convert: add missing data type support Signed-off-by: Aaron Teo <aaron.teo1@ibm.com> --------- Signed-off-by: Aaron Teo <aaron.teo1@ibm.com> --- gguf-py/gguf/scripts/gguf_convert_endian.py | 115 +++++--------------- 1 file changed, 28 insertions(+), 87 deletions(-) diff --git a/gguf-py/gguf/scripts/gguf_convert_endian.py b/gguf-py/gguf/scripts/gguf_convert_endian.py index 31618acfc79e..3b20998c68a7 100755 --- a/gguf-py/gguf/scripts/gguf_convert_endian.py +++ b/gguf-py/gguf/scripts/gguf_convert_endian.py @@ -19,71 +19,27 @@ logger = logging.getLogger("gguf-convert-endian") -def byteswap_noop(tensor, block_offs): - # this function is used when byteswapping is not needed - pass - - -def byteswap_q4_0(tensor, block_offs): - # Each block_q4_0 consists of an f16 delta (scaling factor) followed by 16 int8 quantizations. - - # Byte-Swap f16 sized delta field - delta = tensor.data[block_offs:block_offs + 2].view(dtype=np.uint16) - delta.byteswap(inplace=True) - - -def byteswap_q8_0(tensor, block_offs): - # Each block_q8_0 consists of an f16 delta (scaling factor) followed by 32 int8 quantizations. - - # Byte-Swap f16 sized delta field - delta = tensor.data[block_offs:block_offs + 2].view(dtype=np.uint16) - delta.byteswap(inplace=True) - - -def byteswap_q4_k(tensor, block_offs): - # Each block_q4_k consists of 2 f16 values followed by 140 int8 values. - - # Byte-Swap f16 sized fields - delta = tensor.data[block_offs:block_offs + 2].view(dtype=np.uint16) - delta.byteswap(inplace=True) - - delta = tensor.data[block_offs + 2:block_offs + 4].view(dtype=np.uint16) - delta.byteswap(inplace=True) - - -def byteswap_q6_k(tensor, block_offs): - # Each block_q6_k consists of 208 int8 values followed by 1 f16 value. - - # Byte-Swap f16 sized field - delta = tensor.data[block_offs + 208:block_offs + 210].view(dtype=np.uint16) - delta.byteswap(inplace=True) - - -def byteswap_q1_0(tensor, block_offs): - # Each block_q1_0 consists of an f16 delta followed by 16 int8 quantizations. - - # Byte-Swap f16 sized delta field - delta = tensor.data[block_offs:block_offs + 2].view(dtype=np.uint16) - delta.byteswap(inplace=True) - - -def byteswap_tq2_0(tensor, block_offs): - # Each block_tq2_0 consists of 64 int8 values followed by 1 f16 value. - - # Byte-Swap f16 sized field - delta = tensor.data[block_offs + 64:block_offs + 66].view(dtype=np.uint16) - delta.byteswap(inplace=True) - - -byteswap_tensors = { - gguf.GGMLQuantizationType.Q1_0: byteswap_q1_0, - gguf.GGMLQuantizationType.Q4_0: byteswap_q4_0, - gguf.GGMLQuantizationType.Q8_0: byteswap_q8_0, - gguf.GGMLQuantizationType.Q4_K: byteswap_q4_k, - gguf.GGMLQuantizationType.Q6_K: byteswap_q6_k, - gguf.GGMLQuantizationType.TQ2_0: byteswap_tq2_0, - gguf.GGMLQuantizationType.MXFP4: byteswap_noop, - gguf.GGMLQuantizationType.NVFP4: byteswap_noop, +# byte offset and type of the multi-byte fields in each block, all other fields are single bytes +# use the width ggml reads the field with, e.g. Q5_0 qh is uint8_t[4] but read as uint32_t +byteswap_tensors: dict[gguf.GGMLQuantizationType, tuple[tuple[int, type[np.unsignedinteger]], ...]] = { + gguf.GGMLQuantizationType.Q1_0: ((0, np.uint16),), # d + gguf.GGMLQuantizationType.Q2_0: ((0, np.uint16),), # d + gguf.GGMLQuantizationType.Q4_0: ((0, np.uint16),), # d + gguf.GGMLQuantizationType.Q4_1: ((0, np.uint16), (2, np.uint16)), # d, m + gguf.GGMLQuantizationType.Q5_0: ((0, np.uint16), (2, np.uint32)), # d, qh + gguf.GGMLQuantizationType.Q5_1: ((0, np.uint16), (2, np.uint16), (4, np.uint32)), # d, m, qh + gguf.GGMLQuantizationType.Q8_0: ((0, np.uint16),), # d + gguf.GGMLQuantizationType.Q2_K: ((80, np.uint16), (82, np.uint16)), # d, dmin + gguf.GGMLQuantizationType.Q3_K: ((108, np.uint16),), # d + gguf.GGMLQuantizationType.Q4_K: ((0, np.uint16), (2, np.uint16)), # d, dmin + gguf.GGMLQuantizationType.Q5_K: ((0, np.uint16), (2, np.uint16)), # d, dmin + gguf.GGMLQuantizationType.Q6_K: ((208, np.uint16),), # d + gguf.GGMLQuantizationType.IQ4_NL: ((0, np.uint16),), # d + gguf.GGMLQuantizationType.IQ4_XS: ((0, np.uint16), (2, np.uint16)), # d, scales_h + gguf.GGMLQuantizationType.TQ1_0: ((52, np.uint16),), # d + gguf.GGMLQuantizationType.TQ2_0: ((64, np.uint16),), # d + gguf.GGMLQuantizationType.MXFP4: (), # noop + gguf.GGMLQuantizationType.NVFP4: (), # noop } @@ -140,28 +96,12 @@ def convert_byteorder(reader: gguf.GGUFReader, args: argparse.Namespace) -> None # Byte-swap tensor data if necessary if tensor.tensor_type in byteswap_tensors: - # first flatten structure - oldshape = tensor.data.shape - newshape = 1 - for i in tensor.data.shape: - newshape *= i - - tensor.data.resize(newshape) - - block_size = gguf.constants.GGML_QUANT_SIZES[tensor.tensor_type][1] - byteswap_func = byteswap_tensors[tensor.tensor_type] - - n_blocks = len(tensor.data) // block_size - for block_num in (inner_pbar := tqdm(range(n_blocks), desc="Byte-swapping Blocks", leave=False)): - block_offs = block_num * block_size - - byteswap_func(tensor, block_offs) - - if block_num % 100000 == 0: - inner_pbar.set_description(f"Byte-swapping Blocks [{(n_blocks - block_num) // n_blocks}]") - - # restore old shape in case it's ever used - tensor.data.resize(oldshape) + # one row per block, then swap each field of all blocks at once + type_size = gguf.constants.GGML_QUANT_SIZES[tensor.tensor_type][1] + blocks = tensor.data.reshape(-1, type_size, copy=False) + for offs, dtype in byteswap_tensors[tensor.tensor_type]: + size = np.dtype(dtype).itemsize + blocks[:, offs:offs + size].view(dtype=dtype)[:, 0].byteswap(inplace=True) elif tensor.tensor_type == gguf.GGMLQuantizationType.BF16: # Special case for BF16 # It is 2-bytes data, but by default view loads it as 1-byte data. @@ -173,6 +113,7 @@ def convert_byteorder(reader: gguf.GGUFReader, args: argparse.Namespace) -> None pbar.set_description(log_message) + reader.data.flush() logger.info("* Completion") From 9a16f1dbe4722a93ab75a2d15e29aa71ef764c7b Mon Sep 17 00:00:00 2001 From: Mikhail Podvitskii <podvitskiymichael@gmail.com> Date: Fri, 16 Oct 2026 10:20:03 +0200 Subject: [PATCH 056/423] model : MTP support for Mellum2.1 (#30612) * model : MTP support for Mellum * model : inline rope construction in Mellum --- conversion/mellum.py | 74 ++++++++++++- gguf-py/gguf/constants.py | 5 + src/models/mellum.cpp | 222 +++++++++++++++++++++++++++++++++++--- src/models/models.h | 5 + 4 files changed, 288 insertions(+), 18 deletions(-) diff --git a/conversion/mellum.py b/conversion/mellum.py index 6bcb81df8cfe..9dc2192f3778 100644 --- a/conversion/mellum.py +++ b/conversion/mellum.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Iterable, TYPE_CHECKING +from typing import Callable, Iterable, TYPE_CHECKING import torch @@ -14,6 +14,54 @@ @ModelBase.example("JetBrains/Mellum2-12B-A2.5B-Base") class MellumModel(TextModel): model_arch = gguf.MODEL_ARCH.MELLUM + supports_mtp_export = True + + _n_main_layers: int | None = None + + _experts: list[dict[str, Tensor]] | None = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # number of exported MTP blocks, appended after the trunk layers + self.n_nextn = 0 if self.no_mtp else (self.find_hparam(["num_nextn_predict_layers", "num_mtp_layers"], optional=True) or 0) + if self.mtp_only and self.n_nextn == 0: + raise ValueError("--mtp requested, but the checkpoint has no MTP layers (num_nextn_predict_layers is 0)") + if self.n_nextn > 0: + if any(t != "sparse" for t in self.hparams.get("mtp_mlp_layer_types") or []): + raise ValueError("Mellum MTP layers with dense MLP are not supported") + self.block_count += self.n_nextn + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + def index_tensors(self, remote_hf_model_id: str | None = None): + type(self)._n_main_layers = self.hparams["num_hidden_layers"] + return super().index_tensors(remote_hf_model_id=remote_hf_model_id) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + if (titem := super().filter_tensors(item)) is None: + return None + name, gen = titem + + if name.startswith("mtp."): + if cls.no_mtp: + return None + assert cls._n_main_layers is not None + # mtp.layers.{i}.* -> model.layers.{n_main + i}.*, mtp.{fc,norm,...} -> nextn names of the first MTP block + remapper = { + "fc": "eh_proj", + "pre_fc_norm_embedding": "enorm", + "pre_fc_norm_hidden": "hnorm", + "norm": "shared_head.norm", + } + parts = name.split(".", 3) + if len(parts) == 4 and parts[1] == "layers" and parts[2].isdecimal(): + name = f"model.layers.{cls._n_main_layers + int(parts[2])}.{parts[3]}" + elif len(parts) == 3 and parts[1] in remapper: + name = f"model.layers.{cls._n_main_layers}.{remapper[parts[1]]}.{parts[2]}" + elif cls.mtp_only and name not in ("model.embed_tokens.weight", "model.norm.weight", "lm_head.weight"): + return None + + return name, gen def set_gguf_parameters(self): super().set_gguf_parameters() @@ -26,13 +74,31 @@ def set_gguf_parameters(self): if (use_sliding_window is True or use_sliding_window is None) and sliding_window is not None: self.gguf_writer.add_sliding_window(sliding_window) logger.info(f"gguf: sliding window = {sliding_window}") - self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in self.hparams["layer_types"]]) - logger.info(f"gguf: sliding window pattern length = {len(self.hparams['layer_types'])}") + layer_types = list(self.hparams["layer_types"]) + # MTP blocks default to sliding window attention, as in transformers + if self.n_nextn > 0: + layer_types += self.hparams.get("mtp_layer_types") or ["sliding_attention"] * self.n_nextn + self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in layer_types]) + logger.info(f"gguf: sliding window pattern length = {len(layer_types)}") if (norm_topk_prob := self.hparams.get("norm_topk_prob")) is not None: self.gguf_writer.add_expert_weights_norm(norm_topk_prob) - _experts: list[dict[str, Tensor]] | None = None + if self.n_nextn > 0: + self.gguf_writer.add_nextn_predict_layers(self.n_nextn) + + def prepare_metadata(self, vocab_only: bool): + from_dir = self.fname_out.is_dir() + super().prepare_metadata(vocab_only=vocab_only) + + if not self.mtp_only or not from_dir: + return + + output_type: str = self.ftype.name.partition("_")[2] + fname_default: str = gguf.naming_convention( + self.metadata.name, self.metadata.basename, self.metadata.finetune, + self.metadata.version, size_label=None, output_type=output_type, model_type=None) + self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf" def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: if name.find("experts") != -1: diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 92a69bd57369..00b51165c94b 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -5950,6 +5950,11 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_GATE_EXP, MODEL_TENSOR.FFN_DOWN_EXP, MODEL_TENSOR.FFN_UP_EXP, + # NextN/MTP tensors + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], MODEL_ARCH.NANBEIGE: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/src/models/mellum.cpp b/src/models/mellum.cpp index bc36dfa09cdb..eaa540310c7f 100644 --- a/src/models/mellum.cpp +++ b/src/models/mellum.cpp @@ -29,9 +29,13 @@ void llama_model_mellum::load_arch_hparams(llama_model_loader & ml) { } } -void llama_model_mellum::load_arch_tensors(llama_model_loader &) { +void llama_model_mellum::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; + const auto nf = nextn_flags(ml); + const int trunk_flags = nf.trunk; + const int mtp_flags = nf.mtp; + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); // output @@ -45,31 +49,56 @@ void llama_model_mellum::load_arch_tensors(llama_model_loader &) { throw std::runtime_error("n_expert_used must be > 0 for Mellum"); } - for (int i = 0; i < n_layer; ++i) { + auto load_block = [&](int i, int flags) { auto & layer = layers[i]; - layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); - create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, 0); - layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0); + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, flags); - layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); - layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, flags); + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, flags); - layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); - layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); const int64_t n_ff_exp = hparams.n_ff_exp() ? hparams.n_ff_exp() : n_ff / n_expert_used; - layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, 0); - layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); - layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + }; + + for (int i = 0; i < n_layer; ++i) { + load_block(i, trunk_flags); + } + + // MTP block is a regular Mellum decoder block plus the NextN projection and norms + for (int i = n_layer; i < n_layer_all; ++i) { + auto & layer = layers[i]; + + load_block(i, mtp_flags); + + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, mtp_flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, mtp_flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, mtp_flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags); } } std::unique_ptr<llm_graph_context> llama_model_mellum::build_arch_graph(const llm_graph_params & params) const { - if (hparams.swa_type == LLAMA_SWA_TYPE_STANDARD) { + const bool iswa = hparams.swa_type == LLAMA_SWA_TYPE_STANDARD; + + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + if (iswa) { + return std::make_unique<graph_mtp<true>>(*this, params); + } + return std::make_unique<graph_mtp<false>>(*this, params); + } + + if (iswa) { return std::make_unique<graph<true>>(*this, params); } return std::make_unique<graph<false>>(*this, params); @@ -151,7 +180,7 @@ llama_model_mellum::graph<iswa>::graph(const llama_model & model, const llm_grap model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il); } - if (il == n_layer - 1 && inp_out_ids) { + if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } @@ -198,6 +227,14 @@ llama_model_mellum::graph<iswa>::graph(const llama_model & model, const llm_grap model.output_norm, nullptr, LLM_NORM_RMS, -1); + // post-norm hidden state feeds both the LM head and the MTP head + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (crop_after_nextn(inp_out_ids)) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "result_norm", -1); res->t_embd = cur; @@ -208,5 +245,162 @@ llama_model_mellum::graph<iswa>::graph(const llama_model & model, const llm_grap res->t_logits = cur; } +template <bool iswa> +llama_model_mellum::graph_mtp<iswa>::graph_mtp(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + GGML_ASSERT(hparams.n_layer_nextn == 1 && "MELLUM MTP currently only supports a single MTP block"); + + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + GGML_ASSERT(n_embd_head == n_rot); + + const int il = hparams.n_layer(); + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj"); + GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm"); + GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm"); + GGML_ASSERT(layer.nextn.shared_head_norm && "MTP block missing nextn.shared_head_norm"); + + // TODO: extract in a common llm_graph_context::build_inp_embd_h() + auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd_inp(), hparams.n_embd); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens); + ggml_set_input(inp->embd); + + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + + // TODO: make static using `ggml_build_forward_select()` + // see llm_graph_context::build_inp_embd() for reference + ggml_tensor * tok_embd; + ASSERT_EMBD_OR_TOKEN(ubatch); + if (ubatch.token) { + tok_embd = ggml_get_rows(ctx0, model.tok_embd, inp->tokens); + } else { + tok_embd = inp->embd; + } + cb(tok_embd, "mtp_tok_embd", il); + + ggml_tensor * h_embd = inp->h; + + res->add_input(std::move(inp)); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + using inp_attn_type = std::conditional_t<iswa, llm_graph_input_attn_kv_iswa, llm_graph_input_attn_kv>; + inp_attn_type * inp_attn = nullptr; + + if constexpr (iswa) { + inp_attn = build_attn_inp_kv_iswa(); + } else { + inp_attn = build_attn_inp_kv(); + } + + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + cb(e_norm, "mtp_enorm", il); + + ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + cb(h_norm, "mtp_hnorm", il); + + ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, 0); + cb(concat, "mtp_concat", il); + + ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s); + cb(cur, "mtp_eh_proj", il); + + ggml_tensor * inpSA = cur; + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_norm", il); + + // self_attention + { + auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, n_head, n_head_kv, il); + + Qcur = build_norm(Qcur, layer.attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "mtp_Qcur_normed", il); + + Kcur = build_norm(Kcur, layer.attn_k_norm, nullptr, LLM_NORM_RMS, il); + cb(Kcur, "mtp_Kcur_normed", il); + + const bool is_swa = hparams.is_swa(il); + + const float freq_base_l = model.get_rope_freq_base(cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + const float ext_factor_l = is_swa ? 0.0f : ext_factor; + const float attn_factor_l = is_swa ? 1.0f : attn_factor; + const float beta_fast_l = is_swa ? 32.0f : beta_fast; + const float beta_slow_l = is_swa ? 1.0f : beta_slow; + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l + ); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l + ); + + cb(Qcur, "mtp_Qcur", il); + cb(Kcur, "mtp_Kcur", il); + cb(Vcur, "mtp_Vcur", il); + + cur = build_attn(inp_attn, + layer.wo, layer.wo_b, layer.wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il); + cb(cur, "mtp_attn_out", il); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "mtp_ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_ffn_norm", il); + + cur = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + nullptr, nullptr, + layer.ffn_up_exps_s, + layer.ffn_gate_exps_s, + layer.ffn_down_exps_s); + cb(cur, "mtp_ffn_moe_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "mtp_ffn_out", il); + + cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM_RMS, -1); + + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + cb(cur, "mtp_shared_head_norm", -1); + + cur = build_lora_mm(model.output, cur, model.output_s); + + cb(cur, "result_output", -1); + res->t_logits = cur; +} + template struct llama_model_mellum::graph<false>; template struct llama_model_mellum::graph<true>; +template struct llama_model_mellum::graph_mtp<false>; +template struct llama_model_mellum::graph_mtp<true>; diff --git a/src/models/models.h b/src/models/models.h index 94a7896172f3..ffeed1733ac7 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -466,6 +466,11 @@ struct llama_model_mellum : public llama_model_base { graph(const llama_model & model, const llm_graph_params & params); }; + template <bool iswa> + struct graph_mtp : public llm_graph_context { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override; }; From ceefd16daa2093730bf3340315fa2b84b286ddbc Mon Sep 17 00:00:00 2001 From: Jhen-Jie Hong <iainst0409@gmail.com> Date: Fri, 16 Oct 2026 16:51:14 +0800 Subject: [PATCH 057/423] hexagon: correct HMX FA sinks (#30581) * hexagon: correct HMX FA sinks * hex-fa: vectorize sinks in hmx kernels --------- Co-authored-by: Max Krasnyansky <maxk@qti.qualcomm.com> --- ggml/src/ggml-hexagon/htp/flash-attn-ops.c | 126 ++++++++++++--------- 1 file changed, 73 insertions(+), 53 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 4b273469521e..a0473c13de91 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -755,40 +755,11 @@ static void fa_q_load_thread(unsigned int n, unsigned int i, void * data) { const size_t m_start = i * m_bytes_per_t; const size_t m_end = hex_smin(m_start + m_bytes_per_t, col_vec_bytes); - if (factx->sinks) { - const float * sinks_data = factx->vtcm_sinks; - float * m_vec = (float *) factx->vtcm_m_vec; - const size_t r_start = l_start / sizeof(float); - const size_t r_end = l_end / sizeof(float); - const float scale_factor = EXP_LOG2E_F; - - const HVX_Vector v_scale = hvx_vec_splat_f32(scale_factor); - - for (size_t r = r_start; r < r_end; r += 32) { - HVX_VectorAlias local_m; - for (size_t j = 0; j < 32; ++j) { - size_t curr_r = r + j; - if (curr_r < n_rows_g) { - const size_t h_idx = fastmodulo(curr_r, G, &factx->div_G); - const size_t head = args->kv_head * G + h_idx; - local_m.fp32[j] = sinks_data[head]; - } else { - local_m.fp32[j] = HTP_FA_M_INITIAL_VAL; - } - } - HVX_Vector v_scaled = HVX_OP_MUL_F32(local_m.v, v_scale); - *(HVX_Vector *) (m_vec + r) = v_scaled; - } - if (l_start < col_vec_bytes) { - hvx_splat_u8_a((char *) factx->vtcm_l_vec + l_start, 0, l_end - l_start); - } - } else { - if (l_start < col_vec_bytes) { - hvx_splat_u8_a((char *) factx->vtcm_l_vec + l_start, 0, l_end - l_start); - } - if (m_start < col_vec_bytes) { - hvx_splat_f32_a((char *) factx->vtcm_m_vec + m_start, HTP_FA_M_INITIAL_VAL, (m_end - m_start) / sizeof(float)); - } + if (l_start < col_vec_bytes) { + hvx_splat_u8_a((char *) factx->vtcm_l_vec + l_start, 0, l_end - l_start); + } + if (m_start < col_vec_bytes) { + hvx_splat_f32_a((char *) factx->vtcm_m_vec + m_start, HTP_FA_M_INITIAL_VAL, (m_end - m_start) / sizeof(float)); } // Zero the whole rescale region: vtcm_d_tiles[0], the optional vtcm_d_tiles[1] @@ -1447,21 +1418,10 @@ static inline void fa_softmax_impl( HVX_Vector v_rowsum_acc_f32_0 = Q6_V_lo_W(rowsum_acc_pair); HVX_Vector v_rowsum_acc_f32_1 = Q6_V_hi_W(rowsum_acc_pair); - HVX_Vector v_l_curr0; - HVX_Vector v_l_curr1; - if (args->kv_start == 0 && factx->sinks != NULL) { - // First KV block with sinks: m_prev holds the seeded sink value (not -inf), - // so exp_m_diff = exp2(sink - m_curr) is the sink's contribution to the - // denominator. l_prev is 0 here, so add exp_m_diff directly instead of - // multiplying the (uninitialized) l_prev term. - v_l_curr0 = HVX_OP_ADD_F32(exp_m_diff0, v_rowsum_acc_f32_0); - v_l_curr1 = HVX_OP_ADD_F32(exp_m_diff1, v_rowsum_acc_f32_1); - } else { - HVX_Vector l_prev_v0 = factx->vtcm_l_vec[r_vec_idx * 2 + 0]; - HVX_Vector l_prev_v1 = factx->vtcm_l_vec[r_vec_idx * 2 + 1]; - v_l_curr0 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v0, exp_m_diff0), v_rowsum_acc_f32_0); - v_l_curr1 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v1, exp_m_diff1), v_rowsum_acc_f32_1); - } + HVX_Vector l_prev_v0 = factx->vtcm_l_vec[r_vec_idx * 2 + 0]; + HVX_Vector l_prev_v1 = factx->vtcm_l_vec[r_vec_idx * 2 + 1]; + HVX_Vector v_l_curr0 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v0, exp_m_diff0), v_rowsum_acc_f32_0); + HVX_Vector v_l_curr1 = HVX_OP_ADD_F32(HVX_OP_MUL_F32(l_prev_v1, exp_m_diff1), v_rowsum_acc_f32_1); factx->vtcm_m_vec[r_vec_idx * 2 + 0] = v_m_curr0; factx->vtcm_m_vec[r_vec_idx * 2 + 1] = v_m_curr1; @@ -1532,9 +1492,66 @@ static void fa_softmax_thread(unsigned int n, unsigned int i, void * data) { fa_softmax_impl(n, i, data, has_mask, mask_broadcast, is_g1, has_alibi, has_softcap); } +static inline HVX_Vector hvx_fa_load_sinks(const float * sinks_data, size_t idx, size_t count) { + const HVX_Vector * v_src = (const HVX_Vector *) sinks_data + (idx >> 5); + const size_t lane_off = idx & 31; + const size_t byte_off = lane_off * sizeof(float); + return (lane_off + count > 32) + ? Q6_V_valign_VVR(v_src[1], v_src[0], byte_off) + : Q6_V_vror_VR(v_src[0], byte_off); +} + static __attribute__((noinline)) void fa_build_d_diag_inv_l(struct hmx_fa_context * factx, size_t n_row_tiles, - size_t n_row_tiles_g_br) { + size_t n_row_tiles_g_br, + uint32_t kv_head, + size_t n_rows_g) { + (void) n_row_tiles_g_br; + + // sinks: l += exp(sink - m * ln2) in f32, m and l are in the log2 domain + if (factx->sinks) { + const float * sinks_data = factx->vtcm_sinks; + const size_t G = factx->G; + const size_t head_base = (size_t) kv_head * G; + const HVX_Vector v_ln2 = Q6_V_vsplat_R(EXP_LOGN2); + const HVX_Vector v_head_sinks = (G <= 32) ? hvx_fa_load_sinks(sinks_data, head_base, G) : Q6_V_vzero(); + HVX_Vector v_sink = Q6_V_vzero(); + size_t g_prev = SIZE_MAX; + + for (size_t i = 0; i < n_row_tiles; ++i) { + // the sink pattern of a tile only depends on (i * 32) % G, rebuild it when that changes + const size_t r0 = i * 32; + const size_t g0 = fastmodulo(r0, G, &factx->div_G); + if (g0 != g_prev) { + const size_t rem_g = G - g0; + if (G <= 32) { + const HVX_Vector v_lo = Q6_V_vror_VR(v_head_sinks, g0 * sizeof(float)); + const HVX_Vector v_hi = Q6_V_vror_VR(v_head_sinks, 128 - rem_g * sizeof(float)); + v_sink = Q6_V_vmux_QVV(Q6_Q_vsetq2_R(rem_g * sizeof(float)), v_lo, v_hi); + for (size_t len = G; len < 32; len <<= 1) { + const HVX_VectorPred q_keep = Q6_Q_vsetq_R(len * sizeof(float)); + const HVX_Vector v_shifted = Q6_V_vror_VR(v_sink, 128 - len * sizeof(float)); + v_sink = Q6_V_vmux_QVV(q_keep, v_sink, v_shifted); + } + } else if (rem_g >= 32) { + v_sink = hvx_fa_load_sinks(sinks_data, head_base + g0, 32); + } else { + const HVX_Vector v_lo = hvx_fa_load_sinks(sinks_data, head_base + g0, rem_g); + const HVX_Vector v_hi = hvx_fa_load_sinks(sinks_data, head_base, 32 - rem_g); + v_sink = Q6_V_vmux_QVV(Q6_Q_vsetq_R(rem_g * sizeof(float)), v_lo, Q6_V_vror_VR(v_hi, 128 - rem_g * sizeof(float))); + } + g_prev = g0; + } + + HVX_Vector v_d = HVX_OP_SUB_F32(v_sink, HVX_OP_MUL_F32(factx->vtcm_m_vec[i], v_ln2)); + if (r0 + 32 > n_rows_g) { + const HVX_VectorPred q_valid = Q6_Q_vsetq_R((n_rows_g - r0) * sizeof(float)); + v_d = Q6_V_vmux_QVV(q_valid, v_d, Q6_V_vzero()); + } + factx->vtcm_l_vec[i] = HVX_OP_ADD_F32(factx->vtcm_l_vec[i], hvx_vec_exp_f32(v_d)); + } + } + const HVX_Vector v_offsets = *(const HVX_Vector *) d_tile_scatter_offsets; const HVX_VectorPred q_32_mask = Q6_Q_vsetq_R(32 * sizeof(__fp16)); const HVX_Vector one = hvx_vec_splat_f32(1.0f); @@ -2010,7 +2027,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // ======== Initialize HMX output scales ======== hmx_init_column_scales(factx.vtcm_hmx_scales_id, Q6_V_vsplat_R(0x3c00)); // 1.0 - hmx_init_column_scales(factx.vtcm_hmx_scales_qk, hvx_vec_splat_f16(factx.scale)); + { + union { __fp16 f; uint16_t u; } sc = { .f = factx.scale }; + hmx_init_column_scales(factx.vtcm_hmx_scales_qk, Q6_V_vsplat_R(sc.u)); + } // ======== DMA setup ======== dma_queue * const dma_q = ctx->dma[0]; @@ -2250,7 +2270,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // Overlapped: run HVX build diag inv L while HMX is busy executing the update htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start); - fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br); + fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br, kv_head, n_rows_g); htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start); hmx_queue_pop(hmx_q); @@ -2355,7 +2375,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { if (kv_blk + 1 == factx.n_kv_blocks) { // Overlapped: run HVX build diag inv L while HMX is busy executing the update htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start); - fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br); + fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br, kv_head, n_rows_g); htp_trace_event_stop(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start); } hmx_queue_pop(ctx->hmx_queue); From f3363593f457169cdfbdfe22ea399b1cc139ba18 Mon Sep 17 00:00:00 2001 From: "Piotr Wilkin (ilintar)" <piotr.wilkin@syndatis.com> Date: Fri, 16 Oct 2026 11:52:18 +0200 Subject: [PATCH 058/423] Add more benchmarks to llama-eval (#26174) * llama-eval : add sandboxed execution for code suites Adds a lightweight replacement for the Docker containers usually used to run generated code, so a benchmark can be executed without a daemon, an image or root. A cached venv is provisioned per suite and reused; candidate programs run in a throwaway working directory as their own process group, under the strongest isolation the host offers: bwrap read-only root, tmpfs /tmp and $HOME, no network unshare no network, via --map-current-user rlimit cpu/memory/file-size/process caps only Candidate code must not run as uid 0, which is why "unshare -r" is a last-resort tier that warns: root bypasses file permission bits, so a test asserting that a write to a chmod 0444 file fails would instead see it succeed. Two details worth noting for the bwrap tier. Only paths hidden by the $HOME tmpfs are bound back, because bwrap cannot mount onto a symlink and version-manager alias directories usually are one. And the venv's base interpreter tree is resolved by walking the symlink chain, since a version-managed interpreter is commonly reached through an unversioned alias that startup resolves through. Suites can declare provisioning steps, which run once at build time while the network is still up, so test time never needs it. Assisted-by: Claude Opus 4.8 <noreply@anthropic.com> * llama-eval : add HumanEval and ClassEval code suites Adds two execution-graded suites alongside the existing text-graded ones. The model's reply is parsed for a python code block, combined with the task's own unit tests, and run in the sandbox. --dataset humaneval 164 function-level problems, scored pass@k --dataset classeval 100 class-level tasks, scored at class level (feeding pass@k) and method level, matching upstream pass@k reuses the existing chunk mechanism: --n_cases at a multiple of the dataset size gives that many samples per problem, aggregated with the unbiased estimator from the Codex paper. Program construction and result interpretation are dataset-owned hooks, because the two suites differ: one HumanEval problem is a function judged by a single assert, whereas one ClassEval task is a class judged by many independent test methods, and reducing that to pass/fail would throw away partial credit. Suites also declare their own package set, provisioning steps, interpreter and timeout. The package set is derived from the dataset's own imports rather than a requirements file. HumanEval turns out to need no third-party packages at all; ClassEval needs thirteen, two of which must be declared because no import mentions them -- lxml is selected by string in ClassEval_44, and scipy is a gensim requirement. Notes on ClassEval specifically: * gensim publishes no wheel past cp313, so 3.13 is preferred; on 3.14+ the two pure-Python APIs it uses are provided by a small shim instead. * The three nltk corpora it needs are pre-seeded at venv build time. These are the modern names -- the dataset asks for the since-renamed punkt and averaged_perceptron_tagger. * It defaults to a patched fork, ilintar/ClassEval, because three upstream tests cannot be passed by anything: a hard-coded 2024 date checked against datetime.now(), an assertion on the dataset author's own hostname, and an exact float comparison. Reference solutions go from 96/100 to 100/100 classes and 2182/2196 to 2196/2196 test methods. No skeleton is touched, so what the model is asked to write is unchanged; --dataset-source selects the original. --exec-seed (default 0) seeds the RNG before each test, so suites whose tests draw random inputs score reproducibly. ClassEval_58 fails its own reference solution about 30% of the time without it. The simulator gains both suites, so the harness can be exercised without a model. Assisted-by: Claude Opus 4.8 <noreply@anthropic.com> * llama-eval : fix silent isolation downgrade, generalise the sandbox runner The isolation probe ran the interpreter without binding it back, so whenever this tool is itself started from a venv under $HOME -- the usual case, since it needs `datasets` -- bwrap's tmpfs hid the interpreter, the probe failed, and the sandbox quietly fell back to a tier that does not hide $HOME at all. Nothing in the output looked wrong when that happened, and the existing suites have been running that way. The probe now uses the same interpreter and the same binds the suite will actually get. Two additions the agentic suite needs, neither of which changes the existing suites: * run_argv() runs an arbitrary command in a caller-owned directory that survives the call. run() is now a thin wrapper over it, so HumanEval and ClassEval take exactly the same path as before. A multi-step agent needs its edits to persist across many invocations, and a suite for a language other than Python needs to drive a linter and a test runner rather than a single source string. * RLIMIT_AS is now per-sandbox, defaulting to what it was. It caps reserved address space rather than resident memory, which is the wrong instrument for a runtime that reserves a large region and commits little of it: node's TypeScript type-stripping goes through WebAssembly, reserves gigabytes per instance, and aborts at 4 GiB with an out-of-memory error that reads like a fault in the code under test. HumanEval stays at 164/164 and ClassEval at 100/100. Assisted-by: Claude Opus 5 <noreply@anthropic.com> * llama-eval : add SACB, a multi-turn agentic coding suite `--dataset agentic` is a different shape from the other code suites. Rather than answering a single prompt, the model is given a small but real repository, a narrow tool set and a bug report, drives its own multi-turn conversation, and is graded by running a test suite it was never allowed to see. The corpus is SACB (Simple Agent Coding Benchmark), 60 tasks over three hand-written repositories in Python and TypeScript, hosted at ilintar/SACB. Grading needs no container: a cached venv gives ruff and mypy for Python, and node's own test runner plus tsc for TypeScript, so the whole dependency set is four packages. Three constraints are deliberate: * No shell, and no way to run the tests. The only feedback channel is `lint`. A syntax or type error can be driven out mechanically; a logic error has to be reasoned about. A harness that lets an agent iterate against the test suite measures something closer to search than to understanding. * Two edit tools, one line-addressed and one content-addressed. Which one a model reaches for, and whether it keeps line numbers straight after its own earlier edit, is itself a signal. * The tests never touch the working tree. They are held outside it and copied into a throwaway copy only once the agent has stopped, so they can be neither read nor edited. Every tool path is resolved and confinement-checked against the repository root, so a symlink planted inside the tree cannot reach out of it. A task counts as resolved when every test that was failing passes and nothing that already passed broke. Both sets are derived by running the tests against the defective tree and against the reference fix, never written by hand, so a task whose defect no test exercises cannot enter the corpus. Many tasks carry several independent defects, so the fraction of target tests passed is reported alongside, separating "fixed two of the three faults" from "changed nothing". An episode that arrives with no tool call and no `finish` is nudged rather than ended. Models routinely narrate their analysis mid-task, and treating that as completion ends the episode with no edits and scores a reasoning failure that never happened; about one episode in four exits that way. Nothing here is installed unless this suite is selected, and then only for the language being run: agentic_eval is imported lazily and each toolchain is keyed per language. Assisted-by: Claude Opus 5 <noreply@anthropic.com> * llama-eval : make the agentic hooks typed rather than discovered `ty` flagged nine unresolved attributes, of which the seven in main() predate the agentic suite: `EvalState.dataset` is Optional and was dereferenced without narrowing. * `loaded_dataset` turns the Optional into a value in one place, raising rather than handing back None, so callers do not each have to narrow it. * `is_agentic`, `prepare()` and `run()` are declared on BaseDataset. They were reached through getattr with a default at three call sites, which no checker can see through; `run()` now raises a clear error naming the suite instead of failing as a missing attribute. * The sandbox setup is one branch on `is_agentic` rather than two tests of `grader_type`, which also drops a duplicated condition. No behaviour change: HumanEval stays at 164/164, and `ty check` on examples/llama-eval is clean. Assisted-by: Claude Opus 5 <noreply@anthropic.com> * feat: API key * fix reasoning and response length for agentic runs, add option to save entire transcripts * llama-eval : fix ty unresolved-attribute on loaded_dataset.ae Assisted-by: AI coding agent * llama-eval : trim comments and remove non-ASCII from README Assisted-by: AI coding agent --- examples/llama-eval/README.md | 264 ++++ examples/llama-eval/agentic_eval.py | 918 ++++++++++++ examples/llama-eval/eval_sandbox.py | 488 +++++++ examples/llama-eval/llama-eval.py | 1275 ++++++++++++++++- examples/llama-eval/llama-server-simulator.py | 88 +- 5 files changed, 2993 insertions(+), 40 deletions(-) create mode 100644 examples/llama-eval/agentic_eval.py create mode 100644 examples/llama-eval/eval_sandbox.py diff --git a/examples/llama-eval/README.md b/examples/llama-eval/README.md index 3c5c35f78f6c..1f0b5a09a486 100644 --- a/examples/llama-eval/README.md +++ b/examples/llama-eval/README.md @@ -24,3 +24,267 @@ python3 llama-eval.py \ --dataset aime2025 --n_cases 240 \ --grader-type regex ``` + +### Authentication + +A server started with `--api-key` (or any OpenAI-compatible endpoint behind a +token) needs one here too. `--api-key` is sent as `Authorization: Bearer` on +every eval request; it falls back to `$LLAMA_API_KEY`, which keeps the secret +out of your shell history and off the process list. + +```bash +# One key for every server +python3 llama-eval.py --server http://localhost:8033 --api-key sk-... ... + +# From the environment +export LLAMA_API_KEY=sk-... +python3 llama-eval.py --server http://localhost:8033 ... + +# One key per server, in --server order +python3 llama-eval.py \ + --server http://server1:8033,http://server2:8033 \ + --api-key sk-one,sk-two --threads 16,16 ... +``` + +The LLM grader authenticates too. It reuses the key of the server it runs on, +so `--grader-type llm` without `--grader-server` needs nothing extra; a grader +on a separate endpoint takes `--grader-api-key`. + +Endpoints on an opencode host additionally get an `x-opencode-session` header, +generated fresh for every task. An agentic episode sends the same id on all of +its turns, so opencode sees one session per task instead of one per request; +other servers never see the header. + +## Agentic suite + +`--dataset agentic` is a different shape from the other code suites. It runs +**SACB** (Simple Agent Coding Benchmark), hosted at +[`ilintar/SACB`](https://huggingface.co/datasets/ilintar/SACB) -- 60 tasks over +three hand-written repositories, in Python and TypeScript. Rather than +answering one prompt, the model is given a small but real repository, a +deliberately narrow tool set, and a bug report, and drives its own multi-turn +conversation until it declares itself finished. It is then graded by running a +test suite it was never allowed to see. + +```bash +python3 llama-eval.py \ + --server http://localhost:8033 --model my-model \ + --dataset agentic --threads 1 --temperature 0 +``` + +### The tools + +| tool | purpose | +|---|---| +| `list_files` | enumerate the tree | +| `read_file` | read, with 1-based line numbers and optional line-range narrowing | +| `search` | regular-expression search across the repository | +| `edit_lines` | replace an inclusive line range | +| `edit_replace` | replace an exact snippet, which must be unique | +| `write_file` | write a whole file | +| `lint` | run the language's linter and type checker | +| `finish` | declare the change complete | + +Three constraints are deliberate: + +**There is no shell and no way to run the tests.** The only feedback channel is +`lint`. A syntax or type error can be driven out mechanically; a logic error has +to be reasoned about. A harness that lets an agent iterate against the test +suite measures something closer to search than to understanding. + +**There are two edit tools**, one line-addressed and one content-addressed. +Which one a model reaches for, and whether it keeps line numbers straight after +its own earlier edit, is itself a signal. + +**The tests never touch the working tree.** They are held outside it and copied +into a throwaway copy only after the agent has stopped, so they can be neither +read nor edited. Every tool path is resolved and confinement-checked against the +repository root, so a symlink planted inside the tree cannot reach outside it. + +### Scoring + +A task is **resolved** when every test that was failing now passes *and* nothing +that was already passing broke. Both sets are derived by running the tests twice +when the corpus is built -- once against the defective tree and once against the +reference fix -- never hand-written, so a task whose defect no test exercises is +rejected rather than shipped. + +Because many tasks contain several independent defects, the fraction of target +tests passed is reported alongside, which separates "fixed two of the three +faults" from "changed nothing". Process metrics are reported too: turns, tool +calls, rejected calls, edits, and the peak conversation size. + +### Dependencies + +Nothing here is installed unless you select this suite, and then only for the +language you actually run. `--agentic-lang python` provisions ruff and mypy; +`--agentic-lang typescript` provisions the TypeScript compiler and uses node's +own test runner. Neither touches the other, and no other suite touches either. + +### Useful flags + +| flag | meaning | +|---|---| +| `--agentic-lang` | restrict to one language, and provision only its toolchain | +| `--agentic-max-turns` | give up after this many assistant turns (default 50) | +| `--agentic-max-tokens` | cap generation per turn (default: no cap); bounds a model that falls into a repetition loop, but also cuts reasoning off mid-thought. Truncated turns are counted as `truncated_turns` | +| `--agentic-context-limit` | abandon an episode once its context exceeds this many tokens | +| `--agentic-transcripts DIR` | save each episode's full message history to `DIR/<task_id>.json`, plus a `DIR/<task_id>.html` showing every reasoning block, reply, tool call and tool result; the main report links to it | + + +## Code suites (HumanEval, ClassEval) + +`--dataset humaneval` runs the 164 OpenAI HumanEval problems; `--dataset +classeval` runs the 100 ClassEval class-generation tasks. Unlike the other +suites these are not graded by matching text: the model's reply is parsed for a +```` ```python ```` block, that code is combined with the task's own unit tests, +and the result is **executed**. + +ClassEval asks for a whole class rather than one function, so it is scored the +way upstream does - **class level** (every one of the task's tests passes, which +is what feeds pass@k) and **method level** (what fraction of all 2196 test +methods passed). The second number separates "generated nothing usable" from +"got most of the class right", and both are printed at the end of a run. + +Give it more room to generate than HumanEval - a class with ten methods does not +fit in a short budget: + +```bash +python3 llama-eval.py \ + --server http://localhost:8033 --model my-model \ + --dataset classeval --threads 8 --n_predict 4096 --temperature 0 +``` + +```bash +# pass@1, greedy +python3 llama-eval.py \ + --server http://localhost:8033 --model my-model \ + --dataset humaneval --threads 8 --n_predict 1024 --temperature 0 + +# pass@10 -- each extra multiple of 164 is another sample per problem +python3 llama-eval.py \ + --server http://localhost:8033 --model my-model \ + --dataset humaneval --n_cases 1640 --temperature 0.8 --top-p 0.95 +``` + +`--grader-type exec` is the default for code suites, so neither it nor +`--dataset-source` normally needs to be passed. + +### Execution sandbox - no Docker required + +Generated code is untrusted, but these suites need far less than a container to +run safely. **The package set is derived from the dataset's own imports at +startup**, not hardcoded and not taken from the upstream `requirements.txt` - +which for ClassEval lists `torch`, `openai` and `transformers` that its *tests* +never touch. So each suite installs exactly what it needs: + +| suite | third-party packages | +|---|---| +| HumanEval | **none** - every problem imports only stdlib (`typing, math, random, copy, string, collections, hashlib, re`) | +| ClassEval | 13: `beautifulsoup4, gensim, lxml, netifaces, nltk, numpy, openpyxl, pandas, pillow, pypdf2, python-docx, reportlab, scipy` | + +Two of ClassEval's are *declared* rather than derived, because no import +statement mentions them: `lxml` is selected by string +(`BeautifulSoup(html, 'lxml')`, ClassEval_44 - without it that task silently +drops from 23/23 to 7/23), and `scipy` is a hard gensim requirement. Note also +that `PyPDF2` must **not** be modernised to `pypdf`: ClassEval_69's test imports +`PdfFileReader`, a name pypdf removed. + +The venv is cached under `~/.cache/llama-eval/` (override with +`LLAMA_EVAL_CACHE`) and rebuilt only when the package set, interpreter version +or provisioning steps change; force one with `--rebuild-sandbox`. `uv` is used +when available, which makes even the ClassEval venv build in seconds. + +**Interpreter selection.** A suite can prefer a different Python than the one +running the tool, and declares it: ClassEval asks for 3.13, because gensim +publishes no wheel past cp313 and its generated C no longer compiles on 3.14 +(it still uses `PyLongObject.ob_digit`). That interpreter is found automatically, +including via `uv python`; `--sandbox-python` overrides it. + +3.14 is fully supported anyway. ClassEval uses exactly two gensim APIs +(`utils.decode_htmlentities`, `matutils.unitvec`), both pure Python upstream, so +on 3.14+ the suite drops gensim from the install set and provisions a small shim +instead. Both paths score **100/100**. + +**Offline provisioning.** A suite that would otherwise fetch data at test time +declares a provisioning step that runs once at venv build time, while the +network is still up. ClassEval uses this to pre-seed exactly three nltk corpora +(`punkt_tab`, `averaged_perceptron_tagger_eng`, `wordnet`, ~33 MB), so the tests +themselves run with no network at all. Those are the *current* names - the +dataset asks for `punkt` and `averaged_perceptron_tagger`, which nltk has since +renamed, and seeding the legacy names too is not harmless: with a network +available nltk considers them satisfied and re-fetches the modern ones anyway. +The corpora are written *inside* the venv, not `~/nltk_data`, because `$HOME` is +a tmpfs under the strongest isolation tier. + +Each candidate program runs in a throwaway directory as its own process group, +under the strongest isolation the host offers: + +| tier | requires | isolates | +|---|---|---| +| `bwrap` | bubblewrap installed | read-only root, tmpfs `/tmp` and `$HOME`, **no network** | +| `unshare` | `unshare --map-current-user` | **no network** | +| `unshare-root` | any `unshare -r` | no network, but **runs as uid 0** - see below | +| `rlimit` | always available | cpu/memory/file-size/process caps only | + +**Candidate code must not run as root**, which is why `unshare -r` is a +last-resort tier that prints a warning. Root bypasses file permission bits, so a +test asserting that writing to a `chmod 0444` file *fails* instead sees it +succeed - ClassEval_50 scores 14/16 under uid 0 and 16/16 under uid 1000, with +nothing in the output to suggest the sandbox caused it. Install bubblewrap. + +All tiers additionally cap CPU time, address space (4 GiB), file size and core +dumps, and kill the whole process group on timeout (`--exec-timeout`, default +15 s - raise it for ClassEval, whose tests do real work). Writes to the +throwaway working directory are allowed, which is all the file-handling tasks in +ClassEval need. + +Note that the weakest tier does not contain filesystem access. If you are +scoring output from a model you do not trust, check the tier printed at startup. + +### Reproducibility + +Several tasks draw their test inputs from the *global, unseeded* RNG, so a +verdict is not reproducible run to run - which adds noise exactly when you are +comparing two models or two quantisations. `--exec-seed` (default `0`) seeds the +RNG before each test; pass `--exec-seed -1` for the stock upstream behaviour. +This does not change what the model is asked to compute. + +It matters more than it sounds. HumanEval/38, /50 and /53 are affected, and +ClassEval_58 (duplicate mine placement) fails about 30% of the time on its own +*reference* solution - measured here at 21/30 passes unseeded versus **30/30 +seeded**. + +### Dataset fidelity + +The canonical dataset runs **unmodified on current Python** - all 164 reference +solutions pass on CPython 3.14 with no patches, so no forked or re-hosted copy +is needed. `--dataset-source` accepts a local `.jsonl` or a HuggingFace repo id +if you do want to point at a modified copy. + +HumanEval does contain a handful of long-standing defects where the *prompt* +contradicts the graded test (notably HumanEval/47, /116 and /148). Those +penalise a model that reads the docstring carefully, but correcting them makes +the benchmark marginally easier and the resulting scores no longer comparable +with published HumanEval numbers, so nothing is patched by default. + +ClassEval is different: the upstream data has defects that cap the achievable +score below 100% for reasons unrelated to the model, so `--dataset classeval` +defaults to a **patched fork**, +[`ilintar/ClassEval`](https://huggingface.co/datasets/ilintar/ClassEval). +Six edits over five tasks lift the reference solutions from 96/100 classes and +2182/2196 test methods to **100/100 and 2196/2196**. Three repair tests nothing +could pass: + +- **ClassEval_17** was a time bomb - a hard-coded 2024 date checked against + `datetime.now()`, unpassable since January 2024. +- **ClassEval_48** asserted `get_hostname('0.0.0.0') == 'LAPTOP-2CS86KUM'`, the + dataset author's own machine name. +- **ClassEval_31** compared floats with `assertEqual`, failing on a 4-ULP numpy + difference. + +The other three fix only the *reference* solutions (NumPy 2.0 removals in +ClassEval_51, a 30%-flaky mine generator in ClassEval_58) and cannot affect a +model's score. No `skeleton` was touched, so what the model is asked to write is +byte-identical to upstream. Pass `--dataset-source FudanSELab/ClassEval` for the +unpatched original. diff --git a/examples/llama-eval/agentic_eval.py b/examples/llama-eval/agentic_eval.py new file mode 100644 index 000000000000..96629088a703 --- /dev/null +++ b/examples/llama-eval/agentic_eval.py @@ -0,0 +1,918 @@ +#!/usr/bin/env python3 +"""Multi-turn agentic coding eval: tools over a real working tree. + +The existing code suites are single-shot -- one prompt, one reply, execute, +grade. This one is not. The model is given a small but real repository and a +deliberately narrow tool set, and has to find its way to the defect before it +can fix it. That shape is what puts a run in the tens-of-thousands-of-tokens +range without any artificial padding: the context grows because the model +chose to read things. + +Three deliberate constraints: + +* No shell, and no way to run the tests. The only feedback channel is `lint`. + A syntax or type error can be driven out mechanically; a logic error has to + be reasoned about. Suites that let an agent iterate against the test suite + measure something closer to search than to understanding. +* Two edit tools, one line-addressed and one content-addressed. Which one a + model reaches for, and whether it keeps line numbers straight after its own + earlier edit, is itself a signal. +* Tests never touch the working tree. They are held outside it and copied in + after the agent has stopped, so they can be neither read nor edited. + +Everything here is imported lazily by llama-eval.py, and every dependency it +needs is scoped to the language actually under test, so none of it is +installed for anyone running the other suites. +""" + +import html +import json +import os +import re +import shutil +import tempfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence, Tuple + +MAX_READ_LINES = 400 +MAX_TOOL_CHARS = 12000 +MAX_SEARCH_HITS = 60 + + +@dataclass +class LangSpec: + """Everything language-specific: how to lint, how to test, what to install. + + Kept as data so adding a language is a corpus concern rather than a harness + change, and so a run only ever provisions the toolchain for the language it + is actually evaluating. + """ + name: str + exts: Tuple[str, ...] + packages: Tuple[str, ...] = () + provision: Tuple[str, ...] = () + env: Dict[str, str] = field(default_factory=dict) + address_space_limit: Optional[int] = 4 * 1024 * 1024 * 1024 + + def source_files(self, root: Path) -> List[Path]: + return sorted(p for p in root.rglob("*") + if p.is_file() and p.suffix in self.exts) + + +PYTHON = LangSpec( + name="python", + exts=(".py",), + packages=("ruff==0.14.2", "mypy==1.18.2"), +) + +TYPESCRIPT = LangSpec( + name="typescript", + exts=(".ts",), + address_space_limit=16 * 1024 * 1024 * 1024, + provision=( + "import os, subprocess, shutil, sys\n" + "npm = shutil.which('npm')\n" + "if not npm: sys.exit('npm not found on PATH; needed for the typescript suite')\n" + "dest = os.path.join(os.environ['SANDBOX_VENV'], 'node')\n" + "os.makedirs(dest, exist_ok=True)\n" + "r = subprocess.run([npm, 'install', '--silent', '--no-audit', '--no-fund',\n" + " '--prefix', dest, 'typescript@5.9.3'],\n" + " capture_output=True, text=True)\n" + "sys.exit(0 if r.returncode == 0 else (r.stderr or r.stdout)[-2000:])\n", + ), +) + +LANGS: Dict[str, LangSpec] = {"python": PYTHON, "typescript": TYPESCRIPT} + + +class PathEscape(Exception): + """A tool was handed a path outside the repository.""" + + +class Workspace: + """A private, writable copy of one task's repository. + + Every tool path is resolved through here. Confinement is checked on the + fully-resolved path, so a symlink planted inside the tree cannot be used to + read or write outside it. + """ + + def __init__(self, root: Path): + self.root = root.resolve() + self.escape_attempts = 0 + + def resolve(self, rel: str) -> Path: + if rel is None or not str(rel).strip(): + raise PathEscape("no path given") + raw = Path(str(rel)) + base = raw if raw.is_absolute() else (self.root / raw) + candidate = base.resolve() + if candidate != self.root and self.root not in candidate.parents: + self.escape_attempts += 1 + raise PathEscape(f"path {rel!r} is outside the repository") + return candidate + + def rel(self, path: Path) -> str: + return str(path.relative_to(self.root)) + + def files(self) -> List[Path]: + return sorted(p for p in self.root.rglob("*") + if p.is_file() and "__pycache__" not in p.parts) + + def snapshot(self) -> Dict[str, str]: + out = {} + for p in self.files(): + try: + out[self.rel(p)] = p.read_text() + except (UnicodeDecodeError, OSError): + pass + return out + + +def _clip(text: str, limit: int = MAX_TOOL_CHARS) -> str: + if len(text) <= limit: + return text + return text[:limit] + f"\n... [truncated, {len(text) - limit} more characters]" + + +def _numbered(lines: Sequence[str], start: int) -> str: + width = len(str(start + len(lines) - 1)) + return "\n".join(f"{start + i:>{width}}\t{ln}" for i, ln in enumerate(lines)) + + +def tool_schemas() -> List[Dict[str, Any]]: + """OpenAI-format function schemas for the whole tool set.""" + def fn(name, desc, props, required=()): + return {"type": "function", "function": { + "name": name, "description": desc, + "parameters": {"type": "object", "properties": props, + "required": list(required)}}} + + return [ + fn("list_files", "List files in the repository.", + {"path": {"type": "string", + "description": "Directory relative to the repository root. " + "Omit for the whole tree."}}), + fn("read_file", + "Read a file. Output is prefixed with 1-based line numbers. Give " + "start_line/end_line to read only part of a large file.", + {"path": {"type": "string", "description": "File path relative to the repository root."}, + "start_line": {"type": "integer", "description": "First line to read, 1-based, inclusive."}, + "end_line": {"type": "integer", "description": "Last line to read, 1-based, inclusive."}}, + ["path"]), + fn("search", + "Search the repository with a regular expression. Returns matching " + "lines with their file and line number.", + {"pattern": {"type": "string", "description": "Python regular expression."}, + "path": {"type": "string", "description": "Limit the search to this subdirectory."}, + "glob": {"type": "string", "description": "Limit to files matching this glob, e.g. '*.py'."}}, + ["pattern"]), + fn("edit_lines", + "Replace an inclusive 1-based line range with new text. Line numbers " + "refer to the file as it is now, so re-read it after any edit that " + "changed the number of lines.", + {"path": {"type": "string", "description": "File path relative to the repository root."}, + "start_line": {"type": "integer", "description": "First line to replace, 1-based, inclusive."}, + "end_line": {"type": "integer", "description": "Last line to replace, 1-based, inclusive."}, + "new_text": {"type": "string", "description": "Replacement text. Use an empty string to delete the range."}}, + ["path", "start_line", "end_line", "new_text"]), + fn("edit_replace", + "Replace an exact snippet of text. The snippet must occur exactly " + "once in the file unless replace_all is set.", + {"path": {"type": "string", "description": "File path relative to the repository root."}, + "old_text": {"type": "string", "description": "Exact text to find, including indentation."}, + "new_text": {"type": "string", "description": "Replacement text."}, + "replace_all": {"type": "boolean", "description": "Replace every occurrence instead of requiring exactly one."}}, + ["path", "old_text", "new_text"]), + fn("write_file", + "Write a whole file, creating or overwriting it. Prefer the edit " + "tools for changes to existing files.", + {"path": {"type": "string", "description": "File path relative to the repository root."}, + "content": {"type": "string", "description": "Full new file contents."}}, + ["path", "content"]), + fn("lint", + "Run the language's linter and type checker over the repository and " + "return the diagnostics. This is the only feedback available; the " + "test suite cannot be run.", + {"path": {"type": "string", "description": "Limit to this file or directory."}}), + fn("finish", + "Declare the task complete and end the session. This is the only " + "way to end it; call this once the change is made.", + {"summary": {"type": "string", "description": "Brief description of the change made."}}), + ] + + +class ToolBox: + """Executes tool calls against one workspace. Returns plain text results. + + Failures are reported as ordinary text rather than raised, because how a + model recovers from a rejected edit is part of what is being measured. + """ + + def __init__(self, ws: Workspace, linter=None): + self.ws = ws + self.linter = linter + self.calls = 0 + self.errors = 0 + self.edits = 0 + self.finished = False + self.finish_summary: Optional[str] = None + self.per_tool: Dict[str, int] = {} + + def dispatch(self, name: str, args: Dict[str, Any]) -> str: + self.calls += 1 + self.per_tool[name] = self.per_tool.get(name, 0) + 1 + handler = getattr(self, f"_t_{name}", None) + if handler is None: + result = f"Error: unknown tool {name!r}." + else: + try: + result = handler(args) + except PathEscape as e: + result = f"Error: {e}" + except KeyError as e: + result = f"Error: missing required argument {e}." + except Exception as e: + result = f"Error: {type(e).__name__}: {e}" + if result.startswith("Error:"): + self.errors += 1 + return _clip(result) + + def _t_list_files(self, a): + base = self.ws.resolve(a.get("path") or ".") + if not base.exists(): + return f"Error: {a.get('path')!r} does not exist." + paths = [p for p in self.ws.files() if p == base or base in p.parents] + if not paths: + return "(no files)" + return "\n".join(f"{self.ws.rel(p)}\t{p.stat().st_size} bytes" for p in paths) + + def _t_read_file(self, a): + path = self.ws.resolve(a["path"]) + if not path.is_file(): + return f"Error: {a['path']!r} is not a file." + lines = path.read_text().splitlines() + start = max(1, int(a.get("start_line") or 1)) + end = int(a.get("end_line") or len(lines)) + end = min(end, len(lines)) + if start > len(lines): + return f"Error: {a['path']!r} has only {len(lines)} lines." + window = lines[start - 1:end] + note = "" + if len(window) > MAX_READ_LINES: + window = window[:MAX_READ_LINES] + note = (f"\n... [showing {MAX_READ_LINES} of {end - start + 1} requested " + f"lines; re-read with start_line/end_line for the rest]") + header = f"{a['path']} ({len(lines)} lines total)\n" + return header + _numbered(window, start) + note + + def _t_search(self, a): + try: + rx = re.compile(a["pattern"]) + except re.error as e: + return f"Error: bad regular expression: {e}" + base = self.ws.resolve(a.get("path") or ".") + glob = a.get("glob") + hits, scanned = [], 0 + for p in self.ws.files(): + if not (p == base or base in p.parents): + continue + if glob and not p.match(glob): + continue + try: + text = p.read_text() + except (UnicodeDecodeError, OSError): + continue + scanned += 1 + for i, line in enumerate(text.splitlines(), 1): + if rx.search(line): + hits.append(f"{self.ws.rel(p)}:{i}:{line.strip()[:200]}") + if len(hits) >= MAX_SEARCH_HITS: + return ("\n".join(hits) + + f"\n... [stopped at {MAX_SEARCH_HITS} matches; narrow the pattern]") + if not hits: + return f"No matches for {a['pattern']!r} in {scanned} file(s)." + return "\n".join(hits) + + def _t_edit_lines(self, a): + path = self.ws.resolve(a["path"]) + if not path.is_file(): + return f"Error: {a['path']!r} is not a file." + lines = path.read_text().splitlines() + start, end = int(a["start_line"]), int(a["end_line"]) + if start < 1 or start > len(lines): + return (f"Error: start_line {start} out of range; " + f"{a['path']!r} has {len(lines)} lines.") + if end < start or end > len(lines): + return (f"Error: end_line {end} out of range; must be between " + f"{start} and {len(lines)}.") + new = str(a["new_text"]).splitlines() + path.write_text("\n".join(lines[:start - 1] + new + lines[end:]) + "\n") + self.edits += 1 + delta = len(new) - (end - start + 1) + return (f"Replaced lines {start}-{end} of {a['path']} " + f"({end - start + 1} -> {len(new)} lines, {delta:+d}). " + f"Later line numbers have shifted." if delta else + f"Replaced lines {start}-{end} of {a['path']}.") + + def _t_edit_replace(self, a): + path = self.ws.resolve(a["path"]) + if not path.is_file(): + return f"Error: {a['path']!r} is not a file." + text = path.read_text() + old, new = str(a["old_text"]), str(a["new_text"]) + if not old: + return "Error: old_text must not be empty." + n = text.count(old) + if n == 0: + return (f"Error: old_text not found in {a['path']}. It must match " + f"exactly, including indentation.") + if n > 1 and not a.get("replace_all"): + return (f"Error: old_text occurs {n} times in {a['path']}. Include " + f"more surrounding context to make it unique, or set " + f"replace_all.") + path.write_text(text.replace(old, new)) + self.edits += 1 + return f"Replaced {n} occurrence(s) in {a['path']}." + + def _t_write_file(self, a): + path = self.ws.resolve(a["path"]) + path.parent.mkdir(parents=True, exist_ok=True) + existed = path.is_file() + path.write_text(str(a["content"])) + self.edits += 1 + return f"{'Overwrote' if existed else 'Created'} {a['path']}." + + def _t_lint(self, a): + if self.linter is None: + return "Error: no linter is configured for this task." + target = a.get("path") + if target: + self.ws.resolve(target) + return self.linter(target) + + def _t_finish(self, a): + self.finished = True + self.finish_summary = a.get("summary") + return "Task marked complete." + + +PY_TEST_RUNNER = r''' +import json, os, random, sys, unittest +sys.path.insert(0, os.getcwd()) +random.seed(0) + +class Collect(unittest.TestResult): + def __init__(self): + super().__init__() + self.outcomes = {} + def addSuccess(self, t): self.outcomes[t.id()] = "pass" + def addFailure(self, t, e): self.outcomes[t.id()] = "fail" + def addError(self, t, e): self.outcomes[t.id()] = "error" + def addSkip(self, t, r): self.outcomes[t.id()] = "skip" + def addExpectedFailure(self, t, e): self.outcomes[t.id()] = "pass" + def addUnexpectedSuccess(self, t): self.outcomes[t.id()] = "fail" + +def flatten(s): + for t in s: + if isinstance(t, unittest.TestSuite): + yield from flatten(t) + else: + yield t + +out = {"tests": {}} +try: + suite = unittest.TestLoader().discover(os.environ.get("TEST_DIR", "tests"), + top_level_dir=".") + res = Collect() + for t in flatten(suite): + tid = t.id() + try: + t.run(res) + except Exception: + pass + res.outcomes.setdefault(tid, "error") + out["tests"] = res.outcomes +except Exception as e: + out["error"] = f"{type(e).__name__}: {e}" + +print("###RESULT###") +print(json.dumps(out)) +''' + + +def _node_roots() -> List[Path]: + """Directories that must stay visible for node to start under bwrap.""" + node = shutil.which("node") + if not node: + return [] + real = Path(node).resolve() + return [Path(node).parent.parent, real.parent.parent] + + +class Runner: + """Runs the linter and the hidden tests for one language, in the sandbox.""" + + def __init__(self, sandbox, lang: LangSpec, timeout: float = 90.0): + self.sandbox = sandbox + self.lang = lang + self.timeout = timeout + + def _binds(self) -> List[Path]: + return _node_roots() if self.lang.name == "typescript" else [] + + def _env(self) -> Dict[str, str]: + if self.lang.name != "typescript": + return {} + node = shutil.which("node") + return {"PATH": f"{Path(node).parent}:/usr/bin:/bin"} if node else {} + + def lint(self, workdir: Path, target: Optional[str] = None) -> str: + if self.lang.name == "python": + return self._lint_python(workdir, target) + return self._lint_typescript(workdir, target) + + def _lint_python(self, workdir: Path, target: Optional[str]) -> str: + venv = self.sandbox.venv_dir + tgt = target or "." + chunks = [] + r = self.sandbox.run_argv( + [str(venv / "bin" / "ruff"), "check", "--no-cache", + "--output-format", "concise", tgt], + workdir, timeout=self.timeout) + text = (r.stdout or r.stderr).strip() + if text and "All checks passed" not in text: + chunks.append(text) + r = self.sandbox.run_argv( + [str(venv / "bin" / "mypy"), "--no-color-output", "--no-error-summary", + "--cache-dir", "/dev/null", "--ignore-missing-imports", tgt], + workdir, timeout=self.timeout) + text = (r.stdout or r.stderr).strip() + if text and not text.startswith("Success:"): + chunks.append(text) + return "\n".join(chunks) if chunks else "No problems found." + + def _lint_typescript(self, workdir: Path, target: Optional[str]) -> str: + tsc = self.sandbox.venv_dir / "node" / "node_modules" / "typescript" / "bin" / "tsc" + files = [str(p.relative_to(workdir)) for p in sorted(workdir.rglob("*.ts")) + if "node_modules" not in p.parts] + if target: + t = (workdir / target).resolve() + files = [f for f in files if (workdir / f).resolve() == t + or t in (workdir / f).resolve().parents] + if not files: + return "No TypeScript files to check." + r = self.sandbox.run_argv( + ["node", str(tsc), "--noEmit", "--strict", "--target", "es2022", + "--module", "esnext", "--moduleResolution", "bundler", + "--allowImportingTsExtensions", *files], + workdir, timeout=self.timeout, env=self._env(), ro_binds=self._binds()) + out = (r.stdout or "") + (r.stderr or "") + return out.strip() or "No problems found." + + def test(self, workdir: Path) -> Dict[str, Any]: + if self.lang.name == "python": + return self._test_python(workdir) + return self._test_typescript(workdir) + + def _test_python(self, workdir: Path) -> Dict[str, Any]: + runner = workdir / "__run_tests.py" + runner.write_text(PY_TEST_RUNNER) + r = self.sandbox.run_argv( + [str(self.sandbox.python), "-I", "-B", str(runner)], + workdir, timeout=self.timeout, env={"TEST_DIR": "tests"}) + return self._parse_json_result(r) + + def _parse_json_result(self, r) -> Dict[str, Any]: + if r.status == "timeout": + return {"tests": {}, "error": "timeout"} + marker = "###RESULT###" + if marker not in (r.stdout or ""): + return {"tests": {}, + "error": f"no result: {(r.stderr or r.stdout or '')[-400:]}"} + try: + return json.loads(r.stdout.split(marker, 1)[1].strip()) + except Exception as e: + return {"tests": {}, "error": f"unparseable result: {e}"} + + def _test_typescript(self, workdir: Path) -> Dict[str, Any]: + tests = sorted(p for p in (workdir / "tests").rglob("*.ts")) \ + if (workdir / "tests").is_dir() else [] + if not tests: + return {"tests": {}, "error": "no test files"} + r = self.sandbox.run_argv( + ["node", "--test", "--experimental-strip-types", + "--test-reporter=tap", + *[str(p.relative_to(workdir)) for p in tests]], + workdir, timeout=self.timeout, env=self._env(), ro_binds=self._binds()) + return self._parse_tap(r) + + @staticmethod + def _parse_tap(r) -> Dict[str, Any]: + if r.status == "timeout": + return {"tests": {}, "error": "timeout"} + text = (r.stdout or "") + "\n" + (r.stderr or "") + outcomes: Dict[str, str] = {} + for line in text.splitlines(): + m = re.match(r"\s*(not )?ok\s+\d+\s*-\s*(.+?)\s*$", line) + if not m: + continue + name = m.group(2).strip() + if name.endswith(" # SKIP") or " # SKIP" in name: + outcomes[name.split(" # ")[0]] = "skip" + continue + verdict = "fail" if m.group(1) else "pass" + if outcomes.get(name) != "fail": + outcomes[name] = verdict + if not outcomes: + return {"tests": {}, + "error": f"no TAP output: {text.strip()[-400:]}"} + return {"tests": outcomes} + + +SYSTEM_PROMPT = """You are an expert {lang} engineer working in an existing repository. + +Resolve the user's request by editing files with the tools provided. + +How this session works: +- You are working autonomously. There is no human watching; nobody will answer + questions, approve a plan, or read progress reports. +- Every reply you send must contain at least one tool call until the task is + done. A reply that contains only text does nothing: no work happens and the + turn is wasted. Do not stop to describe what you are about to do; do it. +- The session ends only when you call `finish`. Call it once the change is + made, and not before. Only the files as you leave them are evaluated. + +Important constraints: +- You cannot run the code, the tests, or any shell command. The only feedback + available is the `lint` tool, which reports syntax and type problems. +- The repository already has a test suite that will be run against your work + after you finish, but you cannot see it or run it. Do not create test files. +- Read enough of the code to understand it before editing. A change that only + makes the symptom go away will usually fail the tests. +- Make the smallest correct change. Do not reformat or refactor unrelated code. + +Keep calling tools until the change is complete, then call `finish`.""" + + +@dataclass +class AgenticTask: + task_id: str + lang: str + category: str + difficulty: int + instruction: str + files: Dict[str, str] + tests: Dict[str, str] + gold: Dict[str, str] = field(default_factory=dict) + fail_to_pass: List[str] = field(default_factory=list) + pass_to_pass: List[str] = field(default_factory=list) + + @classmethod + def from_record(cls, rec: Dict[str, Any]) -> "AgenticTask": + def as_map(v): + return v if isinstance(v, dict) else json.loads(v or "{}") + def as_list(v): + return v if isinstance(v, list) else json.loads(v or "[]") + return cls( + task_id=rec["task_id"], lang=rec["lang"], category=rec["category"], + difficulty=int(rec.get("difficulty", 3)), + instruction=rec["instruction"], + files=as_map(rec["files"]), tests=as_map(rec["tests"]), + gold=as_map(rec.get("gold")), + fail_to_pass=as_list(rec.get("fail_to_pass")), + pass_to_pass=as_list(rec.get("pass_to_pass")), + ) + + @property + def spec(self) -> LangSpec: + return LANGS[self.lang] + + +def materialise(files: Dict[str, str], dest: Path) -> Path: + """Write a file map out as a real directory tree.""" + dest.mkdir(parents=True, exist_ok=True) + for rel, content in files.items(): + p = dest / rel + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(content) + return dest + + +def repo_overview(files: Dict[str, str]) -> str: + """A compact tree, so a run is not lost merely to not finding the files.""" + lines = [] + for rel in sorted(files): + n = files[rel].count("\n") + 1 + lines.append(f" {rel} ({n} lines)") + return "\n".join(lines) + + +def user_prompt(task: AgenticTask) -> str: + return (f"Repository files:\n{repo_overview(task.files)}\n\n" + f"Task:\n{task.instruction}") + + +@dataclass +class EpisodeResult: + stop_reason: str = "unknown" + turns: int = 0 + tool_calls: int = 0 + tool_errors: int = 0 + edits: int = 0 + escape_attempts: int = 0 + finished: bool = False + bad_tool_json: int = 0 + nudges: int = 0 + prompt_tokens: int = 0 + completion_tokens: int = 0 + peak_context: int = 0 + truncated_turns: int = 0 + per_tool: Dict[str, int] = field(default_factory=dict) + error: Optional[str] = None + + +NUDGE = ("Your last reply contained no tool call, so nothing happened. There is " + "no human in this session to respond to text. If your change is " + "complete, call `finish`. Otherwise continue working by calling the " + "tools; do not reply with text alone.") + + +def run_episode(task: AgenticTask, toolbox: ToolBox, chat, + max_turns: int = 50, + context_limit: Optional[int] = None, + max_nudges: int = 3) -> Tuple[EpisodeResult, List[Dict]]: + """Drive one multi-turn tool-calling episode. + + `chat(messages, tools) -> response dict` is supplied by the caller so this + module stays free of any HTTP concerns. + """ + tools = tool_schemas() + messages: List[Dict[str, Any]] = [ + {"role": "system", "content": SYSTEM_PROMPT.format(lang=task.spec.name)}, + {"role": "user", "content": user_prompt(task)}, + ] + res = EpisodeResult() + + for turn in range(max_turns): + res.turns = turn + 1 + try: + reply = chat(messages, tools) + except Exception as e: + res.stop_reason, res.error = "request_failed", f"{type(e).__name__}: {e}" + break + + usage = reply.get("usage") or {} + pt = int(usage.get("prompt_tokens") or 0) + ct = int(usage.get("completion_tokens") or 0) + res.prompt_tokens += pt + res.completion_tokens += ct + res.peak_context = max(res.peak_context, pt + ct) + + choices = reply.get("choices") or [] + if not choices: + res.stop_reason, res.error = "no_choices", json.dumps(reply)[:300] + break + msg = choices[0].get("message") or {} + if choices[0].get("finish_reason") == "length": + res.truncated_turns += 1 + calls = msg.get("tool_calls") or [] + + reasoning = msg.get("reasoning_content") + messages.append({ + "role": "assistant", + "content": msg.get("content") or "", + **({"reasoning_content": reasoning} if reasoning else {}), + **({"tool_calls": calls} if calls else {}), + }) + + if not calls: + if toolbox.finished: + res.stop_reason = "finished" + break + if res.nudges >= max_nudges: + res.stop_reason = "no_tool_calls" + break + res.nudges += 1 + messages.append({"role": "user", "content": NUDGE}) + continue + + for call in calls: + fn = call.get("function") or {} + name = fn.get("name") or "" + raw = fn.get("arguments") + if isinstance(raw, dict): + args = raw + else: + try: + args = json.loads(raw or "{}") + except (json.JSONDecodeError, TypeError): + res.bad_tool_json += 1 + args = None + output = (toolbox.dispatch(name, args) if args is not None else + f"Error: arguments for {name!r} were not valid JSON.") + messages.append({"role": "tool", + "tool_call_id": call.get("id") or f"call_{res.tool_calls}", + "content": output}) + + if toolbox.finished: + res.stop_reason = "finished" + break + if context_limit and res.peak_context > context_limit: + res.stop_reason = "context_limit" + break + else: + res.stop_reason = "max_turns" + + res.tool_calls = toolbox.calls + res.tool_errors = toolbox.errors + res.edits = toolbox.edits + res.escape_attempts = toolbox.ws.escape_attempts + res.finished = toolbox.finished + res.per_tool = dict(toolbox.per_tool) + return res, messages + + +def run_tests_against(files: Dict[str, str], task: AgenticTask, runner: Runner, + workroot: Path, label: str) -> Dict[str, str]: + """Materialise a tree, drop the hidden tests in, and run them. + + The tests are only ever written into this throwaway copy, never into the + tree the agent worked in, so there is no window in which they could have + been read or edited. + """ + dest = workroot / f"grade-{label}" + shutil.rmtree(dest, ignore_errors=True) + materialise(files, dest) + materialise(task.tests, dest) + return runner.test(dest).get("tests", {}) + + +def score(task: AgenticTask, outcomes: Dict[str, str]) -> Dict[str, Any]: + """Resolved means every target test passes and nothing previously passing broke.""" + def passing(tid): + return outcomes.get(tid) == "pass" + + f2p_ok = [t for t in task.fail_to_pass if passing(t)] + p2p_ok = [t for t in task.pass_to_pass if passing(t)] + regressions = [t for t in task.pass_to_pass if not passing(t)] + resolved = (len(f2p_ok) == len(task.fail_to_pass) + and not regressions and bool(task.fail_to_pass)) + return { + "resolved": resolved, + "f2p_passed": len(f2p_ok), "f2p_total": len(task.fail_to_pass), + "p2p_passed": len(p2p_ok), "p2p_total": len(task.pass_to_pass), + "regressions": regressions[:10], + "n_regressions": len(regressions), + "partial": (len(f2p_ok) / len(task.fail_to_pass)) if task.fail_to_pass else 0.0, + } + + +def make_sandbox(lang_name: str, sandbox_cls, cache_root=None, base_python=None): + """A sandbox carrying only the toolchain for `lang_name`. + + Keyed per language rather than per suite: running the Python tasks must + never provision node, and vice versa. + """ + lang = LANGS[lang_name] + return sandbox_cls( + name=f"agentic-{lang_name}", + packages=list(lang.packages), + provision=list(lang.provision), + env=dict(lang.env), + address_space_limit=lang.address_space_limit, + cache_root=cache_root, + base_python=base_python, + ) + + +def run_task(task: AgenticTask, chat, sandbox, workroot: Optional[Path] = None, + max_turns: int = 50, context_limit: Optional[int] = None, + test_timeout: float = 90.0) -> Dict[str, Any]: + """One full attempt: fresh workspace, agent episode, then hidden tests.""" + root = Path(tempfile.mkdtemp(prefix=f"agentic-{task.task_id}-", + dir=str(workroot) if workroot else None)) + try: + ws_root = root / "repo" + materialise(task.files, ws_root) + ws = Workspace(ws_root) + runner = Runner(sandbox, task.spec, timeout=test_timeout) + toolbox = ToolBox(ws, linter=lambda t: runner.lint(ws_root, t)) + + episode, messages = run_episode(task, toolbox, chat, + max_turns=max_turns, + context_limit=context_limit) + final = ws.snapshot() + changed = sorted(k for k in final + if final[k] != task.files.get(k)) + \ + sorted(k for k in task.files if k not in final) + + outcomes = run_tests_against(final, task, runner, root, "final") + result = score(task, outcomes) + result.update({ + "task_id": task.task_id, "lang": task.lang, + "category": task.category, "difficulty": task.difficulty, + "changed_files": changed, + "episode": {k: v for k, v in vars(episode).items()}, + "messages": messages, + }) + return result + finally: + shutil.rmtree(root, ignore_errors=True) + + +TRANSCRIPT_CSS = """ +body { font-family: system-ui, sans-serif; margin: 0; padding: 16px; background: #fff; color: #222; max-width: 1100px; } +h1 { font-size: 16px; margin: 0 0 4px; } +.meta { font-size: 13px; color: #555; font-family: 'SF Mono', 'Menlo', 'Consolas', monospace; margin-bottom: 12px; white-space: pre-wrap; } +.msg { border: 1px solid #e1e4e8; border-left-width: 4px; margin: 8px 0; padding: 6px 10px; background: #fafbfc; } +.msg .role { font-size: 12px; font-weight: 600; color: #555; margin-bottom: 4px; } +.system { border-left-color: #888; } +.user { border-left-color: #0969da; } +.nudge { border-left-color: #9a6700; } +.assistant { border-left-color: #1a7f37; background: #fff; } +.tool { border-left-color: #8250df; } +.tool.err { border-left-color: #cf222e; } +pre { font-family: 'SF Mono', 'Menlo', 'Consolas', monospace; font-size: 12px; background: #fff; border: 1px solid #e1e4e8; padding: 6px 8px; margin: 4px 0; white-space: pre-wrap; word-wrap: break-word; } +.reasoning pre { color: #57606a; background: #f6f8fa; } +.call { font-size: 12px; font-weight: 600; color: #8250df; margin-top: 6px; } +details > summary { cursor: pointer; font-size: 12px; color: #555; } +.empty { font-size: 12px; color: #888; font-style: italic; } +""" + + +def _pre(text: str) -> str: + return f"<pre>{html.escape(text)}</pre>" + + +def _render_call(call: Dict[str, Any]) -> str: + fn = call.get("function") or {} + raw = fn.get("arguments") + try: + args = raw if isinstance(raw, dict) else json.loads(raw or "{}") + shown = json.dumps(args, indent=2, ensure_ascii=False) + except (json.JSONDecodeError, TypeError): + shown = f"(invalid JSON)\n{raw}" + return (f'<div class="call">→ {html.escape(fn.get("name") or "?")} ' + f'<span style="font-weight:400;color:#888">{html.escape(call.get("id") or "")}</span></div>' + + _pre(shown)) + + +def render_transcript_html(messages: List[Dict[str, Any]], title: str, + meta: str = "") -> str: + """One page per episode: every turn, with reasoning, tool calls and results.""" + names = {} + for m in messages: + for c in m.get("tool_calls") or []: + names[c.get("id")] = (c.get("function") or {}).get("name") or "?" + + parts = [] + turn = 0 + for m in messages: + role = m.get("role") or "?" + content = m.get("content") or "" + if role == "assistant": + turn += 1 + body = "" + if m.get("reasoning_content"): + body += ('<details class="reasoning" open><summary>reasoning ' + f'({len(m["reasoning_content"])} chars)</summary>' + + _pre(m["reasoning_content"]) + "</details>") + body += _pre(content) if content.strip() else "" + body += "".join(_render_call(c) for c in m.get("tool_calls") or []) + if not body: + body = '<div class="empty">(empty reply)</div>' + parts.append(f'<div class="msg assistant"><div class="role">assistant ' + f'· turn {turn}</div>{body}</div>') + elif role == "tool": + name = names.get(m.get("tool_call_id"), "?") + err = " err" if content.startswith("Error") else "" + parts.append(f'<div class="msg tool{err}"><div class="role">tool result ' + f'· {html.escape(name)}</div>{_pre(content)}</div>') + elif role == "system": + parts.append('<div class="msg system"><details><summary>system prompt' + f'</summary>{_pre(content)}</details></div>') + else: + cls = "nudge" if content == NUDGE else role + label = "user (nudge)" if content == NUDGE else role + parts.append(f'<div class="msg {html.escape(cls)}"><div class="role">' + f'{html.escape(label)}</div>{_pre(content)}</div>') + + return (f'<!DOCTYPE html>\n<html>\n<head>\n<meta charset="UTF-8">\n' + f'<title>{html.escape(title)}\n\n' + f'\n\n

{html.escape(title)}

\n' + f'
{html.escape(meta)}
\n' + + "\n".join(parts) + "\n\n\n") + + +def write_transcript(messages: List[Dict[str, Any]], dest: Path, title: str, + meta: str = "") -> Path: + """Write `.json` (the raw message history) and `.html`.""" + dest.parent.mkdir(parents=True, exist_ok=True) + json_path = dest.with_name(dest.name + ".json") + with open(json_path, "w") as f: + json.dump(messages, f, indent=2, ensure_ascii=False) + with open(dest.with_name(dest.name + ".html"), "w") as f: + f.write(render_transcript_html(messages, title, meta)) + return json_path diff --git a/examples/llama-eval/eval_sandbox.py b/examples/llama-eval/eval_sandbox.py new file mode 100644 index 000000000000..b13bc9d7101b --- /dev/null +++ b/examples/llama-eval/eval_sandbox.py @@ -0,0 +1,488 @@ +#!/usr/bin/env python3 +"""Lightweight sandboxed execution for code-generation evals (HumanEval, ClassEval). + +Replaces the usual Docker-per-suite setup with a cached venv plus whatever +process isolation the host can give us without root: + + tier 1 bubblewrap read-only root, tmpfs /tmp and $HOME, no network + tier 2 unshare -rn no network (home still visible) + tier 3 rlimits only last resort, still capped on cpu/memory/procs/files + +Every tier also applies rlimits and runs the candidate program in a throwaway +working directory as its own process group, so a runaway solution is killed +along with anything it spawned. +""" + +import json +import os +import re +import resource +import shutil +import signal +import subprocess +import sys +import tempfile +import time +import venv +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List, Optional, Sequence + +DEFAULT_CACHE_ROOT = Path( + os.environ.get("LLAMA_EVAL_CACHE", Path.home() / ".cache" / "llama-eval") +) + +LIMIT_AS_BYTES = 4 * 1024 * 1024 * 1024 +LIMIT_CPU_SECONDS_SLACK = 5 +LIMIT_NPROC = 256 +LIMIT_FSIZE_BYTES = 64 * 1024 * 1024 + + +@dataclass +class ExecResult: + passed: bool + status: str + stdout: str = "" + stderr: str = "" + returncode: Optional[int] = None + duration: float = 0.0 + + def summary(self, max_chars: int = 400) -> str: + """Short human-readable reason, for grader logs.""" + if self.passed: + return "ok" + if self.status == "timeout": + return "timeout" + tail = (self.stderr or self.stdout).strip() + if not tail: + return f"{self.status} (rc={self.returncode})" + lines = [ln for ln in tail.splitlines() if ln.strip()] + reason = lines[-1] if lines else tail + return reason[:max_chars] + + +def resolve_python(spec: Optional[str]) -> Optional[str]: + """Turn '3.13' or a path into a usable interpreter, or None if unavailable. + + Suites can need an older interpreter than the one running this tool, because + a dependency has no wheel (or no working source build) for the newest + CPython. Accepts an explicit path, a bare version, or an executable name. + """ + if not spec: + return None + + candidate = Path(spec).expanduser() + if candidate.exists() and candidate.is_file(): + return str(candidate) + + uv = shutil.which("uv") + if uv and re.fullmatch(r"\d+\.\d+(\.\d+)?", spec): + try: + r = subprocess.run([uv, "python", "find", spec], + capture_output=True, text=True, timeout=120) + path = r.stdout.strip() + if r.returncode == 0 and path and Path(path).exists(): + return path + except Exception: + pass + + return shutil.which(spec) or shutil.which(f"python{spec}") + + +def _detect_isolation(probe_python: Optional[str] = None, + ro_binds: Sequence[Path] = ()) -> str: + """Pick the strongest isolation tier this host supports, once. + + The probe must run the *same* interpreter the suite will, with the same + binds. Probing a bare sys.executable silently downgrades the tier whenever + this tool is itself run from a venv under $HOME -- which is the usual case, + since it needs `datasets`. bwrap tmpfs-mounts the home directory, the + interpreter disappears, the probe fails, and the sandbox falls back to a + tier that does not hide $HOME at all. Nothing about the results looks wrong + when that happens, so it has to be got right here. + """ + py = probe_python or sys.executable + probe = "import sys; sys.exit(0)" + + if shutil.which("bwrap"): + try: + r = subprocess.run( + _bwrap_argv(Path.cwd(), [py, "-I", "-c", probe], ro_binds=ro_binds), + capture_output=True, timeout=20, + ) + if r.returncode == 0: + return "bwrap" + except Exception: + pass + + if shutil.which("unshare"): + for tier, argv in ( + ("unshare", ["unshare", "--user", "--map-current-user", "--net"]), + ("unshare-root", ["unshare", "-rn"]), + ): + try: + r = subprocess.run( + [*argv, py, "-I", "-c", probe], + capture_output=True, timeout=20, + ) + if r.returncode == 0: + return tier + except Exception: + continue + + return "rlimit" + + +UNSHARE_ARGV = { + "unshare": ["unshare", "--user", "--map-current-user", "--net"], + "unshare-root": ["unshare", "-rn"], +} + + +def _under(path: Path, parent: str) -> bool: + return bool(parent) and (str(path) == parent + or str(path).startswith(parent.rstrip("/") + os.sep)) + + +def _bwrap_argv(workdir: Path, inner: Sequence[str], + ro_binds: Sequence[Path] = ()) -> List[str]: + home = os.path.expanduser("~") + argv = [ + "bwrap", + "--ro-bind", "/", "/", + "--dev", "/dev", + "--proc", "/proc", + "--tmpfs", "/tmp", + ] + if home and home != "/": + argv += ["--tmpfs", home] + for path in ro_binds: + if _under(path, home): + argv += ["--ro-bind", str(path), str(path)] + argv += [ + "--bind", str(workdir), str(workdir), + "--chdir", str(workdir), + "--unshare-all", + "--die-with-parent", + "--new-session", + "--", + ] + return argv + list(inner) + + +def _nproc_ceiling() -> int: + """A process cap that bounds a fork bomb without tripping on existing load. + + RLIMIT_NPROC is enforced per real UID across the whole system, not per + process tree, so it has to sit above whatever the user is already running. + """ + try: + current = len(os.listdir("/proc")) + except OSError: + current = 512 + return max(current * 2, 1024) + + +def _apply_rlimits(timeout: float, with_nproc: bool, + as_bytes: Optional[int] = LIMIT_AS_BYTES): + def _preexec(): + os.setsid() + cpu = int(timeout) + LIMIT_CPU_SECONDS_SLACK + limits = [ + (resource.RLIMIT_CPU, (cpu, cpu)), + (resource.RLIMIT_FSIZE, (LIMIT_FSIZE_BYTES, LIMIT_FSIZE_BYTES)), + (resource.RLIMIT_CORE, (0, 0)), + ] + if as_bytes: + limits.append((resource.RLIMIT_AS, (as_bytes, as_bytes))) + if with_nproc: + ceiling = _nproc_ceiling() + limits.append((resource.RLIMIT_NPROC, (ceiling, ceiling))) + for res, lim in limits: + try: + resource.setrlimit(res, lim) + except (ValueError, OSError): + pass + return _preexec + + +class Sandbox: + """A cached venv plus an isolated runner for untrusted candidate programs.""" + + def __init__( + self, + name: str, + packages: Optional[Sequence[str]] = None, + cache_root: Optional[Path] = None, + base_python: Optional[str] = None, + isolation: Optional[str] = None, + provision: Optional[Sequence[str]] = None, + env: Optional[Dict[str, str]] = None, + address_space_limit: Optional[int] = LIMIT_AS_BYTES, + ): + self.name = name + self.packages = sorted(set(packages or [])) + self.provision = list(provision or []) + self.env = dict(env or {}) + # None drops the RLIMIT_AS cap entirely -- see _apply_rlimits + self.address_space_limit = address_space_limit + self.cache_root = Path(cache_root or DEFAULT_CACHE_ROOT) + self.base_python = base_python or sys.executable + self.venv_dir = self.cache_root / f"{name}-venv" + self._python: Optional[Path] = None + self._isolation = isolation + self._base_prefix: List[Path] = [] + + @property + def python(self) -> Path: + if self._python is None: + raise RuntimeError("Sandbox.ensure() must be called first") + return self._python + + @property + def isolation(self) -> str: + if self._isolation is None: + if self._python: + py = str(self._python) + roots = [self.venv_dir, *(self._base_prefix or [])] + else: + py = sys.executable + roots = self._detect_base_prefix(Path(py)) + self._isolation = _detect_isolation(py, roots) + return self._isolation + + def _stamp_path(self) -> Path: + return self.venv_dir / ".llama-eval-stamp.json" + + def _stamp_matches(self) -> bool: + try: + stamp = json.loads(self._stamp_path().read_text()) + except Exception: + return False + return ( + stamp.get("packages") == self.packages + and stamp.get("py") == self._base_version() + and stamp.get("provision") == self.provision + ) + + def _resolved_env(self) -> Dict[str, str]: + return {k: v.replace("{venv}", str(self.venv_dir)) for k, v in self.env.items()} + + def _base_version(self) -> str: + r = subprocess.run( + [self.base_python, "-c", + "import sys;print('.'.join(map(str,sys.version_info[:3])))"], + capture_output=True, text=True, + ) + return r.stdout.strip() + + def ensure(self, force: bool = False, quiet: bool = False) -> Path: + """Create/refresh the venv. Idempotent and cheap when already valid.""" + py = self.venv_dir / "bin" / "python" + if force and self.venv_dir.exists(): + shutil.rmtree(self.venv_dir) + + if py.exists() and self._stamp_matches(): + self._python = py + self._base_prefix = self._detect_base_prefix(py) + return py + + if not quiet: + pkgs = ", ".join(self.packages) if self.packages else "stdlib only" + print(f"[sandbox] provisioning venv {self.venv_dir} ({pkgs})") + + self.cache_root.mkdir(parents=True, exist_ok=True) + if self.venv_dir.exists(): + shutil.rmtree(self.venv_dir) + + uv = shutil.which("uv") + if uv: + subprocess.run( + [uv, "venv", "--python", self.base_python, str(self.venv_dir)], + check=True, capture_output=True, + ) + elif Path(self.base_python).resolve() == Path(sys.executable).resolve(): + venv.EnvBuilder(with_pip=bool(self.packages), clear=True).create(self.venv_dir) + else: + subprocess.run( + [self.base_python, "-m", "venv", str(self.venv_dir)], + check=True, capture_output=True, + ) + + if self.packages: + if uv: + cmd = [uv, "pip", "install", "--python", str(py), *self.packages] + else: + cmd = [str(py), "-m", "pip", "install", "--quiet", + "--disable-pip-version-check", *self.packages] + r = subprocess.run(cmd, capture_output=True, text=True) + if r.returncode != 0: + raise RuntimeError( + f"failed to install {len(self.packages)} package(s) into " + f"{self.venv_dir}:\n{(r.stderr or r.stdout)[-3000:]}" + ) + + for i, snippet in enumerate(self.provision, 1): + if not quiet: + print(f"[sandbox] provisioning step {i}/{len(self.provision)}") + prov_env = dict(os.environ) + prov_env["SANDBOX_VENV"] = str(self.venv_dir) + prov_env.update(self._resolved_env()) + r = subprocess.run([str(py), "-c", snippet], env=prov_env, + capture_output=True, text=True, timeout=1800) + if r.returncode != 0: + raise RuntimeError( + f"provisioning step {i} failed:\n{(r.stderr or r.stdout)[-3000:]}" + ) + + self._stamp_path().write_text(json.dumps({ + "packages": self.packages, + "py": self._base_version(), + "provision": self.provision, + })) + self._python = py + self._base_prefix = self._detect_base_prefix(py) + if not quiet: + print(f"[sandbox] ready ({self.isolation} isolation)") + if self.isolation == "unshare-root": + print("Warning: this host's unshare cannot keep the caller's uid, so " + "candidate code runs as root inside the namespace. Tests that " + "assert a permission denial will not see one. Install " + "bubblewrap (bwrap) for correct results.") + return py + + @staticmethod + def _detect_base_prefix(py: Path) -> List[Path]: + """Every directory that must stay visible for the venv to start. + + A venv is only symlinks into its base interpreter, so the whole install + tree has to survive the tmpfs that hides $HOME -- and binding just the + resolved tree is not enough. Version-managed interpreters (uv, pyenv) + are commonly reached through an unversioned alias symlink + (cpython-3.13-... -> cpython-3.13.14-...), and startup resolves through + the alias path, so both spellings need to be mounted. Walk the chain + and keep every install root it passes through. + """ + roots: List[Path] = [] + + def add(path: Path): + if path.exists() and path not in roots: + roots.append(path) + + cur, hops = py, 0 + while hops < 16: + hops += 1 + add(cur.parent.parent) + if cur.is_symlink(): + target = Path(os.readlink(cur)) + cur = target if target.is_absolute() else (cur.parent / target) + else: + break + + try: + r = subprocess.run([str(py), "-c", "import sys; print(sys.base_prefix)"], + capture_output=True, text=True, timeout=60) + if r.returncode == 0 and r.stdout.strip(): + add(Path(r.stdout.strip())) + except Exception: + pass + + return [p for p in roots if str(p) not in ("/", "/usr", "/usr/local")] + + def _isolate(self, argv: Sequence[str], workdir: Path, + ro_binds: Sequence[Path] = ()) -> List[str]: + """Wrap a command in the strongest isolation tier this host supports.""" + if self.isolation == "bwrap": + return _bwrap_argv( + workdir, argv, + ro_binds=[self.venv_dir, *(self._base_prefix or []), *ro_binds], + ) + if self.isolation in UNSHARE_ARGV: + return [*UNSHARE_ARGV[self.isolation], *argv] + return list(argv) + + def _base_env(self, workdir: Path) -> Dict[str, str]: + env = { + "PATH": "/usr/bin:/bin", + "HOME": str(workdir), + "TMPDIR": str(workdir), + "LC_ALL": "C.UTF-8", + "PYTHONHASHSEED": "0", + "PYTHONDONTWRITEBYTECODE": "1", + "OPENBLAS_NUM_THREADS": "1", + "OMP_NUM_THREADS": "1", + } + env.update(self._resolved_env()) + return env + + def run_argv(self, argv: Sequence[str], workdir: Path, + timeout: float = 15.0, env: Optional[Dict[str, str]] = None, + ro_binds: Sequence[Path] = ()) -> ExecResult: + """Run an arbitrary command in `workdir`, which stays writable. + + Unlike run(), the caller owns the directory and it survives the call. + That is what a multi-step agent needs: its edits have to persist across + many tool invocations before anything is graded. Taking an argv rather + than a source string is what lets a suite drive a linter or a test + runner for a language other than Python. + """ + run_env = self._base_env(workdir) + if env: + run_env.update(env) + + started = time.monotonic() + try: + proc = subprocess.Popen( + self._isolate(argv, workdir, ro_binds), + cwd=str(workdir), + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + errors="replace", + env=run_env, + preexec_fn=_apply_rlimits(timeout, + with_nproc=self.isolation == "rlimit", + as_bytes=self.address_space_limit), + ) + except Exception as e: + return ExecResult(False, "error", "", f"{type(e).__name__}: {e}") + + try: + out, err = proc.communicate(timeout=timeout) + except subprocess.TimeoutExpired: + _kill_tree(proc) + out, err = proc.communicate() + return ExecResult(False, "timeout", out or "", err or "", + None, time.monotonic() - started) + + duration = time.monotonic() - started + if proc.returncode == 0: + return ExecResult(True, "ok", out, err, 0, duration) + return ExecResult(False, "failed", out, err, proc.returncode, duration) + + def run(self, code: str, timeout: float = 15.0, + env: Optional[Dict[str, str]] = None) -> ExecResult: + """Run `code` to completion in the sandbox. Exit 0 == passed.""" + workdir = Path(tempfile.mkdtemp(prefix=f"{self.name}-", dir="/tmp")) + try: + prog = workdir / "candidate.py" + prog.write_text(code) + return self.run_argv([str(self.python), "-I", "-B", str(prog)], + workdir, timeout=timeout, env=env) + except Exception as e: + return ExecResult(False, "error", "", f"{type(e).__name__}: {e}") + finally: + shutil.rmtree(workdir, ignore_errors=True) + + +def _kill_tree(proc: subprocess.Popen): + """Kill the candidate and everything it spawned.""" + for sig in (signal.SIGKILL,): + try: + os.killpg(os.getpgid(proc.pid), sig) + except (ProcessLookupError, PermissionError): + try: + proc.kill() + except ProcessLookupError: + pass diff --git a/examples/llama-eval/llama-eval.py b/examples/llama-eval/llama-eval.py index 61bdbddd8b13..3ed68e3b530d 100755 --- a/examples/llama-eval/llama-eval.py +++ b/examples/llama-eval/llama-eval.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import argparse +import ast import json import os import re @@ -8,23 +9,113 @@ import sys import threading import time +import uuid from abc import ABC, abstractmethod from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, asdict, field from pathlib import Path from queue import Queue from typing import Dict, List, Optional, Any, Tuple +from urllib.parse import urlparse import requests from tqdm import tqdm import random from math import sqrt +from eval_sandbox import Sandbox, resolve_python + @dataclass class ServerConfig: url: str threads: int name: str = "" + api_key: Optional[str] = None + + def is_opencode(self) -> bool: + """Whether this endpoint is an opencode gateway rather than a llama-server.""" + return "opencode" in (urlparse(self.url).hostname or "") + + def headers(self, session_id: Optional[str] = None) -> Dict[str, str]: + headers = auth_headers(self.api_key) + if session_id and self.is_opencode(): + headers["x-opencode-session"] = session_id + return headers + + +def auth_headers(api_key: Optional[str]) -> Dict[str, str]: + """Standard request headers, carrying a bearer token when one is configured.""" + headers = {"Content-Type": "application/json"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + return headers + + +SERVER_ERROR = "server_error" + +RETRYABLE_HTTP = {408, 409, 425, 429} + + +class ServerError(Exception): + """A request that still failed, or came back empty, after every retry.""" + + +def null_reply_reason(result: Any) -> Optional[str]: + """Why a chat completion carries no usable reply, or None when it does.""" + if not isinstance(result, dict): + return "reply is not a JSON object" + choices = result.get("choices") + if not choices: + return "reply has no choices" + msg = choices[0].get("message") + if not msg: + return "reply has no message" + if not (msg.get("content") or msg.get("tool_calls") or msg.get("reasoning_content")): + return "reply is empty" + return None + + +def post_with_retry(url: str, headers: Dict[str, str], data: Dict[str, Any], + retries: int, retry_delay: float, + timeout: Optional[float] = None) -> Dict[str, Any]: + """POST a chat completion, resending on errors and null replies. + + The wait before each resend doubles, starting at `retry_delay` seconds. + Raises ServerError once `retries` resends have all failed. + """ + reason = "" + attempts = 0 + for attempt in range(retries + 1): + if attempt: + wait = retry_delay * 2 ** (attempt - 1) + print(f" [{url}] {reason}; retry {attempt}/{retries} in {wait:.0f}s", + file=sys.stderr) + time.sleep(wait) + attempts += 1 + try: + response = requests.post(url, headers=headers, json=data, timeout=timeout) + except requests.RequestException as e: + reason = f"{type(e).__name__}: {e}" + continue + if response.status_code >= 400: + reason = f"HTTP {response.status_code}: {response.text[:300]}" + if response.status_code < 500 and response.status_code not in RETRYABLE_HTTP: + break + continue + try: + result = response.json() + except ValueError: + reason = f"invalid JSON: {response.text[:300]}" + continue + reason = null_reply_reason(result) or "" + if not reason: + return result + raise ServerError(f"{reason} (after {attempts} attempt{'s' if attempts > 1 else ''})") + + +def new_session_id(task_id: str) -> str: + """A fresh session id for one task, unique across tasks, runs and resumes.""" + return f"{task_id}-{uuid.uuid4().hex}" def wilson_interval(correct: int, total: int, z: float = 1.96) -> Tuple[float, float]: """Wilson score confidence interval for a proportion.""" @@ -106,13 +197,118 @@ def wilson_interval(correct: int, total: int, z: float = 1.96) -> Tuple[float, f B) {B} C) {C} D) {D} +""", + "humaneval": """Complete the following Python function. + +```python +{question} +``` + +Reply with the complete function in a single ```python code block, including the \ +signature and any imports it needs. Do not include tests, example usage, or explanation. +""", + "classeval": """Complete the following Python class. + +```python +{question} +``` + +Reply with the complete class in a single ```python code block: keep the class name \ +and every method signature exactly as given, implement every method body, and include \ +any imports the class needs. Do not include tests, example usage, or explanation. """, } +# Suites graded by running the generated code instead of matching its text. +CODE_SUITES = {"humaneval", "classeval", "agentic"} + +AGENTIC_SUITES = {"agentic"} + +# Import name -> pip name, for the cases where they differ. Anything not listed is assumed to install under the name it imports as. +PIP_NAMES = { + "bs4": "beautifulsoup4", + "PIL": "pillow", + "docx": "python-docx", + "fitz": "pymupdf", + "sklearn": "scikit-learn", + "yaml": "PyYAML", + "cv2": "opencv-python-headless", + "PyPDF2": "pypdf2", + "Levenshtein": "levenshtein", + "dateutil": "python-dateutil", + "attr": "attrs", + "OpenSSL": "pyOpenSSL", + "serial": "pyserial", + "Crypto": "pycryptodome", +} + +DATASET_SOURCE: Optional[str] = None + +# Extra construction options for the agentic suite, set from the CLI. +AGENTIC_OPTS: Dict[str, Any] = {} + class BaseDataset(ABC): questions: List[Dict] + is_agentic: bool = False + + def prepare(self, base_python: Optional[str] = None, force: bool = False) -> None: + """Provision whatever the suite needs before any task runs. + + Single-shot suites are provisioned from required_packages() instead, so + this is a no-op for them. + """ + + def run(self, index: int, chat) -> Dict[str, Any]: + """Drive one episode to completion and return its scored result.""" + raise NotImplementedError( + f"{type(self).__name__} does not drive its own conversation") + + def required_packages(self) -> List[str]: + """Third-party pip packages the suite's tests need.""" + return [] + + def provision_steps(self) -> List[str]: + """Python snippets run once at venv build time, with network access.""" + return [] + + def preferred_python(self) -> Optional[str]: + """Interpreter the suite's dependencies need, if not the current one.""" + return None + + def configure_python(self, python: str): + """Told which interpreter the sandbox will use, before it is built. + + Lets a suite adapt its package set to what that version can actually + install. + """ + + def default_exec_timeout(self) -> float: + """Per-task wall-clock limit that suits this suite's workload.""" + return 15.0 + + def sandbox_env(self) -> Dict[str, str]: + """Extra env for every run; '{venv}' expands to the venv path.""" + return {} + + def extract(self, response: str, question: Dict) -> str: + """Pull the candidate code out of a chat reply.""" + return extract_code(response) + + def build_program(self, question: Dict, completion: str, + seed: Optional[int] = None) -> str: + """Assemble the program whose exit status decides pass/fail.""" + raise NotImplementedError + + def interpret(self, result: Any, question: Dict) -> Tuple[bool, Dict[str, Any]]: + """Turn an ExecResult into (passed, detail-for-the-log).""" + return result.passed, { + "status": result.status, + "reason": result.summary(), + "duration": round(result.duration, 3), + } + @abstractmethod def get_question(self, index: int) -> Dict: pass @@ -175,6 +371,13 @@ def __init__( self.total_time: float = 0.0 self._lock = threading.Lock() + @property + def loaded_dataset(self) -> BaseDataset: + """The dataset, once loaded. Raises rather than handing back None.""" + if self.dataset is None: + raise RuntimeError("dataset not loaded; call load_dataset() first") + return self.dataset + def load_dataset(self, seed: int = 1234): if self.dataset_type == "aime": self.dataset = AimeDataset() @@ -186,6 +389,12 @@ def load_dataset(self, seed: int = 1234): self.dataset = Gsm8kDataset() elif self.dataset_type == "gpqa": self.dataset = GpqaDataset(variant="diamond", seed=seed) + elif self.dataset_type == "humaneval": + self.dataset = HumanEvalDataset(source=DATASET_SOURCE) + elif self.dataset_type == "classeval": + self.dataset = ClassEvalDataset(source=DATASET_SOURCE) + elif self.dataset_type == "agentic": + self.dataset = AgenticDataset(source=DATASET_SOURCE, **AGENTIC_OPTS) else: raise ValueError(f"Unknown dataset type: {self.dataset_type}") @@ -263,31 +472,55 @@ def add_result( self.correct = sum(1 for c in self.task_states.get("cases", {}).values() if c.get("correct", False)) + def _display_pair(self, task_state: TaskState) -> Tuple[str, str]: + """The (expected, answer) columns, kept narrow for code suites. + + There the gold value is a whole JSON test harness and the answer is a + whole function, so show the task id and the failure cause instead. + """ + if self.dataset_type not in CODE_SUITES: + return ( + task_state.expected, + task_state.answer if task_state.answer else "N/A", + ) + try: + expected = json.loads(task_state.expected)["task_id"] + except (ValueError, KeyError, TypeError): + expected = self.dataset_type + reason = (task_state.grader_log or {}).get("reason", "") + if task_state.correct: + reason = "ok" + return expected, (reason[:28] or "N/A") + def print_progress(self, task_state: TaskState, total_tasks: int, n_correct: int = 0): - display_answer = task_state.answer if task_state.answer else "N/A" + display_expected, display_answer = self._display_pair(task_state) display_tokens = str(task_state.tokens) if task_state.tokens is not None else "N/A" display_tps = f"{task_state.tps_gen:.1f}" if task_state.tps_gen is not None else "N/A" display_t_gen = f"{task_state.t_gen_ms/1000:.1f}" if task_state.t_gen_ms is not None else "N/A" display_server = task_state.server_name if task_state.server_name else "N/A" success_ratio = n_correct / self.processed if self.processed > 0 else 0.0 - first_line = task_state.question_text.split('\n')[0] + first_line = next( + (ln for ln in task_state.question_text.split('\n') if ln.strip()), "" + ) truncated_question = first_line[:43] if len(first_line) > 43: truncated_question += "..." else: truncated_question = truncated_question.ljust(43) + "..." - print(f"{self.processed:3}/{total_tasks:3} {task_state.task_id:<20} {self.dataset_type.upper()} {truncated_question:<40} {task_state.expected:<10} {display_answer:<10} {display_tokens:<6} {display_tps:<6} {display_t_gen:<8} {'✓' if task_state.correct else '✗'} [{n_correct:3}/{self.processed:3}, {success_ratio:.3f}] {display_server}") + print(f"{self.processed:3}/{total_tasks:3} {task_state.task_id:<20} {self.dataset_type.upper()} {truncated_question:<40} {display_expected:<16} {display_answer:<28} {display_tokens:<6} {display_tps:<6} {display_t_gen:<8} {'✓' if task_state.correct else ('⚠' if task_state.status == SERVER_ERROR else '✗')} [{n_correct:3}/{self.processed:3}, {success_ratio:.3f}] {display_server}") + + def n_server_errors(self) -> int: + cases = self.task_states.get("cases", {}) + return sum(1 for c in cases.values() if c.get("status") == SERVER_ERROR) + + def print_server_errors(self): + n = self.n_server_errors() + if n: + print(f"Server errors: {n} (not counted in accuracy; run with --resume to retry them)") def print_summary(self): - if self.total == 0: - print(f"\n{'='*60}") - print(f"Results: 0/0 correct (0.0%)") - print(f"{'='*60}") - else: - ci_lower, ci_upper = self.accuracy_ci() - print(f"\n{'='*60}") - print(f"Results: {self.correct}/{self.total} correct ({self.correct/self.total*100:.1f}%) [{ci_lower*100:.1f}%, {ci_upper*100:.1f}%]") - print(f"{'='*60}") + print() + self.print_existing_summary() def dump(self): with self._lock: @@ -347,7 +580,8 @@ def dump_html(self, tasks_to_save: List[Tuple[int, str]], all_cases: Dict[str, A completed = {tid: c for tid, c in cases.items() if c.get("status") == "ok"} n_correct = sum(1 for c in completed.values() if c.get("correct", False)) n_incorrect = len(completed) - n_correct - n_pending = len(tasks_to_save) - len(completed) + n_server_errors = sum(1 for c in cases.values() if c.get("status") == SERVER_ERROR) + n_pending = len(tasks_to_save) - len(completed) - n_server_errors accuracy = n_correct / len(completed) * 100 if completed else 0.0 ci_lower, ci_upper = wilson_interval(n_correct, len(completed)) if completed else (0.0, 1.0) @@ -374,6 +608,9 @@ def dump_html(self, tasks_to_save: List[Tuple[int, str]], all_cases: Dict[str, A elif status == "pending": status_class = "pending" status_text = "–" + elif status == SERVER_ERROR: + status_class = "server-error" + status_text = "⚠" else: status_class = "error" status_text = "!" @@ -392,6 +629,11 @@ def dump_html(self, tasks_to_save: List[Tuple[int, str]], all_cases: Dict[str, A escaped_reasoning = self._escape_html(reasoning_content) grader_log_str = self._escape_html(json.dumps(grader_log, indent=2)) escaped_server = self._escape_html(server_name) + transcript = grader_log.get("transcript") if isinstance(grader_log, dict) else None + transcript_link = ( + f'Transcript ' + f'{self._escape_html(transcript[:-5])}.html
' + if transcript and transcript.endswith(".json") else "") answer_class = status_class if status == "ok" else "" rows.append(f""" @@ -407,6 +649,7 @@ def dump_html(self, tasks_to_save: List[Tuple[int, str]], all_cases: Dict[str, A
+ {transcript_link} Prompt
{escaped_prompt}
Response
{escaped_response}
{f'Reasoning
{escaped_reasoning}
' if escaped_reasoning else ''} @@ -483,6 +726,7 @@ def _stat(v, fmt=".1f", avg_fmt=None): .incorrect {{ color: #cf222e; }} .pending {{ color: #888; }} .error {{ color: #9a6700; }} + .server-error {{ color: #8250df; }} .details-row {{ display: none; }} .details-row.open {{ display: table-row; }} .details-content {{ padding: 8px 16px; background: #f6f8fa; font-size: 12px; }} @@ -511,6 +755,7 @@ def _stat(v, fmt=".1f", avg_fmt=None):
Accuracy
{accuracy:.1f}% [{ci_lower*100:.1f}%, {ci_upper*100:.1f}%]
Correct
{n_correct} / {len(completed)}
Pending
{n_pending}
+
Server errors
{n_server_errors}
Time
{self.total_time:.1f}s
Sampling
{sampling_str}
@@ -680,6 +925,7 @@ def print_existing_summary(self): print(f"{'='*60}") print(f"Results: {correct}/{total} correct ({correct/total*100:.1f}%) [{ci_lower*100:.1f}%, {ci_upper*100:.1f}%]") print(f"{'='*60}") + self.print_server_errors() def accuracy_ci(self) -> Tuple[float, float]: """Compute Wilson score confidence interval from completed cases.""" @@ -974,6 +1220,526 @@ def get_prompt(self, question: Dict) -> str: D=question["shuffled_answers"][3] ) +FENCE_RE = re.compile(r"```[ \t]*([A-Za-z0-9_+-]*)[ \t]*\n(.*?)(?:```|\Z)", re.DOTALL) + +RESULT_MARKER = "__LLAMA_EVAL__" + +GENSIM_SHIM = { + "__init__.py": ( + '"""Minimal stand-in for the two gensim entry points ClassEval uses."""\n' + "from . import utils, matutils # noqa: F401\n" + ), + "utils.py": ( + "import re, html\n" + "RE_ENTITY = re.compile(r'&(#?)([xX]?)(\\w{1,8});')\n" + "def decode_htmlentities(text):\n" + " return RE_ENTITY.sub(lambda m: html.unescape(m.group(0)), text)\n" + ), + "matutils.py": ( + "import numpy as np\n" + "def unitvec(vec, norm='l2', return_norm=False):\n" + " v = np.asarray(vec, dtype=float)\n" + " n = float(np.sum(np.abs(v))) if norm == 'l1' " + "else float(np.sqrt(np.sum(v ** 2)))\n" + " if n == 0.0:\n" + " return (v, 1.0) if return_norm else v\n" + " out = v / n\n" + " return (out, n) if return_norm else out\n" + ), +} + +UNITTEST_PROLOGUE = ''' +import unittest as _ut_pre +_ut_pre.main = lambda *a, **k: None +''' + +UNITTEST_RUNNER = f''' + +if True: + import json as _json, os as _os, unittest as _ut + _loader = _ut.TestLoader() + _suite = _ut.TestSuite() + for _c in [v for v in list(globals().values()) + if isinstance(v, type) and issubclass(v, _ut.TestCase) + and v is not _ut.TestCase]: + _suite.addTests(_loader.loadTestsFromTestCase(_c)) + with open(_os.devnull, "w") as _null: + _res = _ut.TextTestRunner(stream=_null, verbosity=0).run(_suite) + _bad = {{t.id() for t, _ in _res.failures}} | {{t.id() for t, _ in _res.errors}} + _total = _res.testsRun + print("{RESULT_MARKER}" + _json.dumps({{ + "total": _total, + "passed": _total - len(_bad), + "failed": sorted(_bad), + "first_error": (_res.failures + _res.errors)[0][1].strip().splitlines()[-1] + if (_res.failures or _res.errors) else None, + }})) +''' + + +def parse_test_report(stdout: str) -> Optional[Dict[str, Any]]: + """Read the JSON line emitted by UNITTEST_RUNNER, if the program got that far.""" + for line in reversed(stdout.splitlines()): + if line.startswith(RESULT_MARKER): + try: + return json.loads(line[len(RESULT_MARKER):]) + except ValueError: + return None + return None + + +def extract_code(response: str, entry_point: Optional[str] = None, + class_name: Optional[str] = None) -> str: + """Pull the candidate program out of a chat reply. + + Prefers a fenced block that actually defines the thing being asked for, + since models like to emit a usage example or a test in a second block. + """ + if not response: + return "" + + if class_name: + wanted = rf"^\s*class\s+{re.escape(class_name)}\s*[:\(]" + elif entry_point: + wanted = rf"^\s*(?:async\s+)?def\s+{re.escape(entry_point)}\s*\(" + else: + wanted = None + + blocks = [(lang.lower(), body) for lang, body in FENCE_RE.findall(response)] + if blocks: + py = [b for lang, b in blocks if lang in ("python", "py", "python3", "")] + candidates = py or [b for _, b in blocks] + if wanted: + defining = [b for b in candidates if re.search(wanted, b, re.M)] + if defining: + return defining[-1].strip("\n") + return candidates[0].strip("\n") + + return response.strip("\n") + + +def _import_preamble(prompt: str) -> str: + """The import lines from the original stub. + + A chat model often returns just the function and drops the imports that the + stub had above it; replaying them keeps that from being scored as a failure. + """ + lines = [] + for line in prompt.splitlines(): + stripped = line.strip() + if stripped.startswith(("import ", "from ")) and not line.startswith((" ", "\t")): + lines.append(stripped) + return "\n".join(lines) + + +def build_program(question: Dict, completion: str, seed: Optional[int] = 0) -> str: + """Assemble the full program that decides pass/fail for one problem.""" + entry_point = question["entry_point"] + code = completion + + if not re.search(rf"^\s*(?:async\s+)?def\s+{re.escape(entry_point)}\s*\(", code, re.M): + body = completion if completion.startswith((" ", "\t")) else \ + "\n".join(" " + ln for ln in completion.splitlines()) + code = question["prompt"] + "\n" + body + + prologue = f"import random; random.seed({seed})" if seed is not None else "" + + return "\n\n".join(part for part in [ + _import_preamble(question["prompt"]), + code, + question["test"], + prologue, + f"check({entry_point})", + ] if part) + + +def pass_at_k(n: int, c: int, k: int) -> float: + """Unbiased pass@k estimator from the Codex paper (arXiv:2107.03374). + + n samples drawn, c of them correct. Computed as 1 - C(n-c,k)/C(n,k) via a + product so it stays stable for large n. + """ + if n - c < k: + return 1.0 + prod = 1.0 + for i in range(n - c + 1, n + 1): + prod *= 1.0 - k / i + return 1.0 - prod + + +class HumanEvalDataset(BaseDataset): + """OpenAI HumanEval, 164 hand-written Python problems. + + Every problem in this suite imports only from the standard library, so the + execution venv needs no pip installs at all -- see `required_packages`. + """ + + HF_REPO = "openai/openai_humaneval" + + def __init__(self, source: Optional[str] = None, split: str = "test"): + self.source = source or self.HF_REPO + self.split = split + self.questions: List[Dict] = [] + self._load_dataset() + + def _load_dataset(self): + print(f"Loading HumanEval dataset (source: {self.source})...") + + local = Path(self.source) + if local.exists(): + rows = [json.loads(ln) for ln in local.read_text().splitlines() if ln.strip()] + else: + from datasets import load_dataset + rows = [dict(r) for r in load_dataset(self.source, split=self.split)] + + required = {"task_id", "prompt", "entry_point", "test"} + for row in rows: + missing = required - row.keys() + if missing: + raise ValueError( + f"{self.source}: record {row.get('task_id', '?')} is missing {sorted(missing)}" + ) + row["dataset_type"] = "humaneval" + self.questions.append(row) + + print(f"HumanEval dataset loaded: {len(self.questions)} problems") + + def required_packages(self) -> List[str]: + """Third-party pip packages this suite needs. Derived, not hardcoded.""" + return sorted(third_party_imports( + "\n".join(q["prompt"] + "\n" + q["test"] for q in self.questions) + )) + + def extract(self, response: str, question: Dict) -> str: + return extract_code(response, question.get("entry_point")) + + def build_program(self, question: Dict, completion: str, + seed: Optional[int] = None) -> str: + return build_program(question, completion, seed=seed) + + def get_question(self, index: int) -> Dict: + return self.questions[index] + + def get_question_text(self, question: Dict) -> str: + return question["prompt"] + + def get_answer(self, question: Dict) -> str: + return json.dumps({ + "task_id": question["task_id"], + "entry_point": question["entry_point"], + "prompt": question["prompt"], + "test": question["test"], + }) + + def get_prompt(self, question: Dict) -> str: + return TEMPLATE_REGISTRY["humaneval"].format( + question=self.get_question_text(question).rstrip(), + ) + + +class ClassEvalDataset(BaseDataset): + """ClassEval: 100 class-level Python generation tasks. + + A task is a whole class judged by several independent unittest classes, so + it is scored twice -- class level (every test passes) and method level + (what fraction of tests passed), matching the upstream protocol. + """ + + HF_REPO = "ilintar/ClassEval" + + def __init__(self, source: Optional[str] = None): + self.source = source or self.HF_REPO + self.questions: List[Dict] = [] + self._gensim_shim = False + self._load_dataset() + + def _load_dataset(self): + print(f"Loading ClassEval dataset (source: {self.source})...") + + local = Path(self.source) + if local.exists(): + text = local.read_text() + rows = (json.loads(text) if local.suffix == ".json" + else [json.loads(ln) for ln in text.splitlines() if ln.strip()]) + else: + from datasets import load_dataset + rows = [dict(r) for r in load_dataset(self.source, split="test")] + + required = {"task_id", "skeleton", "test", "class_name"} + for row in rows: + missing = required - row.keys() + if missing: + raise ValueError( + f"{self.source}: record {row.get('task_id', '?')} is missing {sorted(missing)}" + ) + row["dataset_type"] = "classeval" + row["n_tests"] = _count_test_methods(row["test"]) + self.questions.append(row) + + print(f"ClassEval dataset loaded: {len(self.questions)} classes, " + f"{sum(q['n_tests'] for q in self.questions)} test methods") + + EXTRA_PACKAGES = [ + "lxml", + "scipy", + ] + + def configure_python(self, python: str): + """gensim stops at cp313, so on 3.14+ swap it for a small local shim.""" + try: + r = subprocess.run( + [python, "-c", "import sys; print('%d.%d' % sys.version_info[:2])"], + capture_output=True, text=True, timeout=60, + ) + major, minor = (int(x) for x in r.stdout.strip().split(".")) + except Exception: + return + self._gensim_shim = (major, minor) >= (3, 14) + + def required_packages(self) -> List[str]: + blob = "\n\n".join( + "\n\n".join(_as_text(q.get(f)) for f in + ("skeleton", "test", "solution_code", "import_statement")) + for q in self.questions + ) + found = {PIP_NAMES.get(m, m) for m in third_party_imports(blob)} + found |= set(self.EXTRA_PACKAGES) + if self._gensim_shim: + found.discard("gensim") + return sorted(found) + + NLTK_CORPORA = ("punkt_tab", "averaged_perceptron_tagger_eng", "wordnet") + + def provision_steps(self) -> List[str]: + steps = [] + + if any("nltk" in _as_text(q.get("skeleton")) + _as_text(q.get("test")) + for q in self.questions): + steps.append( + "import os, nltk\n" + "target = os.path.join(os.environ['SANDBOX_VENV'], 'nltk_data')\n" + "os.makedirs(target, exist_ok=True)\n" + f"for corpus in {self.NLTK_CORPORA!r}:\n" + " nltk.download(corpus, download_dir=target, quiet=True)\n" + ) + + if self._gensim_shim: + steps.append( + "import os, sysconfig\n" + "pkg = os.path.join(sysconfig.get_paths()['purelib'], 'gensim')\n" + "os.makedirs(pkg, exist_ok=True)\n" + f"files = {GENSIM_SHIM!r}\n" + "for name, body in files.items():\n" + " open(os.path.join(pkg, name), 'w').write(body)\n" + "import gensim.utils, gensim.matutils\n" + ) + + return steps + + def sandbox_env(self) -> Dict[str, str]: + return {"NLTK_DATA": "{venv}/nltk_data"} + + def preferred_python(self) -> Optional[str]: + return "3.13" + + def default_exec_timeout(self) -> float: + return 30.0 + + def extract(self, response: str, question: Dict) -> str: + return extract_code(response, class_name=question.get("class_name")) + + def build_program(self, question: Dict, completion: str, + seed: Optional[int] = None) -> str: + prologue = f"import random; random.seed({seed})" if seed is not None else "" + return "\n\n".join(part for part in [ + _as_text(question.get("import_statement")), + UNITTEST_PROLOGUE, + completion, + question["test"], + prologue, + UNITTEST_RUNNER, + ] if part) + + def interpret(self, result: Any, question: Dict) -> Tuple[bool, Dict[str, Any]]: + expected = question.get("n_tests", 0) + + report = parse_test_report(result.stdout) + if report is None: + return False, { + "status": result.status if result.status != "ok" else "no-report", + "reason": result.summary(), + "tests_total": expected, + "tests_passed": 0, + "duration": round(result.duration, 3), + } + + total, passed = report.get("total", 0) or expected, report.get("passed", 0) + ok = total > 0 and passed == total + return ok, { + "status": "ok" if ok else "failed", + "reason": "ok" if ok else (report.get("first_error") + or f"{passed}/{total} tests passed")[:200], + "tests_total": total, + "tests_passed": passed, + "duration": round(result.duration, 3), + } + + def get_question(self, index: int) -> Dict: + return self.questions[index] + + def get_question_text(self, question: Dict) -> str: + return question["skeleton"] + + def get_answer(self, question: Dict) -> str: + return json.dumps({ + "task_id": question["task_id"], + "class_name": question["class_name"], + "import_statement": _as_text(question.get("import_statement")), + "test": question["test"], + "n_tests": question["n_tests"], + }) + + def get_prompt(self, question: Dict) -> str: + return TEMPLATE_REGISTRY["classeval"].format( + question=self.get_question_text(question).rstrip(), + ) + + +class AgenticDataset(BaseDataset): + """Multi-turn tool-driven repair tasks over a small real repository. + + The other code suites answer one prompt. Here the model is handed a tool + set and drives its own conversation, so the Processor delegates the whole + episode rather than making a single request. + + Everything language-specific is imported lazily and provisioned per + language, so selecting the Python tasks never installs a node toolchain -- + and running any other suite installs none of it. + """ + + HF_REPO = "ilintar/SACB" + is_agentic = True + ae: Any + + def __init__(self, source: Optional[str] = None, lang: Optional[str] = None, + max_turns: int = 50, context_limit: Optional[int] = None, + max_tokens: Optional[int] = None): + self.source = source or self.HF_REPO + self.lang_filter = lang + self.max_turns = max_turns + self.context_limit = context_limit + self.max_tokens = max_tokens + self._sandboxes: Dict[str, Any] = {} + self._load_dataset() + + def _load_dataset(self): + import agentic_eval as ae + self.ae = ae + + print(f"Loading agentic dataset (source: {self.source})...") + path = Path(self.source) + if path.exists(): + rows = [json.loads(line) for line in path.read_text().splitlines() + if line.strip()] + else: + from datasets import load_dataset + rows = [dict(r) for r in load_dataset(self.source, split="test")] + + if self.lang_filter: + rows = [r for r in rows if r.get("lang") == self.lang_filter] + if not rows: + raise ValueError( + f"no agentic tasks in {self.source}" + + (f" for language {self.lang_filter!r}" if self.lang_filter else "")) + + self.tasks = [ae.AgenticTask.from_record(r) for r in rows] + self.languages = sorted({t.lang for t in self.tasks}) + print(f"Agentic dataset loaded: {len(self.tasks)} tasks " + f"({', '.join(self.languages)})") + + def __len__(self) -> int: + return len(self.tasks) + + def get_question(self, index: int) -> Dict: + return {"index": index} + + def get_question_text(self, question: Dict) -> str: + task = self.tasks[question["index"]] + first = task.instruction.strip().splitlines()[0] + return f"[{task.task_id}] {first}" + + def get_answer(self, question: Dict) -> str: + task = self.tasks[question["index"]] + return json.dumps({"task_id": task.task_id, "index": question["index"]}) + + def get_prompt(self, question: Dict) -> str: + return self.ae.user_prompt(self.tasks[question["index"]]) + + def default_exec_timeout(self) -> float: + return 90.0 + + def required_packages(self) -> List[str]: + return [] + + def prepare(self, base_python: Optional[str] = None, force: bool = False): + """Provision one sandbox per language present in the selected tasks.""" + for lang in self.languages: + sandbox = self.ae.make_sandbox(lang, Sandbox, base_python=base_python) + sandbox.ensure(force=force) + self._sandboxes[lang] = sandbox + return self._sandboxes + + def run(self, index: int, chat) -> Dict[str, Any]: + task = self.tasks[index] + return self.ae.run_task( + task, chat, self._sandboxes[task.lang], + max_turns=self.max_turns, context_limit=self.context_limit, + test_timeout=self.default_exec_timeout()) + + +def _as_text(value: Any) -> str: + """Some dataset fields arrive as a list of lines rather than a blob.""" + if not value: + return "" + if isinstance(value, (list, tuple)): + return "\n".join(str(v) for v in value) + return str(value) + + +def _count_test_methods(test_code: str) -> int: + """How many `def test_*` a task's test suite defines.""" + try: + tree = ast.parse(test_code) + except SyntaxError: + return len(re.findall(r"^\s*def\s+test\w*\s*\(", test_code, re.M)) + return sum( + 1 + for node in ast.walk(tree) + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name.startswith("test") + ) + + +def third_party_imports(source: str) -> set: + """Top-level modules imported by `source` that are not in the stdlib.""" + import ast + + mods = set() + for chunk in source.split("\n\n"): + try: + tree = ast.parse(chunk) + except SyntaxError: + continue + for node in ast.walk(tree): + if isinstance(node, ast.Import): + mods.update(a.name.split(".")[0] for a in node.names) + elif isinstance(node, ast.ImportFrom) and node.module and node.level == 0: + mods.add(node.module.split(".")[0]) + + stdlib = set(sys.stdlib_module_names) + return {m for m in mods if m and m not in stdlib} + + class Grader: def __init__( self, @@ -981,14 +1747,29 @@ def __init__( grader_script: Optional[str] = None, grader_model_name: Optional[str] = None, grader_server_url: str = "", - dataset_type: str = "aime" + grader_api_key: Optional[str] = None, + dataset_type: str = "aime", + sandbox: Optional[Sandbox] = None, + exec_timeout: float = 15.0, + exec_seed: Optional[int] = 0, + dataset: Optional[BaseDataset] = None, ): self.grader_type = grader_type self.grader_script = grader_script self.grader_model_name = grader_model_name self.grader_server_url = grader_server_url + self.grader_api_key = grader_api_key self.dataset_type = dataset_type + self.sandbox = sandbox + self.exec_timeout = exec_timeout + self.exec_seed = exec_seed + self.dataset = dataset self.pattern = self._get_pattern() + self._local = threading.local() + + @property + def last_exec(self) -> Dict[str, Any]: + return getattr(self._local, "exec_info", {}) def _get_pattern(self) -> Optional[str]: if self.grader_type == "regex": @@ -1077,7 +1858,7 @@ def _grade_llm(self, gold: str, pred: str, problem: str) -> Tuple[bool, Optional Please provide only the extracted answer, nothing else. If there is no clear answer that can be extracted from the response, reply with 'no answer'.""" url = f"{self.grader_server_url}/v1/chat/completions" - headers = {"Content-Type": "application/json"} + headers = auth_headers(self.grader_api_key) data = { "model": self.grader_model_name, "messages": [ @@ -1102,6 +1883,28 @@ def _truncate_response(self, response: str, max_lines: int = 6) -> str: lines = response.split('\n') return '\n'.join(lines[-max_lines:]) if len(lines) > max_lines else response + def _grade_exec(self, gold: str, pred: str) -> Tuple[bool, Optional[str]]: + """Grade by running the generated code against the suite's own tests. + + `gold` is the JSON harness produced by the dataset's get_answer(); the + dataset owns how the program is assembled and how its result is read. + """ + if self.sandbox is None or self.dataset is None: + raise ValueError("exec grader requires a sandbox and a dataset") + + question = json.loads(gold) + completion = self.dataset.extract(pred, question) + if not completion.strip(): + self._local.exec_info = {"status": "empty", "reason": "no code in reply"} + return False, None + + program = self.dataset.build_program(question, completion, seed=self.exec_seed) + result = self.sandbox.run(program, timeout=self.exec_timeout) + passed, detail = self.dataset.interpret(result, question) + + self._local.exec_info = detail + return passed, completion + def grade(self, gold: str, pred: str, problem: str = "") -> Tuple[bool, Optional[str]]: """Grade the response""" if self.grader_type == "regex": @@ -1110,6 +1913,8 @@ def grade(self, gold: str, pred: str, problem: str = "") -> Tuple[bool, Optional return self._grade_cli(gold, pred) elif self.grader_type == "llm": return self._grade_llm(gold, pred, problem) + elif self.grader_type == "exec": + return self._grade_exec(gold, pred) else: raise ValueError(f"Unknown grader type: {self.grader_type}") @@ -1119,18 +1924,24 @@ def __init__( server_configs: List[ServerConfig], grader: Grader, model_name: Optional[str] = None, - n_predict: int = -1 + n_predict: int = -1, + retries: int = 5, + retry_delay: float = 10.0, + transcript_dir: Optional[Path] = None, ): self.server_configs = server_configs self.grader = grader self.model_name = model_name self.n_predict = n_predict + self.retries = retries + self.retry_delay = retry_delay + self.transcript_dir = transcript_dir @staticmethod def _check_server(server_config: ServerConfig) -> List[str]: url = f"{server_config.url}/v1/models" try: - response = requests.get(url) + response = requests.get(url, headers=server_config.headers()) response.raise_for_status() models = [m["id"] for m in response.json().get("data", [])] return models @@ -1139,10 +1950,11 @@ def _check_server(server_config: ServerConfig) -> List[str]: sys.exit(1) def _make_request( - self, server_config: ServerConfig, eval_state: EvalState, prompt: str + self, server_config: ServerConfig, eval_state: EvalState, prompt: str, + session_id: Optional[str] = None ) -> Tuple[Dict[str, Any], int, Optional[float], Optional[float], str]: url = f"{server_config.url}/v1/chat/completions" - headers = {"Content-Type": "application/json"} + headers = server_config.headers(session_id) data = { "model": self.model_name if self.model_name else "llama", "messages": [{"role": "user", "content": prompt}], @@ -1156,10 +1968,10 @@ def _make_request( data["top_p"] = eval_state.sampling_config["top_p"] if eval_state.sampling_config.get("min_p") is not None: data["min_p"] = eval_state.sampling_config["min_p"] + if eval_state.sampling_config.get("reasoning_effort") is not None: + data["reasoning_effort"] = eval_state.sampling_config["reasoning_effort"] - response = requests.post(url, headers=headers, json=data) - response.raise_for_status() - result = response.json() + result = post_with_retry(url, headers, data, self.retries, self.retry_delay) tokens = result.get("usage", {}).get("completion_tokens", 0) timings = result.get("timings", {}) tps_gen = timings.get("predicted_per_second") if timings else None @@ -1187,8 +1999,14 @@ def _process_single_case( problem_idx=problem_idx, ) + session_id = new_session_id(task_id) + + if eval_state.loaded_dataset.is_agentic: + return self._process_agentic_case( + server_config, eval_state, i, task_id, task_state, session_id) + try: - response, tokens, tps_gen, t_gen_ms, finish_reason = self._make_request(server_config, eval_state, prompt) + response, tokens, tps_gen, t_gen_ms, finish_reason = self._make_request(server_config, eval_state, prompt, session_id) result = response["choices"][0]["message"]["content"] reasoning_content = response["choices"][0].get("message", {}).get("reasoning_content") task_state.response = result @@ -1208,15 +2026,20 @@ def _process_single_case( eval_state.dump() return task_state - result_truncated = self.grader._truncate_response(result, max_lines=10) - is_correct, answer = self.grader.grade(expected, result_truncated, prompt) + if self.grader.grader_type == "exec": + graded_text = result + else: + graded_text = self.grader._truncate_response(result, max_lines=10) + is_correct, answer = self.grader.grade(expected, graded_text, prompt) grader_log = { - "pred": result_truncated, + "pred": graded_text, "grader_type": self.grader.grader_type } if self.grader.grader_type == "regex" and self.grader.pattern: grader_log["pattern"] = self.grader.pattern + if self.grader.grader_type == "exec": + grader_log.update(self.grader.last_exec) task_state.correct = is_correct task_state.answer = answer @@ -1232,11 +2055,136 @@ def _process_single_case( eval_state.dump() + except ServerError as e: + task_state.status = SERVER_ERROR + task_state.grader_log = {"error": str(e)} + eval_state.add_result( + task_id, prompt, expected, None, None, task_state.grader_log, + False, SERVER_ERROR, None, None, None, None, server_config.name, + chunk_idx, problem_idx, + ) + eval_state.dump() except Exception as e: task_state.status = f"error: {str(e)}" return task_state + def _agentic_chat(self, server_config: ServerConfig, eval_state: EvalState, + session_id: Optional[str] = None): + """A chat callable for agentic_eval: messages + tools in, response out.""" + url = f"{server_config.url}/v1/chat/completions" + + def chat(messages, tools): + data = { + "model": self.model_name if self.model_name else "llama", + "messages": messages, + "tools": tools, + "tool_choice": "auto", + } + per_turn = getattr(eval_state.dataset, "max_tokens", None) + if self.n_predict and self.n_predict > 0: + data["max_tokens"] = self.n_predict + elif per_turn: + data["max_tokens"] = per_turn + for key in ("temperature", "top_k", "top_p", "min_p", "reasoning_effort"): + value = eval_state.sampling_config.get(key) + if value is not None: + data[key] = value + return post_with_retry(url, server_config.headers(session_id), data, + self.retries, self.retry_delay, timeout=1800) + + return chat + + def _process_agentic_case( + self, server_config: ServerConfig, eval_state: EvalState, i: int, + task_id: str, task_state: TaskState, session_id: Optional[str] = None + ) -> TaskState: + """Run one whole tool-driven episode and grade the tree it leaves behind.""" + try: + result = eval_state.loaded_dataset.run( + i, self._agentic_chat(server_config, eval_state, session_id)) + except Exception as e: + task_state.status = f"error: {type(e).__name__}: {e}" + eval_state.add_result( + task_id, task_state.prompt, task_state.expected, + task_state.status, None, + {"error": f"{type(e).__name__}: {e}", "resolved": False}, + False, task_state.status, 0, None, None, None, + server_config.name, task_state.chunk_idx, i, + ) + eval_state.dump() + return task_state + + episode = result.get("episode", {}) + messages = result.pop("messages", None) + if messages is not None and self.transcript_dir is not None: + self._save_transcript(eval_state, task_id, messages, result) + + if episode.get("stop_reason") in ("request_failed", "no_choices"): + task_state.status = SERVER_ERROR + task_state.grader_log = result + task_state.response = f"server error on turn {episode.get('turns')}: {episode.get('error')}" + eval_state.add_result( + task_id, task_state.prompt, task_state.expected, task_state.response, + None, result, False, SERVER_ERROR, + episode.get("completion_tokens", 0), None, None, None, + server_config.name, task_state.chunk_idx, i, + ) + eval_state.dump() + return task_state + + task_state.tokens = episode.get("completion_tokens", 0) + task_state.correct = bool(result["resolved"]) + task_state.answer = ", ".join(result.get("changed_files", [])) or None + task_state.grader_log = result + task_state.status = "ok" + task_state.response = ( + f"stop={episode.get('stop_reason')} turns={episode.get('turns')} " + f"calls={episode.get('tool_calls')} edits={episode.get('edits')} " + f"errors={episode.get('tool_errors')} nudges={episode.get('nudges')} truncated={episode.get('truncated_turns')} " + f"peak_ctx={episode.get('peak_context')}\n" + f"changed: {task_state.answer or '(nothing)'}" + ) + + eval_state.add_result( + task_id, task_state.prompt, task_state.expected, task_state.response, + task_state.answer, result, task_state.correct, "ok", + task_state.tokens, None, None, None, server_config.name, + task_state.chunk_idx, i, + ) + eval_state.dump() + return task_state + + def _save_transcript(self, eval_state: EvalState, task_id: str, + messages: List[Dict[str, Any]], result: Dict[str, Any]): + """Write the episode's messages as .json plus a readable .html.""" + assert self.transcript_dir is not None + ep = result.get("episode", {}) + verdict = "RESOLVED" if result.get("resolved") else "not resolved" + meta = ( + f"{verdict} f2p {result.get('f2p_passed')}/{result.get('f2p_total')} " + f"p2p {result.get('p2p_passed')}/{result.get('p2p_total')}\n" + f"stop={ep.get('stop_reason')} turns={ep.get('turns')} " + f"calls={ep.get('tool_calls')} edits={ep.get('edits')} " + f"errors={ep.get('tool_errors')} nudges={ep.get('nudges')} truncated={ep.get('truncated_turns')} " + f"prompt_tokens={ep.get('prompt_tokens')} " + f"completion_tokens={ep.get('completion_tokens')} " + f"peak_ctx={ep.get('peak_context')}\n" + f"changed: {', '.join(result.get('changed_files') or []) or '(nothing)'}" + + (f"\nerror: {ep.get('error')}" if ep.get("error") else "") + ) + title = f"{task_id} ({result.get('task_id')})" + dataset = eval_state.loaded_dataset + assert isinstance(dataset, AgenticDataset) + try: + path = dataset.ae.write_transcript( + messages, self.transcript_dir / task_id, title, meta) + except OSError as e: + print(f"Warning: could not write transcript for {task_id}: {e}") + return + base = Path(eval_state.output_file).resolve().parent + result["transcript"] = os.path.relpath(path.resolve(), base) + @staticmethod def _worker( server_config: ServerConfig, @@ -1275,7 +2223,7 @@ def evaluate(self, eval_state: EvalState, verbose: bool = False, resume: bool = print(f" {i+1}. {sc.name} — {sc.url} ({sc.threads} threads) [{models_str}]") print(f"Model: {self.model_name}") print(f"Grader: {self.grader.grader_type}") - print(f"Sampling: temp={eval_state.sampling_config.get('temperature', 'skip')}, top-k={eval_state.sampling_config.get('top_k', 'skip')}, top-p={eval_state.sampling_config.get('top_p', 'skip')}, min-p={eval_state.sampling_config.get('min_p', 'skip')}") + print(f"Sampling: temp={eval_state.sampling_config.get('temperature', 'skip')}, top-k={eval_state.sampling_config.get('top_k', 'skip')}, top-p={eval_state.sampling_config.get('top_p', 'skip')}, min-p={eval_state.sampling_config.get('min_p', 'skip')}, reasoning-effort={eval_state.sampling_config.get('reasoning_effort', 'skip')}") print() # Shared task queue: all workers compete for tasks @@ -1342,6 +2290,79 @@ def evaluate(self, eval_state: EvalState, verbose: bool = False, resume: bool = eval_state.print_summary() eval_state.dump() +def preflight_sandbox(sandbox: Sandbox): + """Catch host policies that would fail problems for non-model reasons. + + A FIPS-enforcing OpenSSL makes hashlib.md5() raise, which would sink + HumanEval/162 no matter what the model generated. + """ + probe = sandbox.run("import hashlib; hashlib.md5(b'x').hexdigest()", timeout=15) + if not probe.passed: + print(f"Warning: md5 is unavailable in the execution venv " + f"({probe.summary(120)}). Problems that hash with md5 will fail " + f"regardless of the model -- this is a host policy, not a model result.") + + +def print_pass_at_k(eval_state: "EvalState"): + """Report pass@k plus a breakdown of how the failures failed. + + Each chunk is one full pass over the suite, so running --n_cases at a + multiple of the dataset size gives n samples per problem for free. + """ + cases = eval_state.task_states.get("cases", {}) + if not cases: + return + + by_problem: Dict[int, List[bool]] = {} + reasons: Dict[str, int] = {} + for case in cases.values(): + by_problem.setdefault(case.get("problem_idx", 0), []).append( + bool(case.get("correct", False)) + ) + if not case.get("correct", False): + status = (case.get("grader_log") or {}).get("status", case.get("status", "?")) + reasons[status] = reasons.get(status, 0) + 1 + + n = min(len(v) for v in by_problem.values()) + n_problems = len(by_problem) + + print(f"\n{'=' * 60}") + print(f" {eval_state.dataset_type}: {n_problems} problems, {n} sample(s) each") + print(f"{'=' * 60}") + + for k in sorted({1, 5, 10, 100, n}): + if k > n: + continue + score = sum( + pass_at_k(len(v), sum(v), k) for v in by_problem.values() + ) / n_problems + print(f" pass@{k:<4d} {score:.4f}") + + if n == 1: + solved = sum(1 for v in by_problem.values() if v[0]) + lo, hi = wilson_interval(solved, n_problems) + print(f" {solved}/{n_problems} solved 95% CI [{lo:.4f}, {hi:.4f}]") + + t_total = t_passed = 0 + for case in cases.values(): + log = case.get("grader_log") or {} + if "tests_total" in log: + t_total += log["tests_total"] + else: + try: + t_total += json.loads(case.get("expected") or "{}").get("n_tests", 0) + except ValueError: + pass + t_passed += log.get("tests_passed", 0) + if t_total: + print(f"\n test methods passed {t_passed}/{t_total} ({t_passed / t_total:.4f})") + + if reasons: + print("\n failures by cause:") + for status, count in sorted(reasons.items(), key=lambda kv: -kv[1]): + print(f" {status:<12s} {count}") + + def main(): parser = argparse.ArgumentParser( description="Simplified evaluation tool for llama.cpp" @@ -1358,13 +2379,28 @@ def main(): default="", help="Comma-separated display names for servers (default: use URLs)" ) + parser.add_argument( + "--api-key", + type=str, + default=None, + help="Bearer token for the eval servers: one key for all of them, or a " + "comma-separated key per --server (default: $LLAMA_API_KEY, else none)" + ) parser.add_argument( "--dataset", type=str, default="aime", - choices=["aime", "aime2025", "aime2026", "gsm8k", "gpqa"], + choices=["aime", "aime2025", "aime2026", "gsm8k", "gpqa", "humaneval", + "classeval", "agentic"], help="Dataset type (default: aime)" ) + parser.add_argument( + "--dataset-source", + type=str, + default=None, + help="Override the dataset location: a local .jsonl path or a HuggingFace " + "repo id (e.g. a patched HumanEval). Default: the suite's canonical repo" + ) parser.add_argument( "--n_cases", type=int, @@ -1407,6 +2443,25 @@ def main(): default=None, help="Min P sampling (default: not passed)" ) + parser.add_argument( + "--retries", + type=int, + default=5, + help="Times to resend a request that errors or comes back empty before " + "recording the case as a server error (default: 5)" + ) + parser.add_argument( + "--retry-delay", + type=float, + default=10.0, + help="Seconds to wait before the first resend; doubles on each further one (default: 10)" + ) + parser.add_argument( + "--reasoning-effort", + type=str, + default=None, + help="Reasoning effort sent as reasoning_effort with every request, e.g. low/medium/high (default: not passed)" + ) parser.add_argument( "--threads", type=str, @@ -1433,9 +2488,37 @@ def main(): parser.add_argument( "--grader-type", type=str, - default="llm", - choices=["regex", "cli", "llm"], - help="Grader type: regex, cli, or llm (default: llm)" + default=None, + choices=["regex", "cli", "llm", "exec"], + help="Grader type: regex, cli, llm, or exec (default: exec for code suites, " + "llm otherwise). 'exec' runs the generated code against the suite tests " + "in a sandboxed venv" + ) + parser.add_argument( + "--exec-timeout", + type=float, + default=None, + help="Per-problem wall-clock limit for the exec grader, in seconds " + "(default: per-suite, 15s for humaneval, 30s for classeval)" + ) + parser.add_argument( + "--exec-seed", + type=int, + default=0, + help="Seed the global RNG before running each test, so suites whose tests " + "draw random inputs score reproducibly (default: 0). Use -1 to leave " + "the RNG unseeded, matching the stock upstream harness" + ) + parser.add_argument( + "--sandbox-python", + type=str, + default=None, + help="Base interpreter used to build the execution venv (default: this one)" + ) + parser.add_argument( + "--rebuild-sandbox", + action="store_true", + help="Discard and re-provision the cached execution venv" ) parser.add_argument( "--grader-script", @@ -1455,6 +2538,51 @@ def main(): default="", help="Model name for LLM grader (default: same as main model)" ) + parser.add_argument( + "--grader-api-key", + type=str, + default=None, + help="Bearer token for the LLM grader server (default: the first " + "--api-key when the grader runs on the main server, else none)" + ) + parser.add_argument( + "--agentic-lang", + type=str, + default=None, + choices=["python", "typescript"], + help="Restrict the agentic suite to one language. Only that language's " + "toolchain is then provisioned." + ) + parser.add_argument( + "--agentic-max-turns", + type=int, + default=50, + help="Maximum assistant turns per agentic episode (default: 50)" + ) + parser.add_argument( + "--agentic-max-tokens", + type=int, + default=None, + help="Cap on generated tokens per agentic turn (default: no cap). " + "Bounds a model that falls into a repetition loop, but also cuts " + "off long reasoning; truncated turns are counted in the results." + ) + parser.add_argument( + "--agentic-context-limit", + type=int, + default=None, + help="Abandon an agentic episode once its context exceeds this many " + "tokens (default: no limit)" + ) + parser.add_argument( + "--agentic-transcripts", + type=Path, + default=None, + metavar="DIR", + help="Save each agentic episode's full message history to " + "DIR/.json, with a readable DIR/.html showing " + "every reasoning block, reply, tool call and tool result" + ) parser.add_argument( "--resume", action="store_true", @@ -1463,6 +2591,23 @@ def main(): args = parser.parse_args() + global DATASET_SOURCE, AGENTIC_OPTS + DATASET_SOURCE = args.dataset_source + AGENTIC_OPTS = { + "lang": args.agentic_lang, + "max_turns": args.agentic_max_turns, + "context_limit": args.agentic_context_limit, + "max_tokens": args.agentic_max_tokens, + } + + if args.grader_type is None: + args.grader_type = "exec" if args.dataset in CODE_SUITES else "llm" + + if args.grader_type == "exec" and args.dataset not in CODE_SUITES: + print(f"Error: --grader-type exec is only meaningful for code suites " + f"({', '.join(sorted(CODE_SUITES))}), not '{args.dataset}'") + sys.exit(1) + # Parse server URLs and thread counts server_urls = [u.strip() for u in args.server.split(",") if u.strip()] thread_counts = [int(t.strip()) for t in args.threads.split(",") if t.strip()] @@ -1480,11 +2625,30 @@ def main(): else: server_names = server_urls # fallback to URLs + api_key_arg = args.api_key if args.api_key is not None else os.environ.get("LLAMA_API_KEY") + if api_key_arg: + api_keys = [k.strip() for k in api_key_arg.split(",")] + if len(api_keys) == 1: + api_keys = api_keys * len(server_urls) + elif len(api_keys) != len(server_urls): + print(f"Error: --api-key ({len(api_keys)} keys) must be a single key or " + f"match --server ({len(server_urls)} URLs)") + sys.exit(1) + else: + api_keys = [None] * len(server_urls) + server_configs = [ - ServerConfig(url=url, threads=threads, name=name) - for url, threads, name in zip(server_urls, thread_counts, server_names) + ServerConfig(url=url, threads=threads, name=name, api_key=key) + for url, threads, name, key in zip(server_urls, thread_counts, server_names, api_keys) ] + if args.grader_api_key is not None: + grader_api_key = args.grader_api_key + else: + grader_url = args.grader_server or server_configs[0].url + grader_api_key = next( + (s.api_key for s in server_configs if s.url == grader_url), None) + if args.dataset == "gpqa" and args.grader_type != "llm": print("Error: GPQA dataset requires --grader-type llm") parser.print_help() @@ -1529,6 +2693,7 @@ def main(): grader_script=args.grader_script, grader_model_name=grader_model_name, grader_server_url=grader_server_url, + grader_api_key=grader_api_key, dataset_type=eval_state.dataset_type ) resume = True @@ -1548,6 +2713,7 @@ def main(): grader_script=args.grader_script, grader_model_name=grader_model_name, grader_server_url=grader_server_url, + grader_api_key=grader_api_key, dataset_type=args.dataset ) @@ -1563,6 +2729,8 @@ def main(): sampling_config["top_p"] = args.top_p if args.min_p is not None: sampling_config["min_p"] = args.min_p + if args.reasoning_effort is not None: + sampling_config["reasoning_effort"] = args.reasoning_effort eval_state = EvalState( dataset_type=args.dataset, @@ -1577,14 +2745,53 @@ def main(): eval_state.print_all_tasks() + if args.grader_type == "exec": + dataset = eval_state.loaded_dataset + + if dataset.is_agentic: + dataset.prepare(base_python=args.sandbox_python, + force=args.rebuild_sandbox) + else: + base_python = args.sandbox_python + want = dataset.preferred_python() + if not base_python and want: + base_python = resolve_python(want) + if base_python: + print(f"[sandbox] using Python {want} for " + f"{eval_state.dataset_type} ({base_python})") + + dataset.configure_python(base_python or sys.executable) + + sandbox = Sandbox( + name=eval_state.dataset_type, + packages=dataset.required_packages(), + provision=dataset.provision_steps(), + env=dataset.sandbox_env(), + base_python=base_python, + ) + sandbox.ensure(force=args.rebuild_sandbox) + preflight_sandbox(sandbox) + grader.sandbox = sandbox + grader.dataset = dataset + grader.exec_timeout = (args.exec_timeout if args.exec_timeout is not None + else dataset.default_exec_timeout()) + grader.exec_seed = None if args.exec_seed < 0 else args.exec_seed + processor = Processor( server_configs=server_configs, grader=grader, model_name=args.model, - n_predict=args.n_predict + n_predict=args.n_predict, + retries=args.retries, + retry_delay=args.retry_delay, + transcript_dir=args.agentic_transcripts, ) processor.evaluate(eval_state, verbose=args.verbose, resume=resume) + + if args.grader_type == "exec": + print_pass_at_k(eval_state) + print(f"\nEval state dumped to {args.output}") if __name__ == "__main__": diff --git a/examples/llama-eval/llama-server-simulator.py b/examples/llama-eval/llama-server-simulator.py index e64ba89335d9..dbd46d67f476 100755 --- a/examples/llama-eval/llama-server-simulator.py +++ b/examples/llama-eval/llama-server-simulator.py @@ -13,8 +13,6 @@ from dataclasses import dataclass from pathlib import Path -import datasets - # Set cache directory for HuggingFace datasets cache_dir = Path.home() / ".cache" / "huggingface" / "datasets" cache_dir.mkdir(parents=True, exist_ok=True) @@ -76,6 +74,8 @@ def _get_question_text(self, question: Dict) -> str: return question.get("problem", question.get("question", "")) def _load_dataset(self): + import datasets + if self.dataset_type == "aime": print(f"Loading AIME dataset (split: {self.split})...") cache_path = Path.home() / ".cache" / "huggingface" / "datasets" / "AI-MO___aimo-validation-aime" / "default" / "0.0.0" @@ -164,6 +164,58 @@ def get_answer(self, question: Dict) -> str: return str(normalized) if normalized is not None else answer return str(answer) +class CodeDataset: + """Stand-in for a code model: replays canonical solutions as chat replies. + + Handles both code suites: HumanEval keys off the function stub, ClassEval + off the class skeleton. + """ + + DEFAULT_SOURCE = { + "humaneval": "openai/openai_humaneval", + "classeval": "FudanSELab/ClassEval", + } + + def __init__(self, dataset_type: str, source: Optional[str] = None): + self.dataset_type = dataset_type + self.source = source or self.DEFAULT_SOURCE[dataset_type] + self.questions: List[Dict] = [] + self._load_dataset() + + def _load_dataset(self): + print(f"Loading {self.dataset_type} dataset (source: {self.source})...") + local = Path(self.source) + if local.exists(): + text = local.read_text() + if local.suffix == ".json": + rows = json.loads(text) + else: + rows = [json.loads(ln) for ln in text.splitlines() if ln.strip()] + else: + import datasets + rows = [dict(r) for r in datasets.load_dataset(self.source, split="test")] + self.questions = rows + print(f"{self.dataset_type} dataset loaded: {len(self.questions)} questions") + + def _stub_field(self) -> str: + return "skeleton" if self.dataset_type == "classeval" else "prompt" + + def _get_question_text(self, question: Dict) -> str: + return question[self._stub_field()] + + def find_question(self, request_text: str) -> Optional[Dict]: + for question in self.questions: + stub = question[self._stub_field()].rstrip() + if stub and stub in request_text: + return question + return None + + def get_answer(self, question: Dict) -> str: + if self.dataset_type == "classeval": + return question["solution_code"] + return question["prompt"] + question["canonical_solution"] + + class Simulator: def __init__( self, @@ -171,12 +223,17 @@ def __init__( host: str = "localhost", success_rate: float = 0.8, dataset_split: str = "train", - dataset_type: str = "aime" + dataset_type: str = "aime", + dataset_source: Optional[str] = None ): self.port = port self.host = host self.success_rate = success_rate - self.dataset = AimeDataset(dataset_split, dataset_type) + self.dataset_type = dataset_type + if dataset_type in ("humaneval", "classeval"): + self.dataset = CodeDataset(dataset_type, dataset_source) + else: + self.dataset = AimeDataset(dataset_split, dataset_type) self.eval_state = EvalState( id=dataset_type, tasks=[dataset_type], @@ -196,6 +253,12 @@ def _generate_response( else: response_text = self._generate_wrong_answer(question) + if self.dataset_type in ("humaneval", "classeval"): + response_text = ( + f"Here is the implementation:\n\n```python\n{response_text}\n```\n\n" + "Let me know if you want it optimised further." + ) + comp_tokens = random.randint(10000, 60000) tps_gen = random.uniform(90.0, 110.0) t_gen_ms = comp_tokens / tps_gen * 1000 @@ -229,6 +292,12 @@ def _generate_response( def _generate_wrong_answer(self, question: Dict) -> str: expected_answer = self.dataset.get_answer(question) + if self.dataset_type == "classeval": + return re.sub(r"(\n\s+def [^\n]+:\n)", r"\1 return None\n", + question["skeleton"]) + if self.dataset_type == "humaneval": + return question["prompt"] + " return None\n" + if expected_answer.isdigit(): wrong_answer = str(int(expected_answer) + 1) else: @@ -335,7 +404,7 @@ def main(): "--dataset", type=str, default="aime", - choices=["aime", "aime2025"], + choices=["aime", "aime2025", "humaneval", "classeval"], help="Dataset type (default: aime)" ) parser.add_argument( @@ -344,6 +413,12 @@ def main(): default="train", help="AIME dataset split to use (default: train)" ) + parser.add_argument( + "--dataset-source", + type=str, + default=None, + help="Override the dataset location (local .jsonl path or HuggingFace repo id)" + ) args = parser.parse_args() @@ -353,7 +428,8 @@ def main(): host=args.host, success_rate=args.success_rate, dataset_split=args.dataset_split, - dataset_type=args.dataset + dataset_type=args.dataset, + dataset_source=args.dataset_source ) server = HTTPServer((args.host, args.port), RequestHandler) From 2a667032dc788c3a071d538e0b572e45a3f12731 Mon Sep 17 00:00:00 2001 From: aiqubits Date: Fri, 16 Oct 2026 21:17:30 +0800 Subject: [PATCH 059/423] convert : respect Mistral format for DeepseekV2Model (#30669) Use the instance's is_mistral_format flag when DeepseekV2Model reloads configuration, allowing Mistral MoE checkpoints to read params.json. Fixes #29899 Signed-off-by: aiqubits --- conversion/deepseek.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/conversion/deepseek.py b/conversion/deepseek.py index 817eb76128b6..946de53476ea 100644 --- a/conversion/deepseek.py +++ b/conversion/deepseek.py @@ -243,7 +243,7 @@ class DeepseekV2Model(TextModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - hparams: dict = ModelBase.load_hparams(self.dir_model, is_mistral_format=False) + hparams: dict = ModelBase.load_hparams(self.dir_model, is_mistral_format=self.is_mistral_format) self.origin_hf_arch = hparams.get('architectures', [None])[0] # special handling for Deepseek OCR From 7182ef71e6d5dd0c8ea76486fac134fcd25025d4 Mon Sep 17 00:00:00 2001 From: texasich <101962694+texasich@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:01:37 -0500 Subject: [PATCH 060/423] common: honor HTTP_PROXY/HTTPS_PROXY env vars in http client (#21752) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * common: honor HTTP_PROXY/HTTPS_PROXY env vars in http client common_http_client() was ignoring standard proxy env vars, causing downloads via -hf and similar flags to fail silently in environments behind an HTTP proxy (corporate networks, air-gapped infra, etc.). Checks HTTPS_PROXY/https_proxy for https:// URLs and HTTP_PROXY/http_proxy for http:// URLs, in that order. If a proxy URL is set and parseable, passes it to httplib's set_proxy() and set_proxy_basic_auth() as appropriate. Malformed proxy URLs are silently ignored so existing callers aren't affected. Fixes #21694 * common: drop misleading comment about generic proxy fallback The comment referenced a 'generic one' that doesn't exist in the code — both HTTP_PROXY and HTTPS_PROXY are scheme-specific. Reviewer noted the mismatch. Addresses review feedback from @ServeurpersoCom. * common : warn when the proxy URL from the environment is malformed * common : honor NO_PROXY when selecting the http proxy * ci : retrigger checks (test-thread-safety flake) * common : bypass the proxy for loopback destinations llama cli talks to its spawned server on 127.0.0.1 through common_http_client, so with HTTP_PROXY set and no NO_PROXY the health polling went to the corporate proxy, which rejects loopback targets, and the wait loop never exits ("Loading model..." hangs forever). A remote proxy can never reach the client machine's own loopback, so these destinations are now always contacted directly, before NO_PROXY is consulted: localhost, *.localhost, 127.0.0.0/8, ::1 and 0.0.0.0. Similar to Go's net/http. Proxy handling for downloads and remote --server-base endpoints is unchanged. Signed-off-by: texasich <101962694+texasich@users.noreply.github.com> * docs: move proxy tip from README to server docs per review: the top level README stays lean; the proxy howto now lives with the server docs under Usage Signed-off-by: texasich <101962694+texasich@users.noreply.github.com> --------- Signed-off-by: texasich <101962694+texasich@users.noreply.github.com> Co-authored-by: texasich Co-authored-by: texasich Co-authored-by: marecabo <23156476+marecabo@users.noreply.github.com> --- common/http.h | 141 +++++++++++++++++++++++++++++++++++++++++ tools/server/README.md | 4 ++ 2 files changed, 145 insertions(+) diff --git a/common/http.h b/common/http.h index 878ad1ce28a4..10d017f17ed5 100644 --- a/common/http.h +++ b/common/http.h @@ -1,5 +1,7 @@ #pragma once +#include "log.h" + #include #ifdef _WIN32 @@ -97,6 +99,118 @@ static common_http_url common_http_parse_url(const std::string & url) { return parts; } +// match a loopback destination: a remote proxy can never reach the client machine's own +// loopback, so these hosts are always contacted directly, before NO_PROXY is consulted +// (similar to Go's net/http) +static bool common_http_is_loopback(const std::string & host) { + auto lower = [](std::string s) { + for (auto & c : s) { + if (c >= 'A' && c <= 'Z') { + c += 'a' - 'A'; + } + } + return s; + }; + + const std::string h = lower(host); + + if (h == "localhost" || h == "::1" || h == "0.0.0.0") { + return true; + } + + if (h.size() > 10 && h.compare(h.size() - 10, 10, ".localhost") == 0) { + return true; + } + + // IPv4 literal: exactly four all-digit octets with values 0-255, loopback is 127.0.0.0/8 + int octets[4] = { 0, 0, 0, 0 }; + size_t pos = 0; + int n = 0; + bool complete = false; + while (n < 4 && pos <= h.size()) { + size_t dot = h.find('.', pos); + if (dot == std::string::npos) { + dot = h.size(); + } + if (dot == pos || dot - pos > 3) { + return false; + } + int val = 0; + for (size_t i = pos; i < dot; i++) { + if (h[i] < '0' || h[i] > '9') { + return false; + } + val = val * 10 + (h[i] - '0'); + } + if (val > 255) { + return false; + } + octets[n++] = val; + if (dot == h.size()) { + complete = (n == 4); + break; + } + pos = dot + 1; + } + + return complete && octets[0] == 127; +} + +// match a host against a NO_PROXY list: comma separated entries, "*" bypasses the proxy for every host, +// an entry matches the host itself or any of its subdomains, a leading dot is optional, ports are ignored +static bool common_http_proxy_bypass(const std::string & no_proxy, const std::string & host) { + auto lower = [](std::string s) { + for (auto & c : s) { + if (c >= 'A' && c <= 'Z') { + c += 'a' - 'A'; + } + } + return s; + }; + + const std::string h = lower(host); + + for (size_t pos = 0; pos < no_proxy.size(); ) { + size_t end = no_proxy.find(',', pos); + if (end == std::string::npos) { + end = no_proxy.size(); + } + + std::string entry = lower(no_proxy.substr(pos, end - pos)); + pos = end + 1; + + entry.erase(0, entry.find_first_not_of(" \t")); + auto last = entry.find_last_not_of(" \t"); + if (last == std::string::npos) { + continue; + } + entry.erase(last + 1); + + if (entry == "*") { + return true; + } + + if (entry.front() == '.') { + entry.erase(0, 1); + } + + if (entry.empty()) { + continue; + } + + if (h == entry) { + return true; + } + + if (h.size() > entry.size() && h[h.size() - entry.size() - 1] == '.' && + h.compare(h.size() - entry.size(), entry.size(), entry) == 0) { + return true; + } + } + + return false; +} + static std::pair common_http_client(const std::string & url) { common_http_url parts = common_http_parse_url(url); @@ -123,6 +237,33 @@ static std::pair common_http_client(const std: cli.set_follow_location(true); + // Honor HTTP_PROXY / HTTPS_PROXY environment variables (both lower- and upper-case). + auto getenv_s = [](const char * upper, const char * lower) -> std::string { + const char * val = std::getenv(upper); + if (!val) { + val = std::getenv(lower); + } + return val ? val : ""; + }; + + const std::string proxy_url = (parts.scheme == "https") + ? getenv_s("HTTPS_PROXY", "https_proxy") + : getenv_s("HTTP_PROXY", "http_proxy"); + + if (!proxy_url.empty() && !common_http_is_loopback(parts.host) && + !common_http_proxy_bypass(getenv_s("NO_PROXY", "no_proxy"), parts.host)) { + try { + common_http_url proxy = common_http_parse_url(proxy_url); + cli.set_proxy(proxy.host, proxy.port); + if (!proxy.user.empty()) { + cli.set_proxy_basic_auth(proxy.user, proxy.password); + } + } catch (const std::exception & e) { + // fall back to a direct connection, the URL itself is not logged as it can carry credentials + LOG_WRN("%s: ignoring malformed proxy URL from the environment: %s\n", __func__, e.what()); + } + } + return { std::move(cli), std::move(parts) }; } diff --git a/tools/server/README.md b/tools/server/README.md index 4e4472ec24bd..7d697633edac 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -405,6 +405,10 @@ Recommended `--cors-origins` setting, depending on where the server runs: Related flags: `--cors-origins`, `--cors-methods`, `--cors-headers`, `--cors-credentials` / `--no-cors-credentials`. Background and rationale: [#25655](https://github.com/ggml-org/llama.cpp/pull/25655). +### Running behind a proxy + +If your host reaches the internet through a corporate proxy, export `HTTP_PROXY` and `HTTPS_PROXY` before running commands that download models (e.g. `export HTTPS_PROXY=http://proxy.example.com:8080`). Hosts listed in `NO_PROXY` and loopback addresses are contacted directly. + ## Build `llama-server` is built alongside everything else from the root of the project From 23f7aabc16c72f6f707dba227251c35d7bd1ed96 Mon Sep 17 00:00:00 2001 From: Gabe Goodhart Date: Fri, 16 Oct 2026 12:00:41 -0600 Subject: [PATCH 061/423] model : Granite Switch Decision (#30676) * feat(convert): Add decision.chat to support decision models that can also chat Branch: GraniteSwitchDecision AI-usage: full (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart * feat(convert): Add conversion for Granite Switch as decision model The template follows from the usage sample: https://huggingface.co/ibm-granite/granite-switch-4.1-8b#guardian-and-decision-side-by-side Branch: GraniteSwitchDecision AI-usage: draft (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart * feat(common): Add support for chat+decision models in common Branch: GraniteSwitchDecision AI-usage: full (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart * feat(common): Add granite switch decision model enum Branch: GraniteSwitchDecision AI-usage: full (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart * feat(server): Plumb through decision_chat bool to support decision+chat models Branch: GraniteSwitchDecision AI-usage: full (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart * feat(server): Add Granite Switch decision support Branch: GraniteSwitchDecision AI-usage: draft (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart * docs: Update server README for granite switch decisions / chat+decisions Branch: GraniteSwitchDecision AI-usage: draft (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart * Revert "feat(convert): Add decision.chat to support decision models that can also chat" This reverts commit 952ef567647e53ac9beda7a1f79d6927240d068d. * Revert "feat(common): Add support for chat+decision models in common" This reverts commit 93ed782e3a8b2fd5968732c6525c7ad29ad4bea2. * Revert "feat(server): Plumb through decision_chat bool to support decision+chat models" This reverts commit 97e9af60ddaccbd82f1eb0a488adff43b16ee095. * fix: If Granite Switch supports decisions, it always supports text too Branch: GraniteSwitchDecision AI-usage: none Signed-off-by: Gabe Goodhart * docs: Remove reference to decision.chat in server README Branch: GraniteSwitchDecision AI-usage: none Signed-off-by: Gabe Goodhart --- common/common.cpp | 1 + common/common.h | 1 + conversion/granite.py | 31 ++++++++++++++++++++++++++++++- gguf-py/gguf/constants.py | 1 + tools/server/README.md | 4 ++-- tools/server/server-common.cpp | 2 ++ tools/server/server-decision.cpp | 4 ++-- tools/server/server-decision.h | 1 + 8 files changed, 40 insertions(+), 5 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index c5db228329f5..b59e66e3bd3f 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1174,6 +1174,7 @@ static const std::map COMMON_DECISION_TYPE_NA { COMMON_DECISION_TYPE_VLLM_D3, "vllm-d3" }, { COMMON_DECISION_TYPE_LFM2_D1, "lfm2-d1" }, { COMMON_DECISION_TYPE_LFM2_D1_OMNI, "lfm2-d1-omni" }, + { COMMON_DECISION_TYPE_GRANITE_SWITCH, "granite-switch"}, // for future models, ALWAYS add org name prefix to avoid name collisions, e.g. "orgname-modelname" }; diff --git a/common/common.h b/common/common.h index 2cd6e2ec9858..0c2d5e46f105 100644 --- a/common/common.h +++ b/common/common.h @@ -970,6 +970,7 @@ enum common_decision_type { COMMON_DECISION_TYPE_VLLM_D3, // same as pplx-decider, non-causal attention COMMON_DECISION_TYPE_LFM2_D1, // same as openjev, the labels depend on the question type COMMON_DECISION_TYPE_LFM2_D1_OMNI, // same as laya, other prompt layout + COMMON_DECISION_TYPE_GRANITE_SWITCH, // same as openjev, run by the decision aLoRA of granite switch COMMON_DECISION_TYPE_UNKNOWN, // a decision model of a type that is not supported }; diff --git a/conversion/granite.py b/conversion/granite.py index 4cbeb7617892..d2ef7c836c8d 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -8,7 +8,7 @@ if TYPE_CHECKING: from torch import Tensor -from .base import MmprojModel, ModelBase, gguf, logger +from .base import MmprojModel, ModelBase, gguf, jinja_str_or_json, logger from .llama import LlamaModel from .mamba import Mamba2Model @@ -284,10 +284,39 @@ def set_gguf_parameters(self): self.gguf_writer.add_adapter_lora_rank(self._max_lora_rank) self.gguf_writer.add_adapter_token_ids_activate(self.hparams["adapter_token_ids"]) self.gguf_writer.add_adapter_token_ids_substitute(self.hparams["adapter_substitute_token_ids"]) + if self._has_decision_adapter(): + self.gguf_writer.add_decision_type(gguf.DecisionType.GRANITE_SWITCH) router_gain = float(self.hparams.get("control_token_gain", 15.0)) self.gguf_writer.add_adapter_router_gain(router_gain) logger.info("gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s router_gain=%s", self._n_adapters, self._max_lora_rank, self._n_slots, router_gain) + def _has_decision_adapter(self) -> bool: + return "decision" in self.hparams.get("adapter_names", []) + + def set_vocab(self): + super().set_vocab() + if self._has_decision_adapter(): + self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}]) + + def _systemone_template(self) -> str: + # prompt of the decision aLoRA: a state turn, then a question turn that opens with the control token + # (it replaces the first token of "### DECISION"), the answer is the label token that comes next + description = jinja_str_or_json("o.description") + option = ( + "{% if type != 'noul' %}{{ letters[loop.index0] }}: {% if o.description %}" + description + "{% endif %}" + "{% elif o.key == 'true' %}A: {% if o.description %}" + description + "{% else %}The statement is true.{% endif %}" + "{% else %}B: {% if o.description %}" + description + "{% else %}The statement is false.{% endif %}{% endif %}" + ) + return ( + "{% set letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz' %}" + "<|start_of_role|>user<|end_of_role|>State:\n" + jinja_str_or_json("state") + "<|end_of_text|>\n" + "<|start_of_role|>user<|end_of_role|><|decision|> DECISION\nQuestion:\n" + jinja_str_or_json("instructions") + + "{% if type == 'score' %} Rate along the ordered levels below (lowest first).{% endif %}" + "{{ '\n\nOptions:\n' }}" + "{% for o in options %}" + option + "{{ '\n' }}{% endfor %}" + "{{ '\nReply with exactly one option label.<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>' }}" + ) + def _lora_a(self, data: Tensor) -> Tensor: # on-disk A: [n_adapters, 1, max_rank, in] -> [n_adapters+1, max_rank, in] a = data.squeeze(1) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 00b51165c94b..4261e76e3e80 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -6283,6 +6283,7 @@ class DecisionType: VLLM_D3 = "vllm-d3" # same as pplx-decider, non-causal attention LFM2_D1 = "lfm2-d1" # same as openjev, the labels depend on the question type LFM2_D1_OMNI = "lfm2-d1-omni" # same head as laya on a bidirectional LFM2 trunk, other prompt layout + GRANITE_SWITCH = "granite-switch" # same as openjev, run by the decision aLoRA of granite switch class VisionProjectorType: diff --git a/tools/server/README.md b/tools/server/README.md index 7d697633edac..e4a0c418bf10 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -1723,9 +1723,9 @@ Follows the [TypeSafe API](https://docs.typesafe.ai/api), streaming is not suppo The questions of a request are answered independently, an answer does not depend on the other questions. The exception is clef: it reads all the questions in one prompt and decides them jointly. -The number of options of a `choice` question is limited by the model, for example: 52 for openjev, 255 for laya, clef, pplx-decider, vllm-decision2, vllm-d3, lfm2-d1 and lfm2-d1-omni. For laya, long questions and options are truncated to the token budget the model was trained with. +The number of options of a `choice` question is limited by the model, for example: 52 for openjev and granite-switch, 255 for laya, clef, pplx-decider, vllm-decision2, vllm-d3, lfm2-d1 and lfm2-d1-omni. For laya, long questions and options are truncated to the token budget the model was trained with. -For laya, clef, vllm-d3 and lfm2-d1-omni, the whole prompt is evaluated in one batch: it must fit in `--ubatch-size`. An lfm2-d1-omni prompt is cut to 16384 tokens. For vllm-decision2, the options of a question and the end of the prompt must fit in one batch. A server that runs clef only serves this endpoint, text generation is not available. +For laya, clef, vllm-d3 and lfm2-d1-omni, the whole prompt is evaluated in one batch: it must fit in `--ubatch-size`. An lfm2-d1-omni prompt is cut to 16384 tokens. For vllm-decision2, the options of a question and the end of the prompt must fit in one batch. A server that runs clef only serves this endpoint, text generation is not available. A decision model that is also a text model, such as granite-switch, serves both and reports `["text", "decisions"]`. *Image input:* diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index c32c0d1a5aa7..5577c339dbc0 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -161,6 +161,8 @@ std::vector server_model_output_modalities(common_decision_type dec case COMMON_DECISION_TYPE_LFM2_D1: case COMMON_DECISION_TYPE_LFM2_D1_OMNI: return {"decisions"}; + case COMMON_DECISION_TYPE_GRANITE_SWITCH: + return {"text", "decisions"}; default: // fallback when there is no decision type or the metadata is bad return {"text"}; diff --git a/tools/server/server-decision.cpp b/tools/server/server-decision.cpp index b3c19e1f0304..3905c81b34c3 100644 --- a/tools/server/server-decision.cpp +++ b/tools/server/server-decision.cpp @@ -87,7 +87,7 @@ void server_decision_context::init(const llama_model * model) { bias[level] = std::strtof(val, nullptr); } - if (model_type == COMMON_DECISION_TYPE_OPENJEV) { + if (model_type == COMMON_DECISION_TYPE_OPENJEV || model_type == COMMON_DECISION_TYPE_GRANITE_SWITCH) { // one letter per option, each must be a single token const std::string letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; for (const char c : letters) { @@ -753,7 +753,7 @@ void server_decision_context::fill_task( const std::string prompt = render(state, questions, question, variant, files.size()); - if (type == COMMON_DECISION_TYPE_OPENJEV || type == COMMON_DECISION_TYPE_LEV || type == COMMON_DECISION_TYPE_NIMBLE || type == COMMON_DECISION_TYPE_PPLX_DECIDER || type == COMMON_DECISION_TYPE_VLLM_D3) { + if (type == COMMON_DECISION_TYPE_OPENJEV || type == COMMON_DECISION_TYPE_GRANITE_SWITCH || type == COMMON_DECISION_TYPE_LEV || type == COMMON_DECISION_TYPE_NIMBLE || type == COMMON_DECISION_TYPE_PPLX_DECIDER || type == COMMON_DECISION_TYPE_VLLM_D3) { // lev reads the ratings of a noul question at its first labels, not at the digits task.decision.labels.assign(labels.begin(), labels.begin() + n_outputs(question)); } diff --git a/tools/server/server-decision.h b/tools/server/server-decision.h index bfe9fd0ee518..ed97244293dc 100644 --- a/tools/server/server-decision.h +++ b/tools/server/server-decision.h @@ -39,6 +39,7 @@ struct server_decision_context { bool can_share_prompt() const { switch (type) { case COMMON_DECISION_TYPE_OPENJEV: + case COMMON_DECISION_TYPE_GRANITE_SWITCH: case COMMON_DECISION_TYPE_LEV: case COMMON_DECISION_TYPE_KEV: case COMMON_DECISION_TYPE_NIMBLE: From 01fc3aae846caa644f0cee360f108d4c48e4aaa1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 16 Oct 2026 15:07:55 -0700 Subject: [PATCH 062/423] model: add Beam (Reflection) (#30670) --- common/chat.cpp | 7 + common/jinja/value.cpp | 9 ++ common/parsers/harmony-xml.cpp | 207 ++++++++++++++++++++++++++ common/parsers/parsers.h | 2 + common/parsers/sources.cmake | 1 + conversion/__init__.py | 1 + conversion/beam.py | 135 +++++++++++++++++ gguf-py/gguf/constants.py | 26 ++++ src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/llama-graph.cpp | 26 +++- src/llama-graph.h | 7 +- src/llama-hparams.h | 2 + src/llama-model.cpp | 11 +- src/llama-model.h | 1 + src/llama-vocab.cpp | 5 +- src/models/beam.cpp | 261 +++++++++++++++++++++++++++++++++ src/models/models.h | 13 ++ tests/test-jinja.cpp | 6 + tests/test-llama-archs.cpp | 21 +++ 20 files changed, 732 insertions(+), 11 deletions(-) create mode 100644 common/parsers/harmony-xml.cpp create mode 100644 conversion/beam.py create mode 100644 src/models/beam.cpp diff --git a/common/chat.cpp b/common/chat.cpp index 118072b6b514..99ed74faf7a9 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -1104,6 +1104,13 @@ static common_chat_params_init_fn common_chat_template_detect_params_init(const return common_chat_params_init_llm_jp_harmony; } + // Harmony XML tool calls, checked before the GPT-OSS JSON calls + if (src.find("<|channel|>") != std::string::npos && src.find("<|constrain|>xml") != std::string::npos && + src.find("rfl:parameter") != std::string::npos) { + LOG_DBG("Using specialized template: Harmony XML\n"); + return common_chat_params_init_harmony_xml; + } + // GPT-OSS - has unique channel-based structure that needs dedicated handler if (src.find("<|channel|>") != std::string::npos) { LOG_DBG("Using specialized template: GPT-OSS\n"); diff --git a/common/jinja/value.cpp b/common/jinja/value.cpp index 0ecbbf8a8f7f..a1dc9691b983 100644 --- a/common/jinja/value.cpp +++ b/common/jinja/value.cpp @@ -1299,6 +1299,15 @@ const func_builtins & value_object_t::get_builtins() const { } return result; }}, + {"list", [](const func_args & args) -> value { + args.ensure_vals(); + const auto & obj = args.get_pos(0)->as_ordered_object(); + auto result = mk_val(); + for (const auto & pair : obj) { + result->push_back(pair.first); + } + return result; + }}, {"values", [](const func_args & args) -> value { args.ensure_vals(); const auto & obj = args.get_pos(0)->as_ordered_object(); diff --git a/common/parsers/harmony-xml.cpp b/common/parsers/harmony-xml.cpp new file mode 100644 index 000000000000..bf05fac5d680 --- /dev/null +++ b/common/parsers/harmony-xml.cpp @@ -0,0 +1,207 @@ +#include "parsers.h" + +// Harmony with XML tool calls: <|start|>assistant to=NAME<|channel|>analysis <|constrain|>xml<|message|>\nVALUE\n +common_chat_params common_chat_params_init_harmony_xml(const common_chat_template & tmpl, + const autoparser::generation_params & inputs) { + common_chat_params data; + + data.prompt = common_chat_template_direct_apply_impl(tmpl, inputs); + data.generation_prompt = common_chat_template_generation_prompt_impl(tmpl, inputs); + data.message_delimiters = { + { COMMON_CHAT_ROLE_ASSISTANT, "<|start|>assistant" }, + { COMMON_CHAT_ROLE_USER, "<|start|>user" }, + { COMMON_CHAT_ROLE_SYSTEM, "<|start|>developer" }, + { COMMON_CHAT_ROLE_SYSTEM, "<|start|>system" }, + }; + + data.format = COMMON_CHAT_FORMAT_PEG_NATIVE; + data.supports_thinking = true; + + data.thinking_start_tag = "<|channel|>analysis<|message|>"; + data.thinking_end_tags = {"<|end|>"}; + + // EOG text must reach the parser: an XML call has no closing tag of its own + static const std::vector EOG = { + "<|call|>", "<|return|>", "<|endoftext|>", "<|im_end|>", "<|fim_pad|>", "<|repo_name|>", "<|file_sep|>", + }; + + data.preserved_tokens = { + "<|channel|>", "<|constrain|>", "<|message|>", "<|start|>", "<|end|>", + }; + data.preserved_tokens.insert(data.preserved_tokens.end(), EOG.begin(), EOG.end()); + + if (inputs.has_continuation()) { + const auto & msg = inputs.continue_msg; + + data.generation_prompt = "<|start|>assistant<|channel|>analysis<|message|>" + msg.reasoning_content; + if (inputs.continue_final_message == COMMON_CHAT_CONTINUATION_CONTENT) { + data.generation_prompt += "<|end|><|start|>assistant<|channel|>final<|message|>" + msg.render_content(); + } + + data.prompt.append(data.generation_prompt); + } + + auto has_tools = inputs.tools.is_array() && !inputs.tools.empty(); + auto has_response_format = !inputs.json_schema.is_null() && inputs.json_schema.is_object(); + auto include_grammar = has_response_format || (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE); + auto extract_reasoning = inputs.reasoning_format != COMMON_REASONING_FORMAT_NONE; + + static const std::string PARAM_OPEN = "assistant")); + + auto eog = p.choice(); + for (const auto & tok : EOG) { + eog |= p.literal(tok); + } + auto with_eog = [&](std::vector delims) { + delims.insert(delims.end(), EOG.begin(), EOG.end()); + return delims; + }; + + auto analysis_open = p.literal("<|channel|>analysis<|message|>"); + auto analysis_body = p.until_one_of(with_eog({ "<|end|>", "<|constrain|><|start|>assistant", "<|start|>assistant" })); + auto analysis_close = p.optional(p.literal("<|end|>") | p.literal("<|constrain|>") | eog); + if (extract_reasoning) { + p.rule("analysis", analysis_open + p.reasoning(analysis_body) + analysis_close); + } else { + p.rule("analysis", p.atomic(p.content(analysis_open)) + p.content(analysis_body + analysis_close)); + } + auto analysis = p.ref("analysis"); + + auto final_msg = p.rule("final", + p.literal("<|channel|>final<|message|>") + + p.content(p.until_one_of(with_eog({ "<|end|>" }))) + + p.optional(p.literal("<|end|>") | eog)); + + if (has_response_format) { + auto constrain_type = p.chars("[A-Za-z0-9_-]", 1, -1); + auto constraint = p.optional(p.space() + p.optional(p.literal("<|constrain|>")) + constrain_type); + auto response_format = p.rule("response-format", + p.literal("<|channel|>final") + constraint + p.literal("<|message|>") + + p.content(p.schema(p.json(), "response-format-schema", inputs.json_schema)) + + p.optional(p.literal("<|end|>") | eog)); + + return start + p.zero_or_more(analysis + start) + response_format; + } + + if (has_tools && inputs.tool_choice != COMMON_CHAT_TOOL_CHOICE_NONE) { + // tolerate a missing closing tag + std::vector value_ends = with_eog({ PARAM_OPEN, "<|end|>" }); + std::vector string_ends = { PARAM_CLOSE }; + auto next_open = p.choice(); + for (const auto & e : value_ends) { + string_ends.push_back("\n" + e); + next_open |= p.literal(e); + } + auto arg_close = p.tool_arg_close(p.gbnf( + p.literal(PARAM_CLOSE) | (p.literal("\n") + p.peek(p.optional(p.literal("\n")) + next_open)), + gbnf_format_literal("\n") + " " + gbnf_format_literal(PARAM_CLOSE.substr(1)) + "?")); + auto arg_string = p.rule("rfl-arg-string", + p.ac(p.tool_arg_string_value(p.until_one_of(string_ends)) + arg_close, PARAM_CLOSE)); + auto arg_sep = p.optional(p.literal("\n")); + + std::string call_close_gbnf = gbnf_format_literal("<|end|>"); + for (const auto & tok : EOG) { + call_close_gbnf += " | " + gbnf_format_literal(tok); + } + auto call_close = p.gbnf(p.literal("<|end|>") | eog | p.peek(p.literal("<|start|>assistant")), + "(" + call_close_gbnf + ")?"); + + auto tool_choice = p.choice(); + foreach_function(inputs.tools, [&](size_t tool_index, const json & tool) { + const auto & function = tool.at("function"); + std::string name = function.at("name"); + + std::vector arg_rules; + foreach_parameter(function, [&](size_t param_index, const common_chat_schema_property & param, const common_chat_schema_document_ptr & doc) { + auto rule_name = "tool-" + std::to_string(tool_index) + "-arg-" + std::to_string(param_index); + auto arg_open = p.tool_arg_open(p.literal(PARAM_OPEN) + p.tool_arg_name(p.literal(param.name)) + p.literal("\">\n")); + + auto types = param.schema->value_types(); + auto arg_value = p.eps(); + if (!types.has(common_chat_schema::TYPE_STRING) && + (types.has(common_chat_schema::TYPE_OBJECT) || types.has(common_chat_schema::TYPE_ARRAY))) { + // the schema grammar fixes the key order, and the model writes keys in its own order + arg_value = p.tool_arg_json_value(p.json()) + arg_close; + } else if (!types.has(common_chat_schema::TYPE_STRING)) { + arg_value = p.tool_arg_json_value(p.schema(p.json(), rule_name + "-schema", doc, *param.schema)) + arg_close; + } else if (types.is_only(common_chat_schema::TYPE_STRING)) { + arg_value = arg_string; + } else { + // grammar keeps only the string rule; the parser tries JSON first to type the value + auto json_value = p.choice(); + if (types.has(common_chat_schema::TYPE_OBJECT)) { + json_value |= p.json_object(); + } + if (types.has(common_chat_schema::TYPE_ARRAY)) { + json_value |= p.json_array(); + } + if (types.has(common_chat_schema::TYPE_NUMBER) || types.has(common_chat_schema::TYPE_INTEGER)) { + json_value |= p.json_number(); + } + if (types.has(common_chat_schema::TYPE_BOOLEAN)) { + json_value |= p.json_bool(); + } + if (types.has(common_chat_schema::TYPE_NULL)) { + json_value |= p.json_null(); + } + arg_value = p.gbnf(p.atomic(p.tool_arg_json_value(json_value) + arg_close) | arg_string, "rfl-arg-string"); + } + + arg_rules.push_back(p.rule(rule_name, p.tool_arg(arg_open + arg_value))); + }); + + if (arg_rules.empty()) { + auto any_name = p.chars("[^\"\\n<]", 1, -1); + arg_rules.push_back(p.rule("tool-" + std::to_string(tool_index) + "-arg-any", + p.tool_arg(p.tool_arg_open(p.literal(PARAM_OPEN) + p.tool_arg_name(any_name) + p.literal("\">\n")) + arg_string))); + } + + auto args = p.zero_or_more(p.choice(arg_rules) + arg_sep); + + auto call = p.tool( + p.tool_open(p.literal(" to=") + p.tool_name(p.literal(name)) + + p.literal("<|channel|>analysis <|constrain|>xml<|message|>")) + + p.tool_args(args) + + p.tool_close(call_close)); + + tool_choice |= p.rule("tool-" + std::to_string(tool_index), call); + }); + + auto tool_calls = inputs.parallel_tool_calls + ? tool_choice + p.zero_or_more(start + tool_choice) + : tool_choice; + auto tool_call = p.trigger_rule("tool-call", tool_calls); + + if (inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED) { + return start + p.zero_or_more(analysis + start) + tool_call; + } + + return start + p.zero_or_more(analysis + start) + (tool_call | final_msg | p.content(p.rest())); + } + + auto unsolicited = p.rule("unsolicited", p.atomic(p.literal(" to=") + p.rest())); + + // output without a channel header is returned as content instead of failing the request + return start + p.zero_or_more(analysis + start) + (final_msg | unsolicited | p.content(p.rest())); + }); + + if (include_grammar) { + data.grammar_lazy = !(has_response_format || (has_tools && inputs.tool_choice == COMMON_CHAT_TOOL_CHOICE_REQUIRED)); + data.grammar = build_grammar([&](const common_grammar_builder & builder) { + data.parser.build_grammar(builder, data.grammar_lazy); + }); + + if (data.grammar_lazy) { + data.grammar_triggers = { + { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "^\\s+to$" }, + { COMMON_GRAMMAR_TRIGGER_TYPE_PATTERN, "<\\|start\\|>assistant(\\s+to)" }, + }; + } + } + + return data; +} diff --git a/common/parsers/parsers.h b/common/parsers/parsers.h index fac69c096db1..532a44f014fe 100644 --- a/common/parsers/parsers.h +++ b/common/parsers/parsers.h @@ -59,6 +59,8 @@ common_chat_params common_chat_params_init_gigachat_v3(const common_chat_templat common_chat_params common_chat_params_init_gpt_oss(const common_chat_template & tmpl, const autoparser::generation_params & inputs); +common_chat_params common_chat_params_init_harmony_xml(const common_chat_template & tmpl, const autoparser::generation_params & inputs); + common_chat_params common_chat_params_init_k2_horizon(const common_chat_template & tmpl, const autoparser::generation_params & inputs); common_chat_params common_chat_params_init_kimi_k2(const common_chat_template & tmpl, const autoparser::generation_params & inputs); diff --git a/common/parsers/sources.cmake b/common/parsers/sources.cmake index 5ee65b188a5b..2174f4519a64 100644 --- a/common/parsers/sources.cmake +++ b/common/parsers/sources.cmake @@ -9,6 +9,7 @@ set(LLAMA_CHAT_PARSERS_SOURCES ${CMAKE_CURRENT_LIST_DIR}/gemma4.cpp ${CMAKE_CURRENT_LIST_DIR}/gigachat-v3.cpp ${CMAKE_CURRENT_LIST_DIR}/gpt-oss.cpp + ${CMAKE_CURRENT_LIST_DIR}/harmony-xml.cpp ${CMAKE_CURRENT_LIST_DIR}/k2-horizon.cpp ${CMAKE_CURRENT_LIST_DIR}/kimi-k2.cpp ${CMAKE_CURRENT_LIST_DIR}/kimi-k3.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index c2683aec471c..01c8db29680c 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -24,6 +24,7 @@ "ArcticForCausalLM": "arctic", "AudioFlamingo3ForConditionalGeneration": "qwen", "BaiChuanForCausalLM": "baichuan", + "BeamForCausalLM": "beam", "BaichuanForCausalLM": "baichuan", "BailingMoeForCausalLM": "bailingmoe", "BailingMoeV2ForCausalLM": "bailingmoe", diff --git a/conversion/beam.py b/conversion/beam.py new file mode 100644 index 000000000000..67d0cc93ca63 --- /dev/null +++ b/conversion/beam.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +from typing import Callable, Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import ModelBase, TextModel, gguf, logger + + +@ModelBase.register("BeamForCausalLM") +class BeamModel(TextModel): + model_arch = gguf.MODEL_ARCH.BEAM + + # reference NoPE query-scale constants; not in config.json + NOPE_SCALE_POWER = 0.1 + NOPE_SCALE_LENGTH = 2048 + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + hp = self.hparams + rope = hp.get("rope_parameters") or {} + if rope.get("full_attention") is not None: + raise ValueError("Beam: expected NoPE full-attention layers (rope_parameters.full_attention = null)") + if "sliding_attention" in rope: + swa_rope = rope["sliding_attention"] or {} + else: + # flat rope_theta / rope_scaling, already normalized by TextModel + swa_rope = {k: v for k, v in self.rope_parameters.items() if k not in ("sliding_attention", "full_attention", "type")} + swa_rope.update(self.rope_parameters.get("sliding_attention") or {}) + if swa_rope.get("rope_type", "default") != "default": + raise ValueError(f"Beam: unsupported sliding-attention rope type {swa_rope.get('rope_type')!r}") + # full layers are NoPE, so mirror the sliding rope params for the base class + swa_rope = {k: v for k, v in swa_rope.items() if not (k == "rope_type" and v == "default")} + self.rope_parameters = {"sliding_attention": dict(swa_rope), "full_attention": dict(swa_rope)} + + layer_types = hp["layer_types"] + mlp_types = hp["mlp_layer_types"] + n_layer = hp["num_hidden_layers"] + if len(layer_types) != n_layer or len(mlp_types) != n_layer: + raise ValueError("Beam: layer_types / mlp_layer_types must have num_hidden_layers entries") + if any(t not in ("sliding_attention", "full_attention") for t in layer_types): + raise ValueError(f"Beam: unknown layer type in {layer_types}") + if "sliding_attention" not in layer_types: + raise ValueError("Beam: expected at least one sliding_attention layer") + n_dense = 0 + while n_dense < n_layer and mlp_types[n_dense] == "dense": + n_dense += 1 + if any(t != "sparse" for t in mlp_types[n_dense:]): + raise ValueError("Beam: llama.cpp supports dense MLP layers only as a leading block") + self._n_dense_lead = n_dense + if hp.get("attention_bias", False): + raise ValueError("Beam: attention_bias=True is not supported") + self._experts: dict[tuple[int, str], dict[int, Tensor]] = {} + + def set_vocab(self): + self._set_vocab_gpt2() + # otherwise EOT is guessed from a hash-ordered scan + from transformers import AutoTokenizer + tok = AutoTokenizer.from_pretrained(self.dir_model) + ret_id = tok.convert_tokens_to_ids("<|return|>") # ty: ignore[unresolved-attribute] + if isinstance(ret_id, int) and ret_id >= 0: + self.gguf_writer.add_eot_token_id(ret_id) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hp = self.hparams + + head_dim = hp["head_dim"] + self.gguf_writer.add_rope_dimension_count(head_dim) + + self.gguf_writer.add_sliding_window(hp["sliding_window"]) + self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in hp["layer_types"]]) + + self.gguf_writer.add_attn_temperature_scale(self.NOPE_SCALE_POWER) + self.gguf_writer.add_attn_temperature_length(self.NOPE_SCALE_LENGTH) + + self.gguf_writer.add_leading_dense_block_count(self._n_dense_lead) + self.gguf_writer.add_expert_feed_forward_length(hp["moe_intermediate_size"]) + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) + self.gguf_writer.add_expert_weights_norm(bool(hp.get("norm_topk_prob", True))) + # equivalent to the reference's post-aggregation scaling + self.gguf_writer.add_expert_weights_scale(float(hp["routed_scaling_factor"])) + + if (softcap := hp.get("final_logit_softcapping")) is not None: + self.gguf_writer.add_final_logit_softcapping(float(softcap)) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + if name.endswith(".self_attn.gate_bias"): + name = name.replace(".self_attn.gate_bias", ".self_attn.gate_proj.bias") + return super().filter_tensors((name, gen)) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + hp = self.hparams + if name.endswith(".self_attn.gate_proj.bias"): + assert tuple(data_torch.shape) == (hp["num_attention_heads"],), f"{name}: one gate bias per head expected, got {tuple(data_torch.shape)}" + elif name.endswith(".self_attn.gate_proj.weight"): + assert tuple(data_torch.shape) == (hp["num_attention_heads"] * hp["head_dim"], hp["hidden_size"]), f"{name}: {tuple(data_torch.shape)}" + elif name.endswith(".mlp.gate.e_score_correction.bias"): + assert tuple(data_torch.shape) == (hp["num_experts"],), f"{name}: {tuple(data_torch.shape)}" + elif name.endswith((".self_attn.q_norm.weight", ".self_attn.k_norm.weight")): + assert tuple(data_torch.shape) == (hp["head_dim"],), f"{name}: per-head q/k norm expected, got {tuple(data_torch.shape)}" + elif bid is not None and bid < self._n_dense_lead and ".mlp.experts." in name: + raise ValueError(f"Beam: expert tensor {name} in a dense layer") + if ".mlp.experts." in name: + assert bid is not None + n_experts = self.hparams["num_experts"] + parts = name.split(".") + xid = int(parts[parts.index("experts") + 1]) + proj = parts[parts.index("experts") + 2] + if proj not in ("gate_proj", "up_proj", "down_proj"): + raise ValueError(f"Beam: unexpected expert tensor {name}") + bucket = self._experts.setdefault((bid, proj), {}) + if xid in bucket: + raise ValueError(f"Beam: duplicate expert tensor {name}") + bucket[xid] = data_torch + if len(bucket) < n_experts: + return + del self._experts[(bid, proj)] + data = torch.stack([bucket[e] for e in range(n_experts)], dim=0) + yield from super().modify_tensors(data, f"model.layers.{bid}.mlp.experts.{proj}.weight", bid) + return + + yield from super().modify_tensors(data_torch, name, bid) + + def prepare_tensors(self): + super().prepare_tensors() + if self._experts: + missing = sorted(f"layer {b} {p} ({len(v)} experts)" for (b, p), v in self._experts.items()) + raise ValueError(f"Beam: unprocessed experts: {missing}") + logger.info("Beam: all routed experts stacked") diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 4261e76e3e80..bfeab6adfde6 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -573,6 +573,7 @@ class MODEL_ARCH(IntEnum): QWEN3MOE = auto() QWEN3NEXT = auto() HRM_TEXT = auto() + BEAM = auto() QWEN3VL = auto() QWEN3VLMOE = auto() QWEN35 = auto() @@ -1388,6 +1389,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.QWEN3MOE: "qwen3moe", MODEL_ARCH.QWEN3NEXT: "qwen3next", MODEL_ARCH.HRM_TEXT: "hrm_text", + MODEL_ARCH.BEAM: "beam", MODEL_ARCH.QWEN3VL: "qwen3vl", MODEL_ARCH.QWEN3VLMOE: "qwen3vlmoe", MODEL_ARCH.QWEN35: "qwen35", @@ -2981,6 +2983,30 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.BEAM: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.FFN_PRE_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + ], MODEL_ARCH.HRM_TEXT: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT, diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 1476a079a7cd..bef2d45c0ba2 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -165,6 +165,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_QWEN3TTS, "qwen3tts" }, { LLM_ARCH_POCKETTTS, "pockettts" }, { LLM_ARCH_K2_HORIZON, "k2-horizon" }, + { LLM_ARCH_BEAM, "beam" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; diff --git a/src/llama-arch.h b/src/llama-arch.h index 0c842d312db1..c305c17663cb 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -170,6 +170,7 @@ enum llm_arch { LLM_ARCH_MINIMAX_01, LLM_ARCH_HRM_TEXT, LLM_ARCH_K2_HORIZON, + LLM_ARCH_BEAM, LLM_ARCH_UNKNOWN, }; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 454c21404266..2bc3a291a456 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -218,17 +218,31 @@ void llm_graph_input_attn_temp::set_input(const llama_ubatch * ubatch) { GGML_ASSERT(n_attn_temp_floor_scale != 0); std::vector attn_scale_data(n_tokens, 0.0f); - for (int i = 0; i < n_tokens; ++i) { - const float pos = ubatch->pos[i]; - attn_scale_data[i] = std::log( - std::floor((pos + f_attn_temp_offset) / n_attn_temp_floor_scale) + 1.0 - ) * f_attn_temp_scale + 1.0; + if (attn_temp_pow) { + // (1 + (pos + offset) / length) ^ power + for (int i = 0; i < n_tokens; ++i) { + const double pos = ubatch->pos[i]; + attn_scale_data[i] = (float) std::exp((double) f_attn_temp_scale * + std::log1p((pos + f_attn_temp_offset) / (double) n_attn_temp_floor_scale)); + } + } else { + for (int i = 0; i < n_tokens; ++i) { + const float pos = ubatch->pos[i]; + attn_scale_data[i] = std::log( + std::floor((pos + f_attn_temp_offset) / n_attn_temp_floor_scale) + 1.0 + ) * f_attn_temp_scale + 1.0; + } } ggml_backend_tensor_set(attn_scale, attn_scale_data.data(), 0, n_tokens*ggml_element_size(attn_scale)); } } +bool llm_graph_input_attn_temp::can_reuse(const llm_graph_params & params) { + // only the power form opts in; the log form keeps no reuse + return attn_temp_pow && attn_scale != nullptr && attn_scale->ne[2] == params.ubatch.n_tokens; +} + void llm_graph_input_pos_bucket::set_input(const llama_ubatch * ubatch) { if (pos_bucket) { const int64_t n_tokens = ubatch->n_tokens; @@ -2743,7 +2757,7 @@ ggml_tensor * llm_graph_context::build_inp_pos() const { } ggml_tensor * llm_graph_context::build_inp_attn_scale() const { - auto inp = std::make_unique(hparams.n_attn_temp_floor_scale, hparams.f_attn_temp_scale, hparams.f_attn_temp_offset); + auto inp = std::make_unique(hparams.n_attn_temp_floor_scale, hparams.f_attn_temp_scale, hparams.f_attn_temp_offset, hparams.attn_temp_pow); auto & cur = inp->attn_scale; diff --git a/src/llama-graph.h b/src/llama-graph.h index 816fa07cc7c4..f66bedb3254f 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -188,17 +188,20 @@ class llm_graph_input_pos : public llm_graph_input_i { // temperature tuning, used by llama4 class llm_graph_input_attn_temp : public llm_graph_input_i { public: - llm_graph_input_attn_temp(uint32_t n_attn_temp_floor_scale, float f_attn_temp_scale, float f_attn_temp_offset) - : n_attn_temp_floor_scale(n_attn_temp_floor_scale), f_attn_temp_scale(f_attn_temp_scale), f_attn_temp_offset(f_attn_temp_offset) {} + llm_graph_input_attn_temp(uint32_t n_attn_temp_floor_scale, float f_attn_temp_scale, float f_attn_temp_offset, bool attn_temp_pow = false) + : n_attn_temp_floor_scale(n_attn_temp_floor_scale), f_attn_temp_scale(f_attn_temp_scale), f_attn_temp_offset(f_attn_temp_offset), attn_temp_pow(attn_temp_pow) {} virtual ~llm_graph_input_attn_temp() = default; void set_input(const llama_ubatch * ubatch) override; + bool can_reuse(const llm_graph_params & params) override; + ggml_tensor * attn_scale = nullptr; // F32 [n_batch] const uint32_t n_attn_temp_floor_scale; const float f_attn_temp_scale; const float f_attn_temp_offset; + const bool attn_temp_pow; }; class llm_graph_input_pos_bucket : public llm_graph_input_i { diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 3e025368b91f..719d7081dd21 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -280,6 +280,8 @@ struct llama_hparams { uint32_t n_attn_temp_floor_scale = 0; float f_attn_temp_scale = 0.0f; float f_attn_temp_offset = 0.0f; // offset position index + // use (1 + (pos + offset) / floor_scale) ^ scale instead of the llama4 log form + bool attn_temp_pow = false; // realtime models: number of tokens the output lags behind the audio input uint32_t n_tokens_delay = 0; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 3da32333b10a..099c7a052f46 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -322,6 +322,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_minimax_m3(params); case LLM_ARCH_HRM_TEXT: return new llama_model_hrm_text(params); + case LLM_ARCH_BEAM: + return new llama_model_beam(params); case LLM_ARCH_COGVLM: return new llama_model_cogvlm(params); case LLM_ARCH_PANGU_EMBED: @@ -405,6 +407,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_idx_cache ("cache_idx_(k|v)_l\\d*"); static const std::regex pattern_dsv4_state ("dsv4_(csa|hca|lid)_state_(kv|score)_l\\d*"); static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight"); + static const std::regex pattern_attn_gate_bias ("blk\\.\\d*\\.attn_gate\\.bias"); static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight"); static const std::regex pattern_attn_out_bias ("blk\\.\\d*\\.attn_output.bias"); static const std::regex pattern_attn_out_a_weight("blk\\.\\d*\\.attn_output_a\\.weight"); @@ -550,7 +553,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (std::regex_match(tensor_name, pattern_qk_norm)) { return get_tensor_config_impl(tensor->ne[1] == 1 ? GGML_BACKEND_SPLIT_AXIS_MIRRORED : GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight"); } - if (std::regex_match(tensor_name, pattern_kv_cache) || std::regex_match(tensor_name, pattern_attn_sinks)) { + if (std::regex_match(tensor_name, pattern_kv_cache) || std::regex_match(tensor_name, pattern_attn_sinks) || + std::regex_match(tensor_name, pattern_attn_gate_bias)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output.weight"); } if (std::regex_match(tensor_name, pattern_attn_out_weight)) { @@ -749,6 +753,10 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str const int64_t granularity_q = std::lcm(n_embd_q, blck_size_perf); const int64_t granularity_head = granularity_q / hparams.n_embd_head_k(il); // for tensors with one value per head + if (std::regex_match(tensor_name, pattern_attn_gate_bias)) { + GGML_ASSERT(segments.size() == 1); + return {granularity_head}; + } if (std::regex_match(tensor_name, pattern_attn_sinks)) { GGML_ASSERT(segments.size() == 1); if (is_dsv4) { @@ -3606,6 +3614,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_MELLUM: case LLM_ARCH_MAPLE: case LLM_ARCH_HRM_TEXT: + case LLM_ARCH_BEAM: return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_DFLASH: diff --git a/src/llama-model.h b/src/llama-model.h index 6335dc2bc318..6487b0e9c434 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -310,6 +310,7 @@ struct llama_layer { struct ggml_tensor * wv_enc = nullptr; struct ggml_tensor * wo_enc = nullptr; struct ggml_tensor * wqkv_gate = nullptr; + struct ggml_tensor * wqkv_gate_b = nullptr; // K2 Horizon MoVA struct ggml_tensor * attn_v_gate = nullptr; struct ggml_tensor * attn_v_gate_b = nullptr; diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 9fd5b9ef3f2a..05d2527dc073 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -3142,9 +3142,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { } } - // gemma4 and plamo have a normal token, unlike paddleocr + // gemma4, plamo and beam have a normal token, unlike paddleocr { - bool has_normal_s_marker = false; + // beam's tool call marker is not an EOG token, so it is looked up in the whole vocab + bool has_normal_s_marker = token_to_id.count(" 0; bool has_s = false; llama_token s_id = LLAMA_TOKEN_NULL; diff --git a/src/models/beam.cpp b/src/models/beam.cpp new file mode 100644 index 000000000000..55817dd4839a --- /dev/null +++ b/src/models/beam.cpp @@ -0,0 +1,261 @@ +#include "models.h" + +void llama_model_beam::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_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); + ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm); + hparams.f_final_logit_softcapping = 0.0f; // absent key = no softcap (default is 30) + ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false); + + if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) { + hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID; + } + if (hparams.n_layer_dense_lead > hparams.n_layer()) { + throw std::runtime_error("beam: leading_dense_block_count exceeds the layer count"); + } + + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; // masked iff p1 - p0 >= n_swa + uint32_t n_pattern = 0; + ml.get_arr_n(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, n_pattern); + if (hparams.n_swa == 0 || n_pattern != hparams.n_layer_all) { + throw std::runtime_error("beam: requires a positive sliding window and a sliding window pattern for each layer"); + } + ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl); + if (!hparams.is_swa_any()) { + throw std::runtime_error("beam: sliding window pattern has no sliding window layer"); + } + + // full layers are NoPE + hparams.rope_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); + + // NoPE query scale uses pos + 1 + hparams.f_attn_temp_scale = 0.1f; + hparams.attn_temp_length = 2048; + ml.get_key(LLM_KV_ATTENTION_TEMPERATURE_SCALE, hparams.f_attn_temp_scale, false); + ml.get_key(LLM_KV_ATTENTION_TEMPERATURE_LENGTH, hparams.attn_temp_length, false); + if (hparams.f_attn_temp_scale == 0.0f) { + throw std::runtime_error("beam: attention.temperature_scale must be nonzero"); + } + if (hparams.attn_temp_length == 0) { + throw std::runtime_error("beam: attention.temperature_length must be positive"); + } + hparams.n_attn_temp_floor_scale = hparams.attn_temp_length; + hparams.f_attn_temp_offset = 1.0f; + hparams.attn_temp_pow = true; + + switch (hparams.n_layer()) { + case 52: type = LLM_TYPE_UNKNOWN; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_beam::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 == NULL) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 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); + + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); + + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_embd_head_k * n_head}, 0); + layer.wqkv_gate_b = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "bias", i), {n_head}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0); + + if ((uint32_t) i < hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } else { + const int64_t n_ff_exp = hparams.n_ff_exp(i); + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + } + } +} + +std::unique_ptr llama_model_beam::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +static ggml_tensor * beam_residual_add(ggml_context * ctx, ggml_tensor * h, ggml_tensor * r, int i) { + const float h_scale = (float) (1.0 / std::sqrt((double) i)); + const float r_scale = (float) std::sqrt((double) (i - 1) / (double) i); + return ggml_add(ctx, ggml_scale(ctx, h, h_scale), ggml_scale(ctx, r, r_scale)); +} + +llama_model_beam::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_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + ggml_tensor * inp_pos = build_inp_pos(); + + bool has_nope = false; + for (int il = 0; il < n_layer; ++il) { + has_nope = has_nope || !hparams.is_swa(il); + } + ggml_tensor * inp_attn_scale = has_nope ? build_inp_attn_scale() : nullptr; + + 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) { + const auto & layer = model.layers[il]; + + const bool is_swa = hparams.is_swa(il); + + ggml_tensor * inpSA = inpL; + + cur = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + { + ggml_tensor * attn_inp = cur; + + auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, n_head, n_head_kv, il); + + Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); + Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + cb(Kcur, "Kcur_normed", il); + + if (is_swa) { + const float freq_base_l = model.get_rope_freq_base (cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(Qcur, "Qcur_rope", il); + cb(Kcur, "Kcur_rope", il); + } else { + Qcur = ggml_mul(ctx0, Qcur, inp_attn_scale); + cb(Qcur, "Qcur_scaled", il); + } + + cur = build_attn(inp_attn, + NULL, NULL, NULL, // o_proj is applied after the gate + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "attn_out", il); + + ggml_tensor * gate = build_lora_mm(layer.wqkv_gate, attn_inp); + cb(gate, "attn_gate_proj", il); + gate = ggml_reshape_3d(ctx0, gate, n_embd_head, n_head, n_tokens); + gate = ggml_add(ctx0, gate, ggml_reshape_3d(ctx0, layer.wqkv_gate_b, 1, n_head, 1)); + gate = ggml_sigmoid(ctx0, gate); + gate = ggml_reshape_2d(ctx0, gate, n_embd_head * n_head, n_tokens); + cb(gate, "attn_gate_sig", il); + + cur = ggml_mul(ctx0, cur, gate); + cb(cur, "attn_gated", il); + + cur = build_lora_mm(layer.wo, cur, layer.wo_s); + cb(cur, "attn_o_proj", il); + } + + cur = build_norm(cur, layer.attn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_post_norm", 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); + } + + const int i_attn = 2 * (il + 1); + ggml_tensor * ffn_inp = beam_residual_add(ctx0, cur, inpSA, i_attn); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + cur = build_ffn(cur, + layer.ffn_up, NULL, NULL, + layer.ffn_gate, NULL, NULL, + layer.ffn_down, NULL, NULL, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + } else { + cur = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, + hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + cb(cur, "ffn_moe_out", il); + } + + cur = build_norm(cur, layer.ffn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_post_norm", il); + + cur = beam_residual_add(ctx0, cur, ffn_inp, i_attn + 1); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + inpL = cur; + } + + cur = inpL; + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur, model.output_s); + + if (hparams.f_final_logit_softcapping != 0.0f) { + cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping); + cur = ggml_tanh(ctx0, cur); + cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping); + } + + cb(cur, "result_output", -1); + res->t_logits = cur; +} diff --git a/src/models/models.h b/src/models/models.h index ffeed1733ac7..98f6b6e35dcc 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2041,6 +2041,19 @@ struct llama_model_afmoe : public llama_model_base { }; +struct llama_model_beam : public llama_model_base { + llama_model_beam(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_laguna : public llama_model_base { llama_model_laguna(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tests/test-jinja.cpp b/tests/test-jinja.cpp index 7132c9e22483..1dacf170280e 100644 --- a/tests/test-jinja.cpp +++ b/tests/test-jinja.cpp @@ -2011,6 +2011,12 @@ static void test_object_methods(testing & t) { "a b " ); + test_template(t, "object | list", + "{{ (obj | list) == ['b', 'a'] }} {{ obj | list | join(',') }}", + {{"obj", {{"b", 1}, {"a", 2}}}}, + "True b,a" + ); + test_template(t, "object.values()", "{% for v in obj.values() %}{{ v }} {% endfor %}", {{"obj", {{"a", 1}, {"b", 2}}}}, diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index eb72553b3e81..14ca203be809 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -188,6 +188,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_MUSE_SPARK || arch == LLM_ARCH_AFMOE || arch == LLM_ARCH_MAION_CODER) { n_head_kv = 2; // GQA coverage } + if (arch == LLM_ARCH_BEAM) { + n_head = 4; + n_head_kv = 2; + n_layer = 4; + } const uint32_t n_embd_head = n_embd / n_head; ms.add_kv(LLM_KV_GENERAL_ARCHITECTURE, llm_arch_name(arch)); @@ -294,6 +299,13 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { if (arch == LLM_ARCH_MAION_CODER) { ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, std::vector{1, 1, 0, 0}); + } else if (arch == LLM_ARCH_BEAM) { + std::vector pattern; + pattern.reserve(n_layer); + for (uint32_t il = 0; il < n_layer; il++) { + pattern.push_back(il % 2); + } + ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, pattern); } else if (arch == LLM_ARCH_GEMMA4) { ms.add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, n_embd/2); ms.add_kv(LLM_KV_ATTENTION_SHARED_KV_LAYERS, uint32_t(0)); @@ -470,6 +482,14 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } } + if (arch == LLM_ARCH_BEAM) { + ms.add_kv(LLM_KV_EXPERT_COUNT, uint32_t(8)); + ms.add_kv(LLM_KV_EXPERT_GROUP_COUNT, uint32_t(4)); + ms.add_kv(LLM_KV_EXPERT_GROUP_USED_COUNT, uint32_t(2)); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 2.5f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); + } + if (arch == LLM_ARCH_MAION_CODER) { ms.add_kv(LLM_KV_MOE_LATENT_SIZE, n_embd / 2); ms.add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, std::vector{n_ff, n_ff / 2, n_ff, n_ff}); @@ -712,6 +732,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_DOTS1: case LLM_ARCH_AFMOE: + case LLM_ARCH_BEAM: case LLM_ARCH_ERNIE4_5: case LLM_ARCH_ERNIE4_5_MOE: case LLM_ARCH_HUNYUAN_MOE: From 2f2763db6eaac591e78023d1e2c4a31607c64517 Mon Sep 17 00:00:00 2001 From: Shawn Gu Date: Fri, 16 Oct 2026 15:15:04 -0700 Subject: [PATCH 063/423] opencl: Add MxFP4 MoE bin kernel for Adreno X1 devices (#30691) * opencl: add bin kernel `gemm_moe_mxfp4_q8_1_dp4a_ila_ldck` for X1 * opencl: fix whitespace * opencl: undo unnecessary change * opencl: enable dp4a clc for X1E --------- Co-authored-by: Li He --- ggml/src/ggml-opencl/CMakeLists.txt | 1 + ggml/src/ggml-opencl/ggml-opencl.cpp | 83 ++++++++++++++++++- .../kernels/transpose_activation_q8.cl | 35 ++++++++ 3 files changed, 115 insertions(+), 4 deletions(-) create mode 100644 ggml/src/ggml-opencl/kernels/transpose_activation_q8.cl diff --git a/ggml/src/ggml-opencl/CMakeLists.txt b/ggml/src/ggml-opencl/CMakeLists.txt index 97b862afffd8..5ba63dade180 100644 --- a/ggml/src/ggml-opencl/CMakeLists.txt +++ b/ggml/src/ggml-opencl/CMakeLists.txt @@ -132,6 +132,7 @@ set(GGML_OPENCL_KERNELS gemm_moe_q6_k_q8_1_dp4a gemm_moe_q8_1_dp4a moe_reorder_quant_a_q8_1 + transpose_activation_q8 gemm_noshuffle_q4_k_q8_1_dp4a gemm_noshuffle_q5_k_q8_1_dp4a gemm_noshuffle_q6_k_q8_1_dp4a diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 17ceacc7aecc..de93b0134eb5 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -720,6 +720,10 @@ struct ggml_backend_opencl_context { ggml_cl_buffer prealloc_moe_qa; // int8 quants [tok_slots * ne00] ggml_cl_buffer prealloc_moe_da; // per-block d [tok_slots * ne00/32] (half) ggml_cl_buffer prealloc_moe_sa; // per-block s [tok_slots * ne00/32] (half) + // per-tile transposed qa/da and on-chip scratch for the ldck dp4a MoE GEMM + ggml_cl_buffer prealloc_moe_qa_trans; // int8 quants [tok_slots * ne00] + ggml_cl_buffer prealloc_moe_da_trans; // per-block d [tok_slots * ne00/32] (float) + ggml_cl_buffer prealloc_moe_qda_on_chip; // scratch copy of the router weights to avoid dst aliasing ggml_cl_buffer prealloc_moe_combine_w; ggml_cl_buffer prealloc_splitk_partial; // [ksplit * M] partials for split-K GEMV @@ -1016,6 +1020,7 @@ struct ggml_backend_opencl_context { cl_kernel kernel_gemv_moe_q4_k_f32_ns_wimg = nullptr; // weight-as-texture MoE decode GEMV (opt-in) cl_kernel kernel_gemm_moe_q4_k_q8_1_dp4a = nullptr; // dp4a (int8) prefill GEMM variant cl_kernel kernel_moe_reorder_quant_a_q8_1; // fused reorder + q8_1 quant for the dp4a GEMM + cl_kernel kernel_transpose_activation_q8; // per-tile transpose of q8_1 activations for the ldck dp4a GEMM cl_kernel kernel_gemm_moe_q8_1_dp4a_q80 = nullptr; // generic dp4a MoE GEMM (MOE_QT=80), opt-in cl_kernel kernel_moe_expand_scale_q8_0 = nullptr; // q8_0 per-block d -> uniform scale[16] cl_kernel kernel_gemm_moe_q8_1_dp4a_q50 = nullptr; // generic dp4a MoE GEMM (MOE_QT=50, q5_0), opt-in @@ -1031,6 +1036,7 @@ struct ggml_backend_opencl_context { cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a = nullptr; // dp4a (int8) mxfp4 MoE prefill GEMM cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a = nullptr; // dp4a (int8) q4_0 MoE prefill GEMM cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a_bin = nullptr; // binary dp4a (int8) mxfp4 MoE prefill GEMM + cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a_ldck_bin = nullptr; // binary dp4a (int8) mxfp4 MoE prefill GEMM using ldck cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a_bin = nullptr; // binary dp4a (int8) q4_0 MoE prefill GEMM cl_kernel kernel_gemm_moe_q4_k_q8_1_dp4a_bin = nullptr; // binary dp4a (int8) q4_k MoE prefill GEMM cl_kernel kernel_gemm_moe_q6_k_q8_1_dp4a_bin = nullptr; // binary dp4a (int8) q6_k MoE prefill GEMM @@ -4903,6 +4909,24 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { } } + // gemm_moe_mxfp4_q8_1_dp4a_ldck_bin (dp4a prefill GEMM) + if (backend_ctx->has_integer_dot) { + size_t bin_size = 0; + backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_ldck_bin = nullptr; + + if (use_adreno_bin_kernels(backend_ctx)) { + const char * kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_moe_mxfp4_q8_1_dp4a_ila_ldck", &bin_size); + if (kernel_bin && bin_size > 0) { + cl_program prog = + build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, CL_moe_compile_opts, bin_size); + + CL_CHECK((backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_ldck_bin = clCreateKernel(prog, "kernel_gemm_moe_mxfp4_q8_1_dp4a_ila_ldck", &err), err)); + CL_CHECK(clReleaseProgram(prog)); + GGML_LOG_CONT("."); + } + } + } + // gemm_moe_q4_0_q8_1_dp4a (dp4a prefill GEMM) if (backend_ctx->has_integer_dot) { #ifdef GGML_OPENCL_EMBED_KERNELS @@ -4984,6 +5008,23 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { GGML_LOG_CONT("."); } + // transpose_activation_q8 (per-tile transpose of q8_1 activations for the ldck dp4a GEMM) + { +#ifdef GGML_OPENCL_EMBED_KERNELS + const std::string kernel_src { + #include "transpose_activation_q8.cl.h" + }; +#else + const std::string kernel_src = read_file("transpose_activation_q8.cl"); +#endif + cl_program prog = + build_program_from_source(backend_ctx, kernel_src.c_str(), CL_moe_compile_opts); + + CL_CHECK((backend_ctx->kernel_transpose_activation_q8 = clCreateKernel(prog, "kernel_transpose_activation_q8", &err), err)); + CL_CHECK(clReleaseProgram(prog)); + GGML_LOG_CONT("."); + } + // gemv_moe_q5_k_f32_ns { #ifdef GGML_OPENCL_EMBED_KERNELS @@ -28258,11 +28299,19 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, static const char * mxfp4_moe_dp4a_env = getenv("GGML_OPENCL_MXFP4_MOE_DP4A"); bool use_moe_dp4a = mxfp4_moe_dp4a_env ? (atoi(mxfp4_moe_dp4a_env) != 0) - : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E); + : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E || backend_ctx->adreno_gen == ADRENO_GPU_GEN::X1E); // dot prod has to be available use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a; // bin kernel takes precedence, dp4a bin kernel has higher priority than normal bin kernel - if (backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin == nullptr) { + // ldck dp4a bin kernel (X1) has the same priority as normal bin kernel (X2/A8x), they are not expected to coexist. + const bool use_moe_dp4a_ldck = backend_ctx->has_integer_dot && + (!mxfp4_moe_dp4a_env || atoi(mxfp4_moe_dp4a_env) != 0) && + backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin == nullptr && + backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin == nullptr && + backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_ldck_bin != nullptr; + if (use_moe_dp4a_ldck) { + use_moe_dp4a = true; + } else if (backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin == nullptr) { use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin == nullptr; } @@ -28351,6 +28400,12 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, backend_ctx->prealloc_moe_qa.allocate(backend_ctx->context, tok_slots * ne00 * sizeof(cl_char)); backend_ctx->prealloc_moe_da.allocate(backend_ctx->context, n_blocks * sizeof(cl_half)); backend_ctx->prealloc_moe_sa.allocate(backend_ctx->context, n_blocks * sizeof(cl_half)); + if (use_moe_dp4a_ldck) { + backend_ctx->prealloc_moe_qa_trans.allocate(backend_ctx->context, tok_slots * ne00 * sizeof(cl_char)); + backend_ctx->prealloc_moe_da_trans.allocate(backend_ctx->context, n_blocks * sizeof(cl_float)); + // matches max_constant_size of src1_a_cm in the ldck kernel + backend_ctx->prealloc_moe_qda_on_chip.allocate(backend_ctx->context, 2304); + } // fused reorder + q8_1 quant straight from the original // activations (no intermediate f32 reorder buffer). mxfp4 has no @@ -28375,13 +28430,33 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, cl_kernel dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a; if (backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin) { dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin; + } else if (use_moe_dp4a_ldck) { + dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_ldck_bin; + + // per-tile transpose of qa/da for the ldck kernel + cl_kernel tk = backend_ctx->kernel_transpose_activation_q8; + CL_CHECK(clSetKernelArg(tk, 0, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer)); + CL_CHECK(clSetKernelArg(tk, 1, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa_trans.buffer)); + CL_CHECK(clSetKernelArg(tk, 2, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer)); + CL_CHECK(clSetKernelArg(tk, 3, sizeof(cl_mem), &backend_ctx->prealloc_moe_da_trans.buffer)); + CL_CHECK(clSetKernelArg(tk, 4, sizeof(cl_mem), &(backend_ctx->prealloc_total_tiles.buffer))); + CL_CHECK(clSetKernelArg(tk, 5, sizeof(cl_uint), &ne00)); + size_t tk_local[3] = { 32, 8, 1 }; + size_t tk_global[3] = { 32, (size_t)(ne00 / 4), (size_t)max_post_router_tile }; + backend_ctx->enqueue_ndrange_kernel(tk, 3, tk_global, tk_local, dst); } int aidx = 0; CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->q_img)); CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->e)); - CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer)); - CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer)); + if (use_moe_dp4a_ldck) { + CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa_trans.buffer)); + CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da_trans.buffer)); + CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qda_on_chip.buffer)); + } else { + CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_qa.buffer)); + CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &backend_ctx->prealloc_moe_da.buffer)); + } CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2)); CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_src2_emap)); CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &buf_dst_image)); diff --git a/ggml/src/ggml-opencl/kernels/transpose_activation_q8.cl b/ggml/src/ggml-opencl/kernels/transpose_activation_q8.cl new file mode 100644 index 000000000000..34d54b57e24a --- /dev/null +++ b/ggml/src/ggml-opencl/kernels/transpose_activation_q8.cl @@ -0,0 +1,35 @@ +#pragma OPENCL EXTENSION cl_khr_fp16 : enable +#pragma OPENCL EXTENSION cl_khr_subgroups : enable + +#define tile_size 32 + +kernel void kernel_transpose_activation_q8( + __global uint * src1_qa, + __global uint * src1_qa_trans, + __global half * src1_sa, + __global float * src1_sa_trans, + __global const int * total_tiles, + uint K +) { + uint row_in_tile = get_global_id(0); + uint blk_k = get_group_id(1); + uint col_in_tile = get_local_id(1); + uint tile_n = get_global_id(2); + + if (blk_k >= (K / tile_size) || tile_n >= (uint)total_tiles[0]) { + return; + } + + uint src_idx = (tile_n * tile_size + row_in_tile) * K / 4 + blk_k * 8 + col_in_tile; + uint dst_idx = tile_n * tile_size * K / 4 + (blk_k * 8 + col_in_tile) * tile_size + row_in_tile; + + src1_qa_trans[dst_idx] = src1_qa[src_idx]; + + if (col_in_tile == 0) { + src_idx = (tile_n * tile_size + row_in_tile) * K / 32 + blk_k; + dst_idx = tile_n * tile_size * K / 32 + blk_k * tile_size + row_in_tile; + + src1_sa_trans[dst_idx] = (float)(src1_sa[src_idx]); + } + +} From 27c7a769da1671c97b2e3f6cb88704a6e1e0d0cd Mon Sep 17 00:00:00 2001 From: wes Date: Sat, 17 Oct 2026 06:20:43 +0800 Subject: [PATCH 064/423] hexagon: vectorize softplus with HVX (#30435) softplus_f32 and tile_softplus_f32 ran a scalar logf(1.0f + expf(x)) loop. Replace both with a shared hvx_softplus_f32_aa() helper built from the backend's existing hvx_vec_exp_f32 / hvx_vec_log_f32 primitives, so no new numerical approximation is introduced. hvx_vec_exp_f32 clamps its input to (-88, 88) before exponentiating, so a naive log(1 + exp(x)) composition would saturate to ~88.72 for large x instead of returning x. The Q6_V_vmux_QVV select keeps the scalar reference's x > 20 guard (ggml_compute_softplus_f32 in ggml-impl.h), which is what makes the vector path match the CPU result. Assisted-by: Claude --- ggml/src/ggml-hexagon/htp/unary-ops.c | 51 +++++++++++++++++++-------- 1 file changed, 36 insertions(+), 15 deletions(-) diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.c b/ggml/src/ggml-hexagon/htp/unary-ops.c index c557b0907daa..37df31ddba81 100644 --- a/ggml/src/ggml-hexagon/htp/unary-ops.c +++ b/ggml/src/ggml-hexagon/htp/unary-ops.c @@ -690,22 +690,48 @@ static void tri_f32(const void * restrict src, } } +// softplus(x) = log(1 + exp(x)) +// Match CPU reference: return x for x > 20 to avoid overflow. +static inline void hvx_softplus_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + const HVX_Vector * restrict v_src = (const HVX_Vector *) src; + HVX_Vector * restrict v_dst = (HVX_Vector *) dst; + + const HVX_Vector twenty = hvx_vec_splat_f32(20.0f); + const HVX_Vector one = hvx_vec_splat_f32(1.0f); + + const uint32_t nvec = n / VLEN_FP32; + const uint32_t nloe = n % VLEN_FP32; + + #pragma unroll(4) + for (uint32_t i = 0; i < nvec; i++) { + HVX_Vector x = v_src[i]; + HVX_VectorPred pred_gt20 = Q6_Q_vcmp_gt_VsfVsf(x, twenty); + HVX_Vector expx = hvx_vec_exp_f32(x); + HVX_Vector logv = hvx_vec_log_f32(hvx_vec_add_f32_f32(one, expx)); + v_dst[i] = Q6_V_vmux_QVV(pred_gt20, x, logv); + } + + if (nloe > 0) { + HVX_Vector x = v_src[nvec]; + HVX_VectorPred pred_gt20 = Q6_Q_vcmp_gt_VsfVsf(x, twenty); + HVX_Vector expx = hvx_vec_exp_f32(x); + HVX_Vector logv = hvx_vec_log_f32(hvx_vec_add_f32_f32(one, expx)); + HVX_Vector sp = Q6_V_vmux_QVV(pred_gt20, x, logv); + hvx_vec_store_a((void *) &v_dst[nvec], nloe * sizeof(float), sp); + } +} + static void softplus_f32(const void * restrict src, void * restrict dst, const uint32_t num_rows, const struct htp_unary_context * uctx) { htp_unary_op_preamble; - // softplus(x) = log(1 + exp(x)) - // Match CPU reference: ggml_compute_softplus_f32() in ggml-impl.h + for (uint32_t ir = 0; ir < num_rows; ir++) { - const float * restrict src_f = (const float *)((const uint8_t *)src + (ir * src0_row_size_aligned)); - float * restrict dst_f = (float *)((uint8_t *)dst + (ir * dst_row_size_aligned)); + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned); - for (uint32_t i = 0; i < ne0; i++) { - float x = src_f[i]; - // For x > 20: softplus(x) ~ x (avoids exp overflow) - dst_f[i] = (x > 20.0f) ? x : logf(1.0f + expf(x)); - } + hvx_softplus_f32_aa(dst_local, src_local, ne0); } } @@ -866,12 +892,7 @@ static void tile_gelu_erf_f32(void * restrict dst, const void * restrict src, ui static void tile_softplus_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { (void) uctx; - const float * restrict sf = (const float *) src; - float * restrict df = (float *) dst; - for (uint32_t i = 0; i < tw; i++) { - float x = sf[i]; - df[i] = (x > 20.0f) ? x : logf(1.0f + expf(x)); - } + hvx_softplus_f32_aa((uint8_t *) dst, (const uint8_t *) src, tw); } static void tile_tanh_f32(void * restrict dst, const void * restrict src, uint32_t tw, const struct htp_unary_context * uctx) { From 6420cf5dd588c68e48657a0bbd82bb3bd5ac7209 Mon Sep 17 00:00:00 2001 From: Zack Li <39573601+zhiyuan8@users.noreply.github.com> Date: Fri, 16 Oct 2026 16:25:05 -0700 Subject: [PATCH 065/423] ci: update device matrix for Snapdragon jobs by removing SM8750 (#30620) * ci: update device matrix for Snapdragon jobs by removing SM8750 * ci: increase timeout for QDC device tests and update job timeout parameters --- .github/workflows/build-and-test-snapdragon.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build-and-test-snapdragon.yml b/.github/workflows/build-and-test-snapdragon.yml index d6420b65f141..29be95c8be02 100644 --- a/.github/workflows/build-and-test-snapdragon.yml +++ b/.github/workflows/build-and-test-snapdragon.yml @@ -104,11 +104,11 @@ jobs: name: Test on QDC Device (${{ matrix.device }}) needs: [android-ndk-snapdragon, linux-iot-snapdragon] runs-on: ubuntu-24.04-arm - timeout-minutes: 90 + timeout-minutes: 180 strategy: fail-fast: false matrix: - device: [SM8750, SM8850, QCS9075M] + device: [SM8850, QCS9075M] permissions: actions: read @@ -155,7 +155,7 @@ jobs: --pkg-dir pkg-snapdragon/llama.cpp \ --model-url "https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q4_0.gguf" \ --device ${{ matrix.device }} \ - ${{ startsWith(matrix.device, 'QCS') && '--retries 2 --retry-delay 300' || '' }} + ${{ startsWith(matrix.device, 'QCS') && '--job-timeout 5400 --retries 2 --retry-delay 300' || '' }} env: QDC_API_KEY: ${{ secrets.QDC_API_KEY }} From ea0c4c2c85fb183ec0fbd30a1af4e4275a18a064 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Sat, 17 Oct 2026 02:11:56 +0200 Subject: [PATCH 066/423] skills: add slice-model (#30677) * skills: add slice-model * exclude from py check * add to agents.md --- .flake8 | 2 + AGENTS.md | 3 + skills/slice-model/SKILL.md | 26 ++ skills/slice-model/slice_model.py | 558 ++++++++++++++++++++++++++++++ ty.toml | 1 + 5 files changed, 590 insertions(+) create mode 100644 skills/slice-model/SKILL.md create mode 100644 skills/slice-model/slice_model.py diff --git a/.flake8 b/.flake8 index 669d231f1f63..c7b5cf7f803d 100644 --- a/.flake8 +++ b/.flake8 @@ -5,6 +5,8 @@ exclude = # Do not traverse examples and tools examples, tools, + # Reference scripts for skills + skills, # Do not include package initializers __init__.py, # No need to traverse our git directory diff --git a/AGENTS.md b/AGENTS.md index 3ae583b82bae..36390a6b4ad6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -246,3 +246,6 @@ Chat template and parser: - [PEG parser](docs/development/parsing.md) - alternative to regex that llama.cpp uses to parse model's output - [Auto parser](docs/autoparser.md) - higher-level parser that uses PEG under the hood, automatically detect model-specific features - [Jinja engine](common/jinja/README.md) + +Testing: +- [Slice model](skills/slice-model/SKILL.md) - build a small slice (few layers/experts) of a big HF model without downloading the full weights, to test conversion and inference diff --git a/skills/slice-model/SKILL.md b/skills/slice-model/SKILL.md new file mode 100644 index 000000000000..f19564f90bd3 --- /dev/null +++ b/skills/slice-model/SKILL.md @@ -0,0 +1,26 @@ +--- +name: slice-model +description: Build a small test "slice" of a HF model (fewer layers/experts) by downloading only the needed tensors via HTTP range requests. Use when the user needs a small version of a big model to test conversion or inference. +--- + +# Slice a HF model for testing + +[slice_model.py](slice_model.py) is a **reference** implementation. Read it first: its docstring and `--help` explain how it works and what it does. Its heuristics only cover common layouts, and new models often have tensor names, shapes or config keys it does not handle. Do not run it blindly. If the model needs special handling, copy it to `tmp/slice_.py` and adapt the copy, or write a new script with the same approach. + +## Steps + +1. Ask the user for the repo, what the slice must exercise (e.g. a MoE layer, each layer type of a hybrid model, MTP, the vision encoder), the size budget and the output dir. +2. Inspect the model: run the script with `--inspect`, read `config.json`, and read the HF modeling code if anything is unclear. Check the plan against the pitfalls below. +3. Build the slice, then read the config patches in the log and check each one. +4. Verify: load the slice in transformers (no missing/unexpected weight warnings, no NaN logits), then convert it with `convert_hf_to_gguf.py` and run it. For numerical checks, see `examples/model-conversion/`. + +## Pitfalls + +- Layer counts in nested or unknown config keys, or several towers with the same layer count. +- Hybrid layer patterns (`layer_types`, `full_attention_interval`, ...): keep at least one layer of each type, and check the patched config still matches the kept layers. +- Dense leading layers (`first_k_dense_replace`) and MTP/nextn layers numbered after the regular layers. +- Every per-expert tensor must be sliced to the same expert count: experts, router weight and biases. Tensors with the expert dim not at dim 0 are not contiguous. +- Routing constraints: top-k, expert groups, shared experts, zero-compute experts. +- Quantized weights (FP8, MXFP4): slice the scale tensors in the same way as their weight. +- Sizes derived in modeling code instead of stored in the config: follow the exact formula if you shrink them. +- Slicing the vocab without the tokenizer breaks conversion. diff --git a/skills/slice-model/slice_model.py b/skills/slice-model/slice_model.py new file mode 100644 index 000000000000..88d696a5eda9 --- /dev/null +++ b/skills/slice-model/slice_model.py @@ -0,0 +1,558 @@ +#!/usr/bin/env python3 +""" +REFERENCE ONLY - build a small test "slice" of a HF safetensors model without downloading the full weights. + +The heuristics below (layer/expert detection, config patching) cover common layouts only. New models often use tensor names, shapes or config keys that this script does not know about. Always run --inspect first, check the plan, and adapt the script to the model. + +How it works: each safetensors shard starts with an 8-byte header length + JSON header, which gives the exact byte span of every tensor. We fetch only the headers, then use HTTP range requests to fetch the tensors we keep: + 1. keep only some layers (per numeric "family" in tensor names whose count matches a layer count in config.json), renumbered to 0..n-1 + 2. keep only the first E experts (per-expert tensors are dropped, stacked expert tensors and router weights are row-sliced along dim 0) + 3. optionally keep only the first V rows of vocab-sized tensors (off by default, see --vocab-size) +config.json is patched to match. Other small repo files are copied as-is. + +Usage: + python slice_model.py Qwen/Qwen3-0.6B --inspect + python slice_model.py Qwen/Qwen3-30B-A3B --num-layers 2 --num-experts 8 --dry-run + python slice_model.py Qwen/Qwen3-30B-A3B --keep-layers 0,3 --output ../models/qwen3moe-slice +""" +from __future__ import annotations + +import argparse +import concurrent.futures as cf +import fnmatch +import json +import logging +import os +import re +import struct +import sys +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import requests +from huggingface_hub import HfApi, get_token, hf_hub_url + +logging.basicConfig(level=logging.INFO, format="%(message)s") +logger = logging.getLogger("slice_model") + +INDEX_FILE = "model.safetensors.index.json" +CONFIG_FILE = "config.json" + +# files not worth copying into a slice +AUX_BLACKLIST = [ + # weights in other formats (the slice replaces them) + "*.bin", "*.pth", "*.pt", "*.ckpt", "*.h5", "*.msgpack", + "*.onnx", "*.tflite", "*.mlmodel", "*.mlpackage/*", "*.gguf", "*.nemo", + "onnx/*", "openvino/*", "coreml/*", "original/*", "metal/*", + # media / demo assets + "*.png", "*.jpg", "*.jpeg", "*.gif", "*.webp", "*.avif", "*.svg", + "*.mp4", "*.webm", "*.mov", "*.wav", "*.mp3", "*.flac", "*.ogg", "*.pdf", + ".gitattributes", +] +AUX_MAX_SIZE = 50 * 1024 * 1024 + +LAYER_CONTAINERS = {"layers", "layer", "blocks", "block", "h"} +EXPERT_ROW_KEYWORDS = ("expert", "router", "gate", "score") + +LAYER_KEYS = { + "num_hidden_layers", "num_layers", "n_layer", "n_layers", "num_layer", + "num_decoder_layers", "encoder_layers", "decoder_layers", "depth", + "num_encoder_layers", "num_attention_layers", "layer_count", +} +VOCAB_KEYS = {"vocab_size", "padded_vocab_size"} +NEXTN_KEYS = {"num_nextn_predict_layers", "num_mtp_layers"} +TOPK_SUBSTR = ("per_tok", "topk", "top_k", "moe_k") +EXPERT_KEY_EXCLUDE = ("per_tok", "topk", "top_k", "shared", "group", "intermediate", "size", "dim", "dtype") + + +def parse_size(text: str) -> int: + m = re.fullmatch(r"(\d+(?:\.\d+)?)\s*([kKmMgGtT]?)[bB]?", text.strip()) + if not m: + raise argparse.ArgumentTypeError(f"invalid size: {text!r}") + mult = {"": 1, "k": 1024, "m": 1024**2, "g": 1024**3, "t": 1024**4}[m.group(2).lower()] + return int(float(m.group(1)) * mult) + + +def human(n: float) -> str: + for unit in ("B", "KB", "MB", "GB"): + if n < 1024 or unit == "GB": + return f"{n:.1f}{unit}" if unit != "B" else f"{int(n)}B" + n /= 1024 + return f"{n}GB" + + +# -------------------------------------------------------------------------------------- +# HTTP +# -------------------------------------------------------------------------------------- + +_tls = threading.local() + + +def _session() -> requests.Session: + s = getattr(_tls, "s", None) + if s is None: + s = requests.Session() + _tls.s = s + return s + + +class HubClient: + def __init__(self, repo_id: str, revision: str, token: str | None): + self.repo_id = repo_id + self.revision = revision + self.headers = {"Authorization": f"Bearer {token}"} if token else {} + + def url(self, filename: str) -> str: + return hf_hub_url(self.repo_id, filename, revision=self.revision) + + def _request(self, filename: str, headers: dict, stream: bool = False, retries: int = 5): + last: Exception | None = None + for attempt in range(1, retries + 1): + try: + r = _session().get(self.url(filename), headers=headers, stream=stream, timeout=90, allow_redirects=True) + if r.status_code in (429, 500, 502, 503, 504): + raise requests.HTTPError(f"HTTP {r.status_code}", response=r) + r.raise_for_status() + return r + except Exception as e: # noqa: BLE001 + last = e + if attempt == retries or "404" in str(e) or "401" in str(e) or "403" in str(e): + break + time.sleep(2.0 * attempt) + raise RuntimeError(f"{self.repo_id}/{filename}: {last}") from last + + def get_range(self, filename: str, start: int, end_inclusive: int) -> bytes: + h = dict(self.headers) + h["Range"] = f"bytes={start}-{end_inclusive}" + return self._request(filename, h).content + + def get_json(self, filename: str) -> Any: + return json.loads(self._request(filename, dict(self.headers)).content) + + def download_file(self, filename: str, dest: Path): + r = self._request(filename, dict(self.headers), stream=True) + dest.parent.mkdir(parents=True, exist_ok=True) + with open(dest, "wb") as f: + for chunk in r.iter_content(chunk_size=4 * 1024 * 1024): + if chunk: + f.write(chunk) + + def stream_range_into_fd(self, filename: str, start: int, end_inclusive: int, fd: int, write_offset: int): + h = dict(self.headers) + h["Range"] = f"bytes={start}-{end_inclusive}" + r = self._request(filename, h, stream=True) + pos = write_offset + for chunk in r.iter_content(chunk_size=8 * 1024 * 1024): + if chunk: + os.pwrite(fd, chunk, pos) + pos += len(chunk) + if pos - write_offset != end_inclusive - start + 1: + raise RuntimeError(f"short read: {filename} bytes={start}-{end_inclusive}") + + +def fetch_shard_header(client: HubClient, filename: str) -> tuple[dict, int]: + (header_len,) = struct.unpack(" ConfigInfo: + info = ConfigInfo(cfg=cfg) + if cfg is None: + return info + for _d, key, val in walk_config(cfg): + if not isinstance(val, int) or isinstance(val, bool) or val <= 0: + continue + if key in LAYER_KEYS: + info.layer_counts.add(val) + elif key in NEXTN_KEYS: + info.nextn_counts.add(val) + elif key in VOCAB_KEYS: + info.vocab_sizes.add(val) + elif key == "first_k_dense_replace": + info.first_k_dense = val + elif "expert" in key and val > 1 and not any(x in key for x in EXPERT_KEY_EXCLUDE): + info.expert_counts.add(val) + return info + + +def patch_config(cfg: dict, families: dict[str, Family], new_vocab: int | None, orig_vocabs: set[int]) -> list[str]: + """Patch counts in-place to match the slice. Returns log lines.""" + log = [] + by_card: dict[int, Family] = {} + for fam in families.values(): + if fam.kept is not None: + by_card.setdefault(fam.card, fam) + by_card.setdefault(fam.regular_total, fam) + max_new_experts = max((f.new_count for f in families.values() if f.kind == "expert" and f.kept is not None), default=None) + + for d, key, val in list(walk_config(cfg)): + if isinstance(val, int) and not isinstance(val, bool): + fam = by_card.get(val) + if key in LAYER_KEYS and fam is not None and fam.kind == "layer": + d[key] = fam.new_regular + elif key == "first_k_dense_replace": + lfam = next((f for f in families.values() if f.kind == "layer"), None) + if lfam is None: + continue + d[key] = sum(1 for i in lfam.kept if i < val) + elif "expert" in key and fam is not None and fam.kind == "expert" and not any(x in key for x in EXPERT_KEY_EXCLUDE): + d[key] = fam.new_count + elif "expert" in key and any(x in key for x in TOPK_SUBSTR) and max_new_experts is not None and val > max_new_experts: + d[key] = max_new_experts + elif key in VOCAB_KEYS and new_vocab is not None and val in orig_vocabs: + d[key] = new_vocab + if d[key] != val: + log.append(f"{key}: {val} -> {d[key]}") + elif isinstance(val, list) and val and all(not isinstance(x, (dict, list)) for x in val): + # per-layer lists, e.g. layer_types + fam = by_card.get(len(val)) + if fam is not None and fam.kind == "layer" and len(val) in (fam.card, fam.regular_total): + d[key] = [val[i] for i in fam.kept if i < len(val)] + log.append(f"{key}: list[{len(val)}] -> list[{len(d[key])}]") + return log + + +# -------------------------------------------------------------------------------------- +# tensor name analysis +# -------------------------------------------------------------------------------------- + +@dataclass +class Family: + """One numeric slot in tensor names, e.g. 'model.layers.#'.""" + prefix: str + kind: str # "layer" | "expert" | "other" + indices: set[int] = field(default_factory=set) + kept: list[int] | None = None # None = keep all, no renumbering + renumber: dict[int, int] = field(default_factory=dict) + regular_total: int = 0 # layer count excluding MTP tail + new_regular: int = 0 + + @property + def card(self) -> int: + return len(self.indices) + + @property + def new_count(self) -> int: + return len(self.kept) if self.kept is not None else self.card + + +def name_slots(name: str): + """Yield (slot_pos, family_prefix, index) for each numeric path component.""" + tokens = name.split(".") + norm: list[str] = [] + for i, tok in enumerate(tokens): + if tok.isdigit(): + yield i, ".".join(norm), int(tok) + norm.append("#") + else: + norm.append(tok) + + +def build_families(names: list[str], cfg_info: ConfigInfo) -> dict[str, Family]: + families: dict[str, Family] = {} + for name in names: + tokens = name.split(".") + for i, prefix, idx in name_slots(name): + container = tokens[i - 1] if i > 0 else "" + fam = families.get(prefix) + if fam is None: + fam = Family(prefix=prefix, kind="expert" if "expert" in container else "other") + families[prefix] = fam + fam.indices.add(idx) + + acceptable_layer_cards: dict[int, int] = {} # card -> regular layer count + for c in cfg_info.layer_counts: + acceptable_layer_cards[c] = c + for m in cfg_info.nextn_counts: + acceptable_layer_cards[c + m] = c + + for fam in families.values(): + if fam.kind == "expert": + continue + container = fam.prefix.rsplit(".", 1)[-1] + if cfg_info.cfg is not None: + if fam.card in acceptable_layer_cards: + fam.kind = "layer" + fam.regular_total = acceptable_layer_cards[fam.card] + elif fam.card >= 4 and container in LAYER_CONTAINERS: + fam.kind = "layer" + fam.regular_total = fam.card + return families + + +def plan_families(families: dict[str, Family], num_layers: int, keep_layers: list[int] | None, num_experts: int, first_k_dense: int): + for fam in families.values(): + if fam.kind == "layer": + if keep_layers is not None: + kept = [i for i in keep_layers if i < fam.regular_total] + else: + base = min(num_layers, fam.regular_total) + kept = list(range(base)) + # keep one MoE layer for models whose first k layers are dense + if 0 < first_k_dense < fam.regular_total and first_k_dense >= base >= 1: + kept = list(range(base - 1)) + [first_k_dense] + fam.new_regular = len(kept) + # keep the MTP tail (layer indices beyond the regular count) + kept += [i for i in sorted(fam.indices) if i >= fam.regular_total] + fam.kept = sorted(set(kept)) + elif fam.kind == "expert": + fam.kept = sorted(fam.indices)[:num_experts] + fam.new_regular = len(fam.kept) + else: + fam.kept = None + continue + fam.renumber = {old: new for new, old in enumerate(fam.kept)} + + +def slice_tensor_name(name: str, families: dict[str, Family]) -> str | None: + """Return the renumbered output name, or None if the tensor is dropped.""" + tokens = name.split(".") + for i, prefix, idx in name_slots(name): + fam = families[prefix] + if fam.kept is None: + continue + if idx not in fam.renumber: + return None + tokens[i] = str(fam.renumber[idx]) + return ".".join(tokens) + + +# -------------------------------------------------------------------------------------- +# slicing +# -------------------------------------------------------------------------------------- + +@dataclass +class TensorPlan: + src_name: str + out_name: str + dtype: str + shape: list[int] + orig_shape: list[int] + size: int + src_file: str + seg_start: int # absolute byte range in src_file + seg_end: int # inclusive + + +def pick_weight_files(filenames: list[str]) -> tuple[list[str], bool]: + """Return (weight files, has_index). Preference: index -> model.safetensors -> any root safetensors.""" + if INDEX_FILE in filenames: + return [], True + if "model.safetensors" in filenames: + return ["model.safetensors"], False + root_st = [f for f in filenames if f.endswith(".safetensors") and "/" not in f] + if "consolidated.safetensors" in root_st: + return ["consolidated.safetensors"], False + return root_st[:1], False + + +def row_slice(shape: list[int], size: int, keep_rows: int) -> tuple[list[int], int] | None: + # only a prefix of dim 0 is a single contiguous byte range + if len(shape) < 1 or shape[0] <= keep_rows or size % shape[0] != 0: + return None + row_bytes = size // shape[0] + return [keep_rows] + shape[1:], keep_rows * row_bytes + + +def build_plan(args, families, headers, weight_files, src_file_of, cfg_info) -> list[TensorPlan]: + plan_families(families, args.num_layers, args.keep_layers, args.num_experts, cfg_info.first_k_dense) + plans = [] + for fname in weight_files: + header, data_start = headers[fname] + for name, meta in header.items(): + out_name = slice_tensor_name(name, families) + if out_name is None: + continue + start, end = meta["data_offsets"] + shape, size = list(meta["shape"]), end - start + new_rows = None + if args.vocab_size and shape and shape[0] > args.vocab_size \ + and any(shape[0] >= vs and shape[0] - vs <= 1024 for vs in cfg_info.vocab_sizes): + new_rows = args.vocab_size + elif shape and shape[0] in cfg_info.expert_counts and shape[0] > args.num_experts \ + and any(k in name for k in EXPERT_ROW_KEYWORDS): + new_rows = args.num_experts + if new_rows is not None: + sliced = row_slice(shape, size, new_rows) + if sliced is not None: + shape, size = sliced + plans.append(TensorPlan(name, out_name, meta["dtype"], shape, list(meta["shape"]), size, src_file_of[name], + data_start + start, data_start + start + size - 1)) + return plans + + +def print_inspect(cfg_info: ConfigInfo, families: dict[str, Family], plans: list[TensorPlan], all_names: list[str]): + logger.info("=== config analysis ===") + logger.info(f"layer counts: {sorted(cfg_info.layer_counts)} nextn: {sorted(cfg_info.nextn_counts)} " + f"experts: {sorted(cfg_info.expert_counts)} vocab: {sorted(cfg_info.vocab_sizes)} first_k_dense: {cfg_info.first_k_dense}") + logger.info("\n=== numeric families in tensor names ===") + for fam in sorted(families.values(), key=lambda f: f.prefix): + kept = "all" if fam.kept is None else f"{fam.kept[:16]}{'...' if len(fam.kept) > 16 else ''}" + logger.info(f"{fam.prefix:60s} kind={fam.kind:6s} count={fam.card:4d} kept={kept}") + logger.info("\n=== kept tensors (first index of each family only) ===") + seen = set() + for t in plans: + key = re.sub(r"\.\d+\.", ".#.", t.out_name) + if key in seen and t.shape == t.orig_shape: + continue + seen.add(key) + mark = f" (sliced from {t.orig_shape})" if t.shape != t.orig_shape else "" + logger.info(f"{t.out_name:70s} {t.dtype:5s} {t.shape}{mark}") + kept_src = {t.src_name for t in plans} + dropped = [n for n in all_names if n not in kept_src] + logger.info(f"\n{len(plans)}/{len(all_names)} tensors kept, {len(dropped)} dropped") + + +def write_slice(out_dir: Path, client: HubClient, plans: list[TensorPlan], jobs: int): + header_out: dict[str, Any] = {} + offset = 0 + for t in plans: + header_out[t.out_name] = {"dtype": t.dtype, "shape": t.shape, "data_offsets": [offset, offset + t.size]} + offset += t.size + header_out["__metadata__"] = {"format": "pt"} + hbytes = json.dumps(header_out, separators=(",", ":")).encode() + hbytes += b" " * ((-(len(hbytes) + 8)) % 8) + data_start = 8 + len(hbytes) + + shard_path = out_dir / "model.safetensors" + with open(shard_path, "wb") as f: + f.write(struct.pack(" args.max_size: + logger.error(f"slice is {human(est)} > --max-size {human(args.max_size)}") + return 1 + if args.inspect or args.dry_run: + return 0 + + out_dir = args.output or Path("tmp") / "hf_slices" / args.repo.replace("/", "--") + out_dir.mkdir(parents=True, exist_ok=True) + write_slice(out_dir, client, plans, args.jobs) + + skipped_aux = [] + for fname in filenames: + if fname in weight_files or fname in (INDEX_FILE, CONFIG_FILE) or fname.endswith(".safetensors"): + continue + if any(fnmatch.fnmatch(fname, p) or fnmatch.fnmatch(Path(fname).name, p) for p in AUX_BLACKLIST): + continue + if siblings[fname] > AUX_MAX_SIZE: + skipped_aux.append(f"{fname} ({human(siblings[fname])})") + continue + client.download_file(fname, out_dir / fname) + for f in skipped_aux: + logger.warning(f"skipped large aux file {f}") + + patch_log = [] + if cfg is not None: + new_vocab = args.vocab_size if args.vocab_size and any(vs > args.vocab_size for vs in cfg_info.vocab_sizes) else None + patch_log = patch_config(cfg, families, new_vocab, cfg_info.vocab_sizes) + with open(out_dir / CONFIG_FILE, "w") as f: + json.dump(cfg, f, indent=2) + f.write("\n") + for line in patch_log: + logger.info(f"config: {line}") + + with open(out_dir / ".slice_meta.json", "w") as f: + json.dump({ + "repo": args.repo, "revision": info.sha, + "tensors": len(plans), "total_tensors": len(all_names), "size": est, + "config_patches": patch_log, "skipped_aux": skipped_aux, + }, f, indent=2) + logger.info(f"wrote {out_dir}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/ty.toml b/ty.toml index 340b0649d334..e2036cd2ed05 100644 --- a/ty.toml +++ b/ty.toml @@ -8,6 +8,7 @@ deprecated = "warn" [src] exclude = [ "./tools/mtmd/legacy-models/**", + "./skills/**", ] [[overrides]] From f74e2fa6a31cc33ab4ee4ae516641cfe62df3924 Mon Sep 17 00:00:00 2001 From: Marcus Date: Fri, 16 Oct 2026 19:03:41 -0700 Subject: [PATCH 067/423] sycl: fix prompt-processing regression after Q4_0/Q8_0 weight reorder (#30374) --- ggml/src/ggml-sycl/convert.cpp | 30 ++++++------- ggml/src/ggml-sycl/dequantize.hpp | 74 ++++++++++++------------------- 2 files changed, 43 insertions(+), 61 deletions(-) diff --git a/ggml/src/ggml-sycl/convert.cpp b/ggml/src/ggml-sycl/convert.cpp index 67dfeb9f943e..8ba26d3a9f23 100644 --- a/ggml/src/ggml-sycl/convert.cpp +++ b/ggml/src/ggml-sycl/convert.cpp @@ -158,16 +158,16 @@ static void dequantize_row_q4_0_sycl_reorder(const void *vx, dst_t *y, const int dpct::has_capability_or_fail(stream->get_device(), {sycl::aspect::fp16}); - int constexpr WARP_K = WARP_SIZE * QK4_0; - const int n_warp = (k + WARP_K - 1) / WARP_K; - GGML_ASSERT(k % 2 == 0); - stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, n_warp) * - sycl::range<3>(1, 1, WARP_SIZE), - sycl::range<3>(1, 1, WARP_SIZE)), - [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{ + GGML_ASSERT(k % QK4_0 == 0); + const int64_t n_groups = ceil_div(k / 2, SYCL_DEQUANTIZE_BLOCK_SIZE); + // decrease global range when it exceeds the max int; the kernel loops to cover the rest + const int64_t wg_size = downsample_sycl_global_range(n_groups, SYCL_DEQUANTIZE_BLOCK_SIZE); + stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, n_groups) * + sycl::range<3>(1, 1, wg_size), + sycl::range<3>(1, 1, wg_size)), + [=](sycl::nd_item<3> item_ct1) { dequantize_block_q4_0_reorder(vx, y, k, item_ct1); }); - } template @@ -177,16 +177,16 @@ static void dequantize_row_q8_0_sycl_reorder(const void *vx, dst_t *y, const int dpct::has_capability_or_fail(stream->get_device(), {sycl::aspect::fp16}); - int constexpr WARP_K = WARP_SIZE * QK8_0; - const int n_warp = (k + WARP_K - 1) / WARP_K; GGML_ASSERT(k % QK8_0 == 0); - stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, n_warp) * - sycl::range<3>(1, 1, WARP_SIZE), - sycl::range<3>(1, 1, WARP_SIZE)), - [=](sycl::nd_item<3> item_ct1) [[sycl::reqd_sub_group_size(WARP_SIZE)]]{ + const int64_t n_groups = ceil_div(k / 2, SYCL_DEQUANTIZE_BLOCK_SIZE); + // decrease global range when it exceeds the max int; the kernel loops to cover the rest + const int64_t wg_size = downsample_sycl_global_range(n_groups, SYCL_DEQUANTIZE_BLOCK_SIZE); + stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, n_groups) * + sycl::range<3>(1, 1, wg_size), + sycl::range<3>(1, 1, wg_size)), + [=](sycl::nd_item<3> item_ct1) { dequantize_block_q8_0_reorder(vx, y, k, item_ct1); }); - } template diff --git a/ggml/src/ggml-sycl/dequantize.hpp b/ggml/src/ggml-sycl/dequantize.hpp index 14529a26beb3..48ee651ad68c 100644 --- a/ggml/src/ggml-sycl/dequantize.hpp +++ b/ggml/src/ggml-sycl/dequantize.hpp @@ -818,61 +818,43 @@ static void dequantize_block_q4_0(const void * __restrict__ vx, dst_t * __restri } template -static void dequantize_block_q4_0_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t nb32, +static void dequantize_block_q4_0_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t k, const sycl::nd_item<3> &item_ct1) { - - const int64_t i = item_ct1.get_group(2); - auto k=nb32; - // assume 32 threads - const int64_t tid = item_ct1.get_local_id(2); - const int lane_ib = i * WARP_SIZE + tid; - - if (lane_ib >= k / QK4_0) { - return; + // one thread per qs byte, like dequantize_block; the reorder layout makes the byte reads contiguous + // grid-stride loop because the launcher may shrink the work-group to keep the global range below INT_MAX + // float math, not dequantize_q4_0_reorder: that helper computes in half with GGML_SYCL_F16, which changes f32 output + const sycl::half * d_ptr = (const sycl::half *) ((const uint8_t *) vx + k/2); + const int64_t stride = item_ct1.get_local_range(2) * item_ct1.get_group_range(2); + + for (int64_t iq = item_ct1.get_global_id(2); iq < k/2; iq += stride) { + const int64_t ib = iq / (QK4_0/2); + const int64_t iqs = iq % (QK4_0/2); + const int vq = ((const uint8_t *) vx)[iq]; + const float d = float(d_ptr[ib]); + + dst_t * y = yy + ib*QK4_0 + iqs; + y[0] = d * ((vq & 0xF) - 8); + y[QK4_0/2] = d * ((vq >> 4) - 8); } - - dst_t * y_ptr = yy + lane_ib * QK4_0; - - auto qs = (const uint8_t*)vx + lane_ib * QK4_0 / 2; - auto s_ptr = (const sycl::half*)((const uint8_t*)vx + k / 2) + lane_ib; - - const float d = float(*s_ptr); - -#pragma unroll - for (int l = 0; l < QK4_0 / 2; ++l) { - int vq = qs[l]; - y_ptr[l + 0] = d * ((vq & 0xF) - 8); - y_ptr[l + 16] = d * ((vq >> 4) - 8); - } - } // Dequantize Q8_0 from reorder layout: [all qs (k bytes)][all d values] -// Each thread handles one block of QK8_0 elements. template static void dequantize_block_q8_0_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t k, const sycl::nd_item<3> &item_ct1) { - - const int64_t i = item_ct1.get_group(2); - const int64_t tid = item_ct1.get_local_id(2); - const int lane_ib = i * WARP_SIZE + tid; - - if (lane_ib >= k / QK8_0) { - return; - } - - dst_t * y_ptr = yy + lane_ib * QK8_0; - - auto qs = (const int8_t*)vx + lane_ib * QK8_0; - auto s_ptr = (const sycl::half*)((const uint8_t*)vx + k) + lane_ib; - - const float d = float(*s_ptr); - -#pragma unroll - for (int l = 0; l < QK8_0; ++l) { - y_ptr[l] = d * qs[l]; + // two values per thread, like dequantize_block; the reorder layout keeps all qs contiguous + // grid-stride loop because the launcher may shrink the work-group to keep the global range below INT_MAX + // float math, not dequantize_q8_0_reorder: that helper computes in half with GGML_SYCL_F16, which changes f32 output + const int8_t * qs = (const int8_t *) vx; + const sycl::half * d_ptr = (const sycl::half *) ((const uint8_t *) vx + k); + const int64_t stride = 2 * item_ct1.get_local_range(2) * item_ct1.get_group_range(2); + + for (int64_t i = 2 * item_ct1.get_global_id(2); i < k; i += stride) { + const float d = float(d_ptr[i / QK8_0]); + + yy[i + 0] = d * qs[i + 0]; + yy[i + 1] = d * qs[i + 1]; } - } template From 1268edac579110aeb1407f7928f53193ed4e28d4 Mon Sep 17 00:00:00 2001 From: Gabe Goodhart Date: Sat, 17 Oct 2026 01:14:32 -0600 Subject: [PATCH 068/423] model : Granite switch decision 256 options (#30693) * feat(convert): Attempt to come up with 256 single-token, semantic option markers Each switch model needs to identify a collection of single-token option markers that can be mapped to their respective logits for the output probability distribution. The simple upper/lower/numbers set only yields 66 unique markers, but the model can support up to 255. The challenge is that which combinations of characters result in single-token markers is model- specific, so we do a discovery exercise at conversion time by looking through the specific model's vocab for single-token sequences of 1+ plain letter strings (eg `A`, `AA`, `BB`, etc). Branch: GraniteSwitchDecisionFullOptions AI-usage: none Signed-off-by: Gabe Goodhart * docs: Note granite-switch with 256 option markers Branch: GraniteSwitchDecisionFullOptions AI-usage: none Signed-off-by: Gabe Goodhart * style: Remove extra newline Branch: GraniteSwitchDecisionFullOptions AI-usage: none Signed-off-by: Gabe Goodhart * feat(convert): Use the computed labels in the template so c++ side is source of truth Branch: GraniteSwitchDecisionFullOptions AI-usage: full (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart * feat: Use LEV case for GRANITE_SWITCH w/ computed labels Branch: GraniteSwitchDecisionFullOptions AI-usage: full (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart * docs: Consolidate g-s into 255 bucket Branch: GraniteSwitchDecisionFullOptions AI-usage: full (Claude + Sonnet 5.5) Signed-off-by: Gabe Goodhart --- conversion/granite.py | 7 ++++--- tools/server/README.md | 2 +- tools/server/server-decision.cpp | 7 +++++-- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/conversion/granite.py b/conversion/granite.py index d2ef7c836c8d..4c6c5a0e4cc5 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -296,19 +296,20 @@ def _has_decision_adapter(self) -> bool: def set_vocab(self): super().set_vocab() if self._has_decision_adapter(): - self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}]) + systemone_template = self._systemone_template() + logger.info("System One Template: %s", systemone_template) + self.gguf_writer.add_chat_template([{"name": "systemone", "template": systemone_template}]) def _systemone_template(self) -> str: # prompt of the decision aLoRA: a state turn, then a question turn that opens with the control token # (it replaces the first token of "### DECISION"), the answer is the label token that comes next description = jinja_str_or_json("o.description") option = ( - "{% if type != 'noul' %}{{ letters[loop.index0] }}: {% if o.description %}" + description + "{% endif %}" + "{% if type != 'noul' %}{{ o.label }}: {% if o.description %}" + description + "{% endif %}" "{% elif o.key == 'true' %}A: {% if o.description %}" + description + "{% else %}The statement is true.{% endif %}" "{% else %}B: {% if o.description %}" + description + "{% else %}The statement is false.{% endif %}{% endif %}" ) return ( - "{% set letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz' %}" "<|start_of_role|>user<|end_of_role|>State:\n" + jinja_str_or_json("state") + "<|end_of_text|>\n" "<|start_of_role|>user<|end_of_role|><|decision|> DECISION\nQuestion:\n" + jinja_str_or_json("instructions") + "{% if type == 'score' %} Rate along the ordered levels below (lowest first).{% endif %}" diff --git a/tools/server/README.md b/tools/server/README.md index e4a0c418bf10..c42336c7f217 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -1723,7 +1723,7 @@ Follows the [TypeSafe API](https://docs.typesafe.ai/api), streaming is not suppo The questions of a request are answered independently, an answer does not depend on the other questions. The exception is clef: it reads all the questions in one prompt and decides them jointly. -The number of options of a `choice` question is limited by the model, for example: 52 for openjev and granite-switch, 255 for laya, clef, pplx-decider, vllm-decision2, vllm-d3, lfm2-d1 and lfm2-d1-omni. For laya, long questions and options are truncated to the token budget the model was trained with. +The number of options of a `choice` question is limited by the model, for example: 52 for openjev, 255 for granite-switch, laya, clef, pplx-decider, vllm-decision2, vllm-d3, lfm2-d1 and lfm2-d1-omni. For laya, long questions and options are truncated to the token budget the model was trained with. For laya, clef, vllm-d3 and lfm2-d1-omni, the whole prompt is evaluated in one batch: it must fit in `--ubatch-size`. An lfm2-d1-omni prompt is cut to 16384 tokens. For vllm-decision2, the options of a question and the end of the prompt must fit in one batch. A server that runs clef only serves this endpoint, text generation is not available. A decision model that is also a text model, such as granite-switch, serves both and reports `["text", "decisions"]`. diff --git a/tools/server/server-decision.cpp b/tools/server/server-decision.cpp index 3905c81b34c3..be9bc8ef4912 100644 --- a/tools/server/server-decision.cpp +++ b/tools/server/server-decision.cpp @@ -87,7 +87,7 @@ void server_decision_context::init(const llama_model * model) { bias[level] = std::strtof(val, nullptr); } - if (model_type == COMMON_DECISION_TYPE_OPENJEV || model_type == COMMON_DECISION_TYPE_GRANITE_SWITCH) { + if (model_type == COMMON_DECISION_TYPE_OPENJEV) { // one letter per option, each must be a single token const std::string letters = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; for (const char c : letters) { @@ -99,7 +99,7 @@ void server_decision_context::init(const llama_model * model) { } n_options_max = labels.size(); noul_true_first = true; - } else if (model_type == COMMON_DECISION_TYPE_LEV || model_type == COMMON_DECISION_TYPE_NIMBLE || model_type == COMMON_DECISION_TYPE_PPLX_DECIDER || model_type == COMMON_DECISION_TYPE_VLLM_D3) { + } else if (model_type == COMMON_DECISION_TYPE_LEV || model_type == COMMON_DECISION_TYPE_GRANITE_SWITCH || model_type == COMMON_DECISION_TYPE_NIMBLE || model_type == COMMON_DECISION_TYPE_PPLX_DECIDER || model_type == COMMON_DECISION_TYPE_VLLM_D3) { // label codes are A..Z then AA..ZZ, only the ones that are a single token are used std::vector codes; for (char a = 'A'; a <= 'Z'; a++) { @@ -121,6 +121,9 @@ void server_decision_context::init(const llama_model * model) { if (model_type == COMMON_DECISION_TYPE_VLLM_D3) { canonical_legend = true; } + if (model_type == COMMON_DECISION_TYPE_GRANITE_SWITCH) { + noul_true_first = true; + } } else if (model_type == COMMON_DECISION_TYPE_KEV) { // the hidden state of an option is read at the token that ends it const auto toks = common_tokenize(vocab, "<|box_end|>", false, true); From 83f9389419ef93117d1444f87adfcd9f03ec0509 Mon Sep 17 00:00:00 2001 From: Sandro Hanea <40202887+sandrohanea@users.noreply.github.com> Date: Sat, 17 Oct 2026 11:41:53 +0200 Subject: [PATCH 069/423] ci : added ABI reports as github artifact (#30619) --- .github/workflows/make-release.yml | 9 +++++++++ scripts/check-release-apiabi.sh | 14 +++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/.github/workflows/make-release.yml b/.github/workflows/make-release.yml index 57c3a04a577f..2e5e44caaab8 100644 --- a/.github/workflows/make-release.yml +++ b/.github/workflows/make-release.yml @@ -61,8 +61,17 @@ jobs: RELEASE_BRANCH: ${{ github.ref_name }} SKIP_APIABI_CHECK: ${{ github.event.inputs.skip_apiabi_check }} APIABI_COMPARE_TAG: ${{ github.event.inputs.apiabi_compare_tag }} + APIABI_REPORT_DIR: ${{ runner.temp }}/apiabi-reports REQUIRE_DOCKER: ${{ github.event.inputs.require_docker }} + - name: Upload API/ABI reports + if: ${{ always() && github.event.inputs.skip_apiabi_check != 'true' }} + uses: actions/upload-artifact@v7 + with: + name: apiabi-reports + path: ${{ runner.temp }}/apiabi-reports + if-no-files-found: ignore + - name: Create release tag if: ${{ github.event.inputs.dry_run == 'false' }} run: | diff --git a/scripts/check-release-apiabi.sh b/scripts/check-release-apiabi.sh index f4803adc5155..c624d12c5f10 100755 --- a/scripts/check-release-apiabi.sh +++ b/scripts/check-release-apiabi.sh @@ -15,6 +15,7 @@ # Environment: # SKIP_APIABI_CHECK: set to 1 or true to skip # APIABI_COMPARE_TAG: equivalent to --tag (used by CI) +# APIABI_REPORT_DIR: preserve reports and dumps in this directory, even if the comparison fails set -euo pipefail @@ -113,4 +114,15 @@ cmake --build "$BUILD_NEW" --parallel "$(nproc)" (cd "$WORKTREE_DIR" && "$SCRIPT_DIR/check-apiabi-compat.sh" --include-path ggml/include --generate "$BUILD_OLD" "${LIBS[@]}") (cd "$REPO_ROOT" && "$SCRIPT_DIR/check-apiabi-compat.sh" --include-path ggml/include --generate "$BUILD_NEW" "${LIBS[@]}") -(cd "$REPO_ROOT" && "$SCRIPT_DIR/check-apiabi-compat.sh" "${CHECK_FLAGS[@]}" --check "$BUILD_OLD" "$BUILD_NEW") + +CHECK_RESULT=0 +(cd "$REPO_ROOT" && "$SCRIPT_DIR/check-apiabi-compat.sh" "${CHECK_FLAGS[@]}" --check "$BUILD_OLD" "$BUILD_NEW") || CHECK_RESULT=$? + +if [[ -n "${APIABI_REPORT_DIR:-}" ]]; then + mkdir -p "$APIABI_REPORT_DIR/old" "$APIABI_REPORT_DIR/new" + cp -R "$BUILD_OLD/apiabi/." "$APIABI_REPORT_DIR/old/" + cp -R "$BUILD_NEW/apiabi/." "$APIABI_REPORT_DIR/new/" + echo "API/ABI reports saved to $APIABI_REPORT_DIR" +fi + +exit "$CHECK_RESULT" From 62318cd57dc63900ecb153272c0b10f78550855d Mon Sep 17 00:00:00 2001 From: Xue Zhenliang Date: Sat, 17 Oct 2026 19:00:17 +0800 Subject: [PATCH 070/423] model : add Kolibri-1 (Aleph-Alpha) support (#30673) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * model : add Kolibri-1 (Aleph-Alpha) support Kolibri-1 is a bilingual German/English MoE reasoning model with 78B parameters and 3.46B active per token: 50 layers, 384 routed experts plus one shared expert. Attention interleaves four sliding layers at a 513 window with one full layer. The sliding layers rotate with NeoX RoPE, the full layers are NoPE. GQA runs 48 query heads against 4 KV heads at head_dim 128 with a QK-RMSNorm, and every block carries four RMSNorms, so the attention and the MLP output are normalized before they rejoin the residual stream. Routing picks the experts from logits + expert_bias and weights them with the sigmoid of the unbiased logits, without renorm. The converter reuses AfmoeModel for the per-expert checkpoint layout and the expert selection bias, and writes the shared expert and the per layer sliding window pattern. The released FP8 checkpoint needs no special code, it goes through the generic weight_scale_inv path. Three tensor names are added to the mapping (post_attn_norm, post_ffn_norm, moe.router.expert_bias) and the 128k vocabulary registers the qwen2 pre-tokenizer by hash. The arch is registered in test-llama-archs, with GQA and the sliding window pattern. Assisted-by: DeepSeek Harness (DeepSeek V4.1 Flash) * Update src/models/kolibri1.cpp Co-authored-by: Sigbjørn Skjæret * Add fusion records for kolibri1 * Update src/models/kolibri1.cpp Co-authored-by: Sigbjørn Skjæret --------- Co-authored-by: Sigbjørn Skjæret --- conversion/__init__.py | 1 + conversion/base.py | 3 + conversion/kolibri1.py | 23 ++++ convert_hf_to_gguf_update.py | 2 + gguf-py/gguf/constants.py | 25 +++++ gguf-py/gguf/tensor_mapping.py | 3 + src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/llama-model.cpp | 4 + src/llama-model.h | 1 + src/models/kolibri1.cpp | 187 +++++++++++++++++++++++++++++++++ src/models/models.h | 13 +++ tests/fusion/MTL.csv | 3 + tests/test-llama-archs.cpp | 10 +- 14 files changed, 275 insertions(+), 2 deletions(-) create mode 100644 conversion/kolibri1.py create mode 100644 src/models/kolibri1.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index 01c8db29680c..e3c10d5fc54e 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -146,6 +146,7 @@ "JinaBertModel": "bert", "JinaEmbeddingsV5Model": "bert", "K2HorizonForCausalLM": "k2_horizon", + "Kolibri1ForCausalLM": "kolibri1", "KORMoForCausalLM": "qwen", "KimiK25ForConditionalGeneration": "deepseek", "KimiK3ForConditionalGeneration": "kimi_k3", diff --git a/conversion/base.py b/conversion/base.py index e3d265157add..72c1256ba74f 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1872,6 +1872,9 @@ def get_vocab_base_pre(self, tokenizer) -> str: if chkhsh == "f5c965e4198da6e7525072d87d1ecf4c51d981fd486c7c4ec728ef9cec7e03fc": # ref: https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B res = "gpt-4o" + if chkhsh == "6e040dfe72e4b85855588c53acf4909ac4e98e55bc3a33cd7db499180dc42a78": + # ref: https://huggingface.co/Aleph-Alpha/Kolibri-1 + res = "qwen2" if chkhsh == "b3d1dd861f1d4c5c0d2569ce36baf3f90fe8a102db3de50dd71ff860d91be3df": # ref: https://huggingface.co/aari1995/German_Semantic_V3 res = "jina-v2-de" diff --git a/conversion/kolibri1.py b/conversion/kolibri1.py new file mode 100644 index 000000000000..845501f4ef38 --- /dev/null +++ b/conversion/kolibri1.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from .afmoe import AfmoeModel +from .base import ModelBase, gguf + + +@ModelBase.register("Kolibri1ForCausalLM") +@ModelBase.example("Aleph-Alpha/Kolibri-1") +class Kolibri1Model(AfmoeModel): + model_arch = gguf.MODEL_ARCH.KOLIBRI1 + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # AfmoeModel reads these keys, this config stores the same values under other names + self.hparams["num_shared_experts"] = 1 + self.hparams["route_norm"] = self.hparams["norm_topk_prob"] + self.hparams["intermediate_size"] = self.hparams["moe_intermediate_size"] + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_expert_shared_feed_forward_length(self.hparams["shared_expert_intermediate_size"]) + self.gguf_writer.add_sliding_window_pattern([t == "sliding_attention" for t in self.hparams["layer_types"]]) + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) diff --git a/convert_hf_to_gguf_update.py b/convert_hf_to_gguf_update.py index d49fb8f3ff3e..d3aa3c3807f0 100755 --- a/convert_hf_to_gguf_update.py +++ b/convert_hf_to_gguf_update.py @@ -191,6 +191,8 @@ class TOKENIZER_TYPE(IntEnum): {"name": "grok-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/alvarobartt/grok-2-tokenizer", "chkhsh": "66b8d4e19ab16c3bfd89bce5d785fb7e0155e8648708a1f42077cb9fe002c273"}, # muse-spark uses the l4_202k tokenizer, whose pre-tokenizer regex is gpt-4o's {"name": "gpt-4o", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B", "chkhsh": "f5c965e4198da6e7525072d87d1ecf4c51d981fd486c7c4ec728ef9cec7e03fc"}, + # kolibri1 uses a 128k vocabulary with the qwen2 pre-tokenizer regex + {"name": "qwen2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/Aleph-Alpha/Kolibri-1", "chkhsh": "6e040dfe72e4b85855588c53acf4909ac4e98e55bc3a33cd7db499180dc42a78"}, # jina-v2-de variants {"name": "jina-v2-de", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/aari1995/German_Semantic_V3", "chkhsh": "b3d1dd861f1d4c5c0d2569ce36baf3f90fe8a102db3de50dd71ff860d91be3df"}, {"name": "gpt-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/evilfreelancer/ruGPT3XL", "chkhsh": "0fe1cf6eda062318a1af7270f3331a85c539a01778ff948e24388e949c5282f4"}, diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index bfeab6adfde6..32e088d8820a 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -701,6 +701,7 @@ class MODEL_ARCH(IntEnum): QWEN3TTS = auto() POCKETTTS = auto() K2HORIZON = auto() + KOLIBRI1 = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -1518,6 +1519,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.QWEN3TTS: "qwen3tts", MODEL_ARCH.POCKETTTS: "pockettts", MODEL_ARCH.K2HORIZON: "k2-horizon", + MODEL_ARCH.KOLIBRI1: "kolibri1", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -6053,6 +6055,29 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_UP_SHEXP, MODEL_TENSOR.FFN_DOWN_SHEXP, ], + MODEL_ARCH.KOLIBRI1: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_POST_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + ], } # tensors that will not be serialized diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index ab6abf9cb704..3586d65285ca 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -416,6 +416,7 @@ class TensorNameMap: "model.layers.{bid}.post_attention_layernorm", # gemma2 olmo2 # ge "layers.{bid}.post_attention_layernorm", # embeddinggemma "model.layers.{bid}.post_self_attn_layernorm", # glm-4-0414 + "model.layers.{bid}.post_attn_norm", # kolibri1 "model.layers.layers.{bid}.post_mixer_norm.weight", # plamo2 ), @@ -498,6 +499,7 @@ class TensorNameMap: "model.layers.{bid}.post_feedforward_layernorm", # gemma2 olmo2 "layers.{bid}.post_feedforward_layernorm", # embeddinggemma "model.layers.{bid}.post_mlp_layernorm", # glm-4-0414 + "model.layers.{bid}.post_ffn_norm", # kolibri1 "model.layers.layers.{bid}.post_mlp_norm.weight", # plamo2 "model.layers.{bid}.feed_forward.up_proj", "model.layers.{bid}.post_moe_norm", # grok-2 @@ -542,6 +544,7 @@ class TensorNameMap: "model.layers.{bid}.mlp.moe_statics.e_score_correction", # ernie4.5-moe "model.layers.{bid}.mlp.gate.expert_bias", # bailingmoe2 "model.layers.{bid}.mlp.expert_bias", # afmoe + "model.layers.{bid}.moe.router.expert_bias", # kolibri1 "model.layers.{bid}.feed_forward.expert_bias", # lfm2moe "model.layers.{bid}.block_sparse_moe.e_score_correction", # minimax-m2 "backbone.layers.{bid}.mixer.gate.e_score_correction", # nemotron-h-moe diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index bef2d45c0ba2..ffa35bdea7d4 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -166,6 +166,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_POCKETTTS, "pockettts" }, { LLM_ARCH_K2_HORIZON, "k2-horizon" }, { LLM_ARCH_BEAM, "beam" }, + { LLM_ARCH_KOLIBRI1, "kolibri1" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; diff --git a/src/llama-arch.h b/src/llama-arch.h index c305c17663cb..14ce29ab3238 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -171,6 +171,7 @@ enum llm_arch { LLM_ARCH_HRM_TEXT, LLM_ARCH_K2_HORIZON, LLM_ARCH_BEAM, + LLM_ARCH_KOLIBRI1, LLM_ARCH_UNKNOWN, }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 099c7a052f46..803a8329e9db 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -360,6 +360,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_spark2_5(params); case LLM_ARCH_K2_HORIZON: return new llama_model_k2_horizon(params); + case LLM_ARCH_KOLIBRI1: + return new llama_model_kolibri1(params); default: throw std::runtime_error(std::string("unsupported model architecture: '") + llm_arch_name(arch) + "'"); } @@ -1002,6 +1004,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_10B_128x3_66B: return "10B+128x3.66B"; case LLM_TYPE_57B_A14B: return "57B.A14B"; case LLM_TYPE_816B_A42B: return "816B.A42B"; + case LLM_TYPE_78B_A3_5B: return "78B.A3.5B"; case LLM_TYPE_17B_16E: return "17Bx16E (Scout)"; case LLM_TYPE_17B_128E: return "17Bx128E (Maverick)"; case LLM_TYPE_A13B: return "A13B"; @@ -3615,6 +3618,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_MAPLE: case LLM_ARCH_HRM_TEXT: case LLM_ARCH_BEAM: + case LLM_ARCH_KOLIBRI1: return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_DFLASH: diff --git a/src/llama-model.h b/src/llama-model.h index 6487b0e9c434..a187c274932b 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -115,6 +115,7 @@ enum llm_type { LLM_TYPE_10B_128x3_66B, LLM_TYPE_57B_A14B, LLM_TYPE_816B_A42B, + LLM_TYPE_78B_A3_5B, // Kolibri1 LLM_TYPE_17B_16E, // llama4 Scout LLM_TYPE_17B_128E, // llama4 Maverick LLM_TYPE_A13B, diff --git a/src/models/kolibri1.cpp b/src/models/kolibri1.cpp new file mode 100644 index 000000000000..acc1dd159202 --- /dev/null +++ b/src/models/kolibri1.cpp @@ -0,0 +1,187 @@ +#include "models.h" + +void llama_model_kolibri1::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + + if (hparams.n_swa == 0 || hparams.n_expert_shared != 1 || hparams.n_ff_shexp == 0) { + throw std::runtime_error("Kolibri-1: invalid sliding window or shared expert configuration"); + } + + 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; + + switch (hparams.n_layer()) { + case 50: type = LLM_TYPE_78B_A3_5B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_kolibri1::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t n_ff_exp = hparams.n_ff_exp(); + const int64_t n_ff_shexp = hparams.n_ff_shexp; + + 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}, 0); + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, 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); + + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_post_norm = create_tensor(tn(LLM_TENSOR_FFN_POST_NORM, "weight", i), {n_embd}, 0); + + // Routing is sigmoid, the selection bias is added to the logits before top-k. + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0); + } +} + +llama_model_kolibri1::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_tensor * cur; + 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) { + const auto & layer = model.layers[il]; + + ggml_tensor * residual = inpL; + + cur = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // QK-norm, then RoPE on the sliding layers only: the global layers are NoPE. + auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, n_head, n_head_kv, il); + + Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); + Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + cb(Kcur, "Kcur_normed", il); + + if (hparams.is_swa(il)) { + const float freq_base_l = model.get_rope_freq_base (cparams, il); + const float freq_scale_l = model.get_rope_freq_scale(cparams, il); + + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(Qcur, "Qcur_rope", il); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(Kcur, "Kcur_rope", il); + } + + cur = build_attn(inp_attn, + layer.wo, NULL, layer.wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "attn_out", il); + + cur = build_norm(cur, layer.attn_post_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_post_norm", il); + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + residual = ggml_get_rows(ctx0, residual, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, residual); + cb(ffn_inp, "ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + // The router runs in f32: the selection adds the bias to the logits, the weights + // are the sigmoid of the unbiased logits, so the same logits feed both. + ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur); + ggml_prec_set_acc(logits, GGML_PREC_F32); + ggml_prec_set_src(logits, GGML_PREC_F32, 1); + cb(logits, "ffn_moe_logits", il); + + ggml_tensor * selection = ggml_add(ctx0, logits, layer.ffn_exp_probs_b); + cb(selection, "ffn_moe_logits_biased", il); + + ggml_tensor * selected = ggml_argsort_top_k(ctx0, selection, n_expert_used); + cb(selected, "ffn_moe_topk", il); + + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, 1.0f, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + logits, + nullptr, nullptr, nullptr, nullptr, selected); + cb(moe_out, "ffn_moe_out", il); + + ggml_tensor * shared = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, nullptr, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(shared, "ffn_shexp_out", il); + + cur = ggml_add(ctx0, moe_out, shared); + cur = build_norm(cur, layer.ffn_post_norm, NULL, LLM_NORM_RMS, il); + cur = ggml_add(ctx0, cur, ffn_inp); + + inpL = build_cvec(cur, il); + cb(inpL, "l_out", il); + } + + cur = build_norm(inpL, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur, model.output_s); + cb(cur, "result_output", -1); + res->t_logits = cur; +} + +std::unique_ptr llama_model_kolibri1::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} diff --git a/src/models/models.h b/src/models/models.h index 98f6b6e35dcc..e6cf9d213cee 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2974,3 +2974,16 @@ struct llama_model_k2_horizon : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; + + +struct llama_model_kolibri1 : public llama_model_base { + llama_model_kolibri1(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/tests/fusion/MTL.csv b/tests/fusion/MTL.csv index d31a9f23e5f0..9ab1f1f3f885 100644 --- a/tests/fusion/MTL.csv +++ b/tests/fusion/MTL.csv @@ -139,6 +139,9 @@ kimi-linear ,0 ,any ,GATED_DELTA_NET+CPY , 1 kimi-linear ,0 ,any ,MUL+ADD , 2 kimi-linear ,0 ,any ,RMS_NORM+MUL , 7 kimi-linear ,0 ,any ,RMS_NORM+SCALE , 2 +kolibri1 ,0 ,any ,MUL+ADD , 4 +kolibri1 ,0 ,any ,RMS_NORM+MUL , 10 +kolibri1 ,0 ,any ,RMS_NORM+MUL+ADD , 3 laguna ,0 ,any ,ADD+ADD , 1 laguna ,0 ,any ,MUL+ADD , 2 laguna ,0 ,any ,RMS_NORM+MUL , 9 diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 14ca203be809..143c1335e4ba 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -185,7 +185,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { uint32_t n_head_kv = n_head; if (arch == LLM_ARCH_QWEN3) { n_head_kv = 1; // MQA coverage - } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_MUSE_SPARK || arch == LLM_ARCH_AFMOE || arch == LLM_ARCH_MAION_CODER) { + } else if (arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_MUSE_SPARK || arch == LLM_ARCH_AFMOE || arch == LLM_ARCH_MAION_CODER || + arch == LLM_ARCH_KOLIBRI1) { n_head_kv = 2; // GQA coverage } if (arch == LLM_ARCH_BEAM) { @@ -316,7 +317,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, uint32_t(5)); } else if (arch == LLM_ARCH_COHERE2MOE || arch == LLM_ARCH_MIMO2 || arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_SPARK2_5 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_MUSE_SPARK || arch == LLM_ARCH_GRANITE_SWA || arch == LLM_ARCH_DOTS3NOTE || - arch == LLM_ARCH_MAPLE) { + arch == LLM_ARCH_KOLIBRI1 || arch == LLM_ARCH_MAPLE) { std::vector pattern; pattern.reserve(n_layer); for (uint32_t il = 0; il < n_layer; il++) { @@ -336,6 +337,10 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); } + if (arch == LLM_ARCH_KOLIBRI1) { + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, false); + } + // MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the // indexer head count is independent of the main attention head count. if (arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_GLM5_NEXT) { @@ -758,6 +763,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_LAGUNA: case LLM_ARCH_MAPLE: case LLM_ARCH_MUSE_SPARK: + case LLM_ARCH_KOLIBRI1: return true; default: return false; From f2f667b02dc76d675395ae3b14546c8108d26387 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Sat, 17 Oct 2026 13:55:33 +0200 Subject: [PATCH 071/423] cuda : add grouped experts top-k fusion (#29181) * fuse grouped experts * add tests * hip fix * add to graph_optimize * apply suggestion from am17an * hip fix * fix hip fix * fix fix fix * gate on hip version instead --- ggml/src/ggml-cuda/ggml-cuda.cu | 101 +++++++++++++++++++++++++++----- ggml/src/ggml-cuda/topk-moe.cu | 66 +++++++++++++++++++++ ggml/src/ggml-cuda/topk-moe.cuh | 2 + tests/test-backend-ops.cpp | 38 +++++++++++- 4 files changed, 188 insertions(+), 19 deletions(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 6fa22289f331..cb85321f5762 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3029,6 +3029,65 @@ static int ggml_cuda_try_ssm_scan_cache_fusion( return skip; } +// expert group selection from build_moe_ffn (n_expert_groups > 1) +// src0/src1 index into this table, -1 is the (biased) probs, -2 is not checked +static const struct { + ggml_op op; + int src0; + int src1; +} topk_moe_group_pattern[] = { + { GGML_OP_RESHAPE, -1, -2 }, // 0: [n_exp_per_group, n_expert_groups, n_tokens] + { GGML_OP_RESHAPE, 0, -2 }, // 1: [1, n_exp_per_group, n_expert_groups, n_tokens] + { GGML_OP_ARGSORT, 0, -2 }, // 2 + { GGML_OP_VIEW, 2, -2 }, // 3: top-2 per group + { GGML_OP_GET_ROWS, 1, 3 }, // 4 + { GGML_OP_RESHAPE, 4, -2 }, // 5 + { GGML_OP_SUM_ROWS, 5, -2 }, // 6: group scores + { GGML_OP_RESHAPE, 6, -2 }, // 7 + { GGML_OP_ARGSORT, 7, -2 }, // 8 + { GGML_OP_VIEW, 8, -2 }, // 9: top n_group_used groups + { GGML_OP_GET_ROWS, 0, 9 }, // 10 + { GGML_OP_FILL, 0, -2 }, // 11 + { GGML_OP_SET_ROWS, 10, 9 }, // 12: unused groups -> -inf + { GGML_OP_RESHAPE, 12, -2 }, // 13: [n_expert, n_tokens] +}; + +static bool ggml_cuda_topk_moe_group_fusion(const struct ggml_cgraph * cgraph, int node_idx, const ggml_tensor * selection, ggml_cuda_topk_moe_args & args) { +#if defined(GGML_USE_HIP) && HIP_VERSION < 60300000 + GGML_UNUSED_VARS(cgraph, node_idx, selection, args); + return false; +#else + const int n_ops = (int) std::size(topk_moe_group_pattern); + if (node_idx + n_ops > cgraph->n_nodes) { + return false; + } + + ggml_tensor ** nodes = cgraph->nodes + node_idx; + auto src_ok = [&](const ggml_tensor * src, int idx) { + return idx == -2 || src == (idx == -1 ? selection : nodes[idx]); + }; + + for (int i = 0; i < n_ops; i++) { + const auto & p = topk_moe_group_pattern[i]; + if (nodes[i]->op != p.op || !src_ok(nodes[i]->src[0], p.src0) || !src_ok(nodes[i]->src[1], p.src1)) { + return false; + } + } + + const int n_groups = nodes[0]->ne[1]; + const int n_group_used = nodes[9]->ne[0]; + + if (nodes[3]->ne[0] != 2 || nodes[12]->src[2] != nodes[11] || ggml_get_op_params_f32(nodes[11], 0) != -INFINITY || + n_groups > WARP_SIZE) { + return false; + } + + args.n_groups = n_groups; + args.n_group_used = n_group_used; + return true; +#endif // defined(GGML_USE_HIP) && HIP_VERSION < 60300000 +} + static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) { args.sigmoid = false; args.sqrt_softplus = false; @@ -3036,6 +3095,8 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod args.delayed_softmax = false; args.prob_bias = false; args.norm = false; + args.n_groups = 0; + args.n_group_used = 0; const int n_nodes = cgraph->n_nodes; ggml_tensor ** nodes = cgraph->nodes; @@ -3078,18 +3139,20 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod } // src of bias add is the unreshaped probs (-2 instead of -1) - if (nodes[node_idx]->op == GGML_OP_ADD && nodes[node_idx]->src[0] == nodes[node_idx - 2]) { + const ggml_tensor * selection = nodes[node_idx - 2]; + if (nodes[node_idx]->op == GGML_OP_ADD && nodes[node_idx]->src[0] == selection) { args.prob_bias = true; + selection = nodes[node_idx]; node_idx++; } - // RESHAPE/ADD -> ARGSORT - if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_ARGSORT) { - return false; + + if (ggml_cuda_topk_moe_group_fusion(cgraph, node_idx, selection, args)) { + node_idx += (int) std::size(topk_moe_group_pattern); + selection = nodes[node_idx - 1]; } - if (args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - return false; - } else if (!args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 2]) { + // RESHAPE/ADD/group mask -> ARGSORT + if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_ARGSORT || nodes[node_idx]->src[0] != selection) { return false; } @@ -3100,6 +3163,11 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod nodes[node_idx]->src[0] != nodes[node_idx - 1]) { return false; } + + // kept groups must hold all used experts + if (args.n_groups > 0 && nodes[node_idx]->ne[0] > args.n_group_used * (probs_reshaped->ne[1] / args.n_groups)) { + return false; + } node_idx++; if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_GET_ROWS) { @@ -3754,17 +3822,18 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } else { ops.insert(ops.end(), { GGML_OP_SOFT_MAX }); } - const int i_probs = i + (int) ops.size() - 1; // last node of the gating activation - + ops.push_back(GGML_OP_RESHAPE); if (args.prob_bias) { - bias = cgraph->nodes[i_probs + 2]->src[1]; - ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW, - GGML_OP_GET_ROWS }); - out_nodes[0] = i_probs + 4; - } else { - ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS }); - out_nodes[0] = i_probs + 3; + ops.push_back(GGML_OP_ADD); + bias = cgraph->nodes[i + ops.size() - 1]->src[1]; + } + if (args.n_groups > 0) { + for (const auto & p : topk_moe_group_pattern) { + ops.push_back(p.op); + } } + ops.insert(ops.end(), { GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS }); + out_nodes[0] = i + ops.size() - 2; ids = cgraph->nodes[out_nodes[0]]; if (args.norm) { diff --git a/ggml/src/ggml-cuda/topk-moe.cu b/ggml/src/ggml-cuda/topk-moe.cu index d015e7ddefa5..81d6ec3f7cdc 100644 --- a/ggml/src/ggml-cuda/topk-moe.cu +++ b/ggml/src/ggml-cuda/topk-moe.cu @@ -11,6 +11,8 @@ struct topk_moe_config { bool use_sqrt_softplus; bool with_norm; bool delayed_softmax; + int n_groups; + int n_group_used; }; // Warp-local softmax used for both the pre-top-k logits and the post-top-k delayed path. @@ -78,6 +80,53 @@ __device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], co } } +#if !defined(GGML_USE_HIP) || HIP_VERSION >= 60300000 +// Set -inf for experts outside the top n_group_used groups. Group score is the sum of its top-2 experts (DeepSeek V3). +template +__device__ void group_mask_warp_inplace(float (&vals)[experts_per_thread], const int n_groups, const int n_group_used, const int lane) { + const int n_per_group = n_experts / n_groups; + + float score = -INFINITY; // lane g holds the score of group g + for (int g = 0; g < n_groups; g++) { + float top1 = -INFINITY; + float top2 = -INFINITY; +#pragma unroll + for (int i = 0; i < experts_per_thread; i++) { + if ((i * WARP_SIZE + lane) / n_per_group == g) { + top2 = fmaxf(top2, fminf(top1, vals[i])); + top1 = fmaxf(top1, vals[i]); + } + } +#pragma unroll + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) { + const float o1 = __shfl_xor_sync(0xFFFFFFFF, top1, offset, WARP_SIZE); + const float o2 = __shfl_xor_sync(0xFFFFFFFF, top2, offset, WARP_SIZE); + top2 = fmaxf(fminf(top1, o1), fmaxf(top2, o2)); + top1 = fmaxf(top1, o1); + } + if (lane == g) { + score = top1 + top2; + } + } + + // rank of each group, lower index wins ties + int rank = 0; + for (int g = 0; g < n_groups; g++) { + const float other = __shfl_sync(0xFFFFFFFF, score, g, WARP_SIZE); + rank += other > score || (other == score && g < lane); + } + const auto keep = __ballot_sync(0xFFFFFFFFULL, lane < n_groups && rank < n_group_used); + +#pragma unroll + for (int i = 0; i < experts_per_thread; i++) { + const int g = (i * WARP_SIZE + lane) / n_per_group; + if (g >= n_groups || !((keep >> g) & 1)) { + vals[i] = -INFINITY; + } + } +} +#endif // !defined(GGML_USE_HIP) || HIP_VERSION >= 60300000 + /* This kernel does the following: 1. optionally softmax over the logits per token [n_experts, n_tokens] @@ -171,6 +220,17 @@ __global__ void topk_moe_cuda(const float * logits, } } +#if !defined(GGML_USE_HIP) || HIP_VERSION >= 60300000 + // without bias, masking wt is safe: the matcher ensures the kept groups hold all used experts + if (config.n_groups > 0) { + if constexpr (has_bias) { + group_mask_warp_inplace(selection_wt, config.n_groups, config.n_group_used, threadIdx.x); + } else { + group_mask_warp_inplace(wt, config.n_groups, config.n_group_used, threadIdx.x); + } + } +#endif // !defined(GGML_USE_HIP) || HIP_VERSION >= 60300000 + //at this point, each thread holds either a portion of the softmax distribution //or the raw logits. We do the argmax reduce over n_expert_used, each time marking //the expert weight as -inf to exclude from the next iteration @@ -297,6 +357,10 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx, cudaStream_t stream = ctx.stream(); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, stream); +#if defined(GGML_USE_HIP) && HIP_VERSION < 60300000 + GGML_ASSERT(config.n_groups == 0 && "expert groups fusion is not supported on this HIP version"); +#endif + switch (n_expert) { case 1: ggml_cuda_kernel_launch(topk_moe_cuda<1, has_bias>, launch_params, @@ -389,6 +453,8 @@ void ggml_cuda_op_topk_moe(ggml_backend_cuda_context & ctx, config.use_sqrt_softplus = args.sqrt_softplus; config.with_norm = with_norm; config.delayed_softmax = args.delayed_softmax; + config.n_groups = args.n_groups; + config.n_group_used = args.n_group_used; if (bias) { launch_topk_moe_cuda(ctx, logits_d, weights_d, ids_d, bias_d, n_rows, n_experts, n_expert_used, clamp_val, diff --git a/ggml/src/ggml-cuda/topk-moe.cuh b/ggml/src/ggml-cuda/topk-moe.cuh index 061b37e2971e..3637500ff99d 100644 --- a/ggml/src/ggml-cuda/topk-moe.cuh +++ b/ggml/src/ggml-cuda/topk-moe.cuh @@ -14,6 +14,8 @@ struct ggml_cuda_topk_moe_args { bool prob_bias{}; bool norm{}; bool scale{}; + int n_groups{}; // > 0: pick experts only from the top n_group_used expert groups + int n_group_used{}; }; void ggml_cuda_op_topk_moe(ggml_backend_cuda_context & ctx, diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index b50364660ed9..8826b274b968 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7328,6 +7328,7 @@ struct test_topk_moe : public test_case { const bool bias_probs; const MoeGatingFunc gating_func; const float scale_w; + const int n_expert_groups; ggml_tensor * weights {}; ggml_tensor * selected_experts {}; @@ -7336,17 +7337,21 @@ struct test_topk_moe : public test_case { bool with_norm = false, bool bias_probs = false, MoeGatingFunc gating_func = GATING_FUNC_SOFTMAX, - float scale_w = 0.0f) : + float scale_w = 0.0f, + int n_expert_groups = 1) : ne(ne), n_expert_used(n_expert_used), with_norm(with_norm), bias_probs(bias_probs), gating_func(gating_func), - scale_w(scale_w) { + scale_w(scale_w), + n_expert_groups(n_expert_groups) { GGML_ASSERT(n_expert_used <= ne[0]); + GGML_ASSERT(ne[0] % n_expert_groups == 0); + GGML_ASSERT(n_expert_used <= ne[0] / n_expert_groups); } - std::string vars() override { return VARS_TO_STR6(ne, n_expert_used, with_norm, bias_probs, gating_func, scale_w); } + std::string vars() override { return VARS_TO_STR7(ne, n_expert_used, with_norm, bias_probs, gating_func, scale_w, n_expert_groups); } std::string op_desc(ggml_tensor * t) override { GGML_UNUSED(t); @@ -7358,6 +7363,8 @@ struct test_topk_moe : public test_case { ggml_tensor * build_graph(ggml_context * ctx) override { const int n_expert = ne[0]; const int n_tokens = ne[1]; + const int n_exp_per_group = n_expert / n_expert_groups; + const int n_group_used = n_expert_groups / 2; ggml_tensor * logits = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); ggml_tensor * probs = @@ -7374,6 +7381,24 @@ struct test_topk_moe : public test_case { ggml_set_name(selection_probs, "selection_probs"); } + if (n_expert_groups > 1) { + ggml_tensor * selection_groups = ggml_reshape_3d(ctx, selection_probs, n_exp_per_group, n_expert_groups, n_tokens); // [n_exp_per_group, n_expert_groups, n_tokens] + + ggml_tensor * group_scores = ggml_argsort_top_k(ctx, selection_groups, 2); // [2, n_expert_groups, n_tokens] + group_scores = ggml_get_rows(ctx, ggml_reshape_4d(ctx, selection_groups, 1, selection_groups->ne[0], selection_groups->ne[1], selection_groups->ne[2]), group_scores); // [1, 2, n_expert_groups, n_tokens] + + group_scores = ggml_sum_rows(ctx, ggml_reshape_3d(ctx, group_scores, group_scores->ne[1], group_scores->ne[2], group_scores->ne[3])); // [1, n_expert_groups, n_tokens] + group_scores = ggml_reshape_2d(ctx, group_scores, group_scores->ne[1], group_scores->ne[2]); // [n_expert_groups, n_tokens] + + ggml_tensor * expert_groups = ggml_argsort_top_k(ctx, group_scores, n_group_used); // [n_group_used, n_tokens] + ggml_set_name(expert_groups, "ffn_moe_group_topk"); + + selection_probs = ggml_get_rows(ctx, selection_groups, expert_groups); // [n_exp_per_group, n_group_used, n_tokens] + selection_probs = ggml_set_rows(ctx, ggml_fill(ctx, selection_groups, -INFINITY), selection_probs, expert_groups); // [n_exp_per_group, n_expert_groups, n_tokens] + selection_probs = ggml_reshape_2d(ctx, selection_probs, n_expert, n_tokens); // [n_expert, n_tokens] + ggml_set_name(selection_probs, "ffn_moe_probs_masked"); + } + selected_experts = ggml_argsort_top_k(ctx, selection_probs, n_expert_used); // [n_expert_used, n_tokens] ggml_set_name(selected_experts, "selected_experts"); @@ -11937,6 +11962,13 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_topk_moe({32, 8, 1, 1}, 4, with_norm, bias_probs, gate, scale_w)); test_cases.emplace_back(new test_topk_moe({32, 8, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); test_cases.emplace_back(new test_topk_moe({32, 9, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); + + // grouped experts + test_cases.emplace_back(new test_topk_moe({128, 1, 1, 1}, 8, with_norm, bias_probs, gate, scale_w, 16)); + test_cases.emplace_back(new test_topk_moe({256, 22, 1, 1}, 6, with_norm, bias_probs, gate, scale_w, 32)); + test_cases.emplace_back(new test_topk_moe({256, 7, 1, 1}, 8, with_norm, bias_probs, gate, scale_w, 8)); // DeepSeek V3 + test_cases.emplace_back(new test_topk_moe({128, 5, 1, 1}, 6, with_norm, bias_probs, gate, scale_w, 8)); + test_cases.emplace_back(new test_topk_moe({512, 3, 1, 1}, 8, with_norm, bias_probs, gate, scale_w, 4)); } } } From 9d1e03616efb7520514d65552c7376296ef9208f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrien=20Gallou=C3=ABt?= Date: Sat, 17 Oct 2026 14:36:17 +0200 Subject: [PATCH 072/423] cuda : fix -Wtautological-compare in fattn vec dispatch (#30519) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Adrien Gallouët --- ggml/src/ggml-cuda/fattn.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 5fa3e530840c..1515572fe111 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -417,9 +417,7 @@ static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, gg #define FATTN_VEC_CASE(D, type_K_case, type_V_case) \ if constexpr (GGML_CUDA_FA_##type_K_case##_##type_V_case) { \ - const bool type_K_okay = type_K == GGML_TYPE_##type_K_case || (type_K == GGML_TYPE_F32 && GGML_TYPE_##type_K_case == GGML_TYPE_F16); \ - const bool type_V_okay = type_V == GGML_TYPE_##type_V_case || (type_V == GGML_TYPE_F32 && GGML_TYPE_##type_V_case == GGML_TYPE_F16); \ - if (head_size == (D) && type_K_okay && type_V_okay) { \ + if (head_size == (D) && type_K == GGML_TYPE_##type_K_case && type_V == GGML_TYPE_##type_V_case) { \ return ggml_cuda_flash_attn_ext_vec_case; \ } \ } \ @@ -432,7 +430,9 @@ static void ggml_cuda_flash_attn_ext_mma_f16(ggml_backend_cuda_context & ctx, gg typedef void (* fattn_vec_case_t)(ggml_backend_cuda_context & ctx, ggml_tensor * dst); // Vector kernel for the given head size and K/V types, nullptr if its template instance was not compiled: -static fattn_vec_case_t ggml_cuda_get_fattn_vec_case(const int64_t head_size, const ggml_type type_K, const ggml_type type_V) { +static fattn_vec_case_t ggml_cuda_get_fattn_vec_case(const int64_t head_size, const ggml_type type_K_in, const ggml_type type_V_in) { + const ggml_type type_K = type_K_in == GGML_TYPE_F32 ? GGML_TYPE_F16 : type_K_in; + const ggml_type type_V = type_V_in == GGML_TYPE_F32 ? GGML_TYPE_F16 : type_V_in; FATTN_VEC_CASES_ALL_D(F16, F16) FATTN_VEC_CASES_ALL_D(Q4_0, F16) FATTN_VEC_CASES_ALL_D(Q4_1, F16) From 10b8b7553494f3638da7668cf17cac089229122a Mon Sep 17 00:00:00 2001 From: Pascal Date: Sat, 17 Oct 2026 15:07:55 +0200 Subject: [PATCH 073/423] tests: give the random model generator a per-tensor filler table (#30674) * tests: give the random model generator a per-tensor filler table Muse Spark makes test-llama-context fail on state_load, cp_h and cp_d on about six percent of the random seeds, with an NMSE of 1.0e-05 to 1.4e-05 against a threshold of 1.0e-05. The sliding window is what exposes it: a restored state keeps the window alone, so the attention of an SWA layer reduces over a different cell layout and the continuation is not bit exact. Every SWA model carries that, at 1e-08. Muse Spark carried it at 1e-05. The generator is what turns 1e-08 into 1e-05. A norm preserving residual gate holds alpha^2 + beta^2 == 1 and the conversion resolves both weights from one raw weight, while the generator draws every tensor on its own. The pair then scaled the stream and the branch independently, the stream lost an arbitrary factor at each layer, and the model became an amplifier. The weights the test ran on were not a Muse Spark. Per-tensor rules lived inside the write path, where a flag picked the values of a tensor inside each of the two branches that store them, so a new rule meant a new flag threaded through both. The file asks to be rid of that in a TODO of its own. Give each rule a filler, pick one from a table keyed on a substring of the tensor name, and keep a single write path behind it. Mamba moves into the table and Muse Spark joins it: the gate filler draws the raw weight once per block and resolves alpha and beta the way the conversion does. Over a thousand seeds on a CPU build the worst case drops from 1.36e-05 to 2.48e-06 and the median from 4.10e-06 to 5.48e-07, with no seed above the threshold. The other 131 models come out identical to the bit, since no other architecture carries these tensors. ctest is green on a CPU build and on a CUDA build. * tests: draw every tensor filler from fill_normal The fillers repeated the generator and the distribution each time. The normal draw now takes its key and the other fillers transform its values, so the generator lives in one place. The generated models are identical to the bit. --- tests/test-llama-archs.cpp | 71 ++++++++++++++++++++++++++++---------- 1 file changed, 52 insertions(+), 19 deletions(-) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 143c1335e4ba..94a22ff2f57f 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -53,33 +53,66 @@ struct tensor_data_params { float stdev; }; +// fills the values of one tensor, the name carries whatever the property needs +typedef void (*tensor_filler)(std::vector & dst, const char * name, const tensor_data_params & params); + +// a normal draw keyed on a name, so that tensors which share a draw ask for the same key +static void fill_normal(std::vector & dst, const char * key, const tensor_data_params & params) { + std::mt19937 gen(params.seed ^ std::hash{}(key)); + std::normal_distribution dis(0.0f, params.stdev); + for (float & v : dst) { + v = dis(gen); + } +} + +// Mamba A must be negative (state decay) +static void fill_ssm_a(std::vector & dst, const char * name, const tensor_data_params & params) { + fill_normal(dst, name, params); + for (float & v : dst) { + v = -fabsf(v); + } +} + +// a norm preserving residual gate holds alpha^2 + beta^2 == 1, and the conversion resolves both from one raw weight +// the alpha and the beta of a block share that weight, so the key drops the suffix and keeps the block +static void fill_resid_gate(std::vector & dst, const char * name, const tensor_data_params & params) { + const char * suffix = strstr(name, "_resid_"); + const bool is_beta = strstr(suffix, "_resid_beta") != nullptr; + fill_normal(dst, std::string(name, suffix - name).c_str(), params); + for (float & v : dst) { + const float beta = 1.0f / (1.0f + expf(-v)); + v = is_beta ? beta : sqrtf(1.0f - beta*beta); + } +} + +// a tensor whose values hold a property of their own, matched on a substring of its name, first match wins +static const struct { const char * match; tensor_filler fill; } TENSOR_FILLERS[] = { + { "ssm_a", fill_ssm_a }, + { "_resid_", fill_resid_gate }, +}; + static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { const tensor_data_params & params = *(const tensor_data_params *) userdata; - size_t seed = params.seed; - std::hash hasher; - seed ^= hasher(tensor->name); - std::mt19937 gen(seed); - std::normal_distribution dis(0.0f, params.stdev); - // TODO: refactor per-tensor initialization logic in a cleaner way + tensor_filler fill = fill_normal; + for (const auto & filler : TENSOR_FILLERS) { + if (strstr(tensor->name, filler.match) != nullptr) { + fill = filler.fill; + break; + } + } + + std::vector tmp(ggml_nelements(tensor)); + fill(tmp, tensor->name, params); - // note: Mamba A must be negative (state decay) - const bool is_ssm_a = strstr(tensor->name, "ssm_a") != nullptr; - const int64_t ne = ggml_nelements(tensor); if (tensor->type == GGML_TYPE_F32) { - std::vector tmp(ne); - for (int64_t i = 0; i < ne; i++) { - float val = dis(gen); - tmp[i] = is_ssm_a ? -fabsf(val) : val; - } ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor)); } else if (tensor->type == GGML_TYPE_F16) { - std::vector tmp(ne); - for (int64_t i = 0; i < ne; i++) { - float val = dis(gen); - tmp[i] = ggml_fp32_to_fp16(is_ssm_a ? -fabsf(val) : val); + std::vector tmp16(tmp.size()); + for (size_t i = 0; i < tmp.size(); i++) { + tmp16[i] = ggml_fp32_to_fp16(tmp[i]); } - ggml_backend_tensor_set(tensor, tmp.data(), 0, ggml_nbytes(tensor)); + ggml_backend_tensor_set(tensor, tmp16.data(), 0, ggml_nbytes(tensor)); } else { GGML_ABORT("fatal error"); } From 2744e48d1e88d7416224e92bf58d4ffedda4a205 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Sat, 17 Oct 2026 19:05:30 +0530 Subject: [PATCH 074/423] CUDA: faster MMVQ for low-bit types (#30517) * CUDA: faster MMVQ for low-bit types * ggml-cuda: hip: small nwarps GCN fix on Q3_K(ncols=2,3), Q6_K(ncols=4) --------- Co-authored-by: thelittlefireman --- ggml/src/ggml-cuda/mmvq.cu | 35 ++++++-- ggml/src/ggml-cuda/vecdotq.cuh | 149 +++++++++++++++++++++------------ 2 files changed, 125 insertions(+), 59 deletions(-) diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 8068ba364cce..930f1adc4eee 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -468,10 +468,12 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d } else if (table_id == MMVQ_PARAMETERS_GCN) { switch (ncols_dst) { case 1: + return 2; case 2: case 3: + return type == GGML_TYPE_Q3_K ? 1 : 2; case 4: - return 2; + return type == GGML_TYPE_Q6_K? 1 : 2; case 5: case 6: case 7: @@ -576,11 +578,31 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d return 1; } -static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) { +// Low bit types load few weight bytes per thread, a second row per block shares the q8_1 loads and adds memory parallelism. +static constexpr __host__ __device__ bool mmvq_two_rows_ncols1(ggml_type type) { + switch (type) { + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + return true; + default: + return false; + } +} + +static constexpr __host__ __device__ int calc_rows_per_block(ggml_type type, int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) { if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN || table_id == MMVQ_PARAMETERS_TURING || table_id == MMVQ_PARAMETERS_GB10) { switch (ncols_dst) { case 1: - return small_k ? nwarps : 1; + if (small_k) { + return nwarps; + } + if (table_id == MMVQ_PARAMETERS_GB10) { + return type == GGML_TYPE_IQ2_S || type == GGML_TYPE_IQ3_S ? 2 : 1; + } + return table_id == MMVQ_PARAMETERS_GENERIC && mmvq_two_rows_ncols1(type) ? 2 : 1; case 2: case 3: case 4: @@ -615,7 +637,7 @@ static __global__ void mul_mat_vec_q( constexpr int vdr = get_vdr_mmvq(type); constexpr mmvq_parameter_table_id table_id = get_device_table_id(); constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters); - constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); + constexpr int rows_per_cuda_block = calc_rows_per_block(type, ncols_dst, table_id, small_k, nwarps); constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); @@ -727,7 +749,8 @@ static __global__ void mul_mat_vec_q( #if __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 + // loop iterations, not blocks; q3_K and q6_K keep the distance in blocks they had with vdr 1 + constexpr int pf_dist = type == GGML_TYPE_Q3_K || type == GGML_TYPE_Q6_K ? 1 : 2; const int kbx_pf = kbx + pf_dist*blocks_per_iter; if (kbx_pf < blocks_per_row_x) { #pragma unroll @@ -1013,7 +1036,7 @@ static std::pair calc_launch_params( const int ncols_dst, const int nrows_x, const int nchannels_dst, const int nsamples_or_ntokens, const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false, const bool halve_iters = false) { const int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters); - const int rpb = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); + const int rpb = calc_rows_per_block(type, ncols_dst, table_id, small_k, nwarps); const int64_t nblocks = (nrows_x + rpb - 1) / rpb; const dim3 block_nums(nblocks, nchannels_dst, nsamples_or_ntokens); const dim3 block_dims(warp_size, nwarps, 1); diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index f65fe530eb4f..14c1737c0207 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -360,12 +360,12 @@ static __device__ __forceinline__ float vec_dot_nvfp4_q8_1( return sum; } -#define VDR_Q2_K_Q8_1_MMVQ 1 +#define VDR_Q2_K_Q8_1_MMVQ 2 #define VDR_Q2_K_Q8_1_MMQ 4 // contiguous v/x values static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq( - const int & v, const int * __restrict__ u, const uint8_t * __restrict__ scales, + const int * __restrict__ v, const int * __restrict__ u, const int * __restrict__ scales, const half2 & dm2, const float * __restrict__ d8) { float sumf_d = 0.0f; @@ -373,17 +373,20 @@ static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq( #pragma unroll for (int i = 0; i < QR2_K; ++i) { - const int sc = scales[2*i]; + const int sc = scales[i/2] >> (16*(i%2)); - const int vi = (v >> (2*i)) & 0x03030303; + int sumi_d = 0; + int sumi_m = 0; +#pragma unroll + for (int j = 0; j < VDR_Q2_K_Q8_1_MMVQ; ++j) { + const int vi = (v[j] >> (2*i)) & 0x03030303; - sumf_d += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * (sc & 0xF)); // SIMD dot product + sumi_d = ggml_cuda_dp4a(vi, u[VDR_Q2_K_Q8_1_MMVQ*i + j], sumi_d); // SIMD dot product + sumi_m = ggml_cuda_dp4a(0x01010101, u[VDR_Q2_K_Q8_1_MMVQ*i + j], sumi_m); // sum of q8_1 values + } - // fill int with 4x m - int m = sc >> 4; - m |= m << 8; - m |= m << 16; - sumf_m += d8[i] * ggml_cuda_dp4a(m, u[i], 0); // multiply constant q2_K part with sum of q8_1 values + sumf_d += d8[i] * (sumi_d * (sc & 0xF)); + sumf_m += d8[i] * (sumi_m * ((sc >> 4) & 0xF)); // multiply constant q2_K part with sum of q8_1 values } const float2 dm2f = __half22float2(dm2); @@ -443,37 +446,35 @@ static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmq( return sumf + d8*sumf_d8; } -#define VDR_Q3_K_Q8_1_MMVQ 1 +#define VDR_Q3_K_Q8_1_MMVQ 2 #define VDR_Q3_K_Q8_1_MMQ 2 // contiguous v/x values static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq( - const int & vl, const int & vh, const int * __restrict__ u, const uint8_t * __restrict__ scales, - const int & scale_offset, const float & d3, const float * __restrict__ d8) { + const int * __restrict__ vl, const int * __restrict__ vh, const int * __restrict__ u, const int * __restrict__ scales, + const float & d3, const float * __restrict__ d8) { float sumf = 0.0f; #pragma unroll for (int i = 0; i < QR3_K; ++i) { - const int isc = scale_offset + 2*i; - - const int isc_low = isc % (QK_K/32); - const int sc_shift_low = 4 * (isc / (QK_K/32)); - const int sc_low = (scales[isc_low] >> sc_shift_low) & 0xF; + // 4 low bits from scales[0..7], 2 high bits from scales[8..11] + const int sc_low = (scales[i/2] >> (16*(i%2))) & 0x0F; + const int sc_high = (scales[2] >> (16*(i%2) + 2*(i/2))) & 0x03; + const int sc = (sc_low | (sc_high << 4)) - 32; - const int isc_high = isc % (QK_K/64); - const int sc_shift_high = 2 * (isc / (QK_K/64)); - const int sc_high = ((scales[(QK_K/32) + isc_high] >> sc_shift_high) & 3) << 4; - - const int sc = (sc_low | sc_high) - 32; - - const int vil = (vl >> (2*i)) & 0x03030303; - - const int vih = ((vh >> i) << 2) & 0x04040404; + int sumi = 0; +#pragma unroll + for (int j = 0; j < VDR_Q3_K_Q8_1_MMVQ; ++j) { + const int vil = (vl[j] >> (2*i)) & 0x03030303; + const int vih = ((vh[j] >> i) << 2) & 0x04040404; - const int vi = __vsub4(vil, vih); + // (vil | vih) - 4 for each byte, the - 4 is done by the second dp4a + sumi = ggml_cuda_dp4a(vil | vih, u[VDR_Q3_K_Q8_1_MMVQ*i + j], sumi); // SIMD dot product + sumi = ggml_cuda_dp4a(0xFCFCFCFC, u[VDR_Q3_K_Q8_1_MMVQ*i + j], sumi); + } - sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product + sumf += d8[i] * (sumi * sc); } return d3 * sumf; @@ -620,12 +621,12 @@ static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_mmq( return dm4f.x*sumf_d - dm4f.y*sumf_m; } -#define VDR_Q6_K_Q8_1_MMVQ 1 +#define VDR_Q6_K_Q8_1_MMVQ 2 #define VDR_Q6_K_Q8_1_MMQ 8 // contiguous v/x values static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq( - const int & vl, const int & vh, const int * __restrict__ u, const int8_t * __restrict__ scales, + const int * __restrict__ vl, const int * __restrict__ vh, const int * __restrict__ u, const int8_t * __restrict__ scales, const float & d, const float * __restrict__ d8) { float sumf = 0.0f; @@ -634,13 +635,18 @@ static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq( for (int i = 0; i < QR6_K; ++i) { const int sc = scales[4*i]; - const int vil = (vl >> (4*i)) & 0x0F0F0F0F; - - const int vih = ((vh >> (4*i)) << 4) & 0x30303030; + int sumi = 0; +#pragma unroll + for (int j = 0; j < VDR_Q6_K_Q8_1_MMVQ; ++j) { + const int vil = (vl[j] >> (4*i)) & 0x0F0F0F0F; + const int vih = ((vh[j] >> (4*i)) << 4) & 0x30303030; - const int vi = __vsub4((vil | vih), 0x20202020); // vi = (vil | vih) - 32 + // (vil | vih) - 32 for each byte, the - 32 is done by the second dp4a + sumi = ggml_cuda_dp4a(vil | vih, u[VDR_Q6_K_Q8_1_MMVQ*i + j], sumi); // SIMD dot product + sumi = ggml_cuda_dp4a(0xE0E0E0E0, u[VDR_Q6_K_Q8_1_MMVQ*i + j], sumi); + } - sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product + sumf += d8[i] * (sumi * sc); } return d*sumf; @@ -870,18 +876,31 @@ static __device__ __forceinline__ float vec_dot_q2_K_q8_1( const block_q2_K * bq2_K = (const block_q2_K *) vbq + kbx; + // iqs is in 0,2..14, each 32 bytes of qs hold 4 q8_1 blocks (one per 2 bit shift) const int bq8_offset = QR2_K * (iqs / QI8_1); - const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); + const int iqs8 = iqs % QI8_1; + + int v[VDR_Q2_K_Q8_1_MMVQ]; +#pragma unroll + for (int j = 0; j < VDR_Q2_K_Q8_1_MMVQ; ++j) { + v[j] = get_int_b4(bq2_K->qs, iqs + j); + } - const uint8_t * scales = bq2_K->scales + scale_offset; + // scale of shift i is byte 2*i + (iqs8 >= QI8_1/2) of the 8 bytes for this half + const int sc_shift = 8 * (iqs8 / (QI8_1/2)); + int scales[2]; + scales[0] = get_int_b4(bq2_K->scales, 2*(iqs / QI8_1) + 0) >> sc_shift; + scales[1] = get_int_b4(bq2_K->scales, 2*(iqs / QI8_1) + 1) >> sc_shift; - const int v = get_int_b4(bq2_K->qs, iqs); - int u[QR2_K]; + int u[VDR_Q2_K_Q8_1_MMVQ*QR2_K]; float d8[QR2_K]; #pragma unroll - for (int i = 0; i < QR2_K; ++ i) { - u[i] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + for (int i = 0; i < QR2_K; ++i) { +#pragma unroll + for (int j = 0; j < VDR_Q2_K_Q8_1_MMVQ; ++j) { + u[VDR_Q2_K_Q8_1_MMVQ*i + j] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs8 + j); + } d8[i] = __low2float(bq8_1[bq8_offset + i].ds); } @@ -893,26 +912,42 @@ static __device__ __forceinline__ float vec_dot_q3_K_q8_1( const block_q3_K * bq3_K = (const block_q3_K *) vbq + kbx; + // iqs is in 0,2..14, each 32 bytes of qs hold 4 q8_1 blocks (one per 2 bit shift) const int bq8_offset = QR3_K * (iqs / (QI3_K/2)); - const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); + const int iqs8 = iqs % QI8_1; const float d = bq3_K->d; - const int vl = get_int_b2(bq3_K->qs, iqs); + int vl[VDR_Q3_K_Q8_1_MMVQ]; + int vh[VDR_Q3_K_Q8_1_MMVQ]; +#pragma unroll + for (int j = 0; j < VDR_Q3_K_Q8_1_MMVQ; ++j) { + vl[j] = get_int_b2(bq3_K->qs, iqs + j); + vh[j] = get_int_b2(bq3_K->hmask, iqs8 + j) >> bq8_offset; + } - // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted - const int vh = ~get_int_b2(bq3_K->hmask, iqs % (QI3_K/2)) >> bq8_offset; + // sub-block 8*(iqs/8) + 2*i + (iqs8 >= QI8_1/2) for shift i: + // 4 low bits in nibble iqs/8 of bytes, 2 high bits at bit 4*(iqs/8) + 2*(i/2) of bytes 8..11 + const int sc_shift = 8 * (iqs8 / (QI8_1/2)) + 4 * (iqs / (QI3_K/2)); + int scales[3]; +#pragma unroll + for (int l = 0; l < 3; ++l) { + scales[l] = get_int_b2(bq3_K->scales, l) >> sc_shift; + } - int u[QR3_K]; + int u[VDR_Q3_K_Q8_1_MMVQ*QR3_K]; float d8[QR3_K]; #pragma unroll for (int i = 0; i < QR3_K; ++i) { - u[i] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1); +#pragma unroll + for (int j = 0; j < VDR_Q3_K_Q8_1_MMVQ; ++j) { + u[VDR_Q3_K_Q8_1_MMVQ*i + j] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs8 + j); + } d8[i] = __low2float(bq8_1[bq8_offset + i].ds); } - return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, d8); + return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, scales, d, d8); } static __device__ __forceinline__ float vec_dot_q4_K_q8_1( @@ -1025,17 +1060,25 @@ static __device__ __forceinline__ float vec_dot_q6_K_q8_1( const int scale_offset = (QI6_K/4) * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/8); const int vh_shift = 2 * ((iqs % (QI6_K/2)) / (QI6_K/4)); - const int vl = get_int_b2(bq6_K->ql, iqs); - const int vh = get_int_b2(bq6_K->qh, (QI6_K/4) * (iqs / (QI6_K/2)) + iqs % (QI6_K/4)) >> vh_shift; + int vl[VDR_Q6_K_Q8_1_MMVQ]; + int vh[VDR_Q6_K_Q8_1_MMVQ]; +#pragma unroll + for (int j = 0; j < VDR_Q6_K_Q8_1_MMVQ; ++j) { + vl[j] = get_int_b2(bq6_K->ql, iqs + j); + vh[j] = get_int_b2(bq6_K->qh, (QI6_K/4) * (iqs / (QI6_K/2)) + iqs % (QI6_K/4) + j) >> vh_shift; + } const int8_t * scales = bq6_K->scales + scale_offset; - int u[QR6_K]; + int u[VDR_Q6_K_Q8_1_MMVQ*QR6_K]; float d8[QR6_K]; #pragma unroll for (int i = 0; i < QR6_K; ++i) { - u[i] = get_int_b4(bq8_1[bq8_offset + 2*i].qs, iqs % QI8_1); +#pragma unroll + for (int j = 0; j < VDR_Q6_K_Q8_1_MMVQ; ++j) { + u[VDR_Q6_K_Q8_1_MMVQ*i + j] = get_int_b4(bq8_1[bq8_offset + 2*i].qs, iqs % QI8_1 + j); + } d8[i] = __low2float(bq8_1[bq8_offset + 2*i].ds); } From c0fd78626405c6d9b934efef0ec7291fcc20e3d6 Mon Sep 17 00:00:00 2001 From: Aleksander Grygier Date: Sat, 17 Oct 2026 17:20:51 +0200 Subject: [PATCH 075/423] ui : Discover & Download Models from HF Hub (#29584) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * common : resolve - download tags and list cached sidecars A Q4_0-mtp style tag now resolves the sidecar file when no model file matches it, so a solo draft or mmproj download actually pulls the file. Cached sidecar files list as their own entries so the state survives a restart, and removing such a tag deletes only the sidecar. Assisted-by: pi:zai-org/GLM-5.3 * ui : optional sanitized raw HTML in markdown Add an allowHtml prop to MarkdownContent: raw HTML found in the markdown is rendered after DOMPurify sanitization instead of being escaped as literal text. Default stays escaped. Assisted-by: pi:GLM-5.3-Flash * ui : add the discover models components The Discover surface: the curated-catalog list with its search and skeletons, the details pane with its readme, metadata, Hub stats and download options, and the standalone download progress bar that replaces the plain one. The quant download button sizes its select with a new xs trigger, so the select gains that size and the muted box look. The store stays incomplete when a repo fails, so the next mount retries it. Assisted-by: pi:llama.cpp/DeepSeek-V4.1-Flash * ui : put the discover view in the models manager The manager dialog gains its second view: the Discover Models call to action fades the table out, slides the title with its compass mark, and fades the discover surface in; the arrow button returns to the table. Assisted-by: pi:llama.cpp/DeepSeek-V4.1-Flash # Conflicts: # tools/ui/src/lib/components/ui/button/button.svelte * ui : add the Enable Discover Models setting The flag gates browsing and downloading HuggingFace GGUF models, so it is declared on the discover layer; the repo-tree lookup in the manager layer above reads it from here. Assisted-by: pi:llama.cpp/DeepSeek-V4.1-Flash * ui : resume model downloads from the pause point A stopped download keeps its partial file instead of deleting it, so re-posting the tag resumes from disk. The first progress record now carries the resume offset, the paused snapshot survives into the resumed entry, and a pause settles as soon as it is requested. Assisted-by: pi * ui : lead the discover list with downloads The discover sidebar lists running downloads above the suggested models, and steps aside while a query is active. A paused download stays in the manager table, where it is resumed or dropped. Assisted-by: pi * ui : make resuming a paused download work Re-posting the tag of a paused download hit the router's "already exists" check, because a paused download parks its entry as DOWNLOADED until the next reload: a download entry that is not running is now replaced and resumed. The UI keeps the paused row when the request still fails, instead of dropping it and looking finished. Assisted-by: pi * ui : remember unfinished model downloads A partial file is named after its content hash, so nothing on disk says which model it belongs to once the router restarts. The router now keeps the tags of the downloads it started, clears a tag when the download finishes or is deleted, and lists what is left as still downloading without starting anything. The UI adopts those tags as paused on every list fetch, so a refresh or a fresh open shows an incomplete download and can resume it from the files already on disk. Assisted-by: pi * ui : show a paused download's state on its chip A paused quant chip reads as paused at rest, the way a running one shows its spinner, and offers the resume action on hover. The selector's download row gets a View in Discover action beside its cancel. Assisted-by: pi * ui : collapse the download options by default The copyable download command starts collapsed so the option rows lead; expanding it stays one click away. Assisted-by: pi * ui : keep decision models out of the hub lists Decision models answer with calibrated probabilities instead of text, so llama-ui cannot run them: the hub data layer drops them from the catalog and every search, keyed by the decision-model tag or a text-classification pipeline. Assisted-by: pi * ui : read download tags as quant and sidecar A `:-` tag now parses into its quant and sidecar instead of one opaque quant token, and ModelsService.draftSidecarsByRepo collects the router listing's sidecar-only entries per repo, replacing the inline copy in the models store. Assisted-by: pi * ui : cluster draft downloads under their main model Tracked downloads group by repo: the main weights lead and the draft sidecars ride inside their item as compact sub-lines, on the selector, the discover list and the manager table. The main model's controls pause, resume and cancel the whole cluster, and the manager's download rows select into the configuration pane. Assisted-by: pi * ui : act on tracked downloads from the table rows The table passes its clusters' draft sub-lines to the model rows. A row's actions cascade pause and resume, an in-flight row offers Cancel download while a paused one keeps Delete from disk, and the confirm dialog's copy follows the state. Assisted-by: pi * ui : give downloads a configuration pane of their own Selecting a tracked download opens the pane in a download mode: live progress with pause, resume and cancel, and See in Discover, instead of metadata and load actions that a model without files cannot offer. Assisted-by: pi * ui : make the discover details self-fetching and responsive ModelsDiscoverDetailsContainer owns the details, tree and readme fetches so hosts outside the discover view can render the pane, and the details layout adapts to narrow panes with container queries instead of viewport breakpoints. Assisted-by: pi * ui : download a matching draft with the main model A main quant download now pulls the quant-matched draft of the most preferred kind the repo ships (DFlash before MTP), fire-and-forget, so speculative decoding is ready when the main lands. Manual draft chips stay available; the options row also collapses to one column in narrow panes. Assisted-by: pi * ui : keep discover off by default and couple it with the hub setting Discover Models starts unchecked. Enabling it turns the Hugging Face Hub setting on and disabling it turns it back off, so the checkbox never lies; with discover off the hub setting stays available on its own for the display data. Stored and reset state are corrected the same way on load. Assisted-by: pi * ui : open the manage models dialog on the manager view The dialog now opens on the manager instead of returning to the last view, discover entries still land on the discover view while the setting allows it, and the discover surfaces follow the setting's state. Assisted-by: pi * ui : give the discover search room for its focus ring The sticky search bar insets from the panel edges so the focus ring stays inside the scroll container, and rows scroll behind an opaque bar instead of through it. Assisted-by: pi * common: settle sidecar names on the listing and drop unfinished partials An infix or trailing sidecar token is ambiguous: draft heads use it (Model-MTP-BF16.gguf), so do full models with embedded heads (Model-Q8D-MTP.gguf). Such a name is a sidecar only next to token-free main weights of the same family, while mmproj or imatrix anywhere in the name, or a leading token, decides alone (mmproj_f16.gguf). A quant tag always carries its bit count, so a sidecar named after no quant (mtp-Model-it.gguf) gets the bare tag: the listing, the plan and remove all use it. An imatrix tag resolves into a plan slot of its own, downloaded but never wired into the params. remove resolves the tag the way a download does and deletes the partial files of every planned blob, so cancelling a running download or dropping one left unfinished by an earlier run leaves nothing on disk. The router passes its token and releases its lock around the call, and records the unfinished downloads with an atomic write. llama download succeeds on a sidecar-only plan and prints its paths. * ui : move dialog views and draft states into enums and constants The models dialog's views become the ModelsDialogView enum, the draft subline's states join ModelRowDownloadState, the view fade and the copy-feedback reset move to ui constants, and the dialog's compass marks come from DISCOVER_ICON. Assisted-by: pi * ui : drop the unused discover surface The dead openDiscoverModels twin, the commented-out command divider, the details pane's unused onClose and its header's close button, and a redundant type cast leave; the command's copy feedback joins the shared reset constant. Assisted-by: pi * ui : trim the download option comments and highlight the chat template The comment verbosity the options panel and its chips carried shrinks to the non-obvious why, and the chat template dialog renders through SyntaxHighlightedCode as jinja instead of a plain pre block. Assisted-by: pi * common: only a .gguf file is a sidecar, and mmproj/imatrix as whole words A non-GGUF file named like a sidecar (quantization-mmproj-Q8_0.json) is picked as the mmproj, and a sidecar word inside a longer word of a model name (Phi-2-OmniMatrix.i1-Q4_K_M.gguf) drops the model. * ui: dedupe draft sidecar badges per quant The draft sidecar badges deduped by kind and repo without the quant, so with a draft cached at two quants only one badge survived, and a delete of the other main quant found no badge to remove its sidecar with. Also restores the getDownloadEntries doc lost when the clusters method landed above it. Assisted-by: pi (zai-org/GLM-5.3-Flash) * ui: adopt the incomplete download status and drop the paused-ids registry The store inferred an unfinished download from a downloading entry it held no local state for, which misread a download running in another tab as paused and offered a resume the server rejects. The listing carries the "incomplete" status now, so the adoption keys off it directly. The paused ids no longer need to survive in localStorage: the server remembers them across restarts, and the progress snapshot was stale after a reload anyway. The stale localStorage key is removed. Assisted-by: pi (zai-org/GLM-5.3-Flash) * ui: quant-match the pause and cancel cascades Pause, resume and cancel cascaded to every draft sidecar of the repo while delete only took the quant-matched ones, so canceling a Q4_0 download also killed a Q8_0 draft that belongs to another quant's main. Every cascade acts on the quant-matched drafts now, matching what a delete takes with it. Assisted-by: pi (zai-org/GLM-5.3-Flash) * ui: collapse the duplicated draft subline markup The downloaded and in-flight branches repeated the identical connector, label and badge markup; one block with a conditional tail keeps the two states in sync. Assisted-by: pi (zai-org/GLM-5.3-Flash) * ui: scrub positioned styles out of sanitized model card html The style attribute stays allowed for the KaTeX output, which let a third party card draw a fixed overlay over the app. The sanitize hook drops position:fixed and position:absolute values and leaves the KaTeX layout styles alone. Assisted-by: pi (zai-org/GLM-5.3-Flash) * ui: run the discover mount fetch once and reuse the shared repoOf helper fetch() toggles its loading flag synchronously inside the mount effect, which re-ran the fetch on every toggle and retried failed catalog repos endlessly; the effect now untracks the call. The local repoOf copy and the duplicate services import give way to the shared helpers. Assisted-by: pi (zai-org/GLM-5.3-Flash) * ui: narrow the spec-type map to the draft sidecars SPEC_TYPE carried empty entries for the aux sidecars it is never indexed with; the map and the command component now type the badge as a draft sidecar through a draftBadgeOf helper instead of a cast. Assisted-by: pi (zai-org/GLM-5.3-Flash) * ui: name the decision-model tags and markdown cache keys The decision-model hub tag and pipeline tag move into the huggingface constants next to the other pipeline tags, and the markdown transform cache prefixes get names instead of bare one-letter tokens. Assisted-by: pi (zai-org/GLM-5.3-Flash) * ui: read the discover setting through one truthiness check Three places read the setting with !== false and one with Boolean(); all four read it the same way now. Assisted-by: pi (zai-org/GLM-5.3-Flash) * ui: compare the models dialog view against its enum The template compared the view state against the 'manage' literal in three spots while the script used the enum. Assisted-by: pi (zai-org/GLM-5.3-Flash) * nit: Remove `+` before draft sidecar badge * server: list unfinished downloads as incomplete A download that never finished was reported as "downloading", the same status a live download carries, so a client could not tell a running download from a dead one. The tags remembered in unfinished-downloads now list under a dedicated "incomplete" status, without an instance. Assisted-by: pi (zai-org/GLM-5.3-Flash) * server: download a model with its sidecars as -hf does Resolution goes back to master. POST /models takes the sidecar kinds to pull with the model (mmproj by default) and the download child gets them as -hf flags, so the server picks the files. /models lists real models only, each with the sidecars it has on disk, and incomplete downloads keep their sidecars across restarts. A stopped download keeps its partial files for the next attempt, and DELETE purges them through the same plan the download follows. * ui: download sidecars as parameters of their model The download options keep their chips, but a chip no longer posts a sidecar tag: the store maps it to its model and the sidecar kinds to pull, the quant chip takes its companion draft in the same download, and draft badges come from the sidecars each cached model lists. The server adds the sidecars a cached model lacks when it is posted again, and DELETE removes a single sidecar of a cached model. * ui : mirror the server's sidecar lookup in draft badges The server resolves a sidecar by the case-sensitive `-` token in the same directory as the weights. The UI matched the token loosely, so a badge could offer a draft the server never pulls. Assisted-by: DeepSeek-V4.1-Flash * ui : post the model's sidecar kinds with its download The store keeps one download per model: the chip's kinds ride along with the model's own sidecars, and a download row no longer expands into one row per sidecar. Draft badges carry no quant of their own yet, so a sidecar of a cached model is never labelled with the model's quant. Assisted-by: DeepSeek-V4.1-Flash * ui : join a quant's draft into its download chip One chip per quant, its draft badged inside it and its size covering both files. The draft is the file the server would pick: the same directory, the exact quant, else the closest bit depth, so a 4-bit model can pair with a 16-bit head. Assisted-by: DeepSeek-V4.1-Flash * ui : drop the draft download rows from the manager and selector A download is one row per model again: the drafts ride with it, and the rows read the sidecars on disk off the listing instead of pairing sidecar entries with their model. Pause, resume and cancel act on the model's own download. Assisted-by: DeepSeek-V4.1-Flash * ui : keep the discover download list one row per model The list leads with one row per downloading model, its sidecars as sub-lines. A row badges only the draft kinds the server can resolve, so a badge never outlives the file its lookup would pick. Assisted-by: DeepSeek-V4.1-Flash * ui : keep the discover chips readable in dark mode The metadata chips, the chat template button and the View on Hugging Face link carried a faint border and no surface over the dialog background, so they were near invisible. They take the tertiary button surface now, and their icons size with size-* so the variant's base rule leaves them alone. Assisted-by: DeepSeek-V4.1-Flash * ui : carry the draft through the download options One chip per main quant, its draft badged on it and the chip's downloaded state covering both files. The drafts sit behind an opt-in checkbox with an inline type pick, and each row opens its own command: it pins the row's quant and only names the draft type, so -hfd is gone. Assisted-by: DeepSeek-V4.1-Flash * ui : line the details skeleton up with the loaded pane Same root padding and container context as the loaded view, so the container queries resolve alike. The download options card pads its blocks instead of itself, the readme gets its card, and a description line stands in for the catalog paragraph. Assisted-by: DeepSeek-V4.1-Flash * common : drop the partials of a removed sidecar Removing a sidecar alone left the .downloadInProgress files of the sidecars a download had planned in blobs/. The removal purges them through the same plan the download follows, and the test covers a finished draft beside an unfinished one, whose partial goes while the model keeps its files. Assisted-by: DeepSeek-V4.1-Flash * server : name why a download is refused A model that is already downloading, or that is running, answered "already exists". It now says which: "is already downloading", or "is running, unload it first", both before the resumable check. The router test covers the two. Assisted-by: DeepSeek-V4.1-Flash * ui : post a download once per click Two quick clicks on the same chip posted twice, and the second answer was "already exists". The store tracks the model as pending, returns early while it is pending or in flight, and reads it as in flight before the server answers, so the chip shows the download at once. The new client test pins the flows: one POST per double click, the chip reading downloaded once the list reports the model, a resume keeping its sidecars, a cancel dropping the whole download, and a reload restoring the paused state. Assisted-by: DeepSeek-V4.1-Flash * ui : keep a plain download row the height of the others The sub-lines wrapper of a download row was rendered even with nothing in it, so its padding made the row taller than a plain one and the overlaid progress bar sat in that empty strip. It renders only with sub-lines now, in the selector and in the discover list, and the selector row takes the phone padding its neighbours have. Assisted-by: DeepSeek-V4.1-Flash * ui : badge a repo's draft once The args of a loaded model and the listing of a cached one report the same draft in different words, one with the file's quant and one with the kind alone, so the merge kept both and a repo row showed the same kind twice. The badges merge per kind now, the args one winning as it knows the quant, and the test pins both sources of a real repo. Assisted-by: DeepSeek-V4.1-Flash * server : report the sidecars of a preset or a download on its way A model the preset also names is registered as a preset, so its drafts were never reported and no row could badge them. A download on its way reported none either, though the record of the kinds it pulls is right there: those win over the plan now, the plan resolves through the preset's hf-repo, and a models dir stays out of it. Assisted-by: DeepSeek-V4.1-Flash * ui : badge a download's drafts and open its model from the row A draft still coming down badges the model id now, the way a landed one does, while an on-disk draft keeps its checked sub-line. The row itself opens the model, the status slot beside it pauses or resumes, and the manager resolves a focus on a download row so its panel opens. Assisted-by: DeepSeek-V4.1-Flash * common : report a removed sidecar that only had a partial Removing a sidecar alone dropped the result of the partial purge, so a sidecar that only ever landed as a .downloadInProgress blob answered false and the UI had no reason to refetch. The purge result is folded into the return now, and the test covers a partial-only removal. Assisted-by: DeepSeek-V4.1-Flash * ui : ignore a pause while the download POST is still on its way The chip reads as downloading before the server answers, so a quick second click unloaded a model that had not started and left the store neither in flight nor paused. A pause returns early while the POST is pending, before it touches the feed or the paused marker. The client test holds the POST open and pins it. Assisted-by: DeepSeek-V4.1-Flash * ui : offer a missing draft as its own chip A model already on disk without its draft had no way to add it: the quant chip carried the draft, so it read as a download of both. The quant chip keeps to its own state now, its joined draft shows only while it rides along or is on disk, and a downloaded quant missing its draft gets a chip of its own - badge, size, download - that posts the model with that kind, so the server keeps the files already there and pulls the draft. Assisted-by: DeepSeek-V4.1-Flash * server : list a paused sidecar download by the files on disk A model on disk whose sidecar download stops keeps its unfinished record, and the listing reported the kinds that record plans as if they were on disk: the UI checked a draft that never landed and offered no way to add it. The planned kinds now stand in only while the download runs. * chore: Clean up comments * server : list sidecars only for a model that resolves from a HF repo A preset on a local file has no HF repo, yet the listing resolved its name as one: the plan threw on the invalid repo format and GET /models failed for every model. The sidecars now come from the repo a preset names or a cached model is listed under, and a local file lists without them. * common : fix download cancel regression Signed-off-by: Adrien Gallouët * common : preserve model partials when removing sidecars Signed-off-by: Adrien Gallouët * server : one sidecar table and an unfinished downloads class Address review from @ggerganov. The sidecar kinds live in one table in common/download, mapping each name to its download option and plan file; the server validates, builds its options and lists cached sidecars through it instead of its own lists. The unfinished downloads move into server_unfinished_downloads, a record with its own mutex and a struct per entry, and server_models exposes them through unfinished_downloads_get and unfinished_downloads_remove. * common: one sidecar table, and removals that keep shared files The resolution, the cached listing and the removal follow a single table of sidecar kinds, so a kind is declared once with its file keyword, whether it is a draft, its option and its plan file. The server reads the same table for the kinds it accepts, the default sidecars and the child flags. Removing a tag keeps the files another cached tag still loads, such as an mmproj named after the removed quant. A sidecar removed on its own goes for every tag that resolves it. An unfinished download records the exact paths of its partials when it starts. A removal deletes those without asking the hub, and keeps a partial another recorded download also writes. The record is taken once the child is spawned, and the files go under the models lock, so a download of the same tag cannot start while they are removed. * server: drop stale partials and refuse a tag that resolves nothing Recording a download again deletes the partials of the earlier record that the new plan no longer writes, such as those of an older hub revision or of a sidecar the new request leaves out. A posted tag that resolves no file of the repo is refused before the download starts. Removing a model follows a single path, whether the reload erased its entry during the wait or not. * server: fix the CI race between the router download tests The pytest workers run in parallel on one shared cache, so a download test could find the file another test had just fetched and see no progress, or start from the unfinished downloads record another test wrote. Each download test now runs on an empty cache of its own. * server: refuse a DELETE of an unknown name and decide downloads in load() A DELETE of a name that is neither a model nor a recorded download went to common_download_remove, and without a tag it wiped the whole repo cache, including the files of a loaded quant or of a download on its way. Such a name is now not found, as it was before. A POST that downloads over an existing entry checked its state outside the lock, so a load of the same model in between got its instance killed. The checks now run in load(), under the lock, right before the spawn. The fake load time of the tests applies to model loads only, since a download worker loads none. * server: add incomplete_sidecars to the /models API A sidecar added to a cached model and then paused kept its record, but the listing showed only the complete model, so a fresh UI lost the pause. The model entry now carries incomplete_sidecars, the kinds its unfinished download has not written yet; a POST with the kind resumes one and a DELETE with it drops its partial. * common : drop a kept partial when a resume gets 416 Now that partials survive between attempts, an old or complete partial makes every retry fail with the same 416. Delete it so the next attempt can start over from zero. Signed-off-by: Adrien Gallouët * common: drop a kept partial when a resume gets 200 A server that ignores the range answers 200 to a resume, and the kept partial then fails every attempt the same way a 416 does. Drop it as well, so the next attempt starts over without a range. * server: refuse a download that shares a partial with one on its way Two downloads of quants that share a file, the mmproj or an mtp picked by the closest bits, appended to the same partial and renamed a file with interleaved bytes. A download whose plan writes a partial that a running download also writes is now refused until that one is done. * common : check the saved etag before resuming a partial download Before resuming, compare the etag saved with the partial to the one the server now serves. If it is missing or different, start over, so a resumed download never mixes bytes from two versions of the file. Signed-off-by: Adrien Gallouët --------- Signed-off-by: Adrien Gallouët Co-authored-by: Pascal Co-authored-by: Adrien Gallouët --- common/download.cpp | 236 ++++++--- common/download.h | 39 +- tests/test-model-resolution.cpp | 122 +++++ tools/server/README.md | 37 +- tools/server/server-models.cpp | 421 +++++++++++++++-- tools/server/server-models.h | 65 +++ tools/server/tests/unit/test_router.py | 237 ++++++++++ tools/server/tests/utils.py | 5 +- .../MarkdownContent/MarkdownContent.svelte | 34 +- .../MarkdownContent/markdown-processor.ts | 31 +- .../MarkdownContent/safe-html-config.ts | 138 ++++++ .../app/dialogs/DialogManageModels.svelte | 181 ++++++- .../app/models/ModelDraftSidecars.svelte | 2 - .../app/models/ModelDraftSubline.svelte | 53 +++ .../lib/components/app/models/ModelId.svelte | 15 + .../models/ModelsManager/ModelsManager.svelte | 32 +- .../ModelsManagerDownloadControl.svelte | 5 +- .../ModelsManagerModelConfiguration.svelte | 98 +++- ...elsManagerModelConfigurationActions.svelte | 143 ++++-- ...delsManagerModelConfigurationHeader.svelte | 31 +- .../ModelsManagerModelRow.svelte | 96 ++-- .../ModelsManagerModelsTable.svelte | 8 +- .../app/models/ModelsManager/row-actions.ts | 31 +- .../app/models/ModelsManager/utils.ts | 31 +- .../ModelsSelectorDownloadItem.svelte | 160 +++++-- .../ModelsSelector/ModelsSelectorList.svelte | 9 +- .../app/models/discover/ModelsDiscover.svelte | 210 ++++++++ .../ModelsDiscoverChatTemplateDialog.svelte | 75 +++ .../ModelsDiscoverDetails.svelte | 96 ++++ .../ModelsDiscoverDetailsContainer.svelte | 66 +++ ...odelsDiscoverDetailsDownloadOptions.svelte | 258 ++++++++++ ...tailsDownloadOptionsDownloadCommand.svelte | 131 +++++ ...sDownloadOptionsDraftDownloadButton.svelte | 61 +++ ...sDownloadOptionsQuantDownloadButton.svelte | 272 +++++++++++ ...lsDiscoverDetailsDownloadOptionsRow.svelte | 134 ++++++ .../download-options.utils.ts | 57 +++ .../index.ts | 24 + .../ModelsDiscoverDetailsHeader.svelte | 96 ++++ .../ModelsDiscoverDetailsHfHubStats.svelte | 31 ++ .../ModelsDiscoverDetailsMetadata.svelte | 82 ++++ .../ModelsDiscoverDetailsMetadataItem.svelte | 18 + .../ModelsDiscoverDetailsReadme.svelte | 17 + .../ModelsDiscoverDetailsSkeleton.svelte | 98 ++++ .../discover/ModelsDiscoverDetails/index.ts | 37 ++ .../ModelsDiscoverList.svelte | 41 ++ .../ModelsDiscoverListItem.svelte | 256 ++++++++++ .../ModelsDiscoverListItemSkeleton.svelte | 39 ++ .../ModelsDiscoverListSearch.svelte | 30 ++ .../discover/ModelsDiscoverList/index.ts | 16 + .../components/app/models/discover/index.ts | 16 + .../ui/src/lib/components/app/models/index.ts | 3 + .../settings/SettingsChat/SettingsChat.svelte | 7 + .../SettingsChat/SettingsChatFields.svelte | 14 +- .../ui/select/select-content.svelte | 4 +- .../ui/select/select-trigger.svelte | 14 +- .../lib/constants/huggingface.constants.ts | 6 + tools/ui/src/lib/constants/icons.constants.ts | 4 + tools/ui/src/lib/constants/index.ts | 1 + .../models-discover-download.constants.ts | 41 ++ .../constants/models-discover.constants.ts | 6 + .../lib/constants/settings-keys.constants.ts | 1 + .../src/lib/constants/settings.constants.ts | 13 +- .../ui/src/lib/constants/storage.constants.ts | 2 - tools/ui/src/lib/constants/ui.constants.ts | 6 + tools/ui/src/lib/enums/index.ts | 9 +- tools/ui/src/lib/enums/model.enums.ts | 16 +- tools/ui/src/lib/enums/server.enums.ts | 2 + tools/ui/src/lib/enums/ui.enums.ts | 6 + .../src/lib/services/huggingface.service.ts | 98 ++-- tools/ui/src/lib/services/models.service.ts | 106 +++-- tools/ui/src/lib/stores/index.ts | 3 + .../stores/models-discover/index.svelte.ts | 267 +++++++++++ .../ui/src/lib/stores/models/index.svelte.ts | 72 +-- .../ui/src/lib/stores/models/status.svelte.ts | 447 ++++++++++++------ .../src/lib/stores/settings/index.svelte.ts | 34 +- tools/ui/src/lib/stores/ui.svelte.ts | 10 + tools/ui/src/lib/types/api.d.ts | 5 + tools/ui/src/lib/types/huggingface.d.ts | 2 + tools/ui/src/lib/types/index.ts | 11 + tools/ui/src/lib/types/models-discover.d.ts | 56 +++ tools/ui/src/lib/types/models.d.ts | 2 + tools/ui/src/lib/utils/index.ts | 9 +- tools/ui/src/lib/utils/sidecars.ts | 18 +- .../markdown-content-sanitize.svelte.test.ts | 160 +++++++ .../models-download-flows.svelte.test.ts | 155 ++++++ .../models-manager-utils.svelte.test.ts | 28 +- .../models-sidecar-badges.svelte.test.ts | 99 ++++ tools/ui/tests/unit/hf-helpers.test.ts | 37 ++ tools/ui/tests/unit/model-id-parser.test.ts | 59 +++ .../tests/unit/model-sidecar-grammar.test.ts | 14 +- 90 files changed, 5688 insertions(+), 580 deletions(-) create mode 100644 tools/ui/src/lib/components/app/content/MarkdownContent/safe-html-config.ts create mode 100644 tools/ui/src/lib/components/app/models/ModelDraftSubline.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscover.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverChatTemplateDialog.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetails.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsContainer.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptions.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsDownloadCommand.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsDraftDownloadButton.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsQuantDownloadButton.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsRow.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/download-options.utils.ts create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/index.ts create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsHeader.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsHfHubStats.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsMetadata.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsMetadataItem.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsReadme.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsSkeleton.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/index.ts create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverList.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListItem.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListItemSkeleton.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListSearch.svelte create mode 100644 tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/index.ts create mode 100644 tools/ui/src/lib/components/app/models/discover/index.ts create mode 100644 tools/ui/src/lib/constants/models-discover-download.constants.ts create mode 100644 tools/ui/src/lib/stores/models-discover/index.svelte.ts create mode 100644 tools/ui/src/lib/types/models-discover.d.ts create mode 100644 tools/ui/tests/client/markdown-content-sanitize.svelte.test.ts create mode 100644 tools/ui/tests/client/models-download-flows.svelte.test.ts create mode 100644 tools/ui/tests/client/models-sidecar-badges.svelte.test.ts diff --git a/common/download.cpp b/common/download.cpp index ae0f7ffede42..38f26e1b497a 100644 --- a/common/download.cpp +++ b/common/download.cpp @@ -39,6 +39,9 @@ // downloader // +// the suffix of the file a download writes until it completes +static const std::string PARTIAL_SUFFIX = ".downloadInProgress"; + static void write_etag(const std::string & path, const std::string & etag) { const std::string etag_path = path + ".etag"; fs_write_atomic(std::filesystem::u8path(etag_path), etag); @@ -186,6 +189,13 @@ static bool common_pull_file(httplib::Client & cli, [&](const httplib::Response &response) { if (p.downloaded > 0 && response.status != 206) { LOG_WRN("%s: server did not respond with 206 Partial Content for a resume request. Status: %d\n", func, response.status); + if (response.status == 416 || response.status == 200) { + // the kept partial cannot be resumed: 416 for a complete or stale one, + // 200 from a server that ignores the range; drop it, so the next + // attempt starts over instead of wedging on the same answer + ofs.close(); + remove(path_tmp.c_str()); + } return false; } if (p.downloaded == 0 && response.status != 200) { @@ -214,7 +224,7 @@ static bool common_pull_file(httplib::Client & cli, if (progress_step >= p.total / 1000 || p.downloaded == p.total) { if (callback) { callback->on_update(p); - if (callback->is_cancelled()) { + if (callback->is_cancelled() && (p.total == 0 || p.downloaded != p.total)) { return false; } } @@ -331,10 +341,24 @@ static int common_download_file_single_online(const std::string & url, } bool success = false; - const std::string path_temporary = path + ".downloadInProgress"; + const std::string path_temporary = path + PARTIAL_SUFFIX; int delay = retry_delay_seconds; + if (!skip_etag && !etag.empty()) { + if (std::filesystem::exists(path_temporary) && read_etag(path_temporary) != etag) { + LOG_WRN("%s: partial file does not match the server etag, discarding it: %s\n", __func__, path_temporary.c_str()); + if (remove(path_temporary.c_str()) != 0) { + LOG_ERR("%s: unable to delete file: %s\n", __func__, path_temporary.c_str()); + return -1; + } + } + write_etag(path_temporary, etag); + } + if (opts.callback) { + // report the resume offset, so a resumed file does not start at 0 + p.downloaded = supports_ranges && std::filesystem::exists(path_temporary) ? + std::filesystem::file_size(path_temporary) : 0; opts.callback->on_start(p); } @@ -372,6 +396,7 @@ static int common_download_file_single_online(const std::string & url, } if (!etag.empty() && !skip_etag) { write_etag(path, etag); + remove((path_temporary + ".etag").c_str()); } success = true; break; @@ -381,12 +406,7 @@ static int common_download_file_single_online(const std::string & url, if (opts.callback) { opts.callback->on_done(p, success); } - if (opts.callback && opts.callback->is_cancelled() && - std::filesystem::exists(path_temporary)) { - if (remove(path_temporary.c_str()) != 0) { - LOG_ERR("%s: unable to delete temporary file: %s\n", __func__, path_temporary.c_str()); - } - } + // a stopped download keeps path_temporary, so the next attempt resumes from it if (!success) { LOG_ERR("%s: download failed after %d attempts\n", __func__, max_attempts); return -1; // max attempts reached @@ -594,33 +614,65 @@ static hf_cache::hf_file find_best_sibling(const hf_cache::hf_files & files, return best; } -static hf_cache::hf_file find_best_mmproj(const hf_cache::hf_files & files, - const std::string & model) { - return find_best_sibling(files, model, "mmproj"); +static const std::vector sidecars = { + { "mmproj", "mmproj", false, &common_download_opts::download_mmproj, &common_download_hf_plan::mmproj }, + { "mtp", "mtp-", true, &common_download_opts::download_mtp, &common_download_hf_plan::mtp }, + { "dflash", "dflash-", true, &common_download_opts::download_dflash, &common_download_hf_plan::dflash }, + { "eagle3", "eagle3-", true, &common_download_opts::download_eagle3, &common_download_hf_plan::eagle3 }, + { "dspark", "dspark-", true, &common_download_opts::download_dspark, &common_download_hf_plan::dspark }, +}; + +const std::vector & common_download_sidecars() { + return sidecars; } -static hf_cache::hf_file find_best_mtp(const hf_cache::hf_files & files, - const std::string & model, - const std::string & tag = "") { - return find_best_sibling(files, model, "mtp-", tag); +bool common_download_opts_add_sidecar(common_download_opts & opts, const std::string & kind) { + for (const auto & s : sidecars) { + if (kind == s.name) { + opts.*s.opt = true; + return true; + } + } + return false; } -static hf_cache::hf_file find_best_eagle3(const hf_cache::hf_files & files, - const std::string & model, - const std::string & tag = "") { - return find_best_sibling(files, model, "eagle3-", tag); +std::vector common_download_plan_sidecars(const common_download_hf_plan & plan) { + std::vector out; + for (const auto & s : sidecars) { + if (!(plan.*s.file).path.empty()) { + out.emplace_back(s.name); + } + } + return out; } -static hf_cache::hf_file find_best_dflash(const hf_cache::hf_files & files, - const std::string & model, - const std::string & tag = "") { - return find_best_sibling(files, model, "dflash-", tag); +std::vector common_download_plan_partials(const common_download_hf_plan & plan) { + std::vector out; + + // a file in the cache resolves to its snapshot entry, a file still to download to its blob + auto add = [&](const std::string & kind, const hf_cache::hf_file & f) { + if (!f.local_path.empty() && f.local_path != f.final_path) { + out.push_back({ kind, f.local_path + PARTIAL_SUFFIX }); + } + }; + + for (const auto & f : plan.model_files) { + add("", f); + } + for (const auto & s : sidecars) { + add(s.name, plan.*s.file); + } + return out; } -static hf_cache::hf_file find_best_dspark(const hf_cache::hf_files & files, - const std::string & model, - const std::string & tag = "") { - return find_best_sibling(files, model, "dspark-", tag); +// true when `name` carries the keyword of a sidecar kind +static bool name_is_sidecar(const std::string & name) { + for (const auto & s : sidecars) { + if (name.find(s.keyword) != std::string::npos) { + return true; + } + } + return false; } static bool gguf_filename_is_model(const std::string & filepath) { @@ -633,12 +685,7 @@ static bool gguf_filename_is_model(const std::string & filepath) { filename = filename.substr(pos + 1); } - return filename.find("mmproj") == std::string::npos && - filename.find("imatrix") == std::string::npos && - filename.find("mtp-") == std::string::npos && - filename.find("eagle3-") == std::string::npos && - filename.find("dflash-") == std::string::npos && - filename.find("dspark-") == std::string::npos; + return filename.find("imatrix") == std::string::npos && !name_is_sidecar(filename); } static hf_cache::hf_file find_best_model(const hf_cache::hf_files & files, @@ -730,8 +777,11 @@ common_download_hf_plan common_download_get_hf_plan(const common_params_model & } } else { primary = find_best_model(all, tag); - // a requested sidecar can resolve on its own, without a full model of the same tag - if (primary.path.empty() && !opts.download_mtp && !opts.download_dflash && !opts.download_eagle3 && !opts.download_dspark) { + // a requested draft can resolve on its own, without a full model of the same tag + const bool want_draft = std::any_of(sidecars.begin(), sidecars.end(), [&](const common_download_sidecar & s) { + return s.draft && opts.*s.opt; + }); + if (primary.path.empty() && !want_draft) { LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str()); list_available_gguf_files(all); return plan; @@ -743,24 +793,22 @@ common_download_hf_plan common_download_get_hf_plan(const common_params_model & plan.model_files = get_split_files(all, primary); } - if (opts.download_mmproj && !primary.path.empty()) { - plan.mmproj = find_best_mmproj(all, primary.path); - } - if (opts.download_mtp) { - plan.mtp = find_best_mtp(all, primary.path, tag); - } - if (opts.download_dflash) { - plan.dflash = find_best_dflash(all, primary.path, tag); - } - if (opts.download_eagle3) { - plan.eagle3 = find_best_eagle3(all, primary.path, tag); - } - if (opts.download_dspark) { - plan.dspark = find_best_dspark(all, primary.path, tag); + // a draft resolves at the tag, the mmproj pairs with the model + for (const auto & s : sidecars) { + if (!(opts.*s.opt)) { + continue; + } + if (s.draft) { + plan.*s.file = find_best_sibling(all, primary.path, s.keyword, tag); + } else if (!primary.path.empty()) { + plan.*s.file = find_best_sibling(all, primary.path, s.keyword); + } } - if (primary.path.empty() && - plan.mtp.local_path.empty() && plan.dflash.local_path.empty() && plan.eagle3.local_path.empty() && plan.dspark.local_path.empty()) { + const bool has_draft = std::any_of(sidecars.begin(), sidecars.end(), [&](const common_download_sidecar & s) { + return s.draft && !(plan.*s.file).local_path.empty(); + }); + if (primary.path.empty() && !has_draft) { LOG_ERR("%s: no GGUF files found in repository %s\n", __func__, repo.c_str()); list_available_gguf_files(all); } @@ -926,6 +974,11 @@ std::string common_docker_resolve_model(const std::string & docker) { } } +// a cached file that lists a model under its tag: the first split of a model, not a sidecar +static bool is_listed_model(const gguf_split_info & split) { + return split.index == 1 && !split.tag.empty() && !name_is_sidecar(split.prefix); +} + std::vector common_list_cached_models() { std::unordered_set seen; std::vector result; @@ -934,12 +987,7 @@ std::vector common_list_cached_models() { for (const auto & f : files) { auto split = get_gguf_split_info(f.path); - if (split.index != 1 || split.tag.empty() || - split.prefix.find("mmproj") != std::string::npos || - split.prefix.find("mtp-") != std::string::npos || - split.prefix.find("eagle3-") != std::string::npos || - split.prefix.find("dflash-") != std::string::npos || - split.prefix.find("dspark-") != std::string::npos) { + if (!is_listed_model(split)) { continue; } if (seen.insert(f.repo_id + ":" + split.tag).second) { @@ -970,6 +1018,44 @@ std::string common_download_resolve_path(const std::string & hf_repo_with_tag, c return find_best_model(files, tag).local_path; } +// the files -hf loads from the cache, the model with every sidecar, for each cached +// tag of `repo_id` other than `tag` +static std::unordered_set files_in_use(const std::string & repo_id, const std::string & tag) { + common_download_opts opts; + opts.offline = true; + for (const auto & s : sidecars) { + opts.*s.opt = true; + } + + std::unordered_set tags; + std::unordered_set used; + + for (const auto & f : hf_cache::get_cached_files(repo_id)) { + auto split = get_gguf_split_info(f.path); + if (!is_listed_model(split) || split.tag == tag || !tags.insert(split.tag).second) { + continue; + } + + common_params_model model; + model.hf_repo = repo_id + ":" + split.tag; + + auto plan = common_download_get_hf_plan(model, opts); + + for (const auto & m : plan.model_files) { + used.insert(m.local_path); + } + for (const auto & s : sidecars) { + const auto & sf = plan.*s.file; + if (!sf.local_path.empty()) { + used.insert(sf.local_path); + } + } + } + return used; +} + +static bool remove_snapshot_entries(const std::string & repo_id, const std::vector & to_remove); + bool common_download_remove(const std::string & hf_repo_with_tag) { namespace fs = std::filesystem; @@ -989,11 +1075,13 @@ bool common_download_remove(const std::string & hf_repo_with_tag) { return false; } - // collect snapshot entries whose tag matches + // collect snapshot entries whose tag matches, but a file another tag loads stays + const auto used = files_in_use(repo_id, tag_upper); + std::vector to_remove; for (const auto & f : files) { auto split = get_gguf_split_info(f.path); - if (split.tag == tag_upper) { + if (split.tag == tag_upper && used.find(f.local_path) == used.end()) { to_remove.emplace_back(f.local_path); } } @@ -1002,6 +1090,13 @@ bool common_download_remove(const std::string & hf_repo_with_tag) { return false; } + return remove_snapshot_entries(repo_id, to_remove); +} + +// remove snapshot entries of `repo_id`, then the blobs no remaining entry references +static bool remove_snapshot_entries(const std::string & repo_id, const std::vector & to_remove) { + namespace fs = std::filesystem; + // resolve blob paths from symlinks before deleting snapshot entries std::vector blobs_to_check; for (const auto & p : to_remove) { @@ -1053,3 +1148,26 @@ bool common_download_remove(const std::string & hf_repo_with_tag) { return true; } + +bool common_download_remove_sidecars(const std::string & hf_repo_with_tag, const common_download_opts & opts) { + common_params_model model; + model.hf_repo = hf_repo_with_tag; + + common_download_opts cached = opts; + cached.offline = true; + + auto plan = common_download_get_hf_plan(model, cached); + + std::vector to_remove; + for (const auto & s : sidecars) { + const auto & f = plan.*s.file; + if (!f.local_path.empty()) { + to_remove.emplace_back(std::filesystem::u8path(f.local_path)); + } + } + if (to_remove.empty()) { + return false; + } + + return remove_snapshot_entries(common_download_split_repo_tag(hf_repo_with_tag).first, to_remove); +} diff --git a/common/download.h b/common/download.h index 8c30cfc3eadb..4dfd0c9831ed 100644 --- a/common/download.h +++ b/common/download.h @@ -104,10 +104,17 @@ std::string common_docker_resolve_model(const std::string & docker); // Remove a cached model from disk // input format: "user/model" or "user/model:tag" // - if tag is omitted, removes the entire repo cache directory -// - if tag is present, removes only files matching that tag (and orphaned blobs) +// - if tag is present, removes only files matching that tag (and orphaned blobs), +// except the files another cached tag of the repo still loads // returns true if anything was removed bool common_download_remove(const std::string & hf_repo_with_tag); +// remove the cached sidecar files `opts` selects for "user/model:tag", as the cached +// plan resolves them, and leave the model files in place; a sidecar file another tag +// shares goes for every tag that resolves it +// returns true if anything was removed +bool common_download_remove_sidecars(const std::string & hf_repo_with_tag, const common_download_opts & opts); + struct common_download_hf_plan { hf_cache::hf_file primary; hf_cache::hf_files model_files; @@ -119,3 +126,33 @@ struct common_download_hf_plan { hf_cache::hf_file preset; // if set, only this file is downloaded }; common_download_hf_plan common_download_get_hf_plan(const common_params_model & model, const common_download_opts & opts); + +// a sidecar kind a download can pull next to its model: mmproj, and the draft heads +// of --spec-type draft- +struct common_download_sidecar { + const char * name; // the kind, as the API names it + const char * keyword; // the part of a file name that marks a file of this kind + bool draft; // a draft head, resolved at the tag even without a model + + bool common_download_opts::* opt; // the option that requests it + hf_cache::hf_file common_download_hf_plan::* file; // the plan file it resolves to +}; + +// every sidecar kind, the single list the resolution, the listing and the removal follow +const std::vector & common_download_sidecars(); + +// select the sidecar `kind` in `opts`; returns false for an unknown kind +bool common_download_opts_add_sidecar(common_download_opts & opts, const std::string & kind); + +// the sidecar kinds `plan` resolves a file for +std::vector common_download_plan_sidecars(const common_download_hf_plan & plan); + +// a file `plan` still has to download, at the path its partial takes while it is +// incomplete; kind is the sidecar it belongs to, empty for the model files +struct common_download_partial { + std::string kind; + std::string path; +}; + +// the partials the download of `plan` writes, for the files not in the cache yet +std::vector common_download_plan_partials(const common_download_hf_plan & plan); diff --git a/tests/test-model-resolution.cpp b/tests/test-model-resolution.cpp index 5c793e24a6ca..eb7b30936e29 100644 --- a/tests/test-model-resolution.cpp +++ b/tests/test-model-resolution.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -339,6 +340,124 @@ static void test_plan_resolution() { g_repos.clear(); } +// put a file in the cache the way a finished download leaves it, refs/main included +static void cache_put(const std::string & repo_id, const std::string & path) { + namespace fs = std::filesystem; + + const fs::path local = cached(repo_id, path); + fs::create_directories(local.parent_path()); + { std::ofstream(local) << "gguf"; } + + const fs::path ref = local.parent_path().parent_path().parent_path() / "refs" / "main"; + fs::create_directories(ref.parent_path()); + if (!fs::exists(ref)) { + std::ofstream(ref) << COMMIT << "\n"; + } +} + +// +// partials of a download: the files of its plan not in the cache yet, next to +// the blobs, named after the oids the listing serves +// + +static void test_plan_partials() { + namespace fs = std::filesystem; + + printf("test-model-resolution: plan partials\n"); + + g_context = "plan partials"; + g_repos.clear(); + g_repos["test/partial"] = { "model-Q4_0.gguf", "mmproj-model-F16.gguf", "mtp-model-Q4_0.gguf" }; + + // the mmproj is complete, the model and its mtp are still to download + cache_put("test/partial", "mmproj-model-F16.gguf"); + + common_download_opts opts; + opts.download_mmproj = true; + opts.download_mtp = true; + + const auto parts = common_download_plan_partials(common_download_get_hf_plan(model_ref("test/partial:Q4_0"), opts)); + + const fs::path blobs = cache_dir / "models--test--partial" / "blobs"; + REQUIRE(parts.size() == 2); + REQUIRE(parts[0].kind.empty()); + REQUIRE(fs::path(parts[0].path).lexically_normal() == (blobs / "0000000000000000000000000000000000000001.downloadInProgress").lexically_normal()); + REQUIRE_EQ(parts[1].kind, "mtp"); + REQUIRE(fs::path(parts[1].path).lexically_normal() == (blobs / "0000000000000000000000000000000000000003.downloadInProgress").lexically_normal()); + + fs::remove_all(cache_dir / "models--test--partial"); + g_repos.clear(); +} + +// +// removal of a sidecar alone: its snapshot entry goes, the model keeps its files +// + +static void test_sidecar_removal() { + namespace fs = std::filesystem; + + printf("test-model-resolution: sidecar removal\n"); + + g_context = "sidecar removal"; + g_repos.clear(); + + cache_put("test/sidecar", "model-Q4_0.gguf"); + cache_put("test/sidecar", "dflash-model-Q4_0.gguf"); + cache_put("test/sidecar", "mtp-model-Q4_0.gguf"); + + common_download_opts opts; + opts.download_dflash = true; + + // the selected sidecar goes, the other one and the model stay + REQUIRE(common_download_remove_sidecars("test/sidecar:Q4_0", opts)); + REQUIRE(!fs::exists(cached("test/sidecar", "dflash-model-Q4_0.gguf"))); + REQUIRE(fs::exists(cached("test/sidecar", "mtp-model-Q4_0.gguf"))); + REQUIRE(fs::exists(cached("test/sidecar", "model-Q4_0.gguf"))); + REQUIRE(!common_download_remove_sidecars("test/sidecar:Q4_0", opts)); + + fs::remove_all(cache_dir / "models--test--sidecar"); +} + +// +// removal of a tag keeps the files another cached tag still loads: a sidecar +// named after the removed tag that another quant resolves stays on disk +// + +static void test_shared_removal() { + namespace fs = std::filesystem; + + printf("test-model-resolution: shared removal\n"); + + g_context = "shared removal"; + g_repos.clear(); + + for (const char * f : { "model-Q8_0.gguf", "model-BF16.gguf", "mmproj-model-BF16.gguf", "mtp-model-BF16.gguf" }) { + cache_put("test/shared", f); + } + + // the BF16 model goes, the mmproj and the mtp Q8_0 resolves stay + REQUIRE(common_download_remove("test/shared:BF16")); + REQUIRE(!fs::exists(cached("test/shared", "model-BF16.gguf"))); + REQUIRE(fs::exists(cached("test/shared", "mmproj-model-BF16.gguf"))); + REQUIRE(fs::exists(cached("test/shared", "mtp-model-BF16.gguf"))); + REQUIRE(fs::exists(cached("test/shared", "model-Q8_0.gguf"))); + + fs::remove_all(cache_dir / "models--test--shared"); + + // with no other tag, the files of the removed tag all go, sidecars included + for (const char * f : { "model-BF16.gguf", "mmproj-model-BF16.gguf", "mtp-model-BF16.gguf" }) { + cache_put("test/alone", f); + } + + REQUIRE(common_download_remove("test/alone:BF16")); + REQUIRE(!fs::exists(cached("test/alone", "model-BF16.gguf"))); + REQUIRE(!fs::exists(cached("test/alone", "mmproj-model-BF16.gguf"))); + REQUIRE(!fs::exists(cached("test/alone", "mtp-model-BF16.gguf"))); + REQUIRE(!common_download_remove("test/alone:BF16")); + + fs::remove_all(cache_dir / "models--test--alone"); +} + // // end-to-end assembly: real CLI parsing, real handler init resolving over the // loopback, downloads skipped by flipping offline before apply @@ -509,6 +628,9 @@ int main(void) { common_set_env("MODEL_ENDPOINT", "http://127.0.0.1:" + std::to_string(port) + "/"); test_plan_resolution(); + test_plan_partials(); + test_sidecar_removal(); + test_shared_removal(); test_task_assembly(); server.stop(); diff --git a/tools/server/README.md b/tools/server/README.md index c42336c7f217..ec82db0076ed 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -2106,6 +2106,32 @@ Note: for "downloading" state, there can be multiple files be downloading in par } ``` +A download that never finished (paused or interrupted) is listed under a separate +"incomplete" status, without an instance: posting the model to `/models` again resumes it +from the files already on disk, and `DELETE /models` drops its partial files. + +```json +"status": { + "value": "incomplete", + "args": [] +} +``` + +Cached models and incomplete downloads carry the sidecar kinds they have on disk +(`mmproj`, `mtp`, `dflash`, `eagle3`, `dspark`), resolved the way `-hf` loads them: + +```json +"sidecars": ["mmproj", "mtp"] +``` + +A cached model whose download of more sidecars never finished also lists the kinds not on +disk yet: posting the model with them resumes them, and `DELETE /models` with `&sidecar={kind}` +drops their partial files. + +```json +"incomplete_sidecars": ["mtp"] +``` + ### POST `/models/load`: Load a model Load a model @@ -2225,7 +2251,7 @@ Example events: Trigger a new download (non-blocking), the progress can be tracked via SSE endpoint `/models/sse` -To cancel model downloading, send an event to `/models/unload` +To pause a download, send an event to `/models/unload`: it keeps its partial files and is listed as "incomplete" until it is posted again. To cancel it and drop its files, use `DELETE /models`. Download procedure: - Send POST request to `/models` @@ -2238,9 +2264,14 @@ Payload: ```json { "model": "ggml-org/gemma-3-4b-it-GGUF:Q4_K_M", + "sidecars": ["mmproj", "mtp"] } ``` +`sidecars` lists the sidecar kinds pulled with the model, picked by the server the way `-hf` picks them: `mmproj`, and the draft heads `mtp`, `dflash`, `eagle3` and `dspark`. It defaults to `["mmproj"]`, as `-hf` does. + +Posting a model that is already in the cache and not running downloads the sidecars it lacks and keeps the files already on disk. + Response (download is started in the background): ```json @@ -2267,6 +2298,10 @@ IMPORTANT: only model stored in cache can be deleted. You cannot delete models i Model name must be passed via query param: `?model={name}` +To remove one sidecar of a cached model and keep the model, add `&sidecar={kind}` (`mmproj`, `mtp`, `dflash`, `eagle3` or `dspark`); the model must not be running. + +Deleting a model keeps the files another cached quant of the repo still loads, such as an mmproj named after the deleted quant, and the partial files another download still writes. A sidecar deleted on its own goes for every quant that shares it. + If delete success, it will send an SSE event of type `model_remove` Response: diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index 60353240ac91..d77837c866f8 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -618,6 +619,7 @@ server_models::server_models( LOG_WRN("failed to get server executable path: %s\n", e.what()); LOG_WRN("using original argv[0] as fallback: %s\n", argv[0]); } + unfinished_downloads.load(); load_models(); debug_fake_timing = !common_get_env("LLAMA_SERVER_DEBUG_FAKE_TIMING").empty(); } @@ -686,6 +688,233 @@ void server_models::add_model(server_model_meta && meta) { }; } +static bool has_sidecar(const std::vector & sidecars, const std::string & kind) { + return std::find(sidecars.begin(), sidecars.end(), kind) != sidecars.end(); +} + +// reject a sidecar kind no download can pull +static void check_sidecar(const std::string & kind) { + std::vector kinds; + for (const auto & s : common_download_sidecars()) { + kinds.emplace_back(s.name); + } + if (!has_sidecar(kinds, kind)) { + throw std::invalid_argument("unknown sidecar '" + kind + "', expected one of " + string_join(kinds, ", ")); + } +} + +// the sidecars a download asks for; by default the ones -hf pulls without a draft type +static std::vector parse_sidecars(const json & body) { + std::vector defaults; + for (const auto & s : common_download_sidecars()) { + if (!s.draft) { + defaults.emplace_back(s.name); + } + } + std::vector sidecars = json_value(body, "sidecars", defaults); + for (const auto & kind : sidecars) { + check_sidecar(kind); + } + return sidecars; +} + +// the download options of a model with its sidecars, the plan a removal purges partials from +static common_download_opts sidecar_download_opts(const common_params & params, const std::vector & sidecars) { + common_download_opts opts; + opts.bearer_token = params.hf_token; + for (const auto & kind : sidecars) { + common_download_opts_add_sidecar(opts, kind); + } + return opts; +} + +// the HF repo a model resolves from: a preset names its own, a cached model is +// listed under it, and a local file or a models dir entry has none +static std::string model_hf_repo(const server_model_meta & meta) { + std::string hf_repo; + if (meta.preset.get_option("LLAMA_ARG_HF_REPO", hf_repo) && !hf_repo.empty()) { + return hf_repo; + } + return meta.source == SERVER_MODEL_SOURCE_CACHE ? meta.name : ""; +} + +// the sidecar kinds of a HF repo present on disk, as the offline plan resolves them +static json cached_sidecars(const std::string & hf_repo) { + common_params_model model; + model.hf_repo = hf_repo; + + common_download_opts opts; + opts.offline = true; + for (const auto & s : common_download_sidecars()) { + opts.*s.opt = true; + } + + return common_download_plan_sidecars(common_download_get_hf_plan(model, opts)); +} + +// replace any inherited value of `key` in a child environment +static void set_child_env(std::vector & env, const std::string & key, const std::string & value) { + env.erase(std::remove_if(env.begin(), env.end(), [&](const std::string & e) { + return e.rfind(key + "=", 0) == 0; + }), env.end()); + env.push_back(key + "=" + value); +} + +// where the unfinished downloads are remembered, inside the model cache +static std::string unfinished_downloads_path() { + return hf_cache::get_cache_path() + "/unfinished-downloads.json"; +} + +void server_unfinished_downloads::load() { + std::ifstream in(unfinished_downloads_path()); + if (!in) { + return; + } + std::string raw((std::istreambuf_iterator(in)), std::istreambuf_iterator()); + std::lock_guard lk(mutex); + try { + json data = json::parse(raw); + for (const auto & d : data.value("downloads", json::array())) { + server_unfinished_download download; + download.model = json_value(d, "model", std::string()); + download.sidecars = json_value(d, "sidecars", std::vector()); + for (const auto & p : d.value("partials", json::array())) { + download.partials.push_back({ json_value(p, "kind", std::string()), json_value(p, "path", std::string()) }); + } + if (!download.model.empty()) { + downloads.push_back(std::move(download)); + } + } + } catch (const std::exception & e) { + LOG_WRN("failed to read unfinished downloads: %s\n", e.what()); + } +} + +void server_unfinished_downloads::save() { + json data; + data["downloads"] = json::array(); + for (const auto & d : downloads) { + json partials = json::array(); + for (const auto & p : d.partials) { + partials.push_back({ {"kind", p.kind}, {"path", p.path} }); + } + data["downloads"].push_back({ {"model", d.model}, {"sidecars", d.sidecars}, {"partials", partials} }); + } + try { + fs_write_atomic(std::filesystem::u8path(unfinished_downloads_path()), data.dump()); + } catch (const std::exception & e) { + LOG_WRN("failed to record unfinished downloads: %s\n", e.what()); + } +} + +void server_unfinished_downloads::delete_partials(const std::vector & parts) { + for (const auto & p : parts) { + const bool shared = std::any_of(downloads.begin(), downloads.end(), [&](const server_unfinished_download & d) { + return std::any_of(d.partials.begin(), d.partials.end(), [&](const common_download_partial & q) { + return q.path == p.path; + }); + }); + if (shared) { + continue; + } + std::error_code ec; + std::filesystem::remove(std::filesystem::u8path(p.path), ec); + } +} + +void server_unfinished_downloads::add(const server_unfinished_download & download) { + std::lock_guard lk(mutex); + std::vector stale; + auto it = std::find_if(downloads.begin(), downloads.end(), [&](const auto & d) { return d.model == download.model; }); + if (it != downloads.end()) { + stale = std::move(it->partials); + *it = download; + } else { + downloads.push_back(download); + } + save(); + delete_partials(stale); +} + +void server_unfinished_downloads::forget(const std::string & model) { + std::lock_guard lk(mutex); + auto it = std::find_if(downloads.begin(), downloads.end(), [&](const auto & d) { return d.model == model; }); + if (it != downloads.end()) { + downloads.erase(it); + save(); + } +} + +bool server_unfinished_downloads::remove(const std::string & model) { + std::lock_guard lk(mutex); + auto it = std::find_if(downloads.begin(), downloads.end(), [&](const auto & d) { return d.model == model; }); + if (it == downloads.end()) { + return false; + } + auto partials = std::move(it->partials); + downloads.erase(it); + save(); + delete_partials(partials); + return true; +} + +bool server_unfinished_downloads::remove_sidecar(const std::string & model, const std::string & kind) { + std::lock_guard lk(mutex); + auto it = std::find_if(downloads.begin(), downloads.end(), [&](const auto & d) { return d.model == model; }); + if (it == downloads.end() || !has_sidecar(it->sidecars, kind)) { + return false; + } + it->sidecars.erase(std::remove(it->sidecars.begin(), it->sidecars.end(), kind), it->sidecars.end()); + + std::vector dropped; + auto & partials = it->partials; + auto pos = std::stable_partition(partials.begin(), partials.end(), [&](const common_download_partial & p) { + return p.kind != kind; + }); + dropped.assign(pos, partials.end()); + partials.erase(pos, partials.end()); + + save(); + delete_partials(dropped); + return true; +} + +std::vector server_unfinished_downloads::get_all() { + std::lock_guard lk(mutex); + return downloads; +} + +std::vector server_models::unfinished_downloads_get() { + return unfinished_downloads.get_all(); +} + +void server_models::unfinished_downloads_remove(const std::string & name) { + // under the lock, so a download of the same tag cannot start while its files go + std::lock_guard lk(mutex); + if (mapping.find(name) != mapping.end()) { + throw std::invalid_argument("model '" + name + "' is downloading"); + } + // a name with no record is no model: its files belong to the cached tags of the repo + if (!unfinished_downloads.remove(name)) { + throw std::runtime_error("model name=" + name + " is not found"); + } + bool ok = common_download_remove(name); + SRV_INF("removing unfinished download name=%s (%s)\n", name.c_str(), ok ? "succeeded" : "nothing cached"); +} + +bool server_models::remove_sidecar(const std::string & name, const std::string & kind) { + std::lock_guard lk(mutex); + auto it = mapping.find(name); + if (it != mapping.end() && (it->second.meta.is_running() || it->second.meta.status == SERVER_MODEL_STATUS_DOWNLOADING)) { + throw std::invalid_argument("model '" + name + "' is in use, unload it first"); + } + // a sidecar that only ever landed as a partial is still something removed + bool ok = unfinished_downloads.remove_sidecar(name, kind); + ok = common_download_remove_sidecars(name, sidecar_download_opts(base_params, { kind })) || ok; + SRV_INF("removing sidecar %s of model name=%s (%s)\n", kind.c_str(), name.c_str(), ok ? "succeeded" : "nothing cached"); + return ok; +} + void server_models::notify_sse(const std::string & event, const std::string & model_id, const json & data) { std::unique_ptr result = std::make_unique(); result->data = { @@ -1135,7 +1364,8 @@ void server_models::load(const std::string & name) { } void server_models::load(const std::string & name, const load_options & opts) { - if (debug_fake_timing) { + // the fake load time of a model; a download worker loads none + if (debug_fake_timing && opts.mode == SERVER_CHILD_MODE_NORMAL) { // do not hold the mutex here, other requests must keep making progress std::this_thread::sleep_for(std::chrono::seconds(2)); } @@ -1152,6 +1382,43 @@ void server_models::load(const std::string & name, const load_options & opts) { // against the freshest preset and a consistent mapping state cv.wait(lk, [this]() { return !is_reloading; }); + // a download over an existing entry is decided under the lock, so a load of the + // same model cannot start in between: a paused download or an idle cached model + // is downloaded again with its files kept, a download on its way and a running + // model are refused, and a preset or local model is not a download + if (opts.mode == SERVER_CHILD_MODE_DOWNLOAD) { + auto it = mapping.find(name); + if (it != mapping.end()) { + const auto & cur = it->second.meta; + if (cur.status == SERVER_MODEL_STATUS_DOWNLOADING) { + throw std::invalid_argument("model '" + name + "' is already downloading"); + } + if (cur.is_running()) { + throw std::invalid_argument("model '" + name + "' is running, unload it first"); + } + if (cur.source != SERVER_MODEL_SOURCE_CACHE) { + throw std::invalid_argument("model '" + name + "' already exists"); + } + SRV_INF("resuming download for model '%s'\n", name.c_str()); + } + + // two downloads writing the same partial would interleave their bytes + for (const auto & d : unfinished_downloads.get_all()) { + auto other = mapping.find(d.model); + if (d.model == name || other == mapping.end() || other->second.meta.status != SERVER_MODEL_STATUS_DOWNLOADING) { + continue; + } + for (const auto & p : d.partials) { + const bool shared = std::any_of(opts.partials.begin(), opts.partials.end(), [&](const common_download_partial & q) { + return q.path == p.path; + }); + if (shared) { + throw std::invalid_argument("model '" + name + "' shares files with the download of '" + d.model + "' on its way, retry once it is done"); + } + } + } + } + auto meta = opts.custom_meta.has_value() ? *opts.custom_meta : mapping[name].meta; if (meta.status != SERVER_MODEL_STATUS_UNLOADED) { SRV_INF("model %s is not ready\n", name.c_str()); @@ -1201,6 +1468,23 @@ void server_models::load(const std::string & name, const load_options & opts) { inst.meta.status = SERVER_MODEL_STATUS_DOWNLOADING; child_env.push_back("LLAMA_SERVER_CHILD_MODE=download"); child_env.push_back("LLAMA_ARG_HF_REPO=" + name); + + // the sidecars as -hf would pull them: the draft types, and the mmproj switch + // for the one that is not a draft + bool mmproj = false; + std::string spec_types; + for (const auto & s : common_download_sidecars()) { + if (!has_sidecar(opts.sidecars, s.name)) { + continue; + } + if (s.draft) { + spec_types += (spec_types.empty() ? "" : ",") + std::string("draft-") + s.name; + } else { + mmproj = true; + } + } + set_child_env(child_env, "LLAMA_ARG_MMPROJ_AUTO", mmproj ? "1" : "0"); + set_child_env(child_env, "LLAMA_ARG_SPEC_TYPE", spec_types.empty() ? "none" : spec_types); } SRV_INF("%s", "spawning server instance with args:\n"); @@ -1216,6 +1500,11 @@ void server_models::load(const std::string & name, const load_options & opts) { } } + // remember the download: the partial files carry a content hash, not the name + if (opts.mode == SERVER_CHILD_MODE_DOWNLOAD) { + unfinished_downloads.add({ name, opts.sidecars, opts.partials }); + } + // old process should have exited already, but just in case, we clean it up here { auto it = mapping.find(name); @@ -1390,6 +1679,12 @@ void server_models::update_download_progress(const std::string & name, const com // mark the instance to be erased on next load_models() call it->second.meta.status = SERVER_MODEL_STATUS_DOWNLOADED; need_reload = true; + + // a download that failed keeps its tag: its partial files are still + // there and re-posting the tag resumes them + if (ok) { + unfinished_downloads.forget(name); + } } else { json & info = it->second.meta.loaded_info; if (!info.contains("progress")) { @@ -1445,24 +1740,14 @@ bool server_models::remove(const std::string & name) { || meta.status == SERVER_MODEL_STATUS_DOWNLOADED; }); - // re-find after wait - load_models() may have erased the entry during the wait - it = mapping.find(name); - if (it == mapping.end()) { - // load_models() already erased the entry; we just need to clean up the cached files on disk - lk.unlock(); - bool ok = common_download_remove(name); - SRV_INF("removing model name=%s from cache (%s)\n", name.c_str(), ok ? "succeeded" : "partial"); - notify_sse("model_remove", name, {}); - return true; - } + // load_models() may have erased the entry during the wait; the files go under the + // lock, so a download of the same tag cannot start meanwhile + mapping.erase(name); + unfinished_downloads.remove(name); - // remove from disk (best-effort: cancelled downloads may have no cached files) + // best-effort: a cancelled download may have no cached files bool ok = common_download_remove(name); - mapping.erase(name); - if (!ok) { - SRV_WRN("removing model name=%s from disk returned false (no cached files?)\n", name.c_str()); - } - SRV_INF("removing model name=%s from cache (%s)\n", name.c_str(), ok ? "succeeded" : "partial"); + SRV_INF("removing model name=%s from cache (%s)\n", name.c_str(), ok ? "succeeded" : "nothing cached"); notify_sse("model_remove", name, {}); return true; } @@ -2091,6 +2376,9 @@ void server_models_routes::init_routes() { json models_json = json::array(); auto all_models = models.get_all_meta(); std::time_t t = std::time(0); + // a download on its way has no files to resolve yet, only the kinds it pulls + const auto unfinished = models.unfinished_downloads_get(); + for (const auto & meta : all_models) { if (meta.hidden) { continue; // cache model deduplicated by a preset @@ -2132,6 +2420,35 @@ void server_models_routes::init_routes() { model_info["context_length"] = meta.n_ctx_train; } + // the sidecars a model has on disk, resolved as -hf would load them, for a + // model that resolves from a HF repo. A download on its way reports the + // kinds it pulls, as an unfinished one does + const auto it_unfinished = std::find_if(unfinished.begin(), unfinished.end(), + [&](const server_unfinished_download & d) { return d.model == meta.name; }); + const std::string hf_repo = model_hf_repo(meta); + + if (it_unfinished != unfinished.end() && meta.status == SERVER_MODEL_STATUS_DOWNLOADING) { + model_info["sidecars"] = it_unfinished->sidecars; + } else if (!hf_repo.empty()) { + const json cached = cached_sidecars(hf_repo); + model_info["sidecars"] = cached; + + // the sidecars of an unfinished download of this model that are not on + // disk yet: a POST with the kind resumes one, a DELETE with it drops it + if (it_unfinished != unfinished.end()) { + std::vector incomplete; + for (const auto & p : it_unfinished->partials) { + if (!p.kind.empty() && !has_sidecar(incomplete, p.kind) && + std::find(cached.begin(), cached.end(), p.kind) == cached.end()) { + incomplete.push_back(p.kind); + } + } + if (!incomplete.empty()) { + model_info["incomplete_sidecars"] = incomplete; + } + } + } + // merge with loaded_info from the child process if available if (meta.is_running()) { for (auto it = meta.loaded_info.begin(); it != meta.loaded_info.end(); ++it) { @@ -2142,6 +2459,33 @@ void server_models_routes::init_routes() { } models_json.push_back(model_info); } + + const json architecture_default = json { + {"input_modalities", json::array({"text"})}, + {"output_modalities", json::array({"text"})}, + }; + + // a download that never finished has no instance: report it as incomplete so + // the UI can offer to resume it, without starting anything + for (const auto & d : unfinished) { + if (models.has_model(d.model)) { + continue; // complete, or downloading right now + } + models_json.push_back(json { + {"id", d.model}, + {"aliases", json::array()}, + {"tags", json::array()}, + {"object", "model"}, + {"owned_by", "llamacpp"}, + {"created", t}, + {"status", json { {"value", server_model_status_to_string(SERVER_MODEL_STATUS_INCOMPLETE)}, {"args", json::array()} }}, + {"architecture", architecture_default}, + {"source", server_model_source_to_string(SERVER_MODEL_SOURCE_CACHE)}, + {"can_remove", true}, + {"sidecars", d.sidecars}, + }); + } + res_ok(res, { {"data", models_json}, {"object", "list"}, @@ -2193,32 +2537,27 @@ void server_models_routes::init_routes() { if (name.empty()) { throw std::invalid_argument("model must be a non-empty string"); } + const auto sidecars = parse_sidecars(body); - common_params p; - p.model.hf_repo = name; - p.hf_token = params.hf_token; + common_params_model model; + model.hf_repo = name; - // validate by fetching metadata - bool ok = false; + // validate by fetching metadata, through the plan the download follows + common_download_hf_plan plan; try { - common_models_handler_init(p, LLAMA_EXAMPLE_SERVER); - ok = true; + plan = common_download_get_hf_plan(model, sidecar_download_opts(params, sidecars)); } catch (...) { SRV_ERR("unknown error while validating model '%s'\n", name.c_str()); // other exceptions will be handled by the outer ex_wrapper() throw; } - if (!ok) { - throw std::invalid_argument("model validation failed, unable to download"); - } - - // reject if model already exists - if (models.has_model(name)) { - throw std::invalid_argument("model '" + name + "' already exists"); + // a tag that resolves no file has nothing to download + if (plan.preset.path.empty() && plan.primary.path.empty() && common_download_plan_sidecars(plan).empty()) { + throw std::invalid_argument("model '" + name + "' resolves no file to download"); } - // then, proceed with the actual download + // then, proceed with the actual download; load() decides over an existing entry SRV_INF("starting download for model '%s'\n", name.c_str()); { server_models::load_options load_opts; @@ -2226,6 +2565,8 @@ void server_models_routes::init_routes() { load_opts.custom_meta = server_model_meta{}; load_opts.custom_meta->source = SERVER_MODEL_SOURCE_CACHE; load_opts.custom_meta->name = name; + load_opts.sidecars = sidecars; + load_opts.partials = common_download_plan_partials(plan); models.load(name, load_opts); } @@ -2241,6 +2582,22 @@ void server_models_routes::init_routes() { throw std::invalid_argument("model must be a non-empty string"); } + // a sidecar alone goes, the model keeps its files + std::string sidecar = req.get_param("sidecar"); + if (!sidecar.empty()) { + check_sidecar(sidecar); + res_ok(res, {{"success", models.remove_sidecar(name, sidecar)}}); + return res; + } + + // a download listed from an earlier run has no instance, only its tag and + // its partial files + if (!models.has_model(name)) { + models.unfinished_downloads_remove(name); + res_ok(res, {{"success", true}}); + return res; + } + models.remove(name); // throws on error res_ok(res, {{"success", true}}); diff --git a/tools/server/server-models.h b/tools/server/server-models.h index f415b7c5fd61..b04fc7384d9f 100644 --- a/tools/server/server-models.h +++ b/tools/server/server-models.h @@ -16,6 +16,7 @@ #include #include #include +#include /** * state diagram: @@ -31,6 +32,7 @@ enum server_model_status { // TODO: also add downloading state when the logic is added SERVER_MODEL_STATUS_DOWNLOADING, + SERVER_MODEL_STATUS_INCOMPLETE, // a download that never finished, listed without an instance SERVER_MODEL_STATUS_DOWNLOADED, SERVER_MODEL_STATUS_UNLOADED, SERVER_MODEL_STATUS_LOADING, @@ -52,6 +54,7 @@ enum server_child_mode { static std::string server_model_status_to_string(server_model_status status) { switch (status) { case SERVER_MODEL_STATUS_DOWNLOADING: return "downloading"; + case SERVER_MODEL_STATUS_INCOMPLETE: return "incomplete"; case SERVER_MODEL_STATUS_DOWNLOADED: return "downloaded"; case SERVER_MODEL_STATUS_UNLOADED: return "unloaded"; case SERVER_MODEL_STATUS_LOADING: return "loading"; @@ -108,6 +111,51 @@ struct server_model_meta { void update_caps(const common_params & base); }; +// a download that never finished: the tag it was asked under, the sidecar kinds it +// pulls and the partial files its plan writes +struct server_unfinished_download { + std::string model; + std::vector sidecars; + std::vector partials; +}; + +// the downloads that never finished, recorded next to the model cache. The cache names +// partial files after their content hash, so this record is the only trace of what is +// still incomplete: it lets a restarted router list, resume or remove them +// a partial another recorded download also writes is never deleted +// these functions are thread-safe +struct server_unfinished_downloads { + // read the record an earlier run left + void load(); + + // record a download that starts, replacing an earlier record of the same model, + // and delete the partials of that record the new plan no longer writes + void add(const server_unfinished_download & download); + + // forget a download whose files are complete + void forget(const std::string & model); + + // forget a download and delete its partials; returns false when it is not recorded + bool remove(const std::string & model); + + // drop the sidecar `kind` from a download and delete its partials; returns false + // when the download is not recorded with that kind + bool remove_sidecar(const std::string & model, const std::string & kind); + + // every recorded download + std::vector get_all(); + +private: + std::mutex mutex; + std::vector downloads; + + // write the record back; callers hold the mutex + void save(); + + // delete the partials of `parts` no recorded download writes; callers hold the mutex + void delete_partials(const std::vector & parts); +}; + struct server_models_routes; struct server_lru_sched; // defined in server-models.cpp struct server_monitor; // defined in server-models.cpp @@ -134,6 +182,9 @@ struct server_models { // models asked to stop, still counted as running until the monitor records their exit std::set stopping_models; + // downloads that never finished, with the sidecar kinds they pull + server_unfinished_downloads unfinished_downloads; + // set to true while load_models() is executing a reload; load() will wait until clear bool is_reloading = false; @@ -252,6 +303,16 @@ struct server_models { // check if a model instance exists (thread-safe) bool has_model(const std::string & name); + // downloads that never finished, with their sidecar kinds (thread-safe) + std::vector unfinished_downloads_get(); + + // drop a download that will never be resumed, its record and its partial files; + // throws for a name that is not recorded (thread-safe) + void unfinished_downloads_remove(const std::string & name); + + // remove the sidecar `kind` of a model that is not in use, its files and partials (thread-safe) + bool remove_sidecar(const std::string & name, const std::string & kind); + // return a copy of model metadata (thread-safe) std::optional get_meta(const std::string & name); @@ -262,6 +323,10 @@ struct server_models { server_child_mode mode = SERVER_CHILD_MODE_NORMAL; // used for spawning a downloading child process std::optional custom_meta = std::nullopt; + // download mode: the sidecar kinds pulled with the model, as -hf would, and + // the partial files the download plan writes + std::vector sidecars; + std::vector partials; }; // load and unload model instances diff --git a/tools/server/tests/unit/test_router.py b/tools/server/tests/unit/test_router.py index 5355ab7367ed..6df41f8d63e6 100644 --- a/tools/server/tests/unit/test_router.py +++ b/tools/server/tests/unit/test_router.py @@ -1,3 +1,4 @@ +import shutil import threading import pytest from utils import * @@ -546,6 +547,15 @@ def _wait_for_sse_event(collected: list, event_type: str, model: str, timeout: i return False +def _own_cache(name: str) -> str: + """An empty model cache for one test: the workers run in parallel, and a download + test must not see the files or the record another test writes in the shared cache.""" + path = os.path.join(TMP_DIR, name) + shutil.rmtree(path, ignore_errors=True) + os.makedirs(path) + return path + + def test_router_download_model(): """Case 1: download a model at the model limit, verify SSE events and GET /models.""" global server @@ -626,3 +636,230 @@ def test_router_delete_model(): # Model should no longer appear in GET /models ids = _get_model_ids(is_reload=False) assert MODEL_DOWNLOAD_ID not in ids, f"{MODEL_DOWNLOAD_ID} still present after deletion" + + +def test_router_download_sidecars(): + """Case 3: sidecars are a parameter of the download, never models of their own.""" + global server + server.llama_cache = _own_cache("cache_download_sidecars") + server.start() + + server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}") + + # an unknown sidecar kind is rejected before anything starts + res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID, "sidecars": ["lora"]}) + assert res.status_code == 400 + assert MODEL_DOWNLOAD_ID not in _get_model_ids(is_reload=False) + + sse_events: list = [] + stop = threading.Event() + sse_ready = threading.Event() + threading.Thread( + target=_listen_sse, args=(server, sse_events, stop, sse_ready), daemon=True + ).start() + assert sse_ready.wait(10), "SSE client failed to connect" + + res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID, "sidecars": []}) + assert res.status_code == 200 + finished = _wait_for_sse_event( + sse_events, "download_finished", MODEL_DOWNLOAD_ID, MODEL_DOWNLOAD_TIMEOUT + ) + stop.set() + assert finished, "Never received download_finished SSE event" + + # the model lists the sidecars it has on disk, and no sidecar is listed as a model + res = server.make_request("GET", "/models") + assert res.status_code == 200 + entries = {item["id"]: item for item in res.body.get("data", [])} + assert entries[MODEL_DOWNLOAD_ID]["sidecars"] == [] + for kind in ("mmproj", "mtp", "dflash", "eagle3", "dspark"): + assert not any(i.endswith("-" + kind) for i in entries), f"a {kind} sidecar is listed as a model" + + server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}") + + +def test_router_download_conflicts(): + """A download is refused over one on its way, and over a model that is running.""" + global server + server.models_max = 2 + server.llama_cache = _own_cache("cache_download_conflicts") + server.start() + + server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}") + + # a download on its way is not started again; the entry reports the state before + # the child fetches anything, so the probe lands inside the download + res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID}) + assert res.status_code == 200 + _wait_for_model_status(MODEL_DOWNLOAD_ID, {"downloading"}, timeout=30) + + res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID}) + assert res.status_code == 400 + assert "is already downloading" in res.body["error"]["message"] + + # a loaded model holds its files, so it goes only once it is unloaded; the + # reload that follows the download parks the entry as unloaded + _wait_for_model_status(MODEL_DOWNLOAD_ID, {"downloaded", "unloaded"}, timeout=MODEL_DOWNLOAD_TIMEOUT) + _load_model_and_wait(MODEL_DOWNLOAD_ID, timeout=120) + + res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID}) + assert res.status_code == 400 + assert "is running, unload it first" in res.body["error"]["message"] + + +def test_router_delete_unknown_name(): + """A name that is neither a model nor a recorded download is not found, and the + cached tags of its repo keep their files.""" + global server + server.llama_cache = _own_cache("cache_delete_unknown_name") + server.start() + + res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID, "sidecars": []}) + assert res.status_code == 200 + _wait_for_model_status(MODEL_DOWNLOAD_ID, {"downloaded", "unloaded"}, timeout=MODEL_DOWNLOAD_TIMEOUT) + + repo = MODEL_DOWNLOAD_ID.split(":")[0] + res = server.make_request("DELETE", f"/models?model={repo}") + assert res.status_code != 200 + assert "is not found" in res.body["error"]["message"] + assert MODEL_DOWNLOAD_ID in _get_model_ids(is_reload=True) + + server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}") + + +def test_router_incomplete_sidecars(): + """A cached model lists the sidecars its unfinished download has not written yet, until + they are dropped.""" + global server + cache_dir = _own_cache("cache_incomplete_sidecars") + server.llama_cache = cache_dir + server.start() + + res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID, "sidecars": []}) + assert res.status_code == 200 + _wait_for_model_status(MODEL_DOWNLOAD_ID, {"downloaded", "unloaded"}, timeout=MODEL_DOWNLOAD_TIMEOUT) + server.stop() + + # a paused mtp of the cached model, as the record of an earlier run holds it + partial_path = os.path.abspath(os.path.join(cache_dir, "test-mtp.gguf.downloadInProgress")) + with open(partial_path, "w") as f: + f.write("partial") + with open(os.path.join(cache_dir, "unfinished-downloads.json"), "w") as f: + json.dump({"downloads": [{ + "model": MODEL_DOWNLOAD_ID, + "sidecars": ["mtp"], + "partials": [{"kind": "mtp", "path": partial_path}], + }]}, f) + + server.start() + + def entry() -> dict: + res = server.make_request("GET", "/models") + assert res.status_code == 200 + return {item["id"]: item for item in res.body.get("data", [])}[MODEL_DOWNLOAD_ID] + + assert entry()["sidecars"] == [] + assert entry()["incomplete_sidecars"] == ["mtp"] + + res = server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}&sidecar=mtp") + assert res.status_code == 200 + assert res.body.get("success") is True + assert not os.path.exists(partial_path) + assert "incomplete_sidecars" not in entry() + + server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}") + + +def test_router_preset_model_sidecars(): + """A preset model reports the sidecars its files have on disk, like a cached one.""" + global server + + preset_path = os.path.join(TMP_DIR, "test_preset_sidecars.ini") + + with open(preset_path, "w") as f: + f.write("[preset-sidecars]\nhf-repo = ggml-org/tinygemma3-GGUF:Q8_0\n") + + server.models_preset = preset_path + server.start() + + try: + # the mmproj of that repo lands next to the model, as -hf would pull it + res = server.make_request("POST", "/models", data={"model": MODEL_A, "sidecars": ["mmproj"]}) + assert res.status_code == 200 + _wait_for_model_status(MODEL_A, {"downloaded", "unloaded"}, timeout=MODEL_DOWNLOAD_TIMEOUT) + + res = server.make_request("GET", "/models") + assert res.status_code == 200 + entries = {item["id"]: item for item in res.body.get("data", [])} + + assert entries["preset-sidecars"]["source"] == "preset" + assert entries["preset-sidecars"]["sidecars"] == ["mmproj"] + finally: + os.remove(preset_path) + + +def test_router_local_preset_listing(): + """A preset on a local file resolves from no HF repo: it lists without sidecars.""" + global server + + model_path = os.path.join(TMP_DIR, "test_local_preset.gguf") + preset_path = os.path.join(TMP_DIR, "test_local_preset.ini") + + with open(model_path, "w") as f: + f.write("GGUF") + with open(preset_path, "w") as f: + f.write(f"[local]\nmodel = {model_path}\n") + + server.models_preset = preset_path + server.start() + + try: + res = server.make_request("GET", "/models") + assert res.status_code == 200 + entries = {item["id"]: item for item in res.body.get("data", [])} + + assert entries["local"]["source"] == "preset" + assert "sidecars" not in entries["local"] + finally: + os.remove(preset_path) + os.remove(model_path) + + +def test_router_download_unresolved_tag(): + """A tag that resolves no file of the repo is refused before anything starts.""" + global server + server.start() + + model = MODEL_DOWNLOAD_ID.split(":")[0] + ":Q9_Z" + res = server.make_request("POST", "/models", data={"model": model}) + assert res.status_code == 400 + assert "resolves no file to download" in res.body["error"]["message"] + assert model not in _get_model_ids(is_reload=False) + + +def test_router_download_drops_stale_partials(): + """Posting a recorded download again deletes the partials its new plan no longer writes.""" + global server + + cache_dir = _own_cache("cache_download_drops_stale_partials") + stale_path = os.path.abspath(os.path.join(cache_dir, "test-stale.gguf.downloadInProgress")) + record_path = os.path.join(cache_dir, "unfinished-downloads.json") + + with open(stale_path, "w") as f: + f.write("stale") + with open(record_path, "w") as f: + json.dump({"downloads": [{ + "model": MODEL_DOWNLOAD_ID, + "sidecars": ["mtp"], + "partials": [{"kind": "mtp", "path": stale_path}], + }]}, f) + + server.llama_cache = cache_dir + server.start() + + res = server.make_request("POST", "/models", data={"model": MODEL_DOWNLOAD_ID, "sidecars": []}) + assert res.status_code == 200 + assert not os.path.exists(stale_path), "the stale partial survived the new record" + + _wait_for_model_status(MODEL_DOWNLOAD_ID, {"downloaded", "unloaded"}, timeout=MODEL_DOWNLOAD_TIMEOUT) + server.make_request("DELETE", f"/models?model={MODEL_DOWNLOAD_ID}") diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index 1c84f5121145..288452808093 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -127,6 +127,7 @@ class ServerProcess: mcp_servers_config: str | None = None mcp_servers_json: str | None = None cors_origins: str | None = None + llama_cache: str | None = None # session variables process: subprocess.Popen | None = None @@ -145,7 +146,9 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None: **os.environ, "LLAMA_SERVER_DEBUG_FAKE_TIMING": "1", } - if "LLAMA_CACHE" not in os.environ: + if self.llama_cache is not None: + env["LLAMA_CACHE"] = self.llama_cache + elif "LLAMA_CACHE" not in os.environ: env["LLAMA_CACHE"] = "tmp" if self.external_server: print(f"[external_server]: Assuming external server running on {self.server_host}:{self.server_port}") 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 c217a769a6b4..895e9a80a7b1 100644 --- a/tools/ui/src/lib/components/app/content/MarkdownContent/MarkdownContent.svelte +++ b/tools/ui/src/lib/components/app/content/MarkdownContent/MarkdownContent.svelte @@ -7,6 +7,7 @@ getMdastNodeHash, isAppendMode } from './markdown-utils'; + import { SAFE_HTML_CONFIG } from './safe-html-config'; import { ActionIconCopyToClipboard, CodeBlockActions, @@ -44,6 +45,7 @@ import { detectIncompleteCodeBlock, highlightCode, type IncompleteCodeBlock } from '$lib/utils'; import { sanitizeSvg } from '$lib/utils/sanitize-svg'; import { mountSvgShadow } from '$lib/utils/svg-shadow'; + import DOMPurify from 'dompurify'; import type { Root as HastRoot, RootContent as HastRootContent } from 'hast'; import githubLightCss from 'highlight.js/styles/github.css?inline'; import githubDarkCss from 'highlight.js/styles/github-dark.css?inline'; @@ -57,6 +59,8 @@ content: string; class?: string; disableMath?: boolean; + /** Render raw HTML found in the markdown (sanitized) instead of escaping it. */ + allowHtml?: boolean; } interface MarkdownBlock { @@ -65,7 +69,13 @@ contentHash?: string; } - let { attachments, class: className = '', content, disableMath = false }: Props = $props(); + let { + allowHtml = false, + attachments, + class: className = '', + content, + disableMath = false + }: Props = $props(); let containerRef = $state(); let renderedBlocks = $state([]); @@ -121,6 +131,8 @@ // Per-instance transform cache, avoids re-transforming stable blocks during streaming // Garbage collected when component is destroyed (on conversation change) + const CACHE_KEY_PLAIN = 'plain'; + const CACHE_KEY_ALLOW_HTML = 'allow-html'; const transformCache = new SvelteMap(); let previousContent = ''; @@ -183,7 +195,10 @@ index: number ): Promise<{ html: string; hash: string }> { const hash = getMdastNodeHash(node, index); - const cached = transformCache.get(hash); + // the rendered HTML also depends on allowHtml (sanitized raw vs escaped), + // so the cache is keyed per mode + const cacheKey = `${allowHtml ? CACHE_KEY_ALLOW_HTML : CACHE_KEY_PLAIN}:${hash}`; + const cached = transformCache.get(cacheKey); if (cached) { return { hash, html: cached }; @@ -192,10 +207,13 @@ const singleNodeRoot = { children: [node], type: 'root' }; const transformedRoot = (await processorInstance.run(singleNodeRoot as MdastRoot)) as HastRoot; const html = processorInstance.stringify(transformedRoot); + const safeHtml = allowHtml + ? (DOMPurify.sanitize(html, SAFE_HTML_CONFIG) as unknown as string) + : html; - transformCache.set(hash, html); + transformCache.set(cacheKey, safeHtml); - return { hash, html }; + return { hash, html: safeHtml }; } /** @@ -300,7 +318,7 @@ if (prefixMarkdown.trim()) { const normalizedPrefix = preprocessLaTeX(prefixMarkdown); - const processorInstance = getMarkdownProcessor({ attachments, disableMath }); + const processorInstance = getMarkdownProcessor({ allowHtml, attachments, disableMath }); const ast = processorInstance.parse(normalizedPrefix) as MdastRoot; const mdastChildren = (ast as { children?: unknown[] }).children ?? []; const nextBlocks: MarkdownBlock[] = []; @@ -350,7 +368,7 @@ incompleteCodeBlock = null; const normalized = preprocessLaTeX(markdown); - const processorInstance = getMarkdownProcessor({ attachments, disableMath }); + const processorInstance = getMarkdownProcessor({ allowHtml, attachments, disableMath }); const ast = processorInstance.parse(normalized) as MdastRoot; const mdastChildren = (ast as { children?: unknown[] }).children ?? []; const stableCount = Math.max(mdastChildren.length - 1, 0); @@ -394,6 +412,10 @@ )) as HastRoot; unstableHtml = processorInstance.stringify(transformedRoot); + + if (allowHtml) { + unstableHtml = DOMPurify.sanitize(unstableHtml, SAFE_HTML_CONFIG) as unknown as string; + } } renderedBlocks = nextBlocks; 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 index 57ded3e99ada..26cae5b5c51b 100644 --- a/tools/ui/src/lib/components/app/content/MarkdownContent/markdown-processor.ts +++ b/tools/ui/src/lib/components/app/content/MarkdownContent/markdown-processor.ts @@ -42,12 +42,15 @@ export interface MarkdownProcessor { export interface MarkdownProcessorOptions { attachments?: DatabaseMessageExtra[]; disableMath?: boolean; + /** Render raw HTML found in the markdown instead of escaping it. */ + allowHtml?: boolean; } const sharedPipelines = new Map(); -const attachmentPipelines = new WeakMap(); +const attachmentPipelines = new WeakMap>(); function buildPipeline({ + allowHtml = false, attachments, disableMath = false }: MarkdownProcessorOptions): MarkdownProcessor { @@ -58,11 +61,14 @@ function buildPipeline({ proc = proc.use(remarkMath); // Parse $inline$ and $$block$$ math } - proc = proc - .use(remarkBreaks) // Convert line breaks to
+ proc = proc.use(remarkBreaks); // Convert line breaks to
+ + if (!allowHtml) { // Treat raw HTML as literal text with preserved indentation - .use(remarkLiteralHtml) - .use(remarkRehype); // Convert Markdown AST to rehype + proc = proc.use(remarkLiteralHtml); + } + + proc = proc.use(remarkRehype, allowHtml ? { allowDangerousHtml: true } : undefined); if (!disableMath) { proc = proc.use(rehypeKatex); // Render math using KaTeX @@ -91,17 +97,26 @@ function buildPipeline({ export function getMarkdownProcessor(options: MarkdownProcessorOptions): MarkdownProcessor { if (options.attachments && options.attachments.length > 0) { - let cached = attachmentPipelines.get(options.attachments); + let byOptions = attachmentPipelines.get(options.attachments); + + if (!byOptions) { + byOptions = new Map(); + attachmentPipelines.set(options.attachments, byOptions); + } + + const key = `${Boolean(options.disableMath)}:${Boolean(options.allowHtml)}`; + + let cached = byOptions.get(key); if (!cached) { cached = buildPipeline(options); - attachmentPipelines.set(options.attachments, cached); + byOptions.set(key, cached); } return cached; } - const key = String(Boolean(options.disableMath)); + const key = `${Boolean(options.disableMath)}:${Boolean(options.allowHtml)}`; let cached = sharedPipelines.get(key); diff --git a/tools/ui/src/lib/components/app/content/MarkdownContent/safe-html-config.ts b/tools/ui/src/lib/components/app/content/MarkdownContent/safe-html-config.ts new file mode 100644 index 000000000000..d8dff07ddaca --- /dev/null +++ b/tools/ui/src/lib/components/app/content/MarkdownContent/safe-html-config.ts @@ -0,0 +1,138 @@ +// Sanitization rules for raw HTML in model cards (MarkdownContent's +// allowHtml mode). Kept as its own module so the adversarial test in +// tests/client can import the exact shipped configuration. +import DOMPurify from 'dompurify'; + +// style is allowed for the KaTeX output, which lays its spans out with inline +// styles; a third party card can use the same attribute to draw a fixed +// overlay over the app, so positioned values are scrubbed +DOMPurify.addHook('afterSanitizeAttributes', (node) => { + if (typeof (node as Element).getAttribute !== 'function') return; + + const style = (node as Element).getAttribute('style'); + + if (!style || !/position\s*:\s*(fixed|absolute)/i.test(style)) return; + + (node as Element).setAttribute( + 'style', + style.replace(/position\s*:\s*(fixed|absolute)\s*;?/gi, '') + ); +}); + +/** + * Explicit DOMPurify allow list: the tags and attributes model cards may use + * plus the KaTeX output (spans with inline styles and MathML). Everything + * else - scripts, event handlers, forms, iframes, style tags - is stripped. + */ +export const SAFE_HTML_CONFIG = { + ALLOWED_ATTR: [ + 'class', + 'style', + 'id', + 'title', + 'role', + 'aria-hidden', + 'aria-label', + 'aria-describedby', + 'aria-expanded', + 'href', + 'target', + 'rel', + 'src', + 'srcset', + 'alt', + 'width', + 'height', + 'align', + 'valign', + 'colspan', + 'rowspan', + 'controls', + 'type', + // MathML + 'xmlns', + 'display', + 'encoding', + 'mathvariant' + ], + ALLOWED_TAGS: [ + // text and structure + 'a', + 'b', + 'i', + 'u', + 's', + 'em', + 'strong', + 'code', + 'pre', + 'kbd', + 'sub', + 'sup', + 'br', + 'hr', + 'p', + 'span', + 'div', + 'blockquote', + 'center', + 'ul', + 'ol', + 'li', + 'dl', + 'dt', + 'dd', + 'h1', + 'h2', + 'h3', + 'h4', + 'h5', + 'h6', + 'table', + 'thead', + 'tbody', + 'tr', + 'th', + 'td', + 'colgroup', + 'col', + 'caption', + // media + 'img', + 'picture', + 'source', + 'figure', + 'figcaption', + 'video', + 'audio', + // MathML, for the KaTeX output + 'math', + 'mrow', + 'mi', + 'mo', + 'mn', + 'ms', + 'mtext', + 'mfrac', + 'mroot', + 'msqrt', + 'msub', + 'msup', + 'msubsup', + 'munder', + 'mover', + 'mmultiscripts', + 'mtable', + 'mtr', + 'mtd', + 'mspace', + 'mglyph', + 'maligngroup', + 'malignmark', + 'mpadded', + 'mphantom', + 'mstyle', + 'semantics', + 'annotation' + ] +}; diff --git a/tools/ui/src/lib/components/app/dialogs/DialogManageModels.svelte b/tools/ui/src/lib/components/app/dialogs/DialogManageModels.svelte index 45b73dc72fc9..fbe711c25cd0 100644 --- a/tools/ui/src/lib/components/app/dialogs/DialogManageModels.svelte +++ b/tools/ui/src/lib/components/app/dialogs/DialogManageModels.svelte @@ -1,7 +1,18 @@ {#each badges as badge (badge.kind + (badge.quant ?? ''))} - - + import { Check } from '@lucide/svelte'; + import { MODEL_BADGE_CLASS, MODEL_VARIANT_BADGE_CLASS } from '$lib/constants'; + import { ModelRowDownloadState } from '$lib/enums'; + + interface Props { + /** + * Left inset of the tree connector, tuned per surface so the elbow hangs + * under the main line's avatar and "Draft model" aligns with its id. + */ + connectorClass?: string; + /** Draft sidecar kind, e.g. dflash. */ + kind: string; + /** Quant of the draft, when its tag carries one. */ + quant?: string | null; + /** Live state of the draft download. */ + state: ModelRowDownloadState; + /** Rounded download percent, when the feed reports progress. */ + percent?: number | null; + } + + let { connectorClass = 'ml-6', kind, percent = null, quant = null, state }: Props = $props(); + + + +
+ + + Draft model + + {kind} + + {#if quant} + {quant} + {/if} + + {#if state === ModelRowDownloadState.DOWNLOADED} + + {:else} + + {#if percent !== null} + {percent}% + {:else if state === ModelRowDownloadState.PAUSED} + Paused + {/if} + + {/if} +
diff --git a/tools/ui/src/lib/components/app/models/ModelId.svelte b/tools/ui/src/lib/components/app/models/ModelId.svelte index 26d72f662ff2..9a3464ab861b 100644 --- a/tools/ui/src/lib/components/app/models/ModelId.svelte +++ b/tools/ui/src/lib/components/app/models/ModelId.svelte @@ -5,6 +5,7 @@ import { TruncatedText } from '$lib/components/app'; import * as Tooltip from '$lib/components/ui/tooltip'; import { MODEL_BADGE_CLASS, MODEL_VARIANT_BADGE_CLASS } from '$lib/constants'; + import type { ModelDraftSidecar } from '$lib/enums'; import { HuggingFaceService } from '$lib/services'; import { ModelsService } from '$lib/services/models.service'; import { settingsStore } from '$lib/stores'; @@ -39,6 +40,8 @@ sizeRange?: { min: number; max: number } | null; /** Draft sidecars available for the model, badged with their own quant. */ draftSidecars?: ModelSidecarBadge[]; + /** Draft sidecar kinds the repo offers, one badge per kind and no quant. */ + draftKinds?: ModelDraftSidecar[]; /** Give the id a row of its own, with its badges and icons together under it. */ stackId?: boolean; /** Allow badges to wrap onto new lines instead of truncating. */ @@ -50,6 +53,7 @@ aliases, class: className = '', contextLength, + draftKinds = [], draftSidecars = [], hideCapabilities = false, hideModalities = false, @@ -93,12 +97,17 @@ let uniqueAliases = $derived([...new Set(aliases ?? [])]); let uniqueTags = $derived([...new Set([...(parsed.tags ?? []), ...(tags ?? [])])]); let uniqueDraftSidecars = $derived(draftSidecars.filter((badge) => !isAuxSidecar(badge.kind))); + // the model id can already name one of the kinds + let uniqueDraftKinds = $derived([ + ...new Set(draftKinds.filter((kind) => kind !== parsed.sidecar)) + ]); let primaryAlias = $derived(uniqueAliases.length === 1 ? uniqueAliases[0] : null); let displayName = $derived(primaryAlias ?? parsed.modelName ?? modelId); let hasBadges = $derived( parsed.sidecar || + uniqueDraftKinds.length > 0 || uniqueDraftSidecars.length > 0 || (parsed.params && !hideParameters) || (parsed.quantization && !resolvedHideQuantization) || @@ -144,6 +153,12 @@ {/if} + {#each uniqueDraftKinds as kind (kind)} + + {kind} + + {/each} + {#if primaryAlias} diff --git a/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManager.svelte b/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManager.svelte index 36fd3e22ae2f..9853bafe51a0 100644 --- a/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManager.svelte +++ b/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManager.svelte @@ -32,9 +32,22 @@ let capabilityFilter = $state([]); let selectedId = $state(null); + // tracked downloads get their own section, listed like any other model + let downloads = $derived( + downloadGroups(modelsStore.status.getDownloadEntries(), modelsStore.models, (repo) => + modelsStore.status.getDraftSidecars(repo) + ) + ); + let allModels = $derived(modelsStore.models); - let selected = $derived(allModels.find((option) => option.id === selectedId) ?? null); + // the download groups stand in for rows the listing does not carry yet, so a + // click on them still finds its option and opens the pane + let selected = $derived( + allModels.find((option) => option.id === selectedId) ?? + downloads.find((group) => group.base.id === selectedId)?.base ?? + null + ); // The pane is laid out before it is ever opened, so the first open only slides a // finished panel in. It renders the selection, else the model it last showed. let lastPicked = $state(null); @@ -42,7 +55,11 @@ let shownId = $state(null); let isSwapping = $state(false); let fade = $state<'open' | 'swap'>('open'); - let shownOption = $derived(allModels.find((option) => option.id === shownId) ?? null); + let shownOption = $derived( + allModels.find((option) => option.id === shownId) ?? + downloads.find((group) => group.base.id === shownId)?.base ?? + null + ); $effect(() => { const id = selectedId; @@ -144,10 +161,6 @@ rank.size === 0 ? list : [...list].sort((a, b) => rankOf(a) - rankOf(b)); // one entry per repo, so a model with several quants takes a single table row let entries = $derived(byRecency(groupModelQuants(matching))); - // tracked downloads get their own section, listed like any other model - let downloads = $derived( - downloadGroups(modelsStore.status.getDownloadEntries(), modelsStore.models) - ); let groups = $derived.by(() => { // a loaded quant is a model of its own: its repo keeps the quants left behind const isLoaded = (option: ModelOption) => modelsStore.isModelLoaded(option.model); @@ -216,13 +229,16 @@ return ordered; }); - // a caller can ask for one model to be revealed, the download rows do + // a caller can ask for one model to be revealed, the download rows do; a + // tracked download is not in the listing yet, so its own rows are searched too $effect(() => { const focus = uiStore.manageModelFocus; if (!focus) return; - const option = allModels.find((model) => model.id === focus || model.model === focus); + const matches = (candidate: ModelOption) => candidate.id === focus || candidate.model === focus; + const option = + allModels.find(matches) ?? downloads.find((group) => matches(group.base))?.base ?? null; if (option) selectedId = option.id; diff --git a/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerDownloadControl.svelte b/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerDownloadControl.svelte index 0e81500f2c1f..5e77caed06e8 100644 --- a/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerDownloadControl.svelte +++ b/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerDownloadControl.svelte @@ -38,10 +38,7 @@ class="h-5 w-5 hover:text-foreground" icon={isPaused ? MODEL_DOWNLOAD_ICONS.resume : MODEL_DOWNLOAD_ICONS.pause} iconSize="h-4 w-4" - onclick={() => - void (isPaused - ? modelsStore.status.downloadModel(option.model) - : modelsStore.status.pauseDownload(option.model))} + onclick={() => void modelsStore.status.toggleDownload(option.model)} stopPropagationOnClick tooltip={isPaused ? 'Resume downloading' : 'Pause downloading'} tooltipAsTitle diff --git a/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerModelConfiguration/ModelsManagerModelConfiguration.svelte b/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerModelConfiguration/ModelsManagerModelConfiguration.svelte index 8a109a8db41e..40a7565ed272 100644 --- a/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerModelConfiguration/ModelsManagerModelConfiguration.svelte +++ b/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerModelConfiguration/ModelsManagerModelConfiguration.svelte @@ -1,4 +1,6 @@
- {#if isMobile} - - - {/if} + Pause downloading + {/if} + -
- {#if !isMobile} - +
+ + {#if onSeeInDiscover && discoverEnabled} + {/if} + {:else} + {#if isMobile} + + {/if} -
+ +
+ {#if !isMobile} + + {/if} + + {#if canToggleLoad} + + {/if} + + {#if onSeeInDiscover && discoverEnabled} + + {/if} +
+ {/if} diff --git a/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerModelConfiguration/ModelsManagerModelConfigurationHeader.svelte b/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerModelConfiguration/ModelsManagerModelConfigurationHeader.svelte index 40f5d0d5481b..d2cc2d5b8822 100644 --- a/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerModelConfiguration/ModelsManagerModelConfigurationHeader.svelte +++ b/tools/ui/src/lib/components/app/models/ModelsManager/ModelsManagerModelConfiguration/ModelsManagerModelConfigurationHeader.svelte @@ -8,8 +8,14 @@ import type { ModelOption } from '$lib/types/models'; interface Props { + isDownload?: boolean; + isDownloadPaused?: boolean; isLoaded: boolean; onClose: () => void; + /** Opens the manager dialog on the Discover view, focused on this model's repo. */ + onSeeInDiscover?: () => void; + onRequestCancelDownload?: () => void; + onToggleDownload?: () => void; onToggleLoad: () => void; onUseInChat: () => void; onUseInNewChat: () => void; @@ -18,13 +24,27 @@ status: ServerModelStatus | null; } - let { isLoaded, onClose, onToggleLoad, onUseInChat, onUseInNewChat, option, status }: Props = - $props(); + let { + isDownload = false, + isDownloadPaused = false, + isLoaded, + onClose, + onRequestCancelDownload, + onSeeInDiscover, + onToggleDownload, + onToggleLoad, + onUseInChat, + onUseInNewChat, + option, + status + }: Props = $props(); let isMobile = $derived(deviceStore.isMobile); let isLoading = $derived(status === ServerModelStatus.LOADING); let statusLabel = $derived.by(() => { + if (isDownload) return isDownloadPaused ? 'Paused' : 'Downloading'; + if (status === ServerModelStatus.LOADING) return 'Loading'; if (status === ServerModelStatus.FAILED) return 'Failed to load'; @@ -37,6 +57,8 @@ }); let statusDot = $derived.by(() => { + if (isDownload) return isDownloadPaused ? 'bg-yellow-500' : 'bg-blue-500 animate-pulse'; + if (status === ServerModelStatus.FAILED) return 'bg-red-500'; if (status === ServerModelStatus.LOADING) return 'bg-muted-foreground/50 animate-pulse'; @@ -100,8 +122,13 @@ {#if !isMobile} boolean; @@ -20,11 +21,21 @@ onDelete: (option: ModelOption) => void; onSelect: (option: ModelOption) => void; selected: boolean; + /** Draft sidecars already on disk, when the row stands for a download. */ + sidecarBadges?: ModelSidecarBadge[]; /** Left padding in px, from the nesting depth. */ indent?: number; } - let { indent = 0, isFavorite, onDelete, onSelect, option, selected }: Props = $props(); + let { + indent = 0, + isFavorite, + onDelete, + onSelect, + option, + selected, + sidecarBadges + }: Props = $props(); let favorite = $derived(isFavorite(option)); let isHidden = $derived(modelsStore.isHidden(option.id)); @@ -38,6 +49,10 @@ ? ModelRowDownloadState.DOWNLOADING : null ); + + function activate(): void { + onSelect(option); + }
- - + favorited + {/if} + + + + {#if sidecarBadges?.length} + +
+ {#each sidecarBadges as badge (badge.repo + ':' + badge.kind + ':' + (badge.quant ?? ''))} + + {/each} +
+ {/if} +
(); @@ -129,6 +130,10 @@ // cancel is confirmed once for the whole list, so one dialog serves every row function requestDelete(option: ModelOption): void { pendingDelete = option.model; + // an in-flight row cancels (and drops the partials), a paused one deletes + pendingAction = modelsStore.status.isDownloadInProgress(option.model) + ? ModelDownloadConfirmAction.CANCEL + : ModelDownloadConfirmAction.DELETE; deleteOpen = true; } @@ -248,6 +253,7 @@ {onSelect} option={entry.base} selected={selectedId === entry.base.id} + sidecarBadges={entry.sidecarBadges} /> {/if} {/snippet} @@ -328,7 +334,7 @@ {/snippet} (deleteOpen = false)} open={deleteOpen} repoWithTag={pendingDelete} diff --git a/tools/ui/src/lib/components/app/models/ModelsManager/row-actions.ts b/tools/ui/src/lib/components/app/models/ModelsManager/row-actions.ts index 9931dd3f475f..3407634ec6f1 100644 --- a/tools/ui/src/lib/components/app/models/ModelsManager/row-actions.ts +++ b/tools/ui/src/lib/components/app/models/ModelsManager/row-actions.ts @@ -1,10 +1,10 @@ -import { Eye, EyeOff, Heart, HeartOff, Trash2 } from '@lucide/svelte'; +import { Compass, Eye, EyeOff, Heart, HeartOff, Trash2 } from '@lucide/svelte'; import { MODEL_DOWNLOAD_ICONS } from '$lib/constants'; import { ModelRowDownloadState } from '$lib/enums'; -import { modelsStore } from '$lib/stores'; +import { modelsStore, uiStore } from '$lib/stores'; import type { ModelOption } from '$lib/types/models'; +import { repoOf } from '$lib/utils'; -/** Row actions follow the app's dropdown pattern: icon, label, separators, variants. */ export function modelRowActions( option: ModelOption, favorite: boolean, @@ -13,9 +13,17 @@ export function modelRowActions( /** Download state, when the row stands for a tracked download. */ download?: ModelRowDownloadState | null ) { + const viewInDiscover = { + icon: Compass, + label: 'See in Discover', + // the details pane is keyed by repo, not by `:` + onclick: () => uiStore.openModelsDiscover(repoOf(option.model)), + separator: true + }; + // a tracked download is paused, resumed or dropped: it is not loaded, hidden or - // drafted. Deleting it stops the download and removes what is on disk, which is - // what the same wording offers on a cached model + // drafted. An in-flight row cancels, which stops the download and removes the + // partial files; a paused one deletes, which removes the partials from disk if (download) { const isPaused = download === ModelRowDownloadState.PAUSED; @@ -23,18 +31,16 @@ export function modelRowActions( { icon: isPaused ? MODEL_DOWNLOAD_ICONS.resume : MODEL_DOWNLOAD_ICONS.pause, label: isPaused ? 'Resume downloading' : 'Pause downloading', - onclick: () => - void (isPaused - ? modelsStore.status.downloadModel(option.model) - : modelsStore.status.pauseDownload(option.model)) + onclick: () => void modelsStore.status.toggleDownload(option.model) }, { icon: Trash2, - label: 'Delete from disk', + label: isPaused ? 'Delete from disk' : 'Cancel download', onclick: () => onDelete(option), separator: true, variant: 'destructive' as const - } + }, + viewInDiscover ]; } @@ -56,6 +62,7 @@ export function modelRowActions( label: isHidden ? 'Unhide model' : 'Hide model', onclick: () => modelsStore.toggleHidden(option.id), separator: true - } + }, + viewInDiscover ]; } diff --git a/tools/ui/src/lib/components/app/models/ModelsManager/utils.ts b/tools/ui/src/lib/components/app/models/ModelsManager/utils.ts index 208343664942..435700a10394 100644 --- a/tools/ui/src/lib/components/app/models/ModelsManager/utils.ts +++ b/tools/ui/src/lib/components/app/models/ModelsManager/utils.ts @@ -2,7 +2,12 @@ import { LOCAL_BACKEND_ID, type ModalityKey } from '$lib/constants'; import { ModelCapability, ModelGroupKind, ModelsTableGroupKind } from '$lib/enums'; import { HuggingFaceService, ModelsService } from '$lib/services'; import { modelsStore } from '$lib/stores'; -import type { ModelDownloadEntry, ModelDownloadProgress, ModelOption } from '$lib/types/models'; +import type { + ModelDownloadEntry, + ModelDownloadProgress, + ModelOption, + ModelSidecarBadge +} from '$lib/types/models'; import { detectThinkingSupport, detectToolUseSupport, repoOf } from '$lib/utils'; /** One repo of the table, with the rows it ships as. */ @@ -11,6 +16,8 @@ export interface ModelQuantGroup { key: string; kind: ModelGroupKind; quants: ModelOption[]; + /** Draft sidecars already on disk, shown as checked sub-lines. */ + sidecarBadges?: ModelSidecarBadge[]; } /** True when any of the manager's filters narrows the table. */ @@ -31,28 +38,36 @@ export function downloadProgressFor(repoWithTag: string): ModelDownloadProgress } /** - * One table row per tracked download, so the table lists them the way it lists - * any other model. A paused download is already in the router's listing, so its - * own option is reused; a fresh one stands for its tag alone. + * One table row per tracked download, listed like any other model. A paused download + * is already in the router's listing, so its own option is reused; a fresh one is + * synthesized for its tag alone. */ export function downloadGroups( entries: ModelDownloadEntry[], - models: ModelOption[] + models: ModelOption[], + sidecarsFor: (repo: string) => ModelSidecarBadge[] ): ModelQuantGroup[] { return entries.map((entry) => { - const option = models.find((candidate) => candidate.model === entry.repoWithTag) ?? { + const listed = models.find((candidate) => candidate.model === entry.repoWithTag); + // the synthesized option carries backendId, which ModelOption drops on + // assignment, so the callers spread it instead of typing this as one + const option = listed ?? { backendId: LOCAL_BACKEND_ID, capabilities: [], id: entry.repoWithTag, model: entry.repoWithTag, name: entry.repoWithTag }; + // a draft still coming down badges the id, an on-disk one keeps its sub-line + const sidecars = sidecarsFor(repoOf(entry.repoWithTag)); + const pending = sidecars.filter((badge) => badge.pending); return { - base: option, + base: { ...option, draftSidecars: pending }, key: `download::${entry.repoWithTag}`, kind: ModelGroupKind.QUANTS, - quants: [option] + quants: [{ ...option, draftSidecars: pending }], + sidecarBadges: sidecars.filter((badge) => !badge.pending) }; }); } diff --git a/tools/ui/src/lib/components/app/models/ModelsSelector/ModelsSelectorDownloadItem.svelte b/tools/ui/src/lib/components/app/models/ModelsSelector/ModelsSelectorDownloadItem.svelte index 6ef1433b970f..7cc286f89ddf 100644 --- a/tools/ui/src/lib/components/app/models/ModelsSelector/ModelsSelectorDownloadItem.svelte +++ b/tools/ui/src/lib/components/app/models/ModelsSelector/ModelsSelectorDownloadItem.svelte @@ -1,22 +1,38 @@ - -
- - + when paused; opacity only for the spinner, the spin owns the transform --> + - - - onRequestCancel?.(entry.repoWithTag)} - /> + + + + + // a phone has no manager: its information dialog takes the icon + deviceStore.isMobile + ? uiStore.openModelInformation(option) + : uiStore.openModelsManager(entry.repoWithTag)} + tooltip="Manage model" + tooltipAsTitle + /> + + onRequestCancel?.(entry.repoWithTag)} + /> +
+ + + {#if onDiskDrafts.length} +
+ {#each onDiskDrafts as badge (badge.repo + ':' + badge.kind + ':' + (badge.quant ?? ''))} + + {/each} +
+ {/if} {#if entry.progress && entry.progress.totalBytes > 0} Downloading

{#each downloadEntries as entry (entry.repoWithTag)} - + {/each} {/if} @@ -230,6 +236,7 @@ (cancelOpen = false)} + onConfirm={() => void modelsStore.status.cancelDownload(pendingCancel)} open={cancelOpen} repoWithTag={pendingCancel} /> diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscover.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscover.svelte new file mode 100644 index 000000000000..292938d59c5c --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscover.svelte @@ -0,0 +1,210 @@ + + +{#snippet catalog()} + (selectedId = id)} + showBaseModelAvatar + /> +{/snippet} + + + +
+ {#if selectedId} + + {/if} +
+ + (cancelOpen = false)} + onConfirm={() => void modelsStore.status.cancelDownload(pendingCancel)} + open={cancelOpen} + repoWithTag={pendingCancel} +/> diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverChatTemplateDialog.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverChatTemplateDialog.svelte new file mode 100644 index 000000000000..11240a4a91e1 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverChatTemplateDialog.svelte @@ -0,0 +1,75 @@ + + + + + + + Chat template + + + + + + + + + + + diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetails.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetails.svelte new file mode 100644 index 000000000000..8f7e50da6c24 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetails.svelte @@ -0,0 +1,96 @@ + + +{#if loading} + +{:else if error} +
+

{error}

+
+{:else if details} +
+ + + + + +
+{/if} diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsContainer.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsContainer.svelte new file mode 100644 index 000000000000..7d81ea8abab9 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsContainer.svelte @@ -0,0 +1,66 @@ + + + diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptions.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptions.svelte new file mode 100644 index 000000000000..a1bd05a37eb2 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptions.svelte @@ -0,0 +1,258 @@ + + +{#if bitDepthRows.length} +
+ +
+ {#each rows as row (row.bitDepth)} + (draftPick = kind)} + onRequestCancel={requestCancel} + onRequestDelete={requestDelete} + /> + {/each} +
+ + + {#if draftKinds.length} +
+ + + value && (draftPick = value as ModelDraftSidecar)} + type="single" + value={draftKind} + > + + {draftKindLabel} + + + + {#each draftKinds as kind (kind)} + + {kind.toUpperCase()} + + {/each} + + + + +
+ {/if} +
+{/if} + + (confirmOpen = false)} + open={confirmOpen} + repoWithTag={pending.repoWithTag} +/> diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsDownloadCommand.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsDownloadCommand.svelte new file mode 100644 index 000000000000..eb7f4e55f43b --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsDownloadCommand.svelte @@ -0,0 +1,131 @@ + + +
+ +
+ {SERVE_COMMAND.BIN} + + {SERVE_COMMAND.SUBCOMMAND} + + {SERVE_COMMAND.MODEL_FLAG} + + + {modelId}{quant ? MODEL_ID.QUANTIZATION_SEPARATOR : ''} + + + + {#if quant} + + {quant} + + {/if} + + + {#if draftKind} + {SERVE_COMMAND.SPEC_TYPE_FLAG} + + + {#if draftKinds.length > 1} + value && onDraftKindChange?.(value as ModelDraftSidecar)} + type="single" + value={draftKind} + > + + {SPEC_TYPE[draftKind]} + + + + {#each draftKinds as kind (kind)} + + {SPEC_TYPE[kind]} + + {/each} + + + {:else} + + {SPEC_TYPE[draftKind]} + + {/if} + {/if} +
+ + +
diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsDraftDownloadButton.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsDraftDownloadButton.svelte new file mode 100644 index 000000000000..d7cbd8becd84 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsDraftDownloadButton.svelte @@ -0,0 +1,61 @@ + + + + + + {#snippet child({ props })} + + {/snippet} + + + +

{tooltipText}

+ + diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsQuantDownloadButton.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsQuantDownloadButton.svelte new file mode 100644 index 000000000000..5ccbd31d08c7 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsQuantDownloadButton.svelte @@ -0,0 +1,272 @@ + + +{#snippet chipBody(dividerClass: string)} + + {label} + + {#if showsCompanion && companion} + + + + + + {companion.kind}{companion.quant ? `:${companion.quant}` : ''} + + {#if companion.onDisk} + + {/if} + + {/if} + + +{/snippet} + + +{#snippet tooltipTrigger(tip: string, button: Snippet<[Record]>)} + + + {#snippet child({ props })} + {@render button(props)} + {/snippet} + + + +

{tip}

+ + +{/snippet} + +{#snippet deleteChip(props: Record)} + + +{/snippet} + +{#snippet pauseResumeChip(props: Record)} + +{/snippet} + +{#snippet cancelChip(props: Record)} + + +{/snippet} + +{#snippet downloadChip(props: Record)} + + +{/snippet} + +{#if entry.isDownloaded} + {@render tooltipTrigger(tooltipText, deleteChip)} +{:else if entry.isDownloading || entry.isPaused} + +
+ {@render tooltipTrigger(tooltipText, pauseResumeChip)} + + {@render tooltipTrigger('Cancel downloading', cancelChip)} + + {#if percent !== null} + + {/if} +
+{:else} + {@render tooltipTrigger(tooltipText, downloadChip)} +{/if} diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsRow.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsRow.svelte new file mode 100644 index 000000000000..4527014fc901 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/ModelsDiscoverDetailsDownloadOptionsRow.svelte @@ -0,0 +1,134 @@ + + +
+
+ + + +
+ {#each files as file (file.path)} + + {/each} + + {#each missingDrafts as draft (draft.repoWithTag + ':' + draft.companion.kind)} + + {/each} +
+
+ + {#if commandOpen} +
+ +
+ {/if} +
diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/download-options.utils.ts b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/download-options.utils.ts new file mode 100644 index 000000000000..3c07cbafdf70 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/download-options.utils.ts @@ -0,0 +1,57 @@ +import { DRAFT_FILE_LABEL, MODEL_ID, PATH_SEPARATOR } from '$lib/constants'; +import { ModelDraftSidecar, ModelSelectableFileKind } from '$lib/enums'; +import { HuggingFaceService } from '$lib/services'; +import type { ModelSelectableFile } from '$lib/types'; +import { isAuxSidecar, isDraftSidecar } from '$lib/utils'; + +/** + * Draft kinds in the order a chip lists them. DFlash wins: on the models that + * ship both it drafts a block at a time and stays accurate further out than + * MTP's few tokens; a kind not listed here sorts last. + */ +export const DRAFT_PREFERENCE_ORDER: readonly ModelDraftSidecar[] = [ + ModelDraftSidecar.DFLASH, + ModelDraftSidecar.MTP +]; + +/** Draft sidecar of a file, typed for files already known to be drafts; null when it carries none. */ +export function draftBadgeOf(file: Pick): ModelDraftSidecar | null { + const sidecar = HuggingFaceService.extractQuantMeta(file.path)?.sidecar; + + return sidecar && isDraftSidecar(sidecar) ? sidecar : null; +} + +/** Kind of a file path: the main weights, a draft sidecar, or an aux sidecar (mmproj). */ +export function classify(path: string): ModelSelectableFileKind { + const sidecar = HuggingFaceService.extractQuantMeta(path)?.sidecar; + + if (!sidecar) return ModelSelectableFileKind.MAIN; + + return isAuxSidecar(sidecar) ? ModelSelectableFileKind.AUX : ModelSelectableFileKind.DRAFT; +} + +/** Draft kinds best first, so a chip lists its joined badges in the pull order. */ +export function sortDraftKinds(kinds: readonly ModelDraftSidecar[]): ModelDraftSidecar[] { + const rank = (kind: ModelDraftSidecar) => { + const index = DRAFT_PREFERENCE_ORDER.indexOf(kind); + + return index === -1 ? DRAFT_PREFERENCE_ORDER.length : index; + }; + + return [...kinds].sort((a, b) => rank(a) - rank(b)); +} + +/** Display label of a file: its quant, else the file name without the extension. */ +export function labelFor(path: string): string { + const meta = HuggingFaceService.extractQuantMeta(path); + + if (meta?.quant) return meta.quant; + + // Quantless sidecar files: the chip badge already carries the sidecar type, + // so the label only marks draft files; aux sidecars badge alone. + if (meta?.sidecar) return isDraftSidecar(meta.sidecar) ? DRAFT_FILE_LABEL : ''; + + const basename = path.split(PATH_SEPARATOR).pop() ?? path; + + return basename.replace(MODEL_ID.WEIGHT_EXTENSION_REGEX, ''); +} diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/index.ts b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/index.ts new file mode 100644 index 000000000000..bb82218ccb4b --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsDownloadOptions/index.ts @@ -0,0 +1,24 @@ +/** + * + * MODELS DISCOVER - DETAILS - DOWNLOAD OPTIONS + * + * The download area of the detail pane: one chip per main quant, the drafts it rides with + * badged on it, a draft checkbox and type pick at the bottom, and a terminal command + * preview inside every row. + * + */ + +/** One download action per main quant, the sidecars it pulls, and the row's command. */ +export { default as ModelsDiscoverDetailsDownloadOptions } from './ModelsDiscoverDetailsDownloadOptions.svelte'; + +/** One bit-depth row: the depth label opens its command, the quants of that depth download beside it. */ +export { default as ModelsDiscoverDetailsDownloadOptionsRow } from './ModelsDiscoverDetailsDownloadOptionsRow.svelte'; + +/** A single GGUF file as an independent action chip. */ +export { default as ModelsDiscoverDetailsDownloadOptionsQuantDownloadButton } from './ModelsDiscoverDetailsDownloadOptionsQuantDownloadButton.svelte'; + +/** Chip for a draft a downloaded model is missing; posts the model with that kind, so the server keeps what is on disk and pulls the draft. */ +export { default as ModelsDiscoverDetailsDownloadOptionsDraftDownloadButton } from './ModelsDiscoverDetailsDownloadOptionsDraftDownloadButton.svelte'; + +/** The `llama serve -hf ...` command box of one row, with a draft type pick and a copy button. */ +export { default as ModelsDiscoverDetailsDownloadOptionsDownloadCommand } from './ModelsDiscoverDetailsDownloadOptionsDownloadCommand.svelte'; diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsHeader.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsHeader.svelte new file mode 100644 index 000000000000..8687650796b9 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsHeader.svelte @@ -0,0 +1,96 @@ + + +
+
+
+ + +
+
+

{details.id ?? modelId}

+ + +
+ + {#if baseModels.length} +
+ {baseModels.join(', ')} + + + + +
+ {/if} +
+
+ + +
+ + + + +
diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsHfHubStats.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsHfHubStats.svelte new file mode 100644 index 000000000000..b848a9bfeb35 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsHfHubStats.svelte @@ -0,0 +1,31 @@ + + +
+ {#if typeof details.downloads === 'number'} + + + {HuggingFaceService.formatDownloads(details.downloads)} + + {/if} + + {#if typeof details.likes === 'number'} + + + {HuggingFaceService.formatLikes(details.likes)} + + {/if} + + {#if details.lastModified} + Updated {HuggingFaceService.formatRelativeTime(details.lastModified)} + {/if} +
diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsMetadata.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsMetadata.svelte new file mode 100644 index 000000000000..af7a38089390 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsMetadata.svelte @@ -0,0 +1,82 @@ + + +{#if description} +

{description}

+{/if} + + +
+ {#if gguf?.total} + + {/if} + + {#if gguf?.context_length} + + {/if} + + {#if gguf?.architecture} + + {/if} + + {#if gguf?.chat_template} + + {/if} + + {#if licenseTag} + + {/if} + + {#if details.gated === true} + + gated + + {/if} +
+ +{#if gguf?.chat_template} + +{/if} diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsMetadataItem.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsMetadataItem.svelte new file mode 100644 index 000000000000..4bfc027d8e1e --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsMetadataItem.svelte @@ -0,0 +1,18 @@ + + + + + {label} + + {value} + diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsReadme.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsReadme.svelte new file mode 100644 index 000000000000..fbf043b79712 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsReadme.svelte @@ -0,0 +1,17 @@ + + +{#if readme} +
+ +
+{/if} diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsSkeleton.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsSkeleton.svelte new file mode 100644 index 000000000000..8438a3c2b47d --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/ModelsDiscoverDetailsSkeleton.svelte @@ -0,0 +1,98 @@ + + + +
+ +
+
+
+ + +
+ + + +
+
+ + +
+ + +
+ + + + + +
+ + + + + +
+
+ + + + + + + +
+
+
+ + +
+
+ {#each [0, 1, 2] as _, index (index)} +
+
+ + + +
+ +
+ + + +
+
+ {/each} +
+ +
+ + +
+ + + + + +
+ + +
+
+ + +
+
+ + + + + + + +
+
+
diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/index.ts b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/index.ts new file mode 100644 index 000000000000..91d5c89065d6 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverDetails/index.ts @@ -0,0 +1,37 @@ +/** + * + * MODELS DISCOVER - DETAILS + * + * The right-hand detail pane of the discover view: header (avatar, name, stats, metadata + * chips, capability badges), the download options area and the model-card README. + * + */ + +/** Detail pane of the selected model; fetches its own details and GGUF file list. */ +export { default as ModelsDiscoverDetails } from './ModelsDiscoverDetails.svelte'; + +/** Fetches one repo's details, files and readme for hosts outside Discover. */ +export { default as ModelsDiscoverDetailsContainer } from './ModelsDiscoverDetailsContainer.svelte'; + +/** Header of the detail pane. */ +export { default as ModelsDiscoverDetailsHeader } from './ModelsDiscoverDetailsHeader.svelte'; + +/** Downloads, likes and last updated, formatted the way the Hub shows them. */ +export { default as ModelsDiscoverDetailsHfHubStats } from './ModelsDiscoverDetailsHfHubStats.svelte'; + +/** Static placeholder matching the detail layout. */ +export { default as ModelsDiscoverDetailsSkeleton } from './ModelsDiscoverDetailsSkeleton.svelte'; + +/** Model card README as markdown. */ +export { default as ModelsDiscoverDetailsReadme } from './ModelsDiscoverDetailsReadme.svelte'; + +/** Metadata chips: size, context, architecture, license, chat template. */ +export { default as ModelsDiscoverDetailsMetadata } from './ModelsDiscoverDetailsMetadata.svelte'; + +/** One label | value chip of the metadata row. */ +export { default as ModelsDiscoverDetailsMetadataItem } from './ModelsDiscoverDetailsMetadataItem.svelte'; + +/** Chat template in a scrollable dialog with a copy button. */ +export { default as ModelsDiscoverChatTemplateDialog } from './ModelsDiscoverChatTemplateDialog.svelte'; + +export * from './ModelsDiscoverDetailsDownloadOptions'; diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverList.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverList.svelte new file mode 100644 index 000000000000..5ebff7230d32 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverList.svelte @@ -0,0 +1,41 @@ + + +
    + {#if loading} + {#each Array(loadingSkeletonRowsCount) as _, index (index)} + + {/each} + {:else} + {#each models as model (model.id)} + + {/each} + {/if} +
diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListItem.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListItem.svelte new file mode 100644 index 000000000000..1eee85ea0777 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListItem.svelte @@ -0,0 +1,256 @@ + + +
  • +
    + + + {#if download} + {@const state = download} + + + + + {percent !== null ? `${percent}%` : state.isPaused ? 'Paused' : 'Downloading'} + + + + + {/if} +
    + + + {#if onDiskDrafts.length} +
    + {#each onDiskDrafts as badge (badge.repo + ':' + badge.kind + ':' + (badge.quant ?? ''))} + + {/each} +
    + {/if} + + {#if download?.progress && download.progress.totalBytes > 0} + + {/if} +
    • diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListItemSkeleton.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListItemSkeleton.svelte new file mode 100644 index 000000000000..851abc523860 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListItemSkeleton.svelte @@ -0,0 +1,39 @@ + + + +
  • +
    + + +
    + + +
    + {#each badgeWidths as width, i (i)} + + {/each} +
    +
    +
    +
    • diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListSearch.svelte b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListSearch.svelte new file mode 100644 index 000000000000..458c31648c07 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/ModelsDiscoverListSearch.svelte @@ -0,0 +1,30 @@ + + + +
    + handleInput(v)} {placeholder} size="sm" /> +
    diff --git a/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/index.ts b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/index.ts new file mode 100644 index 000000000000..c3025c0b2601 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/ModelsDiscoverList/index.ts @@ -0,0 +1,16 @@ +/** + * + * MODELS DISCOVER - LIST + * + * The discover sidebar column: a debounced search field above a navigable list of + * HuggingFace GGUF models, with skeleton rows while loading. + * + */ + +export { default as ModelsDiscoverList } from './ModelsDiscoverList.svelte'; +export { default as ModelsDiscoverListSearch } from './ModelsDiscoverListSearch.svelte'; + +/** One model entry; also the row of a tracked download, with its own controls. */ +export { default as ModelsDiscoverListItem } from './ModelsDiscoverListItem.svelte'; + +export { default as ModelsDiscoverListItemSkeleton } from './ModelsDiscoverListItemSkeleton.svelte'; diff --git a/tools/ui/src/lib/components/app/models/discover/index.ts b/tools/ui/src/lib/components/app/models/discover/index.ts new file mode 100644 index 000000000000..bb353c72dfa8 --- /dev/null +++ b/tools/ui/src/lib/components/app/models/discover/index.ts @@ -0,0 +1,16 @@ +/** + * + * MODELS DISCOVER + * + * Components for the Models Discover view: a sidebar search + list of HuggingFace GGUF + * models and a detail view for the selected model, used as the body of the discovery + * dialog. The list and detail trees live in their own subfolders; this barrel re-exports + * them alongside the shared leaves. + * + */ + +/** The whole view: sidebar list on the left, detail pane on the right. */ +export { default as ModelsDiscover } from './ModelsDiscover.svelte'; + +export * from './ModelsDiscoverList'; +export * from './ModelsDiscoverDetails'; diff --git a/tools/ui/src/lib/components/app/models/index.ts b/tools/ui/src/lib/components/app/models/index.ts index dbfb5c60e248..83f076fd11b0 100644 --- a/tools/ui/src/lib/components/app/models/index.ts +++ b/tools/ui/src/lib/components/app/models/index.ts @@ -59,4 +59,7 @@ export { default as ModelCapabilityIcons } from './ModelCapabilityIcons.svelte'; /** Draft sidecar badges, rendered as `+ [KIND] [QUANT]` per sidecar. */ export { default as ModelDraftSidecars } from './ModelDraftSidecars.svelte'; + +/** Compact draft download line, embedded in the item it belongs to. */ +export { default as ModelDraftSubline } from './ModelDraftSubline.svelte'; export * from './ModelsSelector'; diff --git a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte index 879390075988..aebf37251384 100644 --- a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte +++ b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChat.svelte @@ -13,6 +13,7 @@ NUMERIC_FIELDS, POSITIVE_INTEGER_FIELDS, SETTINGS_CHAT_SECTIONS, + SETTINGS_KEYS, SETTINGS_SECTION_SLUGS } from '$lib/constants'; import { ColorMode } from '$lib/enums/ui.enums'; @@ -70,6 +71,12 @@ function handleConfigChange(key: string, value: string | boolean) { localConfig[key] = value; + + // discover needs the hub metadata reads: checking it turns the hub on, + // unchecking it turns the hub back off, so the checkbox never lies + if (key === SETTINGS_KEYS.ENABLE_DISCOVER_MODELS) { + localConfig[SETTINGS_KEYS.USE_HUGGING_FACE_HUB] = value === true; + } } function handleReset() { diff --git a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatFields.svelte b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatFields.svelte index e8795d7c5204..5dfba6bb374c 100644 --- a/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatFields.svelte +++ b/tools/ui/src/lib/components/app/settings/SettingsChat/SettingsChatFields.svelte @@ -22,6 +22,13 @@ let { fields, localConfig, onConfigChange, onThemeChange }: Props = $props(); + // discover needs the hub metadata reads, so its checkbox is read-only while discover is on + let hubLocked = $derived(Boolean(localConfig[SETTINGS_KEYS.ENABLE_DISCOVER_MODELS])); + + function isLockedField(key: string): boolean { + return hubLocked && key === SETTINGS_KEYS.USE_HUGGING_FACE_HUB; + } + let currentModelParams = $derived.by(() => { void modelsStore.props.cacheVersion; @@ -301,13 +308,18 @@ onConfigChange(field.key, checked)} />