ggml: bound views by the bytes their strides reach (#31584)

* ggml: bound views by the bytes their strides reach

A view was checked against its source with the size of a contiguous
tensor of the same shape, computed before its strides were set. Views
whose rows overlap failed that check although they stay inside their
source, and views with wider strides passed it while reaching past
the end. The view functions now pass their strides to the tensor
constructor, which checks the bytes they actually reach, like
ggml_nbytes. The conv windows of delta_net are created with all their
rows at once instead of being widened after the check.

* ggml: assert that a strided view starts from contiguous elements

The view functions set nb[0] to the type size, so the source must have
contiguous elements along dim 0 for the view to read them.

.github/workflows/pr-video-demo-label.yml

@@ -0,0 +1,37 @@
+name: "PR video demo label"
+
+on:
+  pull_request_target:
+    types: [opened, edited, reopened, synchronize]
+
+cache-mode: none
+permissions:
+  pull-requests: write
+  issues: write
+
+jobs:
+  video-demo-label:
+    runs-on: ubuntu-slim
+    steps:
+      - name: Add or remove the video demo label
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          PR_URL: ${{ github.event.pull_request.html_url }}
+          PR_NUMBER: ${{ github.event.pull_request.number }}
+          # passed through env, so the body is only matched, never run, and is safe to read from a fork
+          BODY: ${{ github.event.pull_request.body }}
+          HAS_LABEL: ${{ contains(github.event.pull_request.labels.*.name, 'video demo') }}
+        run: |
+          # sections written by demo:pr (see tools/ui/demo/cli.ts): the managed comment
+          # markers, or a preview linking to a demo video uploaded for this PR
+          has_demo=false
+          if { grep -qF '<!-- demo:start -->' <<< "$BODY" && grep -qF '<!-- demo:end -->' <<< "$BODY"; } ||
+            grep -qE "/resolve/videos/pr/$PR_NUMBER/[0-9a-f]{10}(-[0-9]+)?\.mp4" <<< "$BODY"; then
+            has_demo=true
+          fi
+
+          if [ "$has_demo" = true ] && [ "$HAS_LABEL" = false ]; then
+            gh pr edit "$PR_URL" --add-label 'video demo'
+          elif [ "$has_demo" = false ] && [ "$HAS_LABEL" = true ]; then
+            gh pr edit "$PR_URL" --remove-label 'video demo'
+          fi

common/download.cpp

@@ -748,7 +748,8 @@ static hf_cache::hf_file find_best_model(const hf_cache::hf_files & files,
     }
 
     for (const auto & t : tags) {
-        std::regex pattern(t + "[.-]", std::regex::icase);
+        // match the tag as a full token, so ":F16" does not match a BF16 file
+        std::regex pattern("(^|[-./_])" + t + "[.-]", std::regex::icase);
         for (const auto & f : files) {
             if (gguf_filename_is_model(f.path) &&
                 std::regex_search(f.path, pattern)) {

common/speculative.cpp

@@ -1613,12 +1613,12 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
             }
         }
 
-        llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false);
-        llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
-
         is_mem_shared = llama_get_ctx_other(ctx_dft) == ctx_tgt;
         chain_heads   = n_mtp_layers > 1 && !is_mem_shared;
 
+        llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ is_mem_shared);
+        llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
+
         if (chain_heads) {
             this->params.n_max = std::min(this->params.n_max, n_mtp_layers);
 
@@ -1772,16 +1772,28 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
             }
 
             const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1;
-            verify_h_rows[seq_id] = n_rows;
-            verify_h[seq_id].resize((size_t) n_rows * n_embd);
+            if (is_mem_shared && !batch_in.tokens[i_batch_end[seq_id]].output) {
+                verify_h_rows[seq_id] = 0;
+                continue;
+            }
 
-            for (int32_t i = 0; i < n_rows; ++i) {
-                const float * h = llama_get_embeddings_nextn_ith(ctx_tgt, i_batch_beg[seq_id] + i);
+            const bool all_rows_output = std::all_of(
+                    batch_in.tokens.begin() + i_batch_beg[seq_id],
+                    batch_in.tokens.begin() + i_batch_end[seq_id] + 1,
+                    [](const common_batch::token & token) { return token.output; });
+            const int32_t h_rows = is_mem_shared && !all_rows_output ? 1 : n_rows;
+            const int32_t h_offset = is_mem_shared && !all_rows_output ? i_batch_end[seq_id] : i_batch_beg[seq_id];
+
+            verify_h_rows[seq_id] = h_rows;
+            verify_h[seq_id].resize((size_t) h_rows * n_embd);
+
+            for (int32_t i = 0; i < h_rows; ++i) {
+                const float * h = llama_get_embeddings_nextn_ith(ctx_tgt, h_offset + i);
                 std::memcpy(verify_h[seq_id].data() + (size_t) i * n_embd, h, row_bytes);
             }
 
             std::memcpy(pending_h[seq_id].data(),
-                    verify_h[seq_id].data() + (size_t) (n_rows - 1) * n_embd, row_bytes);
+                    verify_h[seq_id].data() + (size_t) (h_rows - 1) * n_embd, row_bytes);
         }
 
         return true;
@@ -1967,6 +1979,29 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
         const size_t row_bytes = (size_t) n_embd * sizeof(float);
         std::memcpy(pending_h[seq_id].data(), verify_h[seq_id].data() + (size_t) i_h * n_embd, row_bytes);
     }
+
+    bool get_state(llama_seq_id seq_id, std::vector<uint8_t> & data) const override {
+        if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
+            return false;
+        }
+
+        const size_t row_bytes = (size_t) n_embd * sizeof(float);
+        data.resize(row_bytes);
+        std::memcpy(data.data(), pending_h[seq_id].data(), row_bytes);
+        return true;
+    }
+
+    void set_state(llama_seq_id seq_id, const std::vector<uint8_t> & data) override {
+        if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
+            return;
+        }
+
+        const size_t row_bytes = (size_t) n_embd * sizeof(float);
+        if (data.size() != row_bytes) {
+            return;
+        }
+        std::memcpy(pending_h[seq_id].data(), data.data(), row_bytes);
+    }
 };
 
 // state of self-speculation (simple implementation, not ngram-map)

conversion/__init__.py

@@ -45,6 +45,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "ChameleonForConditionalGeneration": "chameleon",
     "ChatGLMForConditionalGeneration": "chatglm",
     "ChatGLMModel": "chatglm",
+    "ChatterboxModel": "chatterbox",
     "ClefModel": "clef",
     "CloverMoeModelLM": "clover",
     "CodeShellForCausalLM": "codeshell",
@@ -197,6 +198,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "MambaLMHeadModel": "mamba",
     "MapleForCausalLM": "maple",
     "MellumForCausalLM": "mellum",
+    "Metis16ForCausalLM": "metis",
     "MiMoV2FlashForCausalLM": "mimo",
     "MiMoV2ForCausalLM": "mimo",
     "MiniCPM3ForCausalLM": "minicpm",
@@ -223,6 +225,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "ModernBertForSequenceClassification": "bert",
     "ModernBertModel": "bert",
     "NanbeigeForCausalLM": "nanbeige",
+    "Nemotron3_5TranscribeForConditionalGeneration": "nemotron",
     "NemotronForCausalLM": "nemotron",
     "NemotronHForCausalLM": "nemotron",
     "NemotronHPuzzleForCausalLM": "nemotron",
@@ -328,6 +331,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
 
 MMPROJ_MODEL_MAP: dict[str, str] = {
     "AudioFlamingo3ForConditionalGeneration": "ultravox",
+    "ChatterboxModel": "chatterbox",
     "ClefModel": "clef",
     "CogVLMForCausalLM": "cogvlm",
     "Cohere2VisionForConditionalGeneration": "command_r",
@@ -380,6 +384,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
     "MiniCPMOcrForConditionalGeneration": "minicpm",
     "Mistral3ForConditionalGeneration": "llava",
     "Mistral3Model": "llava",
+    "Nemotron3_5TranscribeForConditionalGeneration": "nemotron",
     "NemotronH_Omni_Reasoning_V3": "nemotron",
     "NemotronH_Nano_VL_V2": "nemotron",
     "MuseGlimmerForConditionalGeneration": "muse_glimmer",

conversion/base.py

@@ -225,8 +225,6 @@ class ModelBase:
         tensors: dict[str, Callable[[], Tensor]] = {}
 
         if remote_hf_model_id is not None:
-            is_safetensors = True
-
             logger.info(f"Using remote model with HuggingFace id: {remote_hf_model_id}")
             remote_tensors = gguf.utility.SafetensorRemote.get_list_tensors_hf_model(remote_hf_model_id)
             for name, remote_tensor in remote_tensors.items():
@@ -237,9 +235,13 @@ class ModelBase:
 
             return tensors
 
-        prefix = "model" if not self.is_mistral_format else "consolidated"
+        return self.index_local_tensors("model" if not self.is_mistral_format else "consolidated")
+
+    # index the local "<prefix>*.safetensors" files (or pytorch_model*.bin), through filter_tensors
+    def index_local_tensors(self, prefix: str) -> dict[str, Callable[[], Tensor]]:
+        tensors: dict[str, Callable[[], Tensor]] = {}
         part_names: list[str] = ModelBase.get_model_part_names(self.dir_model, prefix, ".safetensors")
-        is_safetensors: bool = len(part_names) > 0 or (not self.is_mistral_format and (self.dir_model / "model.safetensors.index.json").is_file())
+        is_safetensors: bool = len(part_names) > 0 or (not self.is_mistral_format and (self.dir_model / f"{prefix}.safetensors.index.json").is_file())
         if not is_safetensors:
             part_names = ModelBase.get_model_part_names(self.dir_model, "pytorch_model", ".bin")
 
@@ -247,7 +249,7 @@ class ModelBase:
         tensor_names_from_parts: set[str] = set()
 
         if not self.is_mistral_format:
-            index_name = "model.safetensors" if is_safetensors else "pytorch_model.bin"
+            index_name = f"{prefix}.safetensors" if is_safetensors else "pytorch_model.bin"
             index_name += ".index.json"
             index_file = self.dir_model / index_name
 
@@ -1277,6 +1279,7 @@ class ModelBase:
         expert_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}
         expert_input_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}
         expert_shapes: dict[tuple[int, str], list[int]] = {}
+        expert_names: dict[tuple[int, str], str] = {}
         n_experts = self.find_hparam(["num_local_experts", "num_experts", "n_routed_experts"], optional=True) or 0
         consumed: list[str] = []
 
@@ -1325,6 +1328,7 @@ class ModelBase:
                     expert_scales[key] = []
                     expert_input_scales[key] = []
                     expert_shapes[key] = shape
+                    expert_names[key] = name
                 expert_blocks[key].append((expert_id, raw.copy()))
                 # Collect per-expert scale2 (scalar per expert)
                 expert_scales[key].append((expert_id, float(scale2.float().sum())))
@@ -1333,13 +1337,13 @@ class ModelBase:
 
                 # Flush when all experts for this (layer, proj) are collected
                 if n_experts > 0 and len(expert_blocks[key]) >= n_experts:
-                    self._flush_nvfp4_experts(key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type)
+                    self._flush_nvfp4_experts(key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type)
             else:
                 self._repack_nvfp4(name, weight, scale, scale2, input_scale)
 
         # Flush any remaining experts (fallback if n_experts was unknown)
         for bid, proj_type in list(expert_blocks.keys()):
-            self._flush_nvfp4_experts((bid, proj_type), expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type)
+            self._flush_nvfp4_experts((bid, proj_type), expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type)
 
         # Remove consumed tensors so get_tensors/modify_tensors won't see them
         for name in consumed:
@@ -1350,7 +1354,7 @@ class ModelBase:
             if name.endswith((".k_scale", ".v_scale")):
                 del self.model_tensors[name]
 
-    def _flush_nvfp4_experts(self, key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, bid, proj_type):
+    def _flush_nvfp4_experts(self, key, expert_blocks, expert_scales, expert_input_scales, expert_shapes, expert_names, bid, proj_type):
         experts = expert_blocks.pop(key)
         scales = expert_scales.pop(key)
         input_scales = expert_input_scales.pop(key)
@@ -1362,7 +1366,7 @@ class ModelBase:
         new_name = self.map_tensor_name(merged_name)
         logger.info(f"Repacked {new_name} with shape [{len(experts)}, {shape[0]}, {shape[1]}] and quantization NVFP4")
         self.gguf_writer.add_tensor(new_name, merged, raw_dtype=gguf.GGMLQuantizationType.NVFP4)
-        self._tag_prec_a4(merged_name, new_name)
+        self._tag_prec_a4(expert_names.pop(key), new_name)
 
         scales.sort(key=lambda x: x[0])
         self._write_scales_tensor(new_name.replace(".weight", ".scale"), [s[1] for s in scales])
@@ -2114,6 +2118,9 @@ class TextModel(ModelBase):
         if chkhsh == "0fe1cf6eda062318a1af7270f3331a85c539a01778ff948e24388e949c5282f4":
             # ref: https://huggingface.co/evilfreelancer/ruGPT3XL
             res = "gpt-2"
+        if chkhsh == "df1bc34a6366466bcb36a9a1b1cc0afdb08cfe21c9250602c930f159085390a9":
+            # ref: https://huggingface.co/Lernex/Metis-1.6-Logos
+            res = "gpt-2"
         if chkhsh == "9e454714343b69b99b71795c1d27a68c2a1d15dab111f4d353109f966af29da7":
             # ref: https://huggingface.co/LiquidAI/LFM2.5-8B-A1B
             res = "lfm2"
@@ -3384,7 +3391,7 @@ def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> st
     # Kimi-K3's text_config reports "KimiLinearForCausalLM", which is the older
     # Kimi-Linear-48B architecture and cannot load K3 (no attention residuals,
     # latent MoE, situ, ...). Route on the top-level architecture instead.
-    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration"):
+    if model_type == ModelType.TEXT and arch in ("NEOChatModel", "StepVLForConditionalGeneration", "Sarashina2VisionForCausalLM", "Exaone4_5_ForConditionalGeneration", "Step3p7ForConditionalGeneration", "KimiK3ForConditionalGeneration", "Nemotron3_5TranscribeForConditionalGeneration"):
         return arch
 
     # if "architectures" is found in the sub-config, use that instead

conversion/chatterbox.py

@@ -0,0 +1,471 @@
+# Chatterbox (ResembleAI) turbo and multilingual v3, from the raw checkpoints of the official repos
+# the checkpoints ship no config.json, the variant is detected from the talker file and the hparams come from the reference code
+# talker: text_emb and speech_emb are fused into one [text | speech] vocab, speech token i is <|speech_i|>
+# mmproj: the s3gen sidecar, the voice and conditioning encoders, the position and speech tables, and the conds.pt voice
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, Iterable
+
+import torch
+import torch.nn.functional as F
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, TextModel, MmprojModel, LazyTorchTensor, gguf, logger
+from .llama import LlamaModel
+
+# checkpoint file prefixes, the files are "<prefix>.safetensors"
+TURBO_TALKER = "t3_turbo_v1"
+MTL_TALKER   = "t3_mtl23ls_v3"
+TURBO_S3GEN  = "s3gen_meanflow"
+MTL_S3GEN    = "s3gen_v3"
+VE           = "ve"
+
+# relative speech ids shared by both variants (start/stop_speech_token in the reference)
+SPEECH_BOS = 6561
+SPEECH_EOS = 6562
+
+# sampling defaults of the reference, samplers it does not use are disabled so the common defaults do not apply
+TURBO_TOP_K = 1000
+MTL_TOP_K   = 0
+TURBO_SAMPLING = {"min_p": 0.0,  "top_p": 0.95, "temp": 0.8, "penalty_repeat": 1.2}
+MTL_SAMPLING   = {"min_p": 0.05, "top_p": 1.0,  "temp": 0.8, "penalty_repeat": 1.2}
+# the reference repetition penalty covers the whole generated history
+PENALTY_LAST_N = -1
+
+
+def _has_file(dir_model: Path, prefix: str) -> bool:
+    return (dir_model / f"{prefix}.safetensors").is_file()
+
+
+def _is_turbo(dir_model: Path) -> bool:
+    if _has_file(dir_model, TURBO_TALKER):
+        return True
+    if _has_file(dir_model, MTL_TALKER):
+        return False
+    raise FileNotFoundError(f"no chatterbox talker checkpoint in {dir_model}")
+
+
+# talker hparams of the reference (GPT2_MEDIUM_CONFIG / LLAMA_520M_CONFIG_DICT, T3Config)
+TURBO_TEXT_CONFIG = {
+    "n_layer": 24, "n_embd": 1024, "n_inner": 4096, "n_head": 16, "n_ctx": 8196, "layer_norm_epsilon": 1e-05,
+    "vocab_size": 50276, "speech_vocab_size": 6563,
+}
+MTL_TEXT_CONFIG = {
+    "num_hidden_layers": 30, "hidden_size": 1024, "intermediate_size": 4096, "num_attention_heads": 16,
+    "num_key_value_heads": 16, "head_dim": 64, "max_position_embeddings": 131072, "rms_norm_eps": 1e-05,
+    "rope_theta": 500000.0,
+    "rope_scaling": {"factor": 8.0, "high_freq_factor": 4.0, "low_freq_factor": 1.0, "original_max_position_embeddings": 8192},
+    "vocab_size": 2454, "speech_vocab_size": 8194,
+}
+
+
+def _is_chatterbox_checkpoint(dir_model: Path) -> bool:
+    return _has_file(dir_model, VE) and (_has_file(dir_model, TURBO_TALKER) or _has_file(dir_model, MTL_TALKER))
+
+
+@ModelBase.register_hparams_loader(_is_chatterbox_checkpoint)
+def _load_chatterbox_hparams(dir_model: Path) -> dict[str, Any]:
+    turbo = _is_turbo(dir_model)
+    logger.info(f"gguf: detected Chatterbox {'turbo' if turbo else 'multilingual v3'} checkpoint")
+    text_config = dict(TURBO_TEXT_CONFIG if turbo else MTL_TEXT_CONFIG)
+    # the mmproj tensors are named explicitly, they do not go through the block tensor map
+    return {"architectures": ["ChatterboxModel"], "text_config": text_config, "audio_config": {"num_hidden_layers": 0}}
+
+
+@ModelBase.register("ChatterboxModel")
+class ChatterboxTalkerModel(TextModel):
+    model_arch = gguf.MODEL_ARCH.LLAMA  # multilingual; the turbo constructor switches to GPT2
+
+    def __init__(self, dir_model: Path, *args, **kwargs):
+        self.is_turbo = _is_turbo(dir_model)
+        if self.is_turbo:
+            self.model_arch = gguf.MODEL_ARCH.GPT2
+        super().__init__(dir_model, *args, **kwargs)
+        self._text_embd: Tensor | None = None
+        self._speech_embd: Tensor | None = None
+        self._text_head: Tensor | None = None
+        self._speech_head: Tensor | None = None
+
+    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+        return self.index_local_tensors(TURBO_TALKER if self.is_turbo else MTL_TALKER)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, _ = item
+        # the reference feeds inputs_embeds, the transformer input embeddings are unused
+        if name in ("tfmr.wte.weight", "tfmr.embed_tokens.weight"):
+            return None
+        if name.startswith(("cond_enc.", "text_pos_emb.", "speech_pos_emb.")):
+            return None
+        return super().filter_tensors(item)
+
+    def set_vocab(self):
+        if self.is_turbo:
+            self._set_vocab_turbo()
+        else:
+            self._set_vocab_mtl()
+
+    def _speech_token_names(self, n_speech: int) -> list[str]:
+        return [f"<|speech_{i}|>" for i in range(n_speech)]
+
+    def _set_vocab_turbo(self):
+        # stock GPT-2 BPE from the checkpoint dir, extended with the speech tokens
+        tokens, toktypes, tokpre = self.get_vocab_base()
+        n_text = len(tokens)
+        n_speech = self.hparams["speech_vocab_size"]
+        speech = self._speech_token_names(n_speech)
+        tokens += speech
+        toktypes += [gguf.TokenType.CONTROL] * n_speech
+
+        with open(self.dir_model / "merges.txt", "r", encoding="utf-8") as f:
+            merges = [line.rstrip("\n") for line in f if line.strip() and not line.startswith("#version")]
+
+        self._write_vocab(tokens, toktypes, merges, tokpre, n_text)
+
+    def _write_vocab(self, tokens: list[str], toktypes: list[int], merges: list[str], tokpre: str, n_text: int):
+        self.gguf_writer.add_tokenizer_model("gpt2")
+        self.gguf_writer.add_tokenizer_pre(tokpre)
+        self.gguf_writer.add_token_list(tokens)
+        self.gguf_writer.add_token_types(toktypes)
+        self.gguf_writer.add_token_merges(merges)
+
+        self.gguf_writer.add_bos_token_id(n_text + SPEECH_BOS)
+        self.gguf_writer.add_eos_token_id(n_text + SPEECH_EOS)
+        self.gguf_writer.add_add_bos_token(False)
+        self.gguf_writer.add_add_eos_token(False)
+
+        # the reference samples the speech head only
+        self.gguf_writer.add_suppress_tokens(list(range(n_text)))
+
+    def _set_vocab_mtl(self):
+        # the reference BPE is char-level, it is re-encoded to byte-level with extra merges that rebuild each multi-byte char
+        # its lowercasing is folded into the embedding table, so the prompt needs no text preprocessing
+        with open(self.dir_model / "mtl_tokenizer.json", "r", encoding="utf-8") as f:
+            tok = json.load(f)
+
+        byte_map = gguf.vocab.bytes_to_unicode()
+
+        def enc(s: str) -> str:
+            return "".join(byte_map[b] for b in s.encode("utf-8"))
+
+        n_text = self.hparams["vocab_size"]
+        tokens: list[str] = [f"[unused_{i}]" for i in range(n_text)]
+        toktypes = [int(gguf.TokenType.UNUSED)] * n_text
+        char_merges: list[str] = []
+        for t, i in tok["model"]["vocab"].items():
+            if t == " ":
+                # the reference substitutes [SPACE] before encoding, the raw space token is never used
+                continue
+            tokens[i] = enc(t)
+            toktypes[i] = int(gguf.TokenType.NORMAL)
+            if len(t) == 1 and len(t.encode("utf-8")) > 1:
+                parts = [byte_map[b] for b in t.encode("utf-8")]
+                for k in range(1, len(parts)):
+                    char_merges.append("".join(parts[:k]) + " " + parts[k])
+        for entry in tok.get("added_tokens", []):
+            tokens[entry["id"]] = entry["content"]
+            toktypes[entry["id"]] = int(gguf.TokenType.CONTROL)
+
+        # raw spaces tokenize to [SPACE] directly
+        space_id = tok["model"]["vocab"]["[SPACE]"]
+        tokens[space_id] = enc(" ")
+        toktypes[space_id] = int(gguf.TokenType.NORMAL)
+
+        n_speech = self.hparams["speech_vocab_size"]
+        tokens += self._speech_token_names(n_speech)
+        toktypes += [int(gguf.TokenType.CONTROL)] * n_speech
+
+        # chars are atomic in the reference, they must form before any of its merges
+        merges = char_merges
+        for m in tok["model"].get("merges", []):
+            a, b = m if isinstance(m, list) else m.split(" ")
+            merges.append(enc(a) + " " + enc(b))
+
+        self._write_vocab(tokens, toktypes, merges, "default", n_text)
+
+    def _mtl_case_fold_pairs(self) -> list[tuple[int, int]]:
+        # (upper_id, lower_id) of the single-char vocab entries, for every script in the vocab
+        with open(self.dir_model / "mtl_tokenizer.json", "r", encoding="utf-8") as f:
+            vocab = json.load(f)["model"]["vocab"]
+        pairs: list[tuple[int, int]] = []
+        for t, i in vocab.items():
+            if len(t) != 1:
+                continue
+            low = t.lower()
+            if low != t and len(low) == 1 and low in vocab:
+                pairs.append((i, vocab[low]))
+        return pairs
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+
+        sampling = TURBO_SAMPLING if self.is_turbo else MTL_SAMPLING
+        self.gguf_writer.add_sampling_top_k(TURBO_TOP_K if self.is_turbo else MTL_TOP_K)
+        self.gguf_writer.add_sampling_min_p(sampling["min_p"])
+        self.gguf_writer.add_sampling_top_p(sampling["top_p"])
+        self.gguf_writer.add_sampling_temp(sampling["temp"])
+        self.gguf_writer.add_sampling_penalty_repeat(sampling["penalty_repeat"])
+        self.gguf_writer.add_sampling_penalty_last_n(PENALTY_LAST_N)
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        if self.is_turbo:
+            return
+        # llama3 rope scaling baked into the rope_freqs factors tensor
+        rp = self.hparams["rope_scaling"]
+        dim = self.hparams["head_dim"]
+        base = self.hparams["rope_theta"]
+        freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+        factor = rp["factor"]
+        low_freq_wavelen = rp["original_max_position_embeddings"] / rp["low_freq_factor"]
+        high_freq_wavelen = rp["original_max_position_embeddings"] / rp["high_freq_factor"]
+        rope_factors = []
+        for freq in freqs:
+            wavelen = 2 * math.pi / freq
+            if wavelen < high_freq_wavelen:
+                rope_factors.append(1)
+            elif wavelen > low_freq_wavelen:
+                rope_factors.append(factor)
+            else:
+                smooth = (rp["original_max_position_embeddings"] / wavelen - rp["low_freq_factor"]) / (rp["high_freq_factor"] - rp["low_freq_factor"])
+                rope_factors.append(1 / ((1 - smooth) / factor + smooth))
+        yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32))
+
+    def _maybe_emit_fused(self) -> Iterable[tuple[str, Tensor]]:
+        if self._text_embd is not None and self._speech_embd is not None:
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD),
+                   torch.cat([self._text_embd, self._speech_embd], dim=0))
+            self._text_embd = None
+            self._speech_embd = None
+        if self._text_head is not None and self._speech_head is not None:
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT),
+                   torch.cat([self._text_head, self._speech_head], dim=0))
+            self._text_head = None
+            self._speech_head = None
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # generate_extra_tensors output comes back through here with its final name
+        if name.startswith("rope_freqs"):
+            yield (name, data_torch)
+            return
+        # fused [text | speech] vocab: embeddings and output head
+        if name == "text_emb.weight":
+            if not self.is_turbo:
+                # the reference lowercases its input, the uppercase rows take the lowercase ones
+                fold = list(range(data_torch.shape[0]))
+                for upper, lower in self._mtl_case_fold_pairs():
+                    fold[upper] = lower
+                data_torch = data_torch[torch.tensor(fold)]
+            self._text_embd = data_torch
+            yield from self._maybe_emit_fused()
+            return
+        if name == "speech_emb.weight":
+            self._speech_embd = data_torch
+            yield from self._maybe_emit_fused()
+            return
+        if name == "text_head.weight":
+            self._text_head = data_torch
+            yield from self._maybe_emit_fused()
+            return
+        if name == "speech_head.weight":
+            self._speech_head = data_torch
+            yield from self._maybe_emit_fused()
+            return
+        if name == "speech_head.bias":
+            # the text head has no bias, its zone is suppressed at sampling anyway
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT, suffix=".bias"), F.pad(data_torch, (self.hparams["vocab_size"], 0)))
+            return
+
+        assert name.startswith("tfmr.")
+        name = name[len("tfmr."):]
+
+        if self.is_turbo:
+            # HF GPT-2 layout: Conv1D style weights are stored transposed
+            if name.endswith((".c_attn.weight", ".c_proj.weight", ".c_fc.weight")):
+                data_torch = data_torch.transpose(1, 0)
+            yield (self.map_tensor_name(name), data_torch)
+            return
+
+        if name.endswith("q_proj.weight"):
+            data_torch = LlamaModel.permute(data_torch, self.hparams["num_attention_heads"], None)
+        if name.endswith("k_proj.weight"):
+            data_torch = LlamaModel.permute(data_torch, self.hparams["num_attention_heads"], self.hparams["num_key_value_heads"])
+        yield (self.map_tensor_name("model." + name), data_torch)
+
+
+@ModelBase.register("ChatterboxModel")
+class ChatterboxMmprojModel(MmprojModel):
+    has_vision_encoder = False
+    has_audio_encoder = True
+
+    def __init__(self, dir_model: Path, *args, **kwargs):
+        self.is_turbo = _is_turbo(dir_model)
+        super().__init__(dir_model, *args, **kwargs)
+        self._wnorm_g: dict[str, Tensor] = {}
+        self._wnorm_v: dict[str, Tensor] = {}
+
+    def get_audio_config(self) -> dict[str, Any] | None:
+        return self.global_config.get("audio_config")
+
+    def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Callable[[], Tensor]]:
+        tensors: dict[str, Callable[[], Tensor]] = {}
+        for prefix in (TURBO_S3GEN if self.is_turbo else MTL_S3GEN, VE, TURBO_TALKER if self.is_turbo else MTL_TALKER):
+            tensors.update(self.index_local_tensors(prefix))
+        return tensors
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        # the s3gen, ve and talker checkpoints have disjoint top-level names
+        name, gen = item
+        new_name: str | None = None
+        # batchnorm bookkeeping (over the gguf name length cap) and the tokenizer dsp buffers are unused
+        if name.endswith("num_batches_tracked") or name in ("tokenizer.window", "tokenizer._mel_filters"):
+            return None
+        for src, dst in (
+            ("flow.encoder.",           "a.gen.fenc."),
+            ("flow.decoder.estimator.", "a.gen.est."),
+            ("mel2wav.",                "a.gen.hift."),
+            ("speaker_encoder.",        "a.spk."),
+            ("tokenizer.",              "a.s3tok."),
+            ("flow.",                   "a.gen.flow."),  # input_embedding, encoder_proj, spk_embed_affine_layer
+            ("lstm.",                   "a.ve.lstm."),
+            ("proj.",                   "a.ve.proj."),
+            ("cond_enc.",               "a.cenc."),
+        ):
+            if name.startswith(src):
+                new_name = dst + name[len(src):]
+                break
+        if name == "text_pos_emb.emb.weight":
+            new_name = "a.gen.t3.text_pos_emb"
+        elif name == "speech_pos_emb.emb.weight":
+            new_name = "a.gen.t3.speech_pos_emb"
+        elif name == "speech_emb.weight":
+            new_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.A_GEN_CODE_OUT_EMBD] + ".weight"
+        if new_name is None:
+            return None
+        return super().filter_tensors((new_name, gen))
+
+    def set_gguf_parameters(self):
+        self.gguf_writer.add_file_type(self.ftype)
+
+        # speaker encoder, the front-end hparams are hard-coded in clip.cpp
+        self.gguf_writer.add_clip_has_audio_encoder(True)
+        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.CHATTERBOX_SPKENC)
+        self.gguf_writer.add_audio_projection_dim(80)
+        self.gguf_writer.add_audio_num_mel_bins(80)
+        self.gguf_writer.add_audio_block_count(0)
+        self.gguf_writer.add_audio_embedding_length(192)
+        self.gguf_writer.add_audio_head_count(1)
+        self.gguf_writer.add_audio_feed_forward_length(192)
+        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
+
+        # audio generator (s3gen, chatterbox projector); the flow encoder shape
+        self.gguf_writer.add_clip_has_gen_audio_encoder(True)
+        self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.CHATTERBOX)
+        self.gguf_writer.add_gen_audio_projection_dim(80)
+        self.gguf_writer.add_gen_audio_embedding_length(512)
+        self.gguf_writer.add_gen_audio_feed_forward_length(2048)
+        self.gguf_writer.add_gen_audio_block_count(6)
+        self.gguf_writer.add_gen_audio_head_count(8)
+        self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5)
+
+    def _fuse_weight_norm(self, name: str, data_torch: Tensor) -> tuple[str, Tensor] | None:
+        # torch weight_norm parametrization: weight = g * v / |v| over dims 1..n
+        base = name.split(".parametrizations.weight.original")[0]
+        if name.endswith("original0"):
+            self._wnorm_g[base] = data_torch
+        else:
+            self._wnorm_v[base] = data_torch
+        if base in self._wnorm_g and base in self._wnorm_v:
+            g = self._wnorm_g.pop(base)
+            v = self._wnorm_v.pop(base)
+            norm = v.float().norm(dim=tuple(range(1, v.dim())), keepdim=True)
+            return (base + ".weight", g.float() * v.float() / norm)
+        return None
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if ".parametrizations.weight.original" in name:
+            fused = self._fuse_weight_norm(name, data_torch)
+            if fused is not None:
+                yield fused
+            return
+        yield (name, data_torch)
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        del name, bid
+        # tensors read on the host stay F32
+        if new_name.startswith(("a.ve.", "a.cenc.", "voice.", "a.gen.t3.", "a.gen.hift.m_source.")):
+            return gguf.GGMLQuantizationType.F32
+        # the conv ops have no BF16 kernels
+        if n_dims >= 2 and new_name.endswith((".weight", ".weight_v")):
+            return gguf.GGMLQuantizationType.F16
+        return gguf.GGMLQuantizationType.F32
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        conds = torch.load(self.dir_model / "conds.pt", map_location="cpu", weights_only=True)
+        t3c = conds["t3"] if isinstance(conds, dict) else conds.t3
+        genc = conds["gen"] if isinstance(conds, dict) else conds.gen
+        t3c = vars(t3c) if not isinstance(t3c, dict) else t3c
+        genc = vars(genc) if not isinstance(genc, dict) else genc
+
+        def talker_tensor(name: str) -> Tensor:
+            return LazyTorchTensor.to_eager(self.model_tensors[name]())
+
+        # the conds.pt voice: voice.default holds the talker conditioning rows, voice.default.ref_* the flow decoder reference
+        # 1D tensors are written as F32, the speech ids stay exact
+        yield ("voice.default.ref_tokens", genc["prompt_token"][0].float())
+        yield ("voice.default.ref_feat", genc["prompt_feat"][0].float())
+        yield ("voice.default.ref_spk", genc["embedding"][0].float())
+
+        spkr_w = talker_tensor("a.cenc.spkr_enc.weight").float()
+        spkr_b = talker_tensor("a.cenc.spkr_enc.bias").float()
+        spkr_row = spkr_w @ t3c["speaker_emb"][0].float() + spkr_b
+        speech_emb = talker_tensor(self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_OUT_EMBD)).float()
+        cond_tokens = t3c["cond_prompt_speech_tokens"][0]
+
+        if self.is_turbo:
+            # turbo talker conditioning: [spkr, speech rows of the cond tokens]
+            yield ("voice.default", torch.cat([spkr_row[None], speech_emb[cond_tokens]], dim=0))
+            return
+
+        # multilingual talker conditioning [spkr, perceiver x32, emotion], from the reference perceiver over the cond speech tokens
+        with torch.no_grad():
+            speech_pos = talker_tensor("a.gen.t3.speech_pos_emb").float()
+            pse = speech_emb[cond_tokens] + speech_pos[: cond_tokens.shape[0]]
+
+            ln_w = talker_tensor("a.cenc.perceiver.attn.norm.weight").float()
+            ln_b = talker_tensor("a.cenc.perceiver.attn.norm.bias").float()
+            wq = talker_tensor("a.cenc.perceiver.attn.to_q.weight").float()
+            bq = talker_tensor("a.cenc.perceiver.attn.to_q.bias").float()
+            wk = talker_tensor("a.cenc.perceiver.attn.to_k.weight").float()
+            bk = talker_tensor("a.cenc.perceiver.attn.to_k.bias").float()
+            wv = talker_tensor("a.cenc.perceiver.attn.to_v.weight").float()
+            bv = talker_tensor("a.cenc.perceiver.attn.to_v.bias").float()
+            wo = talker_tensor("a.cenc.perceiver.attn.proj_out.weight").float()
+            bo = talker_tensor("a.cenc.perceiver.attn.proj_out.bias").float()
+            query = talker_tensor("a.cenc.perceiver.pre_attention_query")[0].float()
+
+            n_head = 4
+            n_e = query.shape[1]
+
+            def attn_block(x1: Tensor, x2: Tensor) -> Tensor:
+                nx1 = F.layer_norm(x1, (n_e,), ln_w, ln_b)
+                nx2 = F.layer_norm(x2, (n_e,), ln_w, ln_b)
+                q = (nx1 @ wq.T + bq).view(-1, n_head, n_e // n_head).transpose(0, 1)
+                k = (nx2 @ wk.T + bk).view(-1, n_head, n_e // n_head).transpose(0, 1)
+                v = (nx2 @ wv.T + bv).view(-1, n_head, n_e // n_head).transpose(0, 1)
+                ctx = F.scaled_dot_product_attention(q, k, v)
+                ctx = ctx.transpose(0, 1).reshape(-1, n_e)
+                return ctx @ wo.T + bo + x1
+
+            pre = attn_block(query, pse)
+            p32 = attn_block(pre, pre)
+
+            emo = talker_tensor("a.cenc.emotion_adv_fc.weight").float()
+            emo_row = emo[:, 0] * t3c["emotion_adv"].reshape(-1)[0]
+            yield ("voice.default", torch.cat([spkr_row[None], p32, emo_row[None]], dim=0))

conversion/metis.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import math
+from typing import Callable, Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import LazyTorchTensor, ModelBase, TextModel, gguf
+
+
+@ModelBase.register("Metis16ForCausalLM")
+class Metis16Model(TextModel):
+    model_arch = gguf.MODEL_ARCH.METIS16
+
+    # read by generate_extra_tensors() or not used at inference: dropped in modify_tensors()
+    _folded = {
+        "pass_embeddings", "lm_head.weight", "depth_budget.multiplier",
+        "depth_memory.pass_embeddings", "depth_memory.anchor_embeddings",
+        "depth_memory.state_write.bias", "depth_memory.metadata_write.bias",
+    }
+    _folded_layer = {
+        "mixer_connection.controller.weight", "mixer_connection.controller.bias", "mixer_connection.mix_logits",
+        "mixer_connection.read_logits", "mixer_connection.write_logits",
+        "moe_connection.controller.weight", "moe_connection.controller.bias", "moe_connection.mix_logits",
+        "moe_connection.read_logits", "moe_connection.write_logits",
+        "moe.k_router.weight", "moe.k_router.bias", "moe.k_budget.multiplier",
+    }
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        hp = self.hparams
+        self.n_pass = hp["num_passes"]
+        self.attn_layers = list(hp["attention_layer_indices"])
+        self._generated: set[str] = set()
+        # one cache slot per pass and layer, while the tensors exist once per physical layer
+        self.block_count = self.n_pass * hp["num_hidden_layers"]
+
+    def set_vocab(self):
+        self._set_vocab_gpt2()
+        # turns end on the first eos id, so llama.cpp must stop on it too
+        eos = self.hparams.get("eos_token_id")
+        if isinstance(eos, list) and eos:
+            self.gguf_writer.add_eot_token_id(eos[0])
+
+    def set_gguf_parameters(self):
+        hp = self.hparams
+        super().set_gguf_parameters()
+
+        n_layer = hp["num_hidden_layers"]
+        self.gguf_writer.add_head_count_kv([hp["num_key_value_heads"] if il % n_layer in self.attn_layers else 0 for il in range(self.block_count)])
+
+        self.gguf_writer.add_expert_feed_forward_length(hp["expert_intermediate_size"])
+        self.gguf_writer.add_expert_shared_feed_forward_length(hp["shared_expert_intermediate_size"])
+        self.gguf_writer.add_expert_shared_count(1)
+        self.gguf_writer.add_moe_latent_size(hp["latent_dim"])
+
+        d_inner = hp["mamba_expand"] * hp["hidden_size"]
+        self.gguf_writer.add_ssm_conv_kernel(hp["mamba_d_conv"])
+        self.gguf_writer.add_ssm_inner_size(d_inner)
+        self.gguf_writer.add_ssm_state_size(hp["mamba_d_state"])
+        self.gguf_writer.add_ssm_time_step_rank(d_inner // hp["mamba_head_dim"])
+        self.gguf_writer.add_ssm_group_count(hp["mamba_n_groups"])
+        self.gguf_writer.add_ssm_norm_epsilon(hp["mamba_norm_eps"])
+
+        self.gguf_writer.add_hyper_connection_count(hp["n_streams"])
+        self.gguf_writer.add_metis_pass_count(self.n_pass)
+        self.gguf_writer.add_metis_memory_length(hp["memory_dim"])
+        self.gguf_writer.add_metis_route_length(hp["route_feature_dim"])
+        self.gguf_writer.add_metis_continuation_logit_limit(hp["continuation_logit_limit"])
+        self.gguf_writer.add_attention_pass_lora_rank(hp["attention_pass_lora_rank"])
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        # the checkpoint leaves these names suffix-less; add the suffix of the GGUF tensor so the mapper keeps it
+        if name.endswith((".mixer.impl.mixer.dt_bias", ".moe.selection_bias")):
+            name += ".bias"
+        elif (name.endswith((".moe.expert_embeddings", ".mixer.pass_gates"))
+              or name in ("depth_memory.stream_gate", "stream_embeddings")):
+            name += ".weight"
+        elif name == "depth_memory.stream_gate_bias":
+            name = "depth_memory.stream_gate.bias"
+        return super().filter_tensors((name, gen))
+
+    def _load(self, name: str) -> Tensor:
+        return LazyTorchTensor.to_eager(self.model_tensors[name]())
+
+    def _emit(self, name: str, data: Tensor) -> tuple[str, Tensor]:
+        self._generated.add(name)
+        return name, data
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        hp = self.hparams
+        r = hp["route_feature_dim"]
+        n_stream = hp["n_streams"]
+        T = gguf.MODEL_TENSOR
+        passes = self._load("pass_embeddings")[: self.n_pass]
+
+        # the mixes depend only on the weights and the pass, so Sinkhorn runs once here
+        # the controller runs in the checkpoint dtype, as the reference does
+        for il in range(hp["num_hidden_layers"]):
+            for src, (mix, read, write) in (
+                ("mixer_connection", (T.HC_ATTN_MIX, T.HC_ATTN_READ, T.HC_ATTN_WRITE)),
+                ("moe_connection", (T.HC_FFN_MIX, T.HC_FFN_READ, T.HC_FFN_WRITE)),
+            ):
+                prefix = f"layers.{il}.{src}."
+                ctrl_w, ctrl_b = self._load(prefix + "controller.weight"), self._load(prefix + "controller.bias")
+                mixes, reads, writes = [], [], []
+                for p in range(self.n_pass):
+                    delta = torch.nn.functional.linear(passes[p].to(ctrl_w.dtype), ctrl_w, ctrl_b).float()
+                    m_delta, r_delta, w_delta = torch.split(delta, [n_stream * n_stream, n_stream, n_stream])
+                    log_m = self._load(prefix + "mix_logits").float() + m_delta.view(n_stream, n_stream)
+                    for _ in range(hp["mhc_sinkhorn_iterations"]):
+                        log_m = log_m - torch.logsumexp(log_m, dim=-1, keepdim=True)
+                        log_m = log_m - torch.logsumexp(log_m, dim=-2, keepdim=True)
+                    mixes.append(log_m.exp())
+                    reads.append(torch.softmax(self._load(prefix + "read_logits").float() + r_delta, dim=-1))
+                    writes.append(torch.softmax(self._load(prefix + "write_logits").float() + w_delta, dim=-1))
+                act = ctrl_w.dtype
+                yield self._emit(self.format_tensor_name(mix, il), torch.stack(mixes).to(act).float())
+                yield self._emit(self.format_tensor_name(read, il), torch.stack(reads).to(act).float())
+                yield self._emit(self.format_tensor_name(write, il), torch.stack(writes).to(act).float())
+
+        # one bias per memory entry: write biases, pass and anchor embeddings, and the constant routed width
+        meta_w = self._load("depth_memory.metadata_write.weight").float()
+        mean_k = hp["num_experts_per_tok"] / hp["max_routed_k"]
+        base = (self._load("depth_memory.state_write.bias").float() + self._load("depth_memory.metadata_write.bias").float()
+                + mean_k * meta_w[:, r])
+        pass_e = self._load("depth_memory.pass_embeddings").float()
+        anchor_e = self._load("depth_memory.anchor_embeddings").float()
+        n_anchor = len(self.attn_layers) + 1
+        entry_bias = torch.stack([base + pass_e[p] + anchor_e[a] for p in range(self.n_pass) for a in range(n_anchor)])
+        yield self._emit(self.format_tensor_name(T.DEPTH_MEM_ENTRY_BIAS), entry_bias)
+
+    def _split_features(self, data_torch: Tensor, template: str, parts: tuple[str, ...], d: int, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # the graph reads the concatenated features one block at a time; each block is its own tensor
+        for i, part in enumerate(parts):
+            piece = data_torch[:, i * d:(i + 1) * d] if i < 3 else data_torch[:, 3 * d:]
+            yield from super().modify_tensors(piece, template.format(part), bid)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name in self._generated:
+            yield name, data_torch
+            return
+        if name in self._folded:
+            return
+        hp = self.hparams
+        d = hp["hidden_size"]
+
+        if bid is None:
+            if name == "stream_embeddings.weight":
+                # the graph keeps the streams innermost: [n_stream, n_embd, n_tokens]
+                yield from super().modify_tensors(data_torch.T.contiguous(), name, bid)
+                return
+            if name == "depth_memory.metadata_write.weight":
+                # metadata input becomes [summary, entropy, confidence, continuation]; entropy arrives unnormalized
+                r = hp["route_feature_dim"]
+                meta_w = data_torch.float()
+                entropy_col = meta_w[:, r + 1:r + 2] / math.log(hp["n_routed_experts"])
+                data_torch = torch.cat((meta_w[:, :r], entropy_col, meta_w[:, r + 2:]), dim=1)
+            elif name == "depth_memory.stream_gate.weight":
+                data_torch = data_torch.float() / math.sqrt(d)
+            elif name == "depth_memory.route_projection.weight":
+                yield from self._split_features(data_torch, "depth_memory.route_{}.weight", ("state", "retr", "diff", "hist"), d, bid)
+                return
+            elif name == "continuation.hidden.weight":
+                yield from self._split_features(data_torch, "continuation.hidden_{}.weight", ("state", "retr", "diff", "route"), d, bid)
+                return
+            yield from super().modify_tensors(data_torch, name, bid)
+            return
+
+        local = name.removeprefix(f"layers.{bid}.")
+        if local in self._folded_layer:
+            return
+
+        if local == "mixer.impl.mixer.A_log":
+            data_torch = -torch.exp(data_torch).unsqueeze(-1)
+        elif local == "mixer.impl.mixer.D":
+            data_torch = data_torch.unsqueeze(-1)
+        elif local == "mixer.impl.mixer.conv1d.weight":
+            data_torch = data_torch.squeeze(1)
+        elif local == "mixer.impl.mixer.norm.weight":
+            data_torch = data_torch.reshape(hp["mamba_n_groups"], -1)
+        elif local == "mixer.pass_gates.weight":
+            # the reference applies tanh(gate) in the checkpoint dtype
+            data_torch = torch.tanh(data_torch[: self.n_pass]).to(torch.bfloat16).float()
+        elif local == "moe.shared_expert.gate_up.weight":
+            gate, up = data_torch.chunk(2, dim=0)
+            yield from super().modify_tensors(gate, name.replace("gate_up", "gate"), bid)
+            yield from super().modify_tensors(up, name.replace("gate_up", "up"), bid)
+            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:
+        # FP32 in the checkpoint, tiny and read as rows, or computed here at full precision: keep FP32
+        if (new_name.startswith(("cont_", "depth_mem_entry_bias", "depth_mem_stream_gate"))
+                or new_name.endswith(("ffn_exp_embd.weight", "ffn_gate_inp.weight", "hc_attn_mix.weight", "hc_attn_read.weight",
+                                      "hc_attn_write.weight", "hc_ffn_mix.weight", "hc_ffn_read.weight", "hc_ffn_write.weight"))):
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)

conversion/nemotron.py

@@ -428,6 +428,9 @@ class NemotronHModel(GraniteHybridModel):
     is_moe: bool = False
     supports_mtp_export = True
     _experts: list[dict[str, Tensor]] | None = None
+    # tokens past the embedding rows (left out of the vocab), and special tokens that are printed as text
+    _vocab_skip: tuple[str, ...] = ()
+    _user_defined_tokens: re.Pattern[str] | None = None
 
     _SSM_LAYER_TYPES = {"mamba", "linear_attention"}
     _ATTN_LAYER_TYPES = {"attention", "full_attention"}
@@ -616,7 +619,11 @@ class NemotronHModel(GraniteHybridModel):
         vocab_size = -(vocab_size // -pad_vocab) * pad_vocab
         self.hparams["vocab_size"] = vocab_size
 
-        assert max(tokenizer.vocab.values()) < vocab_size  # ty: ignore[unresolved-attribute]
+        # draft-model converters (DFlash, LiLiCorr) call this with their own instance, which has neither hook
+        vocab_skip = getattr(self, "_vocab_skip", ())
+        user_defined_tokens = getattr(self, "_user_defined_tokens", None)
+
+        assert max(id_ for tok, id_ in tokenizer.vocab.items() if tok not in vocab_skip) < vocab_size  # ty: ignore[unresolved-attribute]
 
         tokpre = self.get_vocab_base_pre(tokenizer)
 
@@ -638,7 +645,8 @@ class NemotronHModel(GraniteHybridModel):
                         if previous_token != token:
                             logger.info(f"{repr(previous_token)} is encoded and decoded back to {repr(token)} using AutoTokenizer")
 
-                    if added_tokens_decoder[i].special or self.does_token_look_special(token):
+                    printed = user_defined_tokens is not None and user_defined_tokens.fullmatch(token)
+                    if (added_tokens_decoder[i].special or self.does_token_look_special(token)) and not printed:
                         toktypes.append(gguf.TokenType.CONTROL)
                     else:
                         token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ")  # pre-normalize user-defined spaces
@@ -834,3 +842,140 @@ class NemotronHPuzzleModel(NemotronHModel):
         if item[0].startswith("mtp."):
             return None
         return super().filter_tensors(item)
+
+
+# NeMo speechlm2 SALM checkpoint: a Nemotron-H decoder under llm.*, the speech encoder under perception.*
+@ModelBase.register("Nemotron3_5TranscribeForConditionalGeneration")
+@ModelBase.example("nvidia/NVIDIA-Nemotron-3.5-Transcribe-30B-A4B")
+class NemotronTranscribeModel(NemotronHModel):
+    model_arch = gguf.MODEL_ARCH.NEMOTRON_H
+    # the audio placeholder (id vocab_size) has no embedding row; speaker, event and style tags print with the transcript
+    _vocab_skip = ("<|audio|>",)
+    _user_defined_tokens = re.compile(r"<spk:\d+>|<ev:[a-z_]+>|<sty:(?:stt|end):[a-z_]+>")
+    _renames = {
+        "llm.model.embed_tokens.weight": "backbone.embeddings.weight",
+        "llm.model.norm.weight":         "backbone.norm_f.weight",
+        "llm.lm_head.weight":            "lm_head.weight",
+    }
+    # grouped experts: [n_expert, n_in, n_out] -> [n_expert, n_out, n_in]
+    _grouped_experts = {"gate_and_up_projs": "up_proj", "down_projs": "down_proj"}
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if name.startswith(("perception.", "ctc_timestamp.")) or "_extra_state" in name:
+            return None
+        if not name.startswith("llm."):
+            # the FP8 and NVFP4 releases store the decoder as backbone.*
+            raise ValueError(f"unexpected tensor {name}: only the BF16 checkpoint layout (llm.*) is supported")
+        if name in cls._renames:
+            name = cls._renames[name]
+        elif name.startswith("llm.model.layers."):
+            name = "backbone.layers." + name[len("llm.model.layers."):]
+        elif name.startswith("llm.mtp.layers."):
+            name = "mtp.layers." + name[len("llm.mtp.layers."):]
+        else:
+            raise ValueError(f"unexpected decoder tensor {name}")
+        return super().filter_tensors((name.replace("._fp32_params.", "."), gen))
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        leaf = name.rsplit(".", 1)[-1]
+        if ".mixer.experts." in name and leaf in self._grouped_experts:
+            if name.startswith("mtp."):
+                assert self._mtp_bid is not None, f"{name}: MTP head is not exported"
+                bid = self._mtp_bid
+            assert bid is not None
+            data_torch = data_torch.permute(0, 2, 1).contiguous()
+            name = f"model.layers.{bid}.mlp.experts.{self._grouped_experts[leaf]}.weight"
+            yield from ModelBase.modify_tensors(self, data_torch, name, bid)
+            return
+        yield from super().modify_tensors(data_torch, name, bid)
+
+
+@ModelBase.register("Nemotron3_5TranscribeForConditionalGeneration")
+class NemotronTranscribeAudioModel(MmprojModel):
+    # NeMo ParallelExpertEncoder: ASR and speaker transformers, speaker activity fused into the ASR states
+    has_vision_encoder = False
+    has_audio_encoder = True
+
+    def get_audio_config(self) -> dict[str, Any] | None:
+        asr = self.global_config["pe_encoder_config"]["asr_encoder_cfg"]
+        return {
+            "hidden_size":         asr["d_model"],
+            "intermediate_size":   int(float(asr.get("ff_expansion", 4)) * asr["d_model"]),
+            "num_hidden_layers":   asr["n_layers"],
+            "num_attention_heads": asr["n_heads"],
+        }
+
+    @staticmethod
+    def _rope_dims(cfg: dict[str, Any]) -> int:
+        assert cfg.get("self_attention_model") == "rope" and cfg.get("subsampling", "feature_stacking") == "feature_stacking"
+        assert cfg.get("attn_mode", "full") == "full" and not cfg.get("xscaling", False) and cfg.get("pre_block_norm", True)
+        assert float(cfg.get("rope_base", 10000.0)) == 10000.0
+        n_rot = int(cfg["d_model"] // cfg["n_heads"] * float(cfg.get("rotary_fraction", 1.0)))
+        assert n_rot >= 2 and n_rot % 2 == 0
+        return n_rot
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        pe = self.global_config["pe_encoder_config"]
+        asr = pe["asr_encoder_cfg"]
+        diar = pe["diarization_model_cfg"]
+        spk = diar["encoder"]
+        assert pe.get("asr_encoder_type") == "transformer" and pe.get("speaker_feature_mode", "thresholded") == "thresholded"
+        assert pe.get("asr_normalize_type") == "per_feature" and pe.get("diar_normalize_type") in (None, "NA")
+        assert asr["subsampling_factor"] == spk["subsampling_factor"] == 8 and asr["feat_in"] == spk["feat_in"]
+        assert self._rope_dims(spk) == spk["d_model"] // spk["n_heads"] and not spk.get("qkv_bias", False)
+        assert diar.get("high_resolution", False)
+        # constants in clip: the ParallelExpertEncoder online settings and the Sortformer streaming config
+        assert (pe.get("online_inference_length", 500), pe.get("chunk_left_context", 50), pe.get("chunk_right_context", 50),
+                pe.get("diar_spkcache_len", 188), pe.get("diar_fifo_len", 40), pe.get("diar_spkcache_update_period", 300),
+                pe.get("speaker_activity_threshold", 0.5)) == (500, 50, 50, 188, 40, 300, 0.5), \
+            "streaming settings differ from the constants in tools/mtmd/models/nemotron-transcribe.cpp"
+        sm = diar["sortformer_modules"]
+        assert (sm["spkcache_sil_frames_per_spk"], sm["pred_score_threshold"], sm["scores_boost_latest"], sm["strong_boost_rate"],
+                sm["weak_boost_rate"], sm["min_pos_scores_rate"], sm["use_learnable_sil_emb"]) == (1, 0.25, 0.05, 0.75, 1.5, 0.5, True), \
+            "speaker cache settings differ from the constants in tools/mtmd/models/nemotron-transcribe.cpp"
+
+        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.NEMOTRON_TRANSCRIBE)
+        self.gguf_writer.add_audio_num_mel_bins(asr["feat_in"])
+        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
+        self.gguf_writer.add_audio_subsampling_factor(asr["subsampling_factor"])
+        self.gguf_writer.add_audio_rope_dims(self._rope_dims(asr))
+        self.gguf_writer.add_audio_spk_block_count(spk["n_layers"])
+        self.gguf_writer.add_audio_spk_head_count(spk["n_heads"])
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # ParallelExpertEncoder._build_sinusoid_position_encoding
+        pe = self.global_config["pe_encoder_config"]
+        n_spk = int(pe["diarization_model_cfg"]["sortformer_modules"]["num_spks"])
+        d = int(pe["asr_encoder_cfg"]["d_model"])
+        scale = float(pe.get("spk_kernel_scale", 1.0))
+        pos = torch.arange(n_spk, dtype=torch.float32)[:, None]
+        div = torch.exp(torch.arange(0, d, 2, dtype=torch.float32) * -(np.log(10000.0) / d))
+        kernel = torch.zeros(n_spk, d, dtype=torch.float32)
+        kernel[:, 0::2] = torch.sin(pos * div)
+        kernel[:, 1::2] = torch.cos(pos * div)
+        # stored as [d, n_spk]: the graph multiplies it with the speaker activity directly
+        yield self.format_tensor_name(gguf.MODEL_TENSOR.A_FUSION_SPK_KERNEL, suffix=""), (kernel * scale).T.contiguous()
+
+    def tensor_force_quant(self, name, new_name, bid, n_dims):
+        if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.A_SPK_HEAD_UPSAMPLE, None):
+            return gguf.GGMLQuantizationType.F16  # ggml_conv_1d kernel
+        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
+        if not name.startswith("perception.") or "_extra_state" in name:
+            return None
+        # training-only head; the speaker model's unused front end
+        if name.startswith(("perception.encoder.diarization_model.sortformer_modules.hidden_to_spks.",
+                            "perception.encoder.diarization_model.preprocessor.")):
+            return None
+        return name, gen
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name == "perception.preprocessor.featurizer.fb":
+            data_torch = data_torch.squeeze(0)
+        yield from super().modify_tensors(data_torch, name, bid)

convert_hf_to_gguf_update.py

@@ -198,6 +198,7 @@ pre_computed_hashes = [
     # 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"},
+    {"name": "gpt-2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/Lernex/Metis-1.6-Logos", "chkhsh": "df1bc34a6366466bcb36a9a1b1cc0afdb08cfe21c9250602c930f159085390a9"},
     # lfm2 variants
     {"name": "lfm2", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/LiquidAI/LFM2.5-8B-A1B", "chkhsh": "9e454714343b69b99b71795c1d27a68c2a1d15dab111f4d353109f966af29da7"},
     # hrm-text (DFM Mimir) is SPM-style BPE: normalizer maps ' ' -> '▁', merges

ggml/include/ggml-backend.h

@@ -211,6 +211,18 @@ extern "C" {
 
     // Context management and operations for faster communication between backends, used for tensor parallelism (meta backend)
     typedef void * (*ggml_backend_comm_init_t)(ggml_backend_t * backends, size_t n_backends);
+
+    // Builds the communicator for the case where each rank is a separate process, so a process
+    // owns one backend of the communicator rather than all of them.
+    //
+    // The caller provides bootstrap_callback, which moves bytes between the ranks before the
+    // communicator exists. Each rank passes size bytes in send_buf, and recv_buf receives
+    // world_size * size bytes in rank order. The backend decides what to exchange and how many
+    // times to call it, and every rank must make the same calls in the same order.
+    // comm_init_rank blocks until all world_size ranks have joined.
+    typedef bool   (*ggml_backend_comm_bootstrap_callback)(void * user_data, const void * send_buf, void * recv_buf, size_t size);
+    typedef void * (*ggml_backend_comm_init_rank_t)(ggml_backend_t backend, int rank, int world_size, ggml_backend_comm_bootstrap_callback bootstrap_callback, void * bootstrap_callback_data);
+
     typedef void   (*ggml_backend_comm_free_t)(void * comm_ctx);
     typedef bool   (*ggml_backend_comm_allreduce_tensor_t)(void * comm_ctx, struct ggml_tensor ** tensors);
 

ggml/src/ggml-backend-meta.cpp

@@ -489,6 +489,32 @@ static struct ggml_tensor * ggml_backend_meta_buffer_simple_tensor(const struct
     return it->second[index];
 }
 
+// bytes spanned by all dims except skip_dim
+static size_t ggml_backend_meta_inner_extent(const struct ggml_tensor * tensor, int skip_dim) {
+    size_t extent = 0;
+    for (int i = 0; i < GGML_MAX_DIMS; i++) {
+        if (i != skip_dim && tensor->ne[i] > 1) {
+            extent += (size_t) (tensor->ne[i] - 1) * tensor->nb[i];
+        }
+    }
+    return extent;
+}
+
+// the view dim that holds the split axis of src, or -1
+// all other dims of src and view must stay inside one slab of that axis
+static int ggml_backend_meta_view_slab_dim(const struct ggml_tensor * view, const struct ggml_tensor * src, int axis) {
+    const size_t slab = src->nb[axis];
+    if (ggml_backend_meta_inner_extent(src, axis) >= slab) {
+        return -1;
+    }
+    for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
+        if (view->ne[dim] == src->ne[axis] && view->nb[dim] == slab) {
+            return ggml_backend_meta_inner_extent(view, dim) < slab ? dim : -1;
+        }
+    }
+    return -1;
+}
+
 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync);
 
 static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
@@ -723,6 +749,16 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
         if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED || src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
             return src_ss[0];
         }
+        // a view that only regroups dims inside one slab of the split axis keeps the split,
+        // the per-device strides are only rescaled above nb[axis] (e.g. the KV streams of a permuted q)
+        if (axis >= 0 && axis < GGML_MAX_DIMS && tensor->view_offs == tensor->src[0]->view_offs) {
+            const int dim = ggml_backend_meta_view_slab_dim(tensor, tensor->src[0], axis);
+            if (dim >= 0) {
+                ggml_backend_meta_split_state ret = src_ss[0];
+                ret.axis = ggml_backend_meta_split_axis(dim);
+                return ret;
+            }
+        }
         GGML_ABORT("view of permuted tensor not implemented");
         //return {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1};
     };
@@ -1020,8 +1056,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
                 split_state = handle_pad(src_ss);
             } break;
             case GGML_OP_PAD_REFLECT_1D:
-            case GGML_OP_ROLL:
-            case GGML_OP_ARANGE:
+            case GGML_OP_ROLL: {
+                split_state = handle_generic(src_ss, /*scalar_only =*/ true);
+            } break;
+            case GGML_OP_ARANGE: {
+                // no sources, every device computes the full range
+                split_state = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+            } break;
             case GGML_OP_TIMESTEP_EMBEDDING: {
                 split_state = handle_generic(src_ss, /*scalar_only =*/ true);
             } break;

ggml/src/ggml-cuda/batch-get.cu

@@ -0,0 +1,198 @@
+#include "batch-get.cuh"
+#include "common.cuh"
+#include "ggml-cpp.h"
+
+#include <climits>
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+struct ggml_cuda_batch_get_descriptor {
+    const unsigned char * src;
+    unsigned char * dst;
+    size_t size;
+};
+
+static __global__ void ggml_cuda_batch_get_kernel(const ggml_cuda_batch_get_descriptor * copies, size_t count) {
+    // Blocks along y cooperate on each descriptor.
+    const size_t index = blockIdx.x;
+    if (index >= count) {
+        return;
+    }
+    const size_t first = (size_t) blockIdx.y * blockDim.x + threadIdx.x;
+    const size_t stride = (size_t) gridDim.y * blockDim.x;
+    const auto copy = copies[index];
+    // Use 16-, 8-, 4-, or 2-byte copies when both pointers meet that alignment; threads copy separate chunks.
+    const uintptr_t alignment = (uintptr_t) copy.src | (uintptr_t) copy.dst;
+    size_t tail = 0;
+    if (alignment % 16 == 0) {
+        for (size_t i = first; i < copy.size / 16; i += stride) {
+            ggml_cuda_memcpy_1<16>(copy.dst + 16*i, copy.src + 16*i);
+        }
+        tail = copy.size / 16 * 16;
+    } else if (alignment % 8 == 0) {
+        for (size_t i = first; i < copy.size / 8; i += stride) {
+            ggml_cuda_memcpy_1<8>(copy.dst + 8*i, copy.src + 8*i);
+        }
+        tail = copy.size / 8 * 8;
+    } else if (alignment % 4 == 0) {
+        for (size_t i = first; i < copy.size / 4; i += stride) {
+            ggml_cuda_memcpy_1<4>(copy.dst + 4*i, copy.src + 4*i);
+        }
+        tail = copy.size / 4 * 4;
+    } else if (alignment % 2 == 0) {
+        for (size_t i = first; i < copy.size / 2; i += stride) {
+            ggml_cuda_memcpy_1<2>(copy.dst + 2*i, copy.src + 2*i);
+        }
+        tail = copy.size / 2 * 2;
+    }
+    // Copy trailing bytes, or the entire range when no wider alignment is available.
+    for (size_t i = tail + first; i < copy.size; i += stride) {
+        ggml_cuda_memcpy_1<1>(copy.dst + i, copy.src + i);
+    }
+}
+
+struct ggml_cuda_batch_get_mapping {
+    void * host = nullptr;
+    unsigned char * device = nullptr;
+
+    ggml_cuda_batch_get_mapping(void * host, size_t size) {
+        const cudaError_t err = cudaHostRegister(host, size, cudaHostRegisterMapped | cudaHostRegisterPortable);
+        if (err != cudaSuccess) {
+            (void) cudaGetLastError();
+            GGML_LOG_DEBUG("%s: failed to register %.2f MiB of pinned memory: %s\n", __func__,
+                           size / 1024.0 / 1024.0, cudaGetErrorString(err));
+            return;
+        }
+        this->host = host;
+        if (cudaHostGetDevicePointer((void **) &device, host, 0) != cudaSuccess) {
+            (void) cudaGetLastError();
+            device = nullptr;
+            return;
+        }
+    }
+
+    ggml_cuda_batch_get_mapping(const ggml_cuda_batch_get_mapping &) = delete;
+    ggml_cuda_batch_get_mapping & operator=(const ggml_cuda_batch_get_mapping &) = delete;
+
+    ~ggml_cuda_batch_get_mapping() {
+        if (host) {
+            // Release only the registration acquired by this object.
+            CUDA_CHECK(cudaHostUnregister(host));
+        }
+    }
+};
+#endif
+
+bool ggml_cuda_get_tensor_n(int device,
+        const ggml_tensor * const * tensors, void * const * data, const size_t * offsets,
+        const size_t * sizes, size_t n_tensors) {
+#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
+    GGML_UNUSED(device);
+    GGML_UNUSED(tensors);
+    GGML_UNUSED(data);
+    GGML_UNUSED(offsets);
+    GGML_UNUSED(sizes);
+    GGML_UNUSED(n_tensors);
+    return false;
+#else
+    // Provisional dispatch values based on experiments.
+    constexpr size_t min_batch_reads = 128;
+    constexpr size_t max_gather_read_bytes = 128 * 1024;
+    // Collect non-empty CPU destination ranges.
+    std::vector<std::pair<uintptr_t, uintptr_t>> ranges;
+    ranges.reserve(n_tensors);
+    size_t payload_bytes = 0;
+    // Number of non-empty reads whose size is at most max_gather_read_bytes.
+    size_t n_reads_small = 0;
+    for (size_t i = 0; i < n_tensors; ++i) {
+        if (sizes[i] == 0) {
+            continue;
+        }
+        const uintptr_t address = (uintptr_t) data[i];
+        GGML_ASSERT(address <= UINTPTR_MAX - sizes[i]);
+        GGML_ASSERT(payload_bytes <= SIZE_MAX - sizes[i]);
+        payload_bytes += sizes[i];
+        n_reads_small += sizes[i] <= max_gather_read_bytes;
+        ranges.emplace_back(address, address + sizes[i]);
+    }
+    if (ranges.empty()) {
+        return true;
+    }
+    // Total number of non-empty reads.
+    const size_t n_reads = ranges.size();
+    if (n_reads > SIZE_MAX / sizeof(ggml_cuda_batch_get_descriptor) || n_reads > INT_MAX) {
+        return false;
+    }
+    if (n_reads < 2 || (n_reads < min_batch_reads && n_reads_small == n_reads)) {
+        return false;
+    }
+    std::sort(ranges.begin(), ranges.end());
+    for (size_t i = 1; i < ranges.size(); ++i) {
+        // Preserve the ordered copy path when CPU destinations overlap.
+        if (ranges[i].first < ranges[i - 1].second) {
+            return false;
+        }
+    }
+    // Find the CPU span to register, including gaps between destinations.
+    const uintptr_t begin = ranges.front().first;
+    const uintptr_t end = ranges.back().second;
+    if (end - begin - payload_bytes > payload_bytes) {
+        return false;
+    }
+    // Register and map the existing CPU destination so the kernel can write to it directly.
+    ggml_cuda_batch_get_mapping mapping((void *) begin, end - begin);
+    if (!mapping.host) {
+        return false;
+    }
+    // Reads selected for gather; zero means all reads use DMA.
+    size_t n_reads_gather = mapping.device && n_reads_small >= min_batch_reads ? n_reads_small : 0;
+    // Keep descriptor storage alive until the final stream synchronization.
+    constexpr size_t max_batch_descriptors = 128 * 1024;
+    const size_t bytes = std::min(n_reads_gather, max_batch_descriptors) * sizeof(ggml_cuda_batch_get_descriptor);
+    ggml_backend_buffer_ptr descriptors(n_reads_gather ? ggml_backend_buft_alloc_buffer(ggml_backend_cuda_buffer_type(device), bytes) : nullptr);
+    if (!descriptors) {
+        n_reads_gather = 0;
+    }
+    const auto host_type = ggml_backend_cuda_host_buffer_type();
+    ggml_backend_buffer_ptr host_descriptors(n_reads_gather ? ggml_backend_buft_alloc_buffer(host_type, bytes) : nullptr);
+    if (!host_descriptors || ggml_backend_buffer_get_type(host_descriptors.get()) != host_type) {
+        n_reads_gather = 0;
+    }
+    if (n_reads_gather) {
+        auto * copies = static_cast<ggml_cuda_batch_get_descriptor *>(ggml_backend_buffer_get_base(host_descriptors.get()));
+        auto * device_copies = static_cast<ggml_cuda_batch_get_descriptor *>(ggml_backend_buffer_get_base(descriptors.get()));
+        const int threads = 256;
+        const size_t nsm = ggml_cuda_info().devices[device].nsm;
+        size_t i = 0;
+        for (size_t first = 0; first < n_reads_gather; first += max_batch_descriptors) {
+            const size_t n_batch = std::min(n_reads_gather - first, max_batch_descriptors);
+            size_t count = 0;
+            for (; i < n_tensors && count < n_batch; ++i) {
+                if (sizes[i] == 0 || sizes[i] > max_gather_read_bytes) {
+                    continue;
+                }
+                copies[count++] = {(const unsigned char *) tensors[i]->data + offsets[i],
+                        mapping.device + ((uintptr_t) data[i] - begin), sizes[i]};
+            }
+            const size_t blocks_per_copy = std::min<size_t>(UINT16_MAX, std::max<size_t>(1, (nsm + n_batch - 1) / n_batch));
+            const dim3 blocks((unsigned int) n_batch, (unsigned int) blocks_per_copy);
+            // Stream ordering keeps each kernel ahead of the next descriptor overwrite.
+            CUDA_CHECK(cudaMemcpyAsync(device_copies, copies, n_batch * sizeof(copies[0]), cudaMemcpyHostToDevice, cudaStreamPerThread));
+            ggml_cuda_batch_get_kernel<<<blocks, threads, 0, cudaStreamPerThread>>>(device_copies, n_batch);
+            CUDA_CHECK(cudaGetLastError());
+            if (first + n_batch < n_reads_gather) {
+                // Wait before refilling the pinned host descriptor buffer.
+                CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+            }
+        }
+    }
+    // Keep the destination pinned for DMA too; overlapping destinations already declined above.
+    for (size_t i = 0; i < n_tensors; ++i) {
+        if (sizes[i] && (!n_reads_gather || sizes[i] > max_gather_read_bytes)) {
+            CUDA_CHECK(cudaMemcpyAsync(data[i], (const char *) tensors[i]->data + offsets[i], sizes[i], cudaMemcpyDeviceToHost, cudaStreamPerThread));
+        }
+    }
+    // Wait for GPU writes before releasing the CPU mapping and descriptor buffers.
+    CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread));
+    return true;
+#endif
+}

ggml/src/ggml-cuda/batch-get.cuh

@@ -0,0 +1,7 @@
+#pragma once
+
+#include "ggml-backend.h"
+
+bool ggml_cuda_get_tensor_n(int device,
+        const ggml_tensor * const * tensors, void * const * data, const size_t * offsets,
+        const size_t * sizes, size_t n_tensors);

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -9,6 +9,7 @@
 #include "ggml-cuda/arange.cuh"
 #include "ggml-cuda/argmax.cuh"
 #include "ggml-cuda/argsort.cuh"
+#include "ggml-cuda/batch-get.cuh"
 #include "ggml-cuda/binbcast.cuh"
 #include "ggml-cuda/clamp.cuh"
 #include "ggml-cuda/col2im-1d.cuh"
@@ -886,6 +887,11 @@ static void ggml_backend_cuda_buffer_get_tensor_n(ggml_backend_buffer_t buffer,
 
     ggml_cuda_set_device(ctx->device);
 
+    if (!ggml_cuda_pinned_memory_disabled() &&
+            ggml_cuda_get_tensor_n(ctx->device, tensors, data, offsets, sizes, n_tensors)) {
+        return;
+    }
+
     for (size_t i = 0; i < n_tensors; i++) {
         if (sizes[i] == 0) {
             continue;
@@ -1117,6 +1123,18 @@ struct ggml_backend_cuda_comm_context {
 
     ggml_cuda_ar_pipeline *     ar_pipeline = nullptr;
 
+    // Number of ranks in the communicator when each rank is a separate process. Zero unless
+    // comm_init_rank created this context.
+    int                         world_size = 0;
+
+    // Number of ranks taking part in the reduction. For a communicator from comm_init_rank this is
+    // world_size. Otherwise it is the number of backends this process holds. The NCCL reduction
+    // uses it to choose the tensor size below which it reduces in F32 rather than converting to
+    // BF16 first.
+    size_t n_ranks() const {
+        return world_size > 0 ? (size_t) world_size : backends.size();
+    }
+
 #ifdef GGML_USE_NCCL
     std::vector<ncclComm_t>     comms;
 #endif // GGML_USE_NCCL
@@ -1144,6 +1162,7 @@ static bool ggml_backend_cuda_comm_allreduce_nccl(
     }
 
     const size_t n_backends = comm_ctx->backends.size();
+    const size_t n_ranks    = comm_ctx->n_ranks();
 
     for (size_t i = 0; i < n_backends; ++i) {
         GGML_ASSERT(tensors[i] != nullptr);
@@ -1153,7 +1172,7 @@ static bool ggml_backend_cuda_comm_allreduce_nccl(
 
     // For small tensors, simply reduce them as FP32.
     // The following heuristic for how "small" a tensor should be is based on RTX 4090s connected via 16x PCIe 4.0.
-    if ((n_backends <= 2 && ne < 32768) || (n_backends == 3 && ne < 131072) || (n_backends >= 4 && ne < 262144)) {
+    if ((n_ranks <= 2 && ne < 32768) || (n_ranks == 3 && ne < 131072) || (n_ranks >= 4 && ne < 262144)) {
         for (size_t i = 0; i < n_backends; ++i) {
             if ((tensors[i]->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) {
                 ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) comm_ctx->backends[i]->context;
@@ -1384,6 +1403,90 @@ static void * ggml_backend_cuda_comm_init(ggml_backend_t * backends, size_t n_ba
     return ret;
 }
 
+#ifdef GGML_USE_NCCL
+// Builds the communicator with ncclCommInitRank, for the case where each rank is its own
+// process. Blocks until all world_size ranks have joined.
+static void * ggml_backend_cuda_comm_init_rank(ggml_backend_t backend, int rank, int world_size,
+                                               ggml_backend_comm_bootstrap_callback bootstrap_callback, void * bootstrap_callback_data) {
+    if (world_size < 1 || rank < 0 || rank >= world_size) {
+        GGML_LOG_ERROR("%s: invalid rank %d for world size %d\n", __func__, rank, world_size);
+        return nullptr;
+    }
+
+    // Every rank sends whether its local checks passed, and rank 0 also sends the NCCL unique id.
+    // So all ranks either join together or fail together, and no rank waits in ncclCommInitRank alone.
+    struct bootstrap_msg {
+        uint8_t      ok;
+        ncclUniqueId id;
+    };
+    bootstrap_msg local;
+    memset(&local, 0, sizeof(local));
+    local.ok = 1;
+
+    if (!ggml_backend_is_cuda(backend)) {
+        GGML_LOG_ERROR("%s: backend is not a CUDA backend\n", __func__);
+        local.ok = 0;
+    } else {
+        // NCCL requires one distinct physical GPU per rank, see ggml_backend_cuda_comm_init_nccl.
+        const ggml_cuda_device_info & info = ggml_cuda_info();
+        if (info.device_count > info.physical_device_count) {
+            GGML_LOG_ERROR("%s: CUDA virtual devices are in use, which NCCL does not support\n", __func__);
+            local.ok = 0;
+        }
+    }
+    if (local.ok && rank == 0) {
+        const ncclResult_t rc = ncclGetUniqueId(&local.id);
+        if (rc != ncclSuccess) {
+            GGML_LOG_ERROR("%s: ncclGetUniqueId failed (%s)\n", __func__, ncclGetErrorString(rc));
+            local.ok = 0;
+        }
+    }
+
+    std::vector<bootstrap_msg> all(world_size);
+    if (!bootstrap_callback(bootstrap_callback_data, &local, all.data(), sizeof(local))) {
+        GGML_LOG_ERROR("%s: bootstrap exchange failed on rank %d\n", __func__, rank);
+        return nullptr;
+    }
+    for (int i = 0; i < world_size; i++) {
+        if (!all[i].ok) {
+            GGML_LOG_ERROR("%s: rank %d failed its setup checks\n", __func__, i);
+            return nullptr;
+        }
+    }
+
+    ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context;
+
+    auto * ret = new ggml_backend_cuda_comm_context;
+    ret->backends.push_back(backend);
+    ret->dev_ids.push_back(cuda_ctx->device);
+    ret->world_size = world_size;
+    ret->comms.resize(1);
+    ret->try_allreduce = ggml_backend_cuda_comm_try_allreduce_nccl;
+
+    // On DGX Spark, NCCL picks 64 channels by default, which is slower and uses more memory than 8.
+    // NCCL_MAX_CTAS still overrides this.
+    ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
+    const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc;
+    if (GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_DGX_SPARK) {
+        config.maxCTAs = 8;
+    }
+
+    ggml_cuda_set_device(cuda_ctx->device);
+    const ncclResult_t rc = ncclCommInitRankConfig(&ret->comms[0], world_size, all[0].id, rank, &config);
+    if (rc != ncclSuccess) {
+        GGML_LOG_ERROR("%s: ncclCommInitRankConfig failed for rank %d/%d (%s)\n",
+                       __func__, rank, world_size, ncclGetErrorString(rc));
+        ret->comms.clear();
+        delete ret;
+        return nullptr;
+    }
+
+    GGML_LOG_INFO("%s: joined NCCL communicator as rank %d/%d on device %d\n",
+                  __func__, rank, world_size, cuda_ctx->device);
+    return ret;
+}
+#endif // GGML_USE_NCCL
+
 // Top-level dispatch -- calls the function pointer chosen by comm_init.
 // Returns false to let the meta-backend's butterfly run.
 static bool ggml_backend_cuda_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tensor ** tensors) {
@@ -6384,6 +6487,11 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con
     if (strcmp(name, "ggml_backend_comm_allreduce_tensor") == 0) {
         return (void *)ggml_backend_cuda_comm_allreduce_tensor;
     }
+#ifdef GGML_USE_NCCL
+    if (strcmp(name, "ggml_backend_comm_init_rank") == 0) {
+        return (void *)ggml_backend_cuda_comm_init_rank;
+    }
+#endif // GGML_USE_NCCL
     if (strcmp(name, "ggml_backend_register_host_buffer") == 0) {
         return (void *)ggml_backend_cuda_register_host_buffer;
     }

ggml/src/ggml-cuda/quantize.cu

@@ -83,21 +83,21 @@ static __global__ void quantize_q8_1(
     ggml_cuda_pdl_sync();
     const float xi = i0 < ne00 ? x[i03*s03 + i02*s02 + i01*s01 + i00] : 0.0f;
     float amax = fabsf(xi);
-    float sum = xi;
 
     amax = warp_reduce_max<QK8_1>(amax);
-    sum  = warp_reduce_sum<QK8_1>(sum);
 
     const float  d = amax / 127.0f;
     const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);
 
+    const int sum = warp_reduce_sum<QK8_1>((int) q);
+
     y[ib].qs[iqs] = q;
 
     if (iqs > 0) {
         return;
     }
 
-    y[ib].ds = make_half2(d, sum);
+    y[ib].ds = make_half2(d, d * sum);
 }
 
 __device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) {
@@ -500,17 +500,6 @@ static __global__ void quantize_mmq_q8_1(
         amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE));
     }
 
-    float sum;
-    if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
-        sum = xi.x + xi.y + xi.z + xi.w;
-
-        // Calculate sums across vals_per_sum/4 threads.
-#pragma unroll
-        for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
-            sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE);
-        }
-    }
-
     const float d_inv = 127.0f / amax;
     char4 q;
     q.x = roundf(xi.x*d_inv);
@@ -519,6 +508,19 @@ static __global__ void quantize_mmq_q8_1(
     q.w = roundf(xi.w*d_inv);
     const float d = 1.0f / d_inv;
 
+    float sum = 0.0f;
+    if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) {
+        int sum_q = int(q.x) + int(q.y) + int(q.z) + int(q.w);
+
+        // Calculate sums across vals_per_sum/4 threads.
+#pragma unroll
+        for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) {
+            sum_q += __shfl_xor_sync(0xFFFFFFFF, sum_q, offset, WARP_SIZE);
+        }
+
+        sum = d * float(sum_q);
+    }
+
     // write the block once (normal) or to each of the token's compact rows (scatter)
     const int nwrite = scatter ? n_expert_used : 1;
 #pragma unroll

ggml/src/ggml-hexagon/ggml-hexagon.cpp

@@ -135,6 +135,7 @@ size_t opt_mbuf_total  = 0;                        // total buffer space limit (
 int    opt_mm_select     = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 int    opt_fa_select     = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported)
 int    opt_fa_split      = 1; // 1 = head-split + team-split when n_kv_heads < N_cores (default), 0 = off
+int    opt_li_select     = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported)
 int    opt_gdn_select    = 2; // 2 = HMX -> HVX, 1 = HVX, 0 = CPU (unsupported)
 int    opt_ar_select     = 2; // 2 = fused ALLREDUCE+ADD (default), 1 = unfused ALLREDUCE, 0 = fallback to CPY+FENCE
 int    opt_ar_scatter    = 1; // 1 = reduce-scatter the fused ALLREDUCE+ADD (default), 0 = full reduction
@@ -3409,6 +3410,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
     const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT");
     const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT");
     const char * str_fa_split  = getenv("GGML_HEXAGON_FA_SPLIT");
+    const char * str_li_select = getenv("GGML_HEXAGON_LI_SELECT");
     const char * str_gdn_select = getenv("GGML_HEXAGON_GDN_SELECT");
     const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT");
     const char * str_ar_scatter = getenv("GGML_HEXAGON_AR_SCATTER");
@@ -3475,6 +3477,7 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) {
     opt_mm_select = str_mm_select ? atoi(str_mm_select)                   : opt_mm_select;
     opt_fa_select = str_fa_select ? atoi(str_fa_select)                   : opt_fa_select;
     opt_fa_split  = str_fa_split  ? atoi(str_fa_split)                   : opt_fa_split;
+    opt_li_select = str_li_select ? atoi(str_li_select)                   : opt_li_select;
     opt_gdn_select    = str_gdn_select    ? atoi(str_gdn_select)          : opt_gdn_select;
     opt_ar_select     = str_ar_select     ? atoi(str_ar_select)           : opt_ar_select;
     opt_ar_scatter    = str_ar_scatter    ? atoi(str_ar_scatter)          : opt_ar_scatter;

ggml/src/ggml-hexagon/htp-common.h

@@ -29,6 +29,7 @@ extern size_t opt_mbuf_total;
 extern int    opt_mm_select;
 extern int    opt_fa_select;
 extern int    opt_fa_split;
+extern int    opt_li_select;
 extern int    opt_gdn_select;
 extern int    opt_ar_select;
 extern int    opt_ar_scatter;

ggml/src/ggml-hexagon/htp-kparams.cpp

@@ -50,9 +50,12 @@ htp_op_code htp_op_remap(const struct ggml_tensor * t) {
         case GGML_OP_ROPE:            return HTP_OP_ROPE;
         case GGML_OP_REPEAT:          return HTP_OP_REPEAT;
         case GGML_OP_CUMSUM:          return HTP_OP_CUMSUM;
+        case GGML_OP_TOP_K_MASK:      return HTP_OP_TOP_K_MASK;
         case GGML_OP_FILL:            return HTP_OP_FILL;
         case GGML_OP_DIAG:            return HTP_OP_DIAG;
         case GGML_OP_SOLVE_TRI:       return HTP_OP_SOLVE_TRI;
+        case GGML_OP_DSV4_HC_PRE:     return HTP_OP_DSV4_HC_PRE;
+        case GGML_OP_DSV4_HC_POST:    return HTP_OP_DSV4_HC_POST;
         case GGML_OP_TRI:             return HTP_OP_TRI;
         case GGML_OP_PAD:             return HTP_OP_PAD;
         case GGML_OP_IM2COL:          return HTP_OP_IM2COL;
@@ -61,6 +64,7 @@ htp_op_code htp_op_remap(const struct ggml_tensor * t) {
         case GGML_OP_CONV_2D_DW:      return HTP_OP_CONV_2D_DW;
         case GGML_OP_ROLL:            return HTP_OP_ROLL;
         case GGML_OP_UPSCALE:         return HTP_OP_UPSCALE;
+        case GGML_OP_LIGHTNING_INDEXER: return HTP_OP_LIGHTNING_INDEXER;
 
         case GGML_OP_UNARY:
             switch (ggml_get_unary_op(t)) {
@@ -3400,6 +3404,194 @@ static bool htp_kparams_solve_cumsum(
     return true;
 }
 
+// Shrink until the layout fits: halve tile_cols down to 1024 columns, then drop threads. False below the VTCM floor.
+static bool htp_top_k_mask_search(struct htp_top_k_mask_vtcm_layout * L, const struct htp_top_k_mask_shape * S, uint32_t kernel_type,
+                                  size_t vtcm_size, uint32_t * n_threads, uint32_t * tile_cols) {
+    const uint32_t min_cols = 1024;
+    while (true) {
+        htp_top_k_mask_vtcm_layout_build(L, S, kernel_type, *n_threads, *tile_cols);
+        if (L->total_bytes <= vtcm_size) {
+            return true;
+        }
+        if (*tile_cols > min_cols) {
+            *tile_cols = (std::max)(min_cols, (uint32_t) hex_round_up(*tile_cols / 2, 64));
+        } else if (*n_threads > 1) {
+            (*n_threads)--;
+        } else {
+            return false;
+        }
+    }
+}
+
+// Row r of t (same ne[1..3] as dst) is at r * nb[1]: the NPU needs one stride per input.
+// A mask broadcasts by i mod ne: its rows are dense above nb[1], the divisors and dst ne[1..2] fit 16 bits.
+static bool htp_top_k_mask_mask_ok(const struct ggml_tensor * m, const struct ggml_tensor * dst) {
+    return m->ne[1] <= UINT16_MAX && m->ne[2] <= UINT16_MAX && m->ne[3] <= UINT16_MAX && dst->ne[1] <= UINT16_MAX && dst->ne[2] <= UINT16_MAX &&
+           (m->ne[2] == 1 || m->nb[2] == m->nb[1] * m->ne[1]) &&
+           (m->ne[3] == 1 || m->nb[3] == m->nb[1] * m->ne[1] * m->ne[2]);
+}
+
+static bool htp_top_k_mask_row_flat(const struct ggml_tensor * t, const struct ggml_tensor * dst) {
+    return t->ne[1] == dst->ne[1] && t->ne[2] == dst->ne[2] && t->ne[3] == dst->ne[3] &&
+           (t->ne[2] == 1 || t->nb[2] == t->nb[1] * t->ne[1]) &&
+           (t->ne[3] == 1 || t->nb[3] == t->nb[1] * t->ne[1] * t->ne[2]);
+}
+
+static bool htp_kparams_solve_top_k_mask(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    struct htp_top_k_mask_kernel_params * kparams
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    const struct ggml_tensor * mask   = op->src[0];
+    const struct ggml_tensor * top_k  = op->src[1];
+    const struct ggml_tensor * score  = op->src[2];
+    const struct ggml_tensor * blocks = op->src[3];
+    const struct ggml_tensor * extra  = op->src[4];
+    const struct ggml_tensor * dst    = op;
+
+    if ((mask->type != GGML_TYPE_F16 && mask->type != GGML_TYPE_F32) || dst->type != mask->type) {
+        return false;
+    }
+    const size_t elem_size = ggml_type_size(dst->type);
+    if (mask->nb[0] != elem_size || dst->nb[0] != elem_size || !ggml_is_contiguous(dst)) {
+        return false;
+    }
+    if (top_k->type != GGML_TYPE_I32 || top_k->nb[0] != 4) {
+        return false;
+    }
+    if (score && (score->type != GGML_TYPE_F32 || score->nb[0] != 4)) {
+        return false;
+    }
+    if (extra && (extra->type != GGML_TYPE_I32 || extra->nb[0] != 4)) {
+        return false;
+    }
+    // flat row addressing: inputs follow the dst rows, a mask with more than one row too
+    if (!htp_top_k_mask_row_flat(top_k, dst) || (score && !htp_top_k_mask_row_flat(score, dst)) ||
+        (extra && !htp_top_k_mask_row_flat(extra, dst)) || (ggml_nrows(mask) > 1 && !htp_top_k_mask_mask_ok(mask, dst))) {
+        return false;
+    }
+    // the kernel scales block ids by blocks->nb[1] with a 16-bit vector multiply
+    if (blocks && (blocks->type != GGML_TYPE_I32 || blocks->nb[0] != 4 || blocks->nb[1] != (size_t) blocks->ne[0] * 4 || blocks->nb[1] > INT16_MAX)) {
+        return false;
+    }
+
+    struct htp_top_k_mask_shape S = {};
+    S.n_kv          = (uint32_t) dst->ne[0];
+    S.elem_size     = (uint32_t) elem_size;
+    S.n_top         = (uint32_t) top_k->ne[0];
+    S.n_blk         = blocks ? (uint32_t) blocks->ne[0] : 0;
+    S.n_pick        = (uint32_t) (blocks ? blocks->ne[1] : score ? score->ne[0] : dst->ne[0]);
+    S.n_extra       = extra ? (uint32_t) extra->ne[0] : 0;
+    const uint32_t total_rows = (uint32_t) ggml_nrows(dst);
+    S.has_score     = score  ? 1 : 0;
+    S.has_blocks    = blocks ? 1 : 0;
+    S.mask_resident = ggml_nrows(mask) == 1 ? 1 : 0;
+
+    // empty inputs would give out-of-range gathers and zero-length DMA rows: leave them to the CPU
+    if (S.n_kv == 0 || total_rows == 0 || S.n_pick == 0 || S.n_top == 0 || (blocks && S.n_blk == 0) || (extra && S.n_extra == 0)) {
+        return false;
+    }
+
+    // PICKS for one row with a resident mask when its single row tile fits, ROWS otherwise
+    const uint32_t sess_threads = (std::max)(1u, htp_session_n_threads(sess));
+    const size_t   vtcm_size    = htp_session_vtcm_size(sess);
+
+    uint32_t kernel_type = HTP_TOP_K_MASK_KERNEL_ROWS;
+    uint32_t n_threads   = sess_threads;
+    uint32_t tile_cols   = 0;
+
+    struct htp_top_k_mask_vtcm_layout L;
+    bool found = false;
+    if (total_rows == 1 && S.mask_resident && sess_threads > 1) {
+        kernel_type = HTP_TOP_K_MASK_KERNEL_PICKS;
+        tile_cols   = hex_round_up(S.n_kv, 64);
+        htp_top_k_mask_vtcm_layout_build(&L, &S, kernel_type, n_threads, tile_cols);
+        found       = L.total_bytes <= vtcm_size;
+    }
+    if (!found) {
+        kernel_type = HTP_TOP_K_MASK_KERNEL_ROWS;
+        n_threads   = (std::max)(1u, (std::min)(sess_threads, total_rows));
+        tile_cols   = hex_round_up(S.n_kv, 64);
+        found       = htp_top_k_mask_search(&L, &S, kernel_type, vtcm_size, &n_threads, &tile_cols);
+    }
+    if (!found) {
+        return false;
+    }
+
+    // compact kernel_params fields: sizes and offsets in HTP_TOP_K_MASK_UNIT bytes as u16, all multiples of the unit
+    const uint64_t off_extra = htp_top_k_mask_rowin_off_extra(S.n_top);
+    const uint64_t off_score = htp_top_k_mask_rowin_off_score(S.n_top, S.n_extra);
+    const uint64_t u16_max   = (uint64_t) UINT16_MAX * HTP_TOP_K_MASK_UNIT;
+    if (L.rowin_bytes > u16_max || L.tile_bytes > u16_max || off_extra > u16_max || off_score > u16_max || L.off_scratch > u16_max) {
+        return false;
+    }
+    const auto units = [](uint64_t bytes) { return (uint16_t) (bytes / HTP_TOP_K_MASK_UNIT); };
+
+    kparams->n_threads    = (uint8_t) n_threads;
+    kparams->kernel_type  = (uint8_t) kernel_type;
+    kparams->flags        = (dst->type == GGML_TYPE_F16 ? HTP_TOP_K_MASK_F16 : 0) |
+                            (score  ? HTP_TOP_K_MASK_HAS_SCORE  : 0) |
+                            (blocks ? HTP_TOP_K_MASK_HAS_BLOCKS : 0) |
+                            (extra  ? HTP_TOP_K_MASK_HAS_EXTRA  : 0) |
+                            (S.mask_resident ? HTP_TOP_K_MASK_MASK_RESIDENT : 0);
+    kparams->elem_size    = (uint8_t) elem_size;
+    kparams->vtcm_size    = (uint32_t) L.total_bytes;
+    kparams->n_kv         = S.n_kv;
+    kparams->n_top        = S.n_top;
+    kparams->n_pick       = S.n_pick;
+    kparams->n_extra      = S.n_extra;
+    kparams->ninf         = dst->type == GGML_TYPE_F16 ? 0xfc00fc00u : 0xff800000u;
+    kparams->blocks_bytes = (uint32_t) L.blocks_bytes;
+    kparams->score_bytes  = score ? S.n_pick * 4 : 0;
+    kparams->tile_cols    = tile_cols;
+    kparams->n_col_tiles  = (S.n_kv + tile_cols - 1) / tile_cols;
+    kparams->n_blk        = (uint16_t) S.n_blk;
+    kparams->blocks_nb1   = (uint16_t) (blocks ? blocks->nb[1] : 0);
+    kparams->off_blocks_u  = units(L.off_blocks);
+    kparams->off_mask_u    = units(L.off_mask);
+    kparams->off_rowin_u   = units(L.off_rowin);
+    kparams->off_tiles_u   = units(L.off_tiles);
+    kparams->off_mtiles_u  = units(L.off_mtiles);
+    kparams->off_scratch_u = units(L.off_scratch);
+    kparams->rowin_units     = units(L.rowin_bytes);
+    kparams->tile_units      = units(L.tile_bytes);
+    kparams->off_extra_units = units(off_extra);
+    kparams->off_score_units = units(off_score);
+
+    if (kernel_type == HTP_TOP_K_MASK_KERNEL_PICKS) {
+        // per thread slices of the static copies and of the flag clear, a multiple of the unit
+        const auto slice = [&](uint64_t n) { return units(hex_round_up((n + n_threads - 1) / n_threads, HTP_TOP_K_MASK_UNIT)); };
+        kparams->u.picks.picks_per_thread = (uint32_t) hex_round_up((S.n_top + n_threads - 1) / n_threads, 32);
+        kparams->u.picks.blocks_slice_u   = slice(L.blocks_bytes);
+        kparams->u.picks.top_slice_u      = slice((uint64_t) S.n_top * 4);
+        kparams->u.picks.extra_slice_u    = slice((uint64_t) S.n_extra * 4);
+        kparams->u.picks.score_slice_u    = slice(kparams->score_bytes);
+        kparams->u.picks.flags_slice_u    = slice(L.tile_bytes);
+    } else {
+        kparams->u.rows.total_rows       = total_rows;
+        kparams->u.rows.rs_top           = (uint32_t) top_k->nb[1];
+        kparams->u.rows.rs_extra         = extra ? (uint32_t) extra->nb[1] : 0;
+        kparams->u.rows.rs_score         = score ? (uint32_t) score->nb[1] : 0;
+        kparams->u.rows.rs_dst           = (uint32_t) dst->nb[1];
+        kparams->u.rows.mask_nb1         = (uint32_t) mask->nb[1];
+        if (S.mask_resident) {
+            kparams->u.rows.m.mask_bytes = (uint32_t) L.mask_bytes;
+        } else {
+            const int64_t d[5] = { dst->ne[1], dst->ne[2], mask->ne[1], mask->ne[2], mask->ne[3] };
+            for (int k = 0; k < 5; k++) {
+                const struct fastdiv_values fd = init_fastdiv_values((uint32_t) d[k]);
+                kparams->u.rows.m.dec.mp[k] = fd.mp;
+                kparams->u.rows.m.dec.l[k]  = (uint8_t) fd.l;
+                kparams->u.rows.m.dec.ne[k] = (uint16_t) d[k];
+            }
+        }
+    }
+
+    return true;
+}
+
 static bool htp_kparams_solve_fill(
     const struct ggml_hexagon_session * sess,
     const struct ggml_tensor * op,
@@ -4141,6 +4333,609 @@ static bool htp_kparams_solve_conv2d_dw(
     return true;
 }
 
+static bool htp_dsv4_hc_pre_can_collapse(
+    const struct ggml_tensor * x,
+    const struct ggml_tensor * weights,
+    const struct ggml_tensor * dst
+) {
+    if (!x || !weights || !dst) {
+        return false;
+    }
+    if (x->type != GGML_TYPE_F32 || weights->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (x->nb[0] != sizeof(float) || weights->nb[0] != sizeof(float) || dst->nb[0] != sizeof(float)) {
+        return false;
+    }
+    if (x->ne[1] != 4 || dst->ne[0] != x->ne[0]) {
+        return false;
+    }
+
+    const bool gated = ggml_get_op_params_i32(dst, 1) != 0;
+    const int64_t x_tokens = x->ne[2] * x->ne[3];
+    const int64_t d_tokens = dst->ne[1] * dst->ne[2] * dst->ne[3];
+    const int64_t w_tokens = gated ? (weights->ne[2] * weights->ne[3]) : (weights->ne[1] * weights->ne[2] * weights->ne[3]);
+
+    if (x_tokens != d_tokens || x_tokens != w_tokens) {
+        return false;
+    }
+
+    if (x->ne[3] == 1 && dst->ne[2] == 1 && dst->ne[3] == 1 &&
+        (gated ? (weights->ne[3] == 1) : (weights->ne[2] == 1 && weights->ne[3] == 1))) {
+        return false;
+    }
+
+    if (x->ne[3] > 1 && x->nb[3] != x->ne[2] * x->nb[2]) return false;
+
+    if (dst->ne[2] > 1 && dst->nb[2] != dst->ne[1] * dst->nb[1]) return false;
+    if (dst->ne[3] > 1 && dst->nb[3] != dst->ne[2] * dst->nb[2]) return false;
+
+    if (gated) {
+        if (weights->ne[3] > 1 && weights->nb[3] != weights->ne[2] * weights->nb[2]) return false;
+    } else {
+        if (weights->ne[2] > 1 && weights->nb[2] != weights->ne[1] * weights->nb[1]) return false;
+        if (weights->ne[3] > 1 && weights->nb[3] != weights->ne[2] * weights->nb[2]) return false;
+    }
+
+    return true;
+}
+
+static void htp_dsv4_hc_pre_collapse_tensors(
+    const struct ggml_tensor * x,
+    const struct ggml_tensor * weights,
+    const struct ggml_tensor * dst,
+    struct ggml_tensor * x_c,
+    struct ggml_tensor * w_c,
+    struct ggml_tensor * d_c
+) {
+    const int64_t n_tokens = x->ne[2] * x->ne[3];
+    const bool gated = ggml_get_op_params_i32(dst, 1) != 0;
+
+    *x_c = *x;
+    x_c->ne[2] = n_tokens;
+    x_c->ne[3] = 1;
+    x_c->nb[3] = x_c->ne[2] * x_c->nb[2];
+
+    *d_c = *dst;
+    d_c->ne[1] = n_tokens;
+    d_c->ne[2] = 1;
+    d_c->ne[3] = 1;
+    d_c->nb[2] = d_c->ne[1] * d_c->nb[1];
+    d_c->nb[3] = d_c->nb[2];
+
+    *w_c = *weights;
+    if (gated) {
+        w_c->ne[2] = n_tokens;
+        w_c->ne[3] = 1;
+        w_c->nb[3] = w_c->ne[2] * w_c->nb[2];
+    } else {
+        w_c->ne[1] = n_tokens;
+        w_c->ne[2] = 1;
+        w_c->ne[3] = 1;
+        w_c->nb[2] = w_c->ne[1] * w_c->nb[1];
+        w_c->nb[3] = w_c->nb[2];
+    }
+}
+
+static bool htp_kparams_solve_dsv4_hc_pre_impl(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * x,
+    const struct ggml_tensor * weights,
+    const struct ggml_tensor * dst,
+    struct htp_dsv4_hc_pre_kernel_params * kparams
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    if (x->type != GGML_TYPE_F32 || weights->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (x->ne[1] != 4 || x->ne[3] != 1) {
+        return false;
+    }
+    if (x->ne[0] % HTP_DSV4_HC_CHUNK_ELEMS != 0) {
+        return false;
+    }
+    if (dst->ne[0] != x->ne[0] || dst->ne[1] != x->ne[2] || dst->ne[2] != 1 || dst->ne[3] != 1) {
+        return false;
+    }
+    if (x->nb[0] != sizeof(float) || weights->nb[0] != sizeof(float) || dst->nb[0] != sizeof(float)) {
+        return false;
+    }
+    if (ggml_is_permuted(x) || ggml_is_permuted(weights) || ggml_is_permuted(dst)) {
+        return false;
+    }
+
+    const bool gated = ggml_get_op_params_i32(dst, 1) != 0;
+    if (gated) {
+        if (weights->ne[0] != x->ne[0] || weights->ne[1] != x->ne[1] || weights->ne[2] != x->ne[2] || weights->ne[3] != 1) {
+            return false;
+        }
+    } else {
+        if (weights->ne[0] != x->ne[1] || weights->ne[1] != x->ne[2] || weights->ne[2] != 1 || weights->ne[3] != 1) {
+            return false;
+        }
+    }
+
+    const uint32_t n_tokens = (uint32_t) x->ne[2];
+    kparams->kernel_type    = gated ? HTP_DSV4_HC_PRE_KERNEL_GATED : HTP_DSV4_HC_PRE_KERNEL_PLAIN;
+
+    kparams->n_chunks   = ((uint32_t) x->ne[0] + HTP_DSV4_HC_CHUNK_ELEMS - 1) / HTP_DSV4_HC_CHUNK_ELEMS;
+    kparams->total_jobs = n_tokens * kparams->n_chunks;
+    if (kparams->total_jobs == 0) {
+        return true;
+    }
+    kparams->n_threads       = (std::min)(htp_session_n_threads(sess), kparams->total_jobs);
+    kparams->jobs_per_thread = (kparams->total_jobs + kparams->n_threads - 1) / kparams->n_threads;
+
+    const size_t vtcm_budget = htp_session_vtcm_size(sess);
+    const uint32_t n_rows    = gated ? HTP_DSV4_HC_PRE_GATED_ROWS : HTP_DSV4_HC_PRE_PLAIN_ROWS;
+    uint32_t block_chunks    = (std::min)((uint32_t) HTP_DSV4_HC_MAX_CHUNKS, kparams->n_chunks);
+    while (block_chunks > 16) {
+        const uint32_t stride = block_chunks * 128;
+        const uint32_t buf_sz = n_rows * stride;
+        if (2 * buf_sz * kparams->n_threads <= vtcm_budget) {
+            break;
+        }
+        block_chunks /= 2;
+    }
+
+    kparams->dma_block_chunks = block_chunks;
+    kparams->dma_row_stride   = block_chunks * 128;
+    kparams->dma_buffer_size  = n_rows * kparams->dma_row_stride;
+    kparams->vtcm_size        = 2 * kparams->dma_buffer_size * kparams->n_threads;
+    if (kparams->vtcm_size > vtcm_budget) {
+        return false;
+    }
+
+    kparams->div_n_chunks  = init_fastdiv_values(kparams->n_chunks);
+    kparams->div_n_threads = init_fastdiv_values(kparams->n_threads);
+
+    return true;
+}
+
+static inline bool htp_kparams_solve_dsv4_hc_pre(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    struct htp_dsv4_hc_pre_kernel_params * kparams
+) {
+    if (htp_dsv4_hc_pre_can_collapse(op->src[0], op->src[1], op)) {
+        struct ggml_tensor x_c, w_c, d_c;
+        htp_dsv4_hc_pre_collapse_tensors(op->src[0], op->src[1], op, &x_c, &w_c, &d_c);
+        return htp_kparams_solve_dsv4_hc_pre_impl(sess, &x_c, &w_c, &d_c, kparams);
+    }
+    return htp_kparams_solve_dsv4_hc_pre_impl(sess, op->src[0], op->src[1], op, kparams);
+}
+
+static bool htp_dsv4_hc_post_can_collapse(
+    const struct ggml_tensor * x,
+    const struct ggml_tensor * residual,
+    const struct ggml_tensor * post,
+    const struct ggml_tensor * comb,
+    const struct ggml_tensor * dst
+) {
+    if (!x || !residual || !post || !dst) {
+        return false;
+    }
+    if (x->type != GGML_TYPE_F32 || residual->type != GGML_TYPE_F32 || post->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (comb && comb->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (x->nb[0] != sizeof(float) || residual->nb[0] != sizeof(float) || post->nb[0] != sizeof(float) || dst->nb[0] != sizeof(float)) {
+        return false;
+    }
+    if (comb && comb->nb[0] != sizeof(float)) {
+        return false;
+    }
+    if (residual->ne[0] != x->ne[0] || residual->ne[1] != 4 || dst->ne[0] != x->ne[0] || dst->ne[1] != 4 || post->ne[0] != 4) {
+        return false;
+    }
+    if (comb && (comb->ne[0] != 4 || comb->ne[1] != 4)) {
+        return false;
+    }
+
+    const int64_t x_tokens = x->ne[1] * x->ne[2] * x->ne[3];
+    const int64_t r_tokens = residual->ne[2] * residual->ne[3];
+    const int64_t p_tokens = post->ne[1] * post->ne[2] * post->ne[3];
+    const int64_t d_tokens = dst->ne[2] * dst->ne[3];
+    if (x_tokens != r_tokens || x_tokens != p_tokens || x_tokens != d_tokens) {
+        return false;
+    }
+    if (comb && (comb->ne[2] * comb->ne[3] != x_tokens)) {
+        return false;
+    }
+
+    if (x->ne[2] == 1 && x->ne[3] == 1 &&
+        residual->ne[3] == 1 &&
+        post->ne[2] == 1 && post->ne[3] == 1 &&
+        dst->ne[3] == 1 &&
+        (!comb || comb->ne[3] == 1)) {
+        return false;
+    }
+
+    if (x->ne[2] > 1 && x->nb[2] != x->ne[1] * x->nb[1]) return false;
+    if (x->ne[3] > 1 && x->nb[3] != x->ne[2] * x->nb[2]) return false;
+
+    if (residual->ne[3] > 1 && residual->nb[3] != residual->ne[2] * residual->nb[2]) return false;
+    if (dst->ne[3] > 1 && dst->nb[3] != dst->ne[2] * dst->nb[2]) return false;
+
+    if (post->ne[2] > 1 && post->nb[2] != post->ne[1] * post->nb[1]) return false;
+    if (post->ne[3] > 1 && post->nb[3] != post->ne[2] * post->nb[2]) return false;
+
+    if (comb && comb->ne[3] > 1 && comb->nb[3] != comb->ne[2] * comb->nb[2]) return false;
+
+    return true;
+}
+
+static void htp_dsv4_hc_post_collapse_tensors(
+    const struct ggml_tensor * x,
+    const struct ggml_tensor * residual,
+    const struct ggml_tensor * post,
+    const struct ggml_tensor * comb,
+    const struct ggml_tensor * dst,
+    struct ggml_tensor * x_c,
+    struct ggml_tensor * r_c,
+    struct ggml_tensor * p_c,
+    struct ggml_tensor * c_c,
+    struct ggml_tensor * d_c
+) {
+    const int64_t n_tokens = x->ne[1] * x->ne[2] * x->ne[3];
+
+    *x_c = *x;
+    x_c->ne[1] = n_tokens;
+    x_c->ne[2] = 1;
+    x_c->ne[3] = 1;
+    x_c->nb[2] = x_c->ne[1] * x_c->nb[1];
+    x_c->nb[3] = x_c->nb[2];
+
+    *r_c = *residual;
+    r_c->ne[2] = n_tokens;
+    r_c->ne[3] = 1;
+    r_c->nb[3] = r_c->ne[2] * r_c->nb[2];
+
+    *d_c = *dst;
+    d_c->ne[2] = n_tokens;
+    d_c->ne[3] = 1;
+    d_c->nb[3] = d_c->ne[2] * d_c->nb[2];
+
+    *p_c = *post;
+    p_c->ne[1] = n_tokens;
+    p_c->ne[2] = 1;
+    p_c->ne[3] = 1;
+    p_c->nb[2] = p_c->ne[1] * p_c->nb[1];
+    p_c->nb[3] = p_c->nb[2];
+
+    if (comb && c_c) {
+        *c_c = *comb;
+        c_c->ne[2] = n_tokens;
+        c_c->ne[3] = 1;
+        c_c->nb[3] = c_c->ne[2] * c_c->nb[2];
+    }
+}
+
+static bool htp_kparams_solve_dsv4_hc_post_impl(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * x,
+    const struct ggml_tensor * residual,
+    const struct ggml_tensor * post,
+    const struct ggml_tensor * comb,
+    const struct ggml_tensor * dst,
+    struct htp_dsv4_hc_post_kernel_params * kparams
+) {
+    memset(kparams, 0, sizeof(*kparams));
+
+    if (x->type != GGML_TYPE_F32 || residual->type != GGML_TYPE_F32 || post->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (comb && comb->type != GGML_TYPE_F32) {
+        return false;
+    }
+    if (x->ne[2] != 1 || x->ne[3] != 1) {
+        return false;
+    }
+    if (x->ne[0] % HTP_DSV4_HC_CHUNK_ELEMS != 0) {
+        return false;
+    }
+    if (residual->ne[0] != x->ne[0] || residual->ne[1] != 4 || residual->ne[2] != x->ne[1] || residual->ne[3] != 1) {
+        return false;
+    }
+    if (post->ne[0] != 4 || post->ne[1] != x->ne[1] || post->ne[2] != 1 || post->ne[3] != 1) {
+        return false;
+    }
+    if (comb && (comb->ne[0] != 4 || comb->ne[1] != 4 || comb->ne[2] != x->ne[1] || comb->ne[3] != 1)) {
+        return false;
+    }
+    if (dst->ne[0] != x->ne[0] || dst->ne[1] != 4 || dst->ne[2] != x->ne[1] || dst->ne[3] != 1) {
+        return false;
+    }
+    if (x->nb[0] != sizeof(float) || residual->nb[0] != sizeof(float) || post->nb[0] != sizeof(float) || dst->nb[0] != sizeof(float)) {
+        return false;
+    }
+    if (comb && comb->nb[0] != sizeof(float)) {
+        return false;
+    }
+    if (ggml_is_permuted(x) || ggml_is_permuted(residual) || ggml_is_permuted(post) || ggml_is_permuted(dst) || (comb && ggml_is_permuted(comb))) {
+        return false;
+    }
+
+    const uint32_t n_tokens = (uint32_t) x->ne[1];
+
+    kparams->has_comb    = comb != nullptr;
+    kparams->kernel_type = kparams->has_comb ? HTP_DSV4_HC_POST_KERNEL_COMB : HTP_DSV4_HC_POST_KERNEL_IDENTITY;
+    kparams->n_chunks    = ((uint32_t) x->ne[0] + HTP_DSV4_HC_CHUNK_ELEMS - 1) / HTP_DSV4_HC_CHUNK_ELEMS;
+    kparams->total_jobs  = n_tokens * kparams->n_chunks;
+    if (kparams->total_jobs == 0) {
+        return true;
+    }
+    kparams->n_threads       = (std::min)(htp_session_n_threads(sess), kparams->total_jobs);
+    kparams->jobs_per_thread = (kparams->total_jobs + kparams->n_threads - 1) / kparams->n_threads;
+
+    const size_t vtcm_budget = htp_session_vtcm_size(sess);
+    uint32_t block_chunks    = (std::min)((uint32_t) HTP_DSV4_HC_MAX_CHUNKS, kparams->n_chunks);
+    while (block_chunks > 16) {
+        const uint32_t stride = block_chunks * 128;
+        const uint32_t buf_sz = HTP_DSV4_HC_POST_ROWS * stride;
+        if (2 * buf_sz * kparams->n_threads <= vtcm_budget) {
+            break;
+        }
+        block_chunks /= 2;
+    }
+
+    kparams->dma_block_chunks = block_chunks;
+    kparams->dma_row_stride   = block_chunks * 128;
+    kparams->dma_buffer_size  = HTP_DSV4_HC_POST_ROWS * kparams->dma_row_stride;
+    kparams->vtcm_size        = 2 * kparams->dma_buffer_size * kparams->n_threads;
+    if (kparams->vtcm_size > vtcm_budget) {
+        return false;
+    }
+
+    kparams->div_n_chunks  = init_fastdiv_values(kparams->n_chunks);
+    kparams->div_n_threads = init_fastdiv_values(kparams->n_threads);
+
+    return true;
+}
+
+static inline bool htp_kparams_solve_dsv4_hc_post(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    struct htp_dsv4_hc_post_kernel_params * kparams
+) {
+    if (htp_dsv4_hc_post_can_collapse(op->src[0], op->src[1], op->src[2], op->src[3], op)) {
+        struct ggml_tensor x_c, r_c, p_c, c_c, d_c;
+        htp_dsv4_hc_post_collapse_tensors(op->src[0], op->src[1], op->src[2], op->src[3], op, &x_c, &r_c, &p_c, op->src[3] ? &c_c : nullptr, &d_c);
+        return htp_kparams_solve_dsv4_hc_post_impl(sess, &x_c, &r_c, &p_c, op->src[3] ? &c_c : nullptr, &d_c, kparams);
+    }
+    return htp_kparams_solve_dsv4_hc_post_impl(sess, op->src[0], op->src[1], op->src[2], op->src[3], op, kparams);
+}
+
+static bool htp_lightning_indexer_can_collapse(
+    const struct ggml_tensor * q,
+    const struct ggml_tensor * k,
+    const struct ggml_tensor * w,
+    const struct ggml_tensor * m,
+    const struct ggml_tensor * dst
+) {
+    if (!q || !k || !w || !m || !dst) {
+        return false;
+    }
+
+    // 2D K normalization: ne[1] is n_kv, ne[2] is 1
+    const bool can_collapse_k = (k->ne[1] > 1 && k->ne[2] == 1);
+
+    // Stream-to-token collapse across q, w, dst (and m)
+    bool can_collapse_stream = false;
+    if (q->ne[3] > 1 && k->ne[3] == 1) {
+        const bool q_ok = (q->nb[3] == (size_t) q->ne[2] * q->nb[2]);
+        const bool dst_ok = (dst->ne[3] == q->ne[3] && dst->ne[2] == 1 && dst->ne[1] == q->ne[2] &&
+                             dst->nb[3] == (size_t) dst->ne[1] * dst->nb[1]);
+        const bool w_ok = (w->ne[3] == q->ne[3] && w->ne[2] == 1 && w->ne[1] == q->ne[2] &&
+                           w->nb[3] == (size_t) w->ne[1] * w->nb[1]) ||
+                          (w->ne[3] == 1 && w->ne[1] == q->ne[2] * q->ne[3]);
+        const bool m_ok = (m->ne[3] == q->ne[3] && m->ne[2] == 1 && m->ne[1] == q->ne[2] &&
+                           m->nb[3] == (size_t) m->ne[1] * m->nb[1]) ||
+                          (m->ne[3] == 1 && m->ne[1] == q->ne[2] * q->ne[3]);
+        can_collapse_stream = q_ok && dst_ok && w_ok && m_ok;
+    }
+
+    return can_collapse_k || can_collapse_stream;
+}
+
+static void htp_lightning_indexer_collapse_tensors(
+    const struct ggml_tensor * q,
+    const struct ggml_tensor * k,
+    const struct ggml_tensor * w,
+    const struct ggml_tensor * m,
+    const struct ggml_tensor * dst,
+    struct ggml_tensor * q_c,
+    struct ggml_tensor * k_c,
+    struct ggml_tensor * w_c,
+    struct ggml_tensor * m_c,
+    struct ggml_tensor * dst_c
+) {
+    *q_c   = *q;
+    *k_c   = *k;
+    *w_c   = *w;
+    *m_c   = *m;
+    *dst_c = *dst;
+
+    if (k->ne[1] > 1 && k->ne[2] == 1) {
+        k_c->ne[2] = k->ne[1];
+        k_c->ne[1] = 1;
+        k_c->nb[1] = k->nb[1];
+        k_c->nb[2] = k->nb[1];
+        k_c->nb[3] = k->nb[2];
+    }
+
+    if (q->ne[3] > 1 && k->ne[3] == 1) {
+        const int64_t n_tokens_collapsed = (int64_t) q->ne[2] * q->ne[3];
+
+        q_c->ne[2] = n_tokens_collapsed;
+        q_c->ne[3] = 1;
+        q_c->nb[3] = (size_t) q_c->ne[2] * q_c->nb[2];
+
+        if (w->ne[3] == q->ne[3]) {
+            w_c->ne[1] = n_tokens_collapsed;
+            w_c->ne[3] = 1;
+            w_c->nb[2] = (size_t) w_c->ne[1] * w_c->nb[1];
+            w_c->nb[3] = w_c->nb[2];
+        }
+
+        dst_c->ne[1] = n_tokens_collapsed;
+        dst_c->ne[3] = 1;
+        dst_c->nb[2] = (size_t) dst_c->ne[1] * dst_c->nb[1];
+        dst_c->nb[3] = dst_c->nb[2];
+
+        if (m->ne[3] == q->ne[3] && m->ne[1] == q->ne[2]) {
+            m_c->ne[1] = n_tokens_collapsed;
+            m_c->ne[3] = 1;
+            m_c->nb[2] = (size_t) m_c->ne[1] * m_c->nb[1];
+            m_c->nb[3] = m_c->nb[2];
+        }
+    }
+}
+
+static bool htp_kparams_solve_lightning_indexer_impl(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * q,
+    const struct ggml_tensor * k,
+    const struct ggml_tensor * w,
+    const struct ggml_tensor * m,
+    const struct ggml_tensor * op,
+    struct htp_lid_kernel_params * kparams
+) {
+    if (opt_li_select < 1) {
+        return false;
+    }
+
+    if (k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_Q8_0 && k->type != GGML_TYPE_Q4_0) {
+        return false;
+    }
+
+    if (q->ne[0] % 32 != 0 ||
+        !ggml_is_contiguous_rows(q) || !ggml_is_contiguous_rows(k) ||
+        !ggml_is_contiguous_rows(w) || !ggml_is_contiguous_rows(m) ||
+        !ggml_is_contiguous_rows(op)) {
+        return false;
+    }
+
+    const uint32_t n_kv_blocks = (k->ne[2] + HTP_LID_KV_BLOCK - 1) / HTP_LID_KV_BLOCK;
+    if (n_kv_blocks == 0 || q->ne[2] == 0) {
+        return false;
+    }
+
+    const uint32_t D          = q->ne[0];
+    const uint32_t n_head     = q->ne[1];
+    const uint32_t n_tokens   = q->ne[2];
+    const size_t   k_row_size = ggml_row_size(k->type, D);
+
+    kparams->k_row_size = (uint32_t) k_row_size;
+    kparams->div_n_head = init_fastdiv_values(n_head);
+    kparams->div_m_ne3  = init_fastdiv_values((uint32_t) m->ne[3]);
+
+    if (opt_li_select >= 2 && htp_session_n_hmx(sess) > 0 && D % HMX_FP16_TILE_N_COLS == 0) {
+        const uint32_t n_threads = htp_session_n_threads(sess);
+        const uint32_t n_stream  = q->ne[3];
+        const uint32_t n_dot     = D / HMX_FP16_TILE_N_COLS;
+
+        auto vtcm_usage = [&](uint32_t n_tok, uint32_t kv_block, uint32_t n_qsets) {
+            return htp_lid_hmx_compute_vtcm_usage(D, n_head, n_tok, kv_block, n_qsets, n_threads, k_row_size);
+        };
+
+        uint32_t n_tok = (std::min)(n_tokens, (uint32_t) HTP_LID_HMX_TOKENS_MAX);
+        while (n_tok > 32 && vtcm_usage(n_tok, 64, 1) > htp_session_vtcm_size(sess)) {
+            n_tok /= 2;
+        }
+
+        uint32_t n_qsets = n_tok == n_tokens ? (std::min)(n_stream, (uint32_t) HTP_LID_HMX_MAX_QSETS) : 1;
+        while (n_qsets > 1 && vtcm_usage(n_tok, 64, n_qsets) > htp_session_vtcm_size(sess)) {
+            n_qsets /= 2;
+        }
+
+        uint32_t n_qbuf = 1;
+        if (n_qsets == 1 && n_head >= 16 && n_stream * ((n_tokens + n_tok - 1) / n_tok) > 1) {
+            uint32_t nt2 = n_tok;
+            while (nt2 > 32 && vtcm_usage(nt2, 64, 2) > htp_session_vtcm_size(sess)) {
+                nt2 /= 2;
+            }
+            const size_t q_bytes = (size_t) nt2 * n_head * D * sizeof(float);
+            const size_t k_bytes = (size_t) k->ne[2] * ggml_row_size(k->type, D);
+            if (vtcm_usage(nt2, 64, 2) <= htp_session_vtcm_size(sess) && q_bytes >= k_bytes) {
+                n_tok  = nt2;
+                n_qbuf = 2;
+            }
+        }
+
+        const uint32_t n_ch   = (n_tok * n_head + 31) / 32;
+        const uint32_t kv_seg = hex_round_up(k->ne[2], 64);
+        const uint32_t n_rows = n_qsets * kv_seg;
+        const uint32_t bc_max = n_rows <= 512 ? n_rows : (std::max)(256u, n_rows / 4);
+        uint32_t kv_block = 64;
+        while (kv_block < HTP_LID_HMX_KV_BLOCK_MAX && kv_block < bc_max && (kv_block / 32) * n_ch * (n_dot + 1) < 256 &&
+               vtcm_usage(n_tok, kv_block * 2, n_qsets * n_qbuf) <= htp_session_vtcm_size(sess)) {
+            kv_block *= 2;
+        }
+
+        const size_t vtcm_size = vtcm_usage(n_tok, kv_block, n_qsets * n_qbuf);
+        if (vtcm_size <= htp_session_vtcm_size(sess)) {
+            const uint32_t n_tt = (n_tokens + n_tok - 1) / n_tok;
+            kparams->kernel_type       = HTP_LID_KERNEL_HMX;
+            kparams->n_threads         = n_threads;
+            kparams->n_tokens_tile     = n_tok;
+            kparams->n_kv_blocks       = (n_rows + kv_block - 1) / kv_block;
+            kparams->blocks_per_thread = 0;
+            kparams->kv_block          = kv_block;
+            kparams->n_qsets           = n_qsets;
+            kparams->n_qbuf            = n_qbuf;
+            kparams->kv_seg            = kv_seg;
+            kparams->n_tt              = n_tt;
+            kparams->vtcm_size         = (uint32_t) vtcm_size;
+            kparams->div_kv_seg        = init_fastdiv_values(kv_seg);
+            kparams->div_n_tt          = init_fastdiv_values(n_tt);
+            kparams->div_n_ch          = init_fastdiv_values(n_ch);
+            return true;
+        }
+    }
+
+    const uint32_t n_threads = (std::min)(htp_session_n_threads(sess), n_kv_blocks);
+    uint32_t n_tok = (std::min)((uint32_t) q->ne[2], (uint32_t) HTP_LID_TOKENS_MAX);
+    while (n_tok > 1 && htp_lid_compute_vtcm_usage(D, n_head, n_tok, n_threads, k_row_size) > htp_session_vtcm_size(sess)) {
+        n_tok /= 2;
+    }
+
+    const size_t vtcm_size = htp_lid_compute_vtcm_usage(D, n_head, n_tok, n_threads, k_row_size);
+    if (vtcm_size > htp_session_vtcm_size(sess)) {
+        return false;
+    }
+
+    kparams->kernel_type       = HTP_LID_KERNEL_HVX;
+    kparams->n_threads         = n_threads;
+    kparams->n_tokens_tile     = n_tok;
+    kparams->n_kv_blocks       = n_kv_blocks;
+    kparams->blocks_per_thread = (n_kv_blocks + n_threads - 1) / n_threads;
+    kparams->kv_block          = HTP_LID_KV_BLOCK;
+    kparams->n_qsets           = 1;
+    kparams->n_qbuf            = 1;
+    kparams->kv_seg            = 0;
+    kparams->n_tt              = 0;
+    kparams->vtcm_size         = (uint32_t) vtcm_size;
+
+    return true;
+}
+
+static bool htp_kparams_solve_lightning_indexer(
+    const struct ggml_hexagon_session * sess,
+    const struct ggml_tensor * op,
+    struct htp_lid_kernel_params * kparams
+) {
+    if (htp_lightning_indexer_can_collapse(op->src[0], op->src[1], op->src[2], op->src[3], op)) {
+        struct ggml_tensor q_c, k_c, w_c, m_c, d_c;
+        htp_lightning_indexer_collapse_tensors(op->src[0], op->src[1], op->src[2], op->src[3], op,
+                                               &q_c, &k_c, &w_c, &m_c, &d_c);
+        return htp_kparams_solve_lightning_indexer_impl(sess, &q_c, &k_c, &w_c, &m_c, &d_c, kparams);
+    }
+    return htp_kparams_solve_lightning_indexer_impl(sess, op->src[0], op->src[1], op->src[2], op->src[3], op, kparams);
+}
 template <auto SolveFn>
 static inline bool htp_kparams_invoke(
     const struct ggml_hexagon_session * sess,
@@ -4252,6 +5047,70 @@ void htp_opnode_collapse(const struct ggml_hexagon_session * sess, htp_opnode &
         }
         return;
     }
+
+    if (node.opcode == HTP_OP_DSV4_HC_PRE) {
+        if (node.inputs.size() < 2 || !node.inputs[0] || !node.inputs[1] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        const struct ggml_tensor * x       = node.inputs[0];
+        const struct ggml_tensor * weights = node.inputs[1];
+        const struct ggml_tensor * dst     = node.outputs[0];
+
+        if (htp_dsv4_hc_pre_can_collapse(x, weights, dst)) {
+            struct ggml_tensor x_c, w_c, d_c;
+            htp_dsv4_hc_pre_collapse_tensors(x, weights, dst, &x_c, &w_c, &d_c);
+            node.inputs[0]  = node.add_dummy(x_c);
+            node.inputs[1]  = node.add_dummy(w_c);
+            node.outputs[0] = node.add_dummy(d_c);
+        }
+        return;
+    }
+
+    if (node.opcode == HTP_OP_DSV4_HC_POST) {
+        if (node.inputs.size() < 3 || !node.inputs[0] || !node.inputs[1] || !node.inputs[2] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        const struct ggml_tensor * x        = node.inputs[0];
+        const struct ggml_tensor * residual = node.inputs[1];
+        const struct ggml_tensor * post     = node.inputs[2];
+        const struct ggml_tensor * comb     = node.inputs.size() > 3 ? node.inputs[3] : nullptr;
+        const struct ggml_tensor * dst      = node.outputs[0];
+
+        if (htp_dsv4_hc_post_can_collapse(x, residual, post, comb, dst)) {
+            struct ggml_tensor x_c, r_c, p_c, c_c, d_c;
+            htp_dsv4_hc_post_collapse_tensors(x, residual, post, comb, dst, &x_c, &r_c, &p_c, comb ? &c_c : nullptr, &d_c);
+            node.inputs[0]  = node.add_dummy(x_c);
+            node.inputs[1]  = node.add_dummy(r_c);
+            node.inputs[2]  = node.add_dummy(p_c);
+            if (comb) {
+                node.inputs[3] = node.add_dummy(c_c);
+            }
+            node.outputs[0] = node.add_dummy(d_c);
+        }
+        return;
+    }
+
+    if (node.opcode == HTP_OP_LIGHTNING_INDEXER) {
+        if (node.inputs.size() < 4 || !node.inputs[0] || !node.inputs[1] || !node.inputs[2] || !node.inputs[3] || node.outputs.empty() || !node.outputs[0]) {
+            return;
+        }
+        const struct ggml_tensor * q   = node.inputs[0];
+        const struct ggml_tensor * k   = node.inputs[1];
+        const struct ggml_tensor * w   = node.inputs[2];
+        const struct ggml_tensor * m   = node.inputs[3];
+        const struct ggml_tensor * dst = node.outputs[0];
+
+        if (htp_lightning_indexer_can_collapse(q, k, w, m, dst)) {
+            struct ggml_tensor q_c, k_c, w_c, m_c, d_c;
+            htp_lightning_indexer_collapse_tensors(q, k, w, m, dst, &q_c, &k_c, &w_c, &m_c, &d_c);
+            node.inputs[0]  = node.add_dummy(q_c);
+            node.inputs[1]  = node.add_dummy(k_c);
+            node.inputs[2]  = node.add_dummy(w_c);
+            node.inputs[3]  = node.add_dummy(m_c);
+            node.outputs[0] = node.add_dummy(d_c);
+        }
+        return;
+    }
 }
 
 bool htp_kparams_solve(
@@ -4284,6 +5143,19 @@ bool htp_kparams_solve(
     if (node.opcode == HTP_OP_UPSCALE) {
         return htp_kparams_solve_upscale(sess, src0, dst, node.kparams<htp_upscale_kernel_params>());
     }
+    if (node.opcode == HTP_OP_DSV4_HC_PRE) {
+        return htp_kparams_solve_dsv4_hc_pre_impl(sess, src0, src1, dst, node.kparams<htp_dsv4_hc_pre_kernel_params>());
+    }
+    if (node.opcode == HTP_OP_DSV4_HC_POST) {
+        const ggml_tensor * src2 = node.inputs.size() > 2 ? node.inputs[2] : nullptr;
+        const ggml_tensor * src3 = node.inputs.size() > 3 ? node.inputs[3] : nullptr;
+        return htp_kparams_solve_dsv4_hc_post_impl(sess, src0, src1, src2, src3, dst, node.kparams<htp_dsv4_hc_post_kernel_params>());
+    }
+    if (node.opcode == HTP_OP_LIGHTNING_INDEXER) {
+        const ggml_tensor * src2 = node.inputs.size() > 2 ? node.inputs[2] : nullptr;
+        const ggml_tensor * src3 = node.inputs.size() > 3 ? node.inputs[3] : nullptr;
+        return htp_kparams_solve_lightning_indexer_impl(sess, src0, src1, src2, src3, dst, node.kparams<htp_lid_kernel_params>());
+    }
 
     return htp_kparams_solve(sess, node.node, node.kernel_params);
 }
@@ -4379,18 +5251,30 @@ bool htp_kparams_solve(
         case HTP_OP_CUMSUM:
             return htp_kparams_invoke<htp_kparams_solve_cumsum>(sess, op, kparams);
 
+        case HTP_OP_TOP_K_MASK:
+            return htp_kparams_invoke<htp_kparams_solve_top_k_mask>(sess, op, kparams);
+
         case HTP_OP_FILL:
             return htp_kparams_invoke<htp_kparams_solve_fill>(sess, op, kparams);
 
         case HTP_OP_ROLL:
             return htp_kparams_invoke<htp_kparams_solve_roll>(sess, op, kparams);
 
+        case HTP_OP_LIGHTNING_INDEXER:
+            return htp_kparams_invoke<htp_kparams_solve_lightning_indexer>(sess, op, kparams);
+
         case HTP_OP_DIAG:
             return htp_kparams_invoke<htp_kparams_solve_diag>(sess, op, kparams);
 
         case HTP_OP_SOLVE_TRI:
             return htp_kparams_invoke<htp_kparams_solve_solve_tri>(sess, op, kparams);
 
+        case HTP_OP_DSV4_HC_PRE:
+            return htp_kparams_invoke<htp_kparams_solve_dsv4_hc_pre>(sess, op, kparams);
+
+        case HTP_OP_DSV4_HC_POST:
+            return htp_kparams_invoke<htp_kparams_solve_dsv4_hc_post>(sess, op, kparams);
+
         default:
             if (htp_op_is_unary(opcode)) {
                 return htp_kparams_invoke<htp_kparams_solve_unary>(sess, op, kparams);

ggml/src/ggml-hexagon/htp-kparams.h

@@ -11,12 +11,14 @@
 #include "htp/cpy-ops.h"
 #include "htp/cumsum-ops.h"
 #include "htp/diag-ops.h"
+#include "htp/dsv4-hc-ops.h"
 #include "htp/fill-ops.h"
 #include "htp/flash-attn-ops.h"
 #include "htp/gated-delta-net-ops.h"
 #include "htp/get-rows-ops.h"
 #include "htp/glu-ops.h"
 #include "htp/im2col-ops.h"
+#include "htp/lightning-indexer-ops.h"
 #include "htp/matmul-ops.h"
 #include "htp/pad-ops.h"
 #include "htp/pool-ops.h"
@@ -28,6 +30,7 @@
 #include "htp/solve-tri-ops.h"
 #include "htp/ssm-conv.h"
 #include "htp/sum-rows-ops.h"
+#include "htp/top-k-mask-ops.h"
 #include "htp/unary-ops.h"
 #include "htp/upscale-ops.h"
 

ggml/src/ggml-hexagon/htp-opnode.cpp

@@ -5,11 +5,13 @@
 #include "htp/binary-ops.h"
 #include "htp/cumsum-ops.h"
 #include "htp/diag-ops.h"
+#include "htp/dsv4-hc-ops.h"
 #include "htp/fill-ops.h"
 #include "htp/flash-attn-ops.h"
 #include "htp/gated-delta-net-ops.h"
 #include "htp/get-rows-ops.h"
 #include "htp/glu-ops.h"
+#include "htp/lightning-indexer-ops.h"
 #include "htp/matmul-ops.h"
 #include "htp/pad-ops.h"
 #include "htp/repeat-ops.h"
@@ -18,6 +20,7 @@
 #include "htp/solve-tri-ops.h"
 #include "htp/ssm-conv.h"
 #include "htp/sum-rows-ops.h"
+#include "htp/top-k-mask-ops.h"
 #include "htp/unary-ops.h"
 #include "htp/upscale-ops.h"
 
@@ -185,6 +188,22 @@ static std::string format_kernel_params(const htp_opnode & node) {
         const char * path = (kparams->kernel_type == HTP_GDN_KERNEL_HMX_CHUNKED) ? "hmx-chunked" : "hvx-recurrent";
         snprintf(str, max_size, "%s-%s vtcm %u", path, kparams->kda ? "kda" : "scalar",
                  (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_DSV4_HC_PRE) {
+        const auto * kparams = node.kparams<htp_dsv4_hc_pre_kernel_params>();
+        const char * ktype_str = kparams->kernel_type == HTP_DSV4_HC_PRE_KERNEL_GATED ? "gated" : "plain";
+        snprintf(str, max_size, "%s nth %u vtcm %u", ktype_str,
+                 (unsigned int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_DSV4_HC_POST) {
+        const auto * kparams = node.kparams<htp_dsv4_hc_post_kernel_params>();
+        const char * ktype_str = kparams->kernel_type == HTP_DSV4_HC_POST_KERNEL_COMB ? "comb" : "identity";
+        snprintf(str, max_size, "%s nth %u vtcm %u", ktype_str,
+                 (unsigned int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_LIGHTNING_INDEXER) {
+        const auto * kparams = node.kparams<htp_lid_kernel_params>();
+        snprintf(str, max_size, "%s tok %u bc %u nth %u vtcm %u",
+                 kparams->kernel_type == HTP_LID_KERNEL_HMX ? "hmx" : "hvx",
+                 (unsigned int) kparams->n_tokens_tile, (unsigned int) kparams->kv_block,
+                 (unsigned int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_MUL || node.opcode == HTP_OP_ADD || node.opcode == HTP_OP_ADD_ID ||
                node.opcode == HTP_OP_SUB || node.opcode == HTP_OP_DIV) {
         const auto * kparams = node.kparams<htp_binary_kernel_params>();
@@ -213,6 +232,14 @@ static std::string format_kernel_params(const htp_opnode & node) {
     } else if (node.opcode == HTP_OP_CUMSUM) {
         const auto * kparams = node.kparams<htp_cumsum_kernel_params>();
         snprintf(str, max_size, "%s nth %d vtcm %u", kparams->use_dma ? "dma" : "hvx", (int) kparams->n_threads, (unsigned int) kparams->vtcm_size);
+    } else if (node.opcode == HTP_OP_TOP_K_MASK) {
+        const auto * kparams = node.kparams<htp_top_k_mask_kernel_params>();
+        const char * ktype_str = "?";
+        switch (kparams->kernel_type) {
+            case HTP_TOP_K_MASK_KERNEL_ROWS:  ktype_str = "rows"; break;
+            case HTP_TOP_K_MASK_KERNEL_PICKS: ktype_str = "picks"; break;
+        }
+        snprintf(str, max_size, "%s nth %d tile %u x%u vtcm %u", ktype_str, (int) kparams->n_threads, (unsigned int) kparams->tile_cols, (unsigned int) kparams->n_col_tiles, (unsigned int) kparams->vtcm_size);
     } else if (node.opcode == HTP_OP_FILL) {
         const auto * kparams = node.kparams<htp_fill_kernel_params>();
         snprintf(str, max_size, "%s nth %d", kparams->kernel_type == HTP_FILL_KERNEL_CONTIGUOUS ? "contig" : "strided", (int) kparams->n_threads);

ggml/src/ggml-hexagon/htp/CMakeLists.txt

@@ -37,6 +37,7 @@ add_library(${HTP_LIB} SHARED
     repeat-ops.c
     ssm-conv.c
     cumsum-ops.c
+    top-k-mask-ops.c
     fill-ops.c
     concat-ops.c
     diag-ops.c
@@ -49,6 +50,8 @@ add_library(${HTP_LIB} SHARED
     pool-ops.c
     upscale-ops.c
     allreduce-ops.c
+    dsv4-hc-ops.c
+    lightning-indexer-ops.c
 )
 
 target_compile_definitions(${HTP_LIB} PRIVATE

ggml/src/ggml-hexagon/htp/argsort-ops.c

@@ -315,55 +315,61 @@ static void bitonic_sort_vtcm_desc(uint8_t * values, uint8_t * indices, uint32_t
     }
 }
 
-static void top_k_select_tiled(dma_queue * dma_q, struct htp_thread_trace * tr,
-                               dma_addr_t src_addr, uint32_t n, uint32_t k,
-                               float * values_buf, int32_t * indices_buf,
-                               float * out_values, int32_t * out_indices,
-                               uint32_t chunk_base) {
-    const uint32_t tile_elems = 1024;
-    uint32_t n_tiles            = (n + tile_elems - 1) / tile_elems;
-    uint32_t candidate_count    = n_tiles * k;
-    uint32_t merge_n_vec        = hmx_ceil_div(candidate_count, 32);
-    uint32_t merge_n_vec_pow2   = 1;
-    while (merge_n_vec_pow2 < merge_n_vec) merge_n_vec_pow2 <<= 1;
-    uint32_t merge_elems        = merge_n_vec_pow2 * 32;
-    float * candidate_values    = values_buf + tile_elems;
-    int32_t * candidate_indices = indices_buf + tile_elems;
-    uint32_t candidate_pos      = 0;
+// top-k of one row held in VTCM, k <= 64. Keeps a sorted list of the k best, rejects 256-float blocks with one max-reduce.
+// Writes k values and k indices (ties: first seen wins). Slots with no finite value get -INF and index 0.
+static void top_k_scan(const float * vals, uint32_t n, uint32_t k, float * out_v, int32_t * out_i) {
+    float   best_v[64];
+    int32_t best_i[64];
+    uint32_t cnt = 0;
+    float    thr = -INFINITY;
 
-    for (uint32_t offset = 0; offset < n; offset += tile_elems) {
-        uint32_t tile_count = MIN(tile_elems, n - offset);
-        dma_sync_read(dma_q, values_buf, src_addr + (size_t) offset * sizeof(float), tile_count * sizeof(float));
-
-        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) offset);
-        if (tile_count < tile_elems) {
-            hvx_splat_f32_u((uint8_t *) (values_buf + tile_count), -INFINITY, tile_elems - tile_count);
+    const uint32_t blk = 256;
+    for (uint32_t off = 0; off < n; off += blk) {
+        const uint32_t bn = MIN(blk, n - off);
+        if (hvx_reduce_max_f32((const uint8_t *) (vals + off), bn) <= thr) {
+            continue;
         }
-
-        bitonic_sort_vtcm_desc((uint8_t *) values_buf, (uint8_t *) indices_buf, tile_elems / 32, true);
-        uint32_t tile_k = MIN(k, tile_count);
-        for (uint32_t j = 0; j < tile_k; j++) {
-            candidate_values[candidate_pos] = values_buf[j];
-            candidate_indices[candidate_pos] = indices_buf[j] + (int32_t) offset;
-            candidate_pos++;
-        }
-        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) offset);
-    }
-
-    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, 0);
-    if (merge_elems > candidate_pos) {
-        hvx_splat_f32_u((uint8_t *) (candidate_values + candidate_pos), -INFINITY, merge_elems - candidate_pos);
-        for (uint32_t j = candidate_pos; j < merge_elems; j++) {
-            candidate_indices[j] = 0;
+        for (uint32_t vi = 0; vi < bn; vi += 32) {
+            const uint32_t vn = MIN(32, bn - vi);
+            HVX_VectorAlias v_local;
+            v_local.v = hvx_vmem(vals + off + vi);
+            for (uint32_t i = 0; i < vn; i++) {
+                const float v = v_local.fp32[i];
+                if (!(v > thr)) continue;
+                uint32_t pos = cnt < k ? cnt++ : k - 1;
+                while (pos > 0 && best_v[pos - 1] < v) {
+                    best_v[pos] = best_v[pos - 1];
+                    best_i[pos] = best_i[pos - 1];
+                    pos--;
+                }
+                best_v[pos] = v;
+                best_i[pos] = (int32_t) (off + vi + i);
+                if (cnt == k) thr = best_v[k - 1];
+            }
         }
     }
 
-    bitonic_sort_vtcm_desc((uint8_t *) candidate_values, (uint8_t *) candidate_indices, merge_n_vec_pow2, false);
-    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, 0);
+    // fewer than k values above -INF: take -INF or NaN elements, then pad with index 0
+    for (uint32_t off = 0; off < n && cnt < k; off += 32) {
+        const uint32_t vn = MIN(32, n - off);
+        HVX_VectorAlias v_local;
+        v_local.v = hvx_vmem(vals + off);
+        for (uint32_t j = 0; j < vn && cnt < k; j++) {
+            if (!(v_local.fp32[j] > -INFINITY)) {
+                best_v[cnt] = -INFINITY;
+                best_i[cnt] = (int32_t) (off + j);
+                cnt++;
+            }
+        }
+    }
+    for (; cnt < k; cnt++) {
+        best_v[cnt] = -INFINITY;
+        best_i[cnt] = 0;
+    }
 
     for (uint32_t j = 0; j < k; j++) {
-        out_values[j]  = candidate_values[j];
-        out_indices[j] = candidate_indices[j] + (int32_t) chunk_base;
+        out_v[j] = best_v[j];
+        out_i[j] = best_i[j];
     }
 }
 
@@ -509,10 +515,14 @@ static void htp_sort_chunk_job(unsigned int n, unsigned int i, void * data) {
         }
 
         if (local_k > 1 && local_k <= 64 && chunk_elems > 1024) {
-            top_k_select_tiled(octx->ctx->dma[i], tr, src_addr, real_count, local_k,
-                               values_buf, indices_buf,
-                               &merge_values[i * local_k], &merge_indices[i * local_k],
-                               chunk_base);
+            dma_sync_read(octx->ctx->dma[i], values_buf, src_addr, real_count * sizeof(float));
+
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
+            top_k_scan(values_buf, real_count, local_k, &merge_values[i * local_k], &merge_indices[i * local_k]);
+            for (uint32_t j = 0; j < local_k; j++) {
+                merge_indices[i * local_k + j] += (int32_t) chunk_base;
+            }
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) i);
             return;
         }
 
@@ -668,10 +678,14 @@ static inline void compute_row_sort(
     uint32_t ne00,
     uint32_t chunk_elems,
     bool is_top_k,
-    enum ggml_sort_order order
+    enum ggml_sort_order order,
+    uint32_t k
 ) {
     if (is_top_k) {
-        if (ne00 == 32 || ne00 == 64 || ne00 == 128 || ne00 == 256 || ne00 == 512 || ne00 == 1024) {
+        if (k > 1 && k <= 64 && ne00 > 1024) {
+            float out_v[64];
+            top_k_scan(cur_values, ne00, k, out_v, cur_indices);
+        } else if (ne00 == 32 || ne00 == 64 || ne00 == 128 || ne00 == 256 || ne00 == 512 || ne00 == 1024) {
             switch (ne00) {
                 case 32:   sort32_f32_hvx((uint8_t *) cur_values, (uint8_t *) cur_indices, GGML_SORT_ORDER_DESC); break;
                 case 64:   sort64_f32_hvx((uint8_t *) cur_values, (uint8_t *) cur_indices, GGML_SORT_ORDER_DESC); break;
@@ -788,7 +802,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
             dma_queue_pop(dma_q);
 
             htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);
-            compute_row_sort(cur_values, cur_indices, ne00, chunk_elems, is_top_k, order);
+            compute_row_sort(cur_values, cur_indices, ne00, chunk_elems, is_top_k, order, k);
             htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);
 
             // Push writeback of current slot
@@ -817,7 +831,7 @@ static void htp_sort_multi_row_job(unsigned int n, unsigned int ith, void * data
             dma_sync_read(dma_q, values_buf, src_addr, src_bytes);
 
             htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);
-            compute_row_sort(values_buf, indices_buf, ne00, chunk_elems, is_top_k, order);
+            compute_row_sort(values_buf, indices_buf, ne00, chunk_elems, is_top_k, order, k);
             htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r);
 
             dma_sync_write(dma_q, dst_addr, indices_buf, dst_bytes);

ggml/src/ggml-hexagon/htp/dsv4-hc-ops.c

@@ -0,0 +1,675 @@
+#pragma clang diagnostic ignored "-Wunused-variable"
+#pragma clang diagnostic ignored "-Wunused-function"
+
+#include <HAP_farf.h>
+
+#include <stdint.h>
+#include <string.h>
+#include <math.h>
+
+#include "dma-queue.h"
+#include "hex-common.h"
+#include "hex-profile.h"
+#include "htp-ctx.h"
+#include "htp-tensor.h"
+#include "hvx-utils.h"
+#include "dsv4-hc-ops.h"
+
+// ---------------------------------------------------------------------------
+// DSV4_HC_PRE
+// ---------------------------------------------------------------------------
+
+// hvx_vec_fast_sigmoid_f32_guard_2it with fewer live constants and predicates
+static inline HVX_Vector dsv4_hc_sigmoid_f32(HVX_Vector v, HVX_Vector max_exp, HVX_Vector min_exp) {
+    v = Q6_Vsf_vmin_VsfVsf(Q6_Vsf_vmax_VsfVsf(v, min_exp), max_exp);
+    v = Q6_Vqf32_vmpy_VsfVsf(v, hvx_vec_splat_f32(0.5f * 1.442695022f));
+
+    HVX_Vector in_int = Q6_Vw_equals_Vsf(Q6_Vsf_equals_Vqf32(v));
+    HVX_Vector x      = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
+    HVX_Vector xx     = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x);
+
+    HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2));
+    v1            = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
+
+    HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1));
+    v2            = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx);
+    v2            = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x);
+
+    HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
+    v3            = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
+
+    HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
+    HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
+
+    // Newton-Raphson with 2 iterations
+    HVX_Vector r = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5);
+    for (int i = 0; i < 2; ++i) {
+        const HVX_Vector rv = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(r, v5));
+        r = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(Q6_Vqf32_vadd_VsfVsf(r, r), Q6_Vqf32_vmpy_VsfVsf(r, rv)));
+    }
+
+    return Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, r));
+}
+
+struct dsv4_hc_pre_context {
+    struct htp_ops_context * octx;
+    const struct htp_dsv4_hc_pre_kernel_params * kparams;
+    uint32_t token_start;
+    uint32_t n_chunks;
+    uint32_t total_jobs;
+    uint32_t jobs_per_thread;
+    float scale;
+    uint8_t * vtcm;
+    uint32_t dma_block_chunks;
+    uint32_t dma_row_stride;
+    uint32_t dma_buffer_size;
+};
+
+struct dsv4_hc_block {
+    uint32_t it;
+    uint32_t chunk;
+    uint32_t count;
+};
+
+static inline struct dsv4_hc_block dsv4_hc_pre_block_make(
+        const struct dsv4_hc_pre_context * dctx,
+        uint32_t job,
+        uint32_t end) {
+    const uint32_t local_token = fastdiv(job, &dctx->kparams->div_n_chunks);
+    const uint32_t chunk       = job - local_token*dctx->n_chunks;
+    const uint32_t count       = hex_smin(dctx->dma_block_chunks, hex_smin(dctx->n_chunks - chunk, end - job));
+    return (struct dsv4_hc_block) {
+        .it    = dctx->token_start + local_token,
+        .chunk = chunk,
+        .count = count,
+    };
+}
+
+static inline void dsv4_hc_pre_push_plain(
+        dma_queue * dma_q,
+        const struct dsv4_hc_pre_context * dctx,
+        const struct dsv4_hc_block * block,
+        uint8_t * spad) {
+    const struct htp_tensor * x       = dctx->octx->src[0];
+    const struct htp_tensor * weights = dctx->octx->src[1];
+    const uint32_t row_stride         = dctx->dma_row_stride;
+    const uint32_t row_size           = block->count * VLEN;
+
+    dma_queue_push(dma_q, dma_make_data(spad, x->data + (block->it*x->nb[2] + block->chunk*VLEN)), row_stride, x->nb[1], row_size, 4);
+    dma_queue_push(dma_q, dma_make_data(spad + 5*row_stride, weights->data + (block->it*weights->nb[1])), 16, 16, 16, 1);
+}
+
+static inline void dsv4_hc_pre_push_gated(
+        dma_queue * dma_q,
+        const struct dsv4_hc_pre_context * dctx,
+        const struct dsv4_hc_block * block,
+        uint8_t * spad) {
+    const struct htp_tensor * x       = dctx->octx->src[0];
+    const struct htp_tensor * g       = dctx->octx->src[1];
+    const uint32_t row_stride         = dctx->dma_row_stride;
+    const uint32_t row_size           = block->count * VLEN;
+
+    dma_queue_push(dma_q, dma_make_data(spad, x->data + (block->it*x->nb[2] + block->chunk*VLEN)), row_stride, x->nb[1], row_size, 4);
+    dma_queue_push(dma_q, dma_make_data(spad + 4*row_stride, g->data + (block->it*g->nb[2] + block->chunk*VLEN)), row_stride, g->nb[1], row_size, 4);
+}
+
+static __attribute__((noinline)) void dsv4_hc_pre_plain_block_f32(
+        const struct dsv4_hc_pre_context * dctx,
+        const struct dsv4_hc_block * block,
+        const uint8_t * x_spad,
+        uint8_t * d_spad,
+        const float * wb) {
+    const uint32_t row_stride = dctx->dma_row_stride;
+    const HVX_Vector scale_v  = hvx_vec_splat_f32(dctx->scale);
+    const HVX_Vector w0       = hvx_vec_splat_f32(wb[0]);
+    const HVX_Vector w1       = hvx_vec_splat_f32(wb[1]);
+    const HVX_Vector w2       = hvx_vec_splat_f32(wb[2]);
+    const HVX_Vector w3       = hvx_vec_splat_f32(wb[3]);
+
+    for (uint32_t c = 0; c < block->count; ++c) {
+        const uint8_t * xb = x_spad + c*VLEN;
+
+        HVX_Vector x0 = *(const HVX_Vector *) xb;
+        HVX_Vector x1 = *(const HVX_Vector *) (xb + row_stride);
+        HVX_Vector x2 = *(const HVX_Vector *) (xb + 2*row_stride);
+        HVX_Vector x3 = *(const HVX_Vector *) (xb + 3*row_stride);
+
+        HVX_Vector acc = hvx_vec_mul_f32_f32(x0, w0);
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(x1, w1));
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(x2, w2));
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(x3, w3));
+
+        *(HVX_Vector *) (d_spad + c*VLEN) = hvx_vec_mul_f32_f32(acc, scale_v);
+    }
+}
+
+static __attribute__((noinline)) void dsv4_hc_pre_gated_block_f32(
+        const struct dsv4_hc_pre_context * dctx,
+        const struct dsv4_hc_block * block,
+        const uint8_t * x_spad,
+        const uint8_t * g_spad,
+        uint8_t * d_spad) {
+    const uint32_t row_stride = dctx->dma_row_stride;
+    const HVX_Vector scale_v = hvx_vec_splat_f32(dctx->scale);
+    const HVX_Vector max_exp = hvx_vec_splat_f32(87.0f);
+    const HVX_Vector min_exp = hvx_vec_splat_f32(-87.0f);
+
+    for (uint32_t c = 0; c < block->count; ++c) {
+        const uint8_t * xb = x_spad + c*VLEN;
+        const uint8_t * gb = g_spad + c*VLEN;
+
+        HVX_Vector xv = *(const HVX_Vector *) xb;
+        HVX_Vector gv = dsv4_hc_sigmoid_f32(*(const HVX_Vector *) gb, max_exp, min_exp);
+        HVX_Vector acc = hvx_vec_mul_f32_f32(xv, gv);
+
+        xv = *(const HVX_Vector *) (xb + row_stride);
+        gv = dsv4_hc_sigmoid_f32(*(const HVX_Vector *) (gb + row_stride), max_exp, min_exp);
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(xv, gv));
+
+        xv = *(const HVX_Vector *) (xb + 2*row_stride);
+        gv = dsv4_hc_sigmoid_f32(*(const HVX_Vector *) (gb + 2*row_stride), max_exp, min_exp);
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(xv, gv));
+
+        xv = *(const HVX_Vector *) (xb + 3*row_stride);
+        gv = dsv4_hc_sigmoid_f32(*(const HVX_Vector *) (gb + 3*row_stride), max_exp, min_exp);
+        acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(xv, gv));
+
+        *(HVX_Vector *) (d_spad + c*VLEN) = hvx_vec_mul_f32_f32(acc, scale_v);
+    }
+}
+
+static void dsv4_hc_pre_thread_plain_f32(unsigned int nth, unsigned int ith, void * data) {
+    struct dsv4_hc_pre_context * dctx = (struct dsv4_hc_pre_context *) data;
+    struct htp_ops_context *     octx = dctx->octx;
+
+    const struct htp_tensor * dst    = octx->dst;
+    const uint64_t            d_data = dst->data;
+
+    const uint32_t row_stride  = dctx->dma_row_stride;
+    const uint32_t buffer_size = dctx->dma_buffer_size;
+    const uint32_t start       = dctx->jobs_per_thread * ith;
+    const uint32_t end         = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    if (start >= end) {
+        return;
+    }
+
+    uint8_t * spad[2] = {
+        dctx->vtcm + 2*buffer_size*ith,
+        dctx->vtcm + 2*buffer_size*ith + buffer_size,
+    };
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    uint32_t prefetch_job = start;
+    for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
+        const struct dsv4_hc_block block = dsv4_hc_pre_block_make(dctx, prefetch_job, end);
+        // empty dst writeback keeps pop order dst, x, weights in loop
+        dma_queue_push(dma_q, dma_make_data(d_data, spad[i] + 4*row_stride), block.count*VLEN, block.count*VLEN, 0, 0);
+        dsv4_hc_pre_push_plain(dma_q, dctx, &block, spad[i]);
+        prefetch_job += block.count;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    uint32_t job  = start;
+    uint32_t curr = 0;
+    while (job < end) {
+        const struct dsv4_hc_block block = dsv4_hc_pre_block_make(dctx, job, end);
+
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+
+        const float * wb = (const float *) (spad[curr] + 5*row_stride);
+        uint8_t * d_spad = spad[curr] + 4*row_stride;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+        dsv4_hc_pre_plain_block_f32(dctx, &block, spad[curr], d_spad, wb);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+
+        const uint64_t d_dst = d_data + (block.it*dst->nb[1] + block.chunk*VLEN);
+        dma_queue_push(dma_q, dma_make_data(d_dst, d_spad), block.count*VLEN, block.count*VLEN, block.count*VLEN, 1);
+
+        if (prefetch_job < end) {
+            const struct dsv4_hc_block next = dsv4_hc_pre_block_make(dctx, prefetch_job, end);
+            dsv4_hc_pre_push_plain(dma_q, dctx, &next, spad[curr]);
+            prefetch_job += next.count;
+        }
+
+        job += block.count;
+        curr ^= 1;
+    }
+
+    dma_queue_flush(dma_q);
+
+    (void) nth;
+}
+
+static void dsv4_hc_pre_thread_gated_f32(unsigned int nth, unsigned int ith, void * data) {
+    struct dsv4_hc_pre_context * dctx = (struct dsv4_hc_pre_context *) data;
+    struct htp_ops_context *     octx = dctx->octx;
+
+    const struct htp_tensor * dst    = octx->dst;
+    const uint64_t            d_data = dst->data;
+
+    const uint32_t row_stride  = dctx->dma_row_stride;
+    const uint32_t buffer_size = dctx->dma_buffer_size;
+    const uint32_t start       = dctx->jobs_per_thread * ith;
+    const uint32_t end         = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    if (start >= end) {
+        return;
+    }
+
+    uint8_t * spad[2] = {
+        dctx->vtcm + 2*buffer_size*ith,
+        dctx->vtcm + 2*buffer_size*ith + buffer_size,
+    };
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    uint32_t prefetch_job = start;
+    for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
+        const struct dsv4_hc_block block = dsv4_hc_pre_block_make(dctx, prefetch_job, end);
+        // empty dst writeback keeps pop order dst, x, g in loop
+        dma_queue_push(dma_q, dma_make_data(d_data, spad[i] + 8*row_stride), block.count*VLEN, block.count*VLEN, 0, 0);
+        dsv4_hc_pre_push_gated(dma_q, dctx, &block, spad[i]);
+        prefetch_job += block.count;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    uint32_t job  = start;
+    uint32_t curr = 0;
+    while (job < end) {
+        const struct dsv4_hc_block block = dsv4_hc_pre_block_make(dctx, job, end);
+
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+
+        uint8_t * d_spad = spad[curr] + 8*row_stride;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+        dsv4_hc_pre_gated_block_f32(dctx, &block, spad[curr], spad[curr] + 4*row_stride, d_spad);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+
+        const uint64_t d_dst = d_data + (block.it*dst->nb[1] + block.chunk*VLEN);
+        dma_queue_push(dma_q, dma_make_data(d_dst, d_spad), block.count*VLEN, block.count*VLEN, block.count*VLEN, 1);
+
+        if (prefetch_job < end) {
+            const struct dsv4_hc_block next = dsv4_hc_pre_block_make(dctx, prefetch_job, end);
+            dsv4_hc_pre_push_gated(dma_q, dctx, &next, spad[curr]);
+            prefetch_job += next.count;
+        }
+
+        job += block.count;
+        curr ^= 1;
+    }
+
+    dma_queue_flush(dma_q);
+
+    (void) nth;
+}
+
+int op_dsv4_hc_pre(struct htp_ops_context * octx) {
+    const struct htp_tensor * x       = octx->src[0];
+    const struct htp_tensor * weights = octx->src[1];
+    const struct htp_tensor * dst     = octx->dst;
+
+    const struct htp_dsv4_hc_pre_kernel_params * kparams =
+        (const struct htp_dsv4_hc_pre_kernel_params *) octx->kernel_params;
+
+    const uint32_t n_threads = kparams->n_threads;
+    if (!htp_ops_context_set_n_threads_div(octx, n_threads, kparams->div_n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    uint32_t token_start     = 0;
+    uint32_t n_tokens        = x->ne[2];
+    uint32_t total_jobs      = kparams->total_jobs;
+    uint32_t jobs_per_thread = kparams->jobs_per_thread;
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_mdev_data_aligned(dst) && ((dst->nb[1] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_tokens, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        token_start     = range.start;
+        n_tokens        = range.count;
+        total_jobs      = n_tokens * kparams->n_chunks;
+        jobs_per_thread = fastdiv(total_jobs + n_threads - 1, &kparams->div_n_threads);
+    }
+
+    if (n_tokens == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    float scale;
+    memcpy(&scale, octx->op_params, sizeof(scale));
+
+    struct dsv4_hc_pre_context dctx = {
+        .octx             = octx,
+        .kparams          = kparams,
+        .token_start      = token_start,
+        .n_chunks         = kparams->n_chunks,
+        .total_jobs       = total_jobs,
+        .jobs_per_thread  = jobs_per_thread,
+        .scale            = scale,
+        .vtcm             = (uint8_t *) octx->ctx->vtcm_base,
+        .dma_block_chunks = kparams->dma_block_chunks,
+        .dma_row_stride   = kparams->dma_row_stride,
+        .dma_buffer_size  = kparams->dma_buffer_size,
+    };
+
+    work_queue_func_t worker = NULL;
+    switch (kparams->kernel_type) {
+        case HTP_DSV4_HC_PRE_KERNEL_PLAIN: worker = dsv4_hc_pre_thread_plain_f32; break;
+        case HTP_DSV4_HC_PRE_KERNEL_GATED: worker = dsv4_hc_pre_thread_gated_f32; break;
+        default: return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    FARF(HIGH, "dsv4-hc-pre: (%ux%ux%u) -> (%ux%u) : ktype %u tokens %u:%u nth %u\n",
+         x->ne[0], x->ne[1], x->ne[2], dst->ne[0], dst->ne[1], kparams->kernel_type,
+         token_start, token_start + n_tokens, n_threads);
+
+    work_queue_run(octx->ctx->work_queue, worker, &dctx, n_threads);
+    return HTP_STATUS_OK;
+}
+
+// ---------------------------------------------------------------------------
+// DSV4_HC_POST
+// ---------------------------------------------------------------------------
+
+struct dsv4_hc_post_context {
+    struct htp_ops_context * octx;
+    const struct htp_dsv4_hc_post_kernel_params * kparams;
+    uint32_t token_start;
+    uint32_t total_jobs;
+    uint32_t jobs_per_thread;
+    uint8_t * vtcm;
+};
+
+static inline struct dsv4_hc_block dsv4_hc_post_block_make(
+        const struct dsv4_hc_post_context * dctx,
+        uint32_t job,
+        uint32_t end,
+        uint32_t block_chunks) {
+    const uint32_t local_token = fastdiv(job, &dctx->kparams->div_n_chunks);
+    const uint32_t chunk       = job - local_token*dctx->kparams->n_chunks;
+    const uint32_t count       = hex_smin(block_chunks, hex_smin(dctx->kparams->n_chunks - chunk, end - job));
+    return (struct dsv4_hc_block) {
+        .it    = dctx->token_start + local_token,
+        .chunk = chunk,
+        .count = count,
+    };
+}
+
+static __attribute__((noinline)) void dsv4_hc_post_identity_block_f32(
+        const uint8_t * xb,
+        const uint8_t * rb,
+        size_t r_stride,
+        uint8_t * db,
+        size_t d_stride,
+        uint32_t count,
+        const float * p) {
+    for (uint32_t id = 0; id < 4; ++id) {
+        const HVX_Vector pv = hvx_vec_splat_f32(p[id]);
+        const uint8_t *  rr = rb + id*r_stride;
+        uint8_t *        dd = db + id*d_stride;
+
+        for (uint32_t i = 0; i < count; ++i) {
+            const HVX_Vector xv = *(const HVX_Vector *) (xb + i*VLEN);
+            const HVX_Vector rv = *(const HVX_Vector *) (rr + i*VLEN);
+            *(HVX_Vector *) (dd + i*VLEN) = hvx_vec_add_f32_f32(hvx_vec_mul_f32_f32(xv, pv), rv);
+        }
+    }
+}
+
+static __attribute__((noinline)) void dsv4_hc_post_comb_block_f32(
+        const uint8_t * xb,
+        const uint8_t * rb,
+        size_t r_stride,
+        uint8_t * db,
+        size_t d_stride,
+        uint32_t count,
+        const float * p,
+        const float * c) {
+    for (uint32_t id = 0; id < 4; ++id) {
+        const HVX_Vector pv = hvx_vec_splat_f32(p[id]);
+        const HVX_Vector c0 = hvx_vec_splat_f32(c[id]);
+        const HVX_Vector c1 = hvx_vec_splat_f32(c[4 + id]);
+        const HVX_Vector c2 = hvx_vec_splat_f32(c[8 + id]);
+        const HVX_Vector c3 = hvx_vec_splat_f32(c[12 + id]);
+        uint8_t * dd = db + id*d_stride;
+
+        for (uint32_t i = 0; i < count; ++i) {
+            const uint8_t * r = rb + i*VLEN;
+            HVX_Vector acc = hvx_vec_mul_f32_f32(*(const HVX_Vector *) (xb + i*VLEN), pv);
+            acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(*(const HVX_Vector *) r, c0));
+            acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(*(const HVX_Vector *) (r + r_stride), c1));
+            acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(*(const HVX_Vector *) (r + 2*r_stride), c2));
+            acc = hvx_vec_add_f32_f32(acc, hvx_vec_mul_f32_f32(*(const HVX_Vector *) (r + 3*r_stride), c3));
+            *(HVX_Vector *) (dd + i*VLEN) = acc;
+        }
+    }
+}
+
+static inline void dsv4_hc_post_push_identity(
+        dma_queue * dma_q,
+        const struct dsv4_hc_post_context * dctx,
+        const struct dsv4_hc_block * block,
+        uint8_t * spad) {
+    const struct htp_tensor * x        = dctx->octx->src[0];
+    const struct htp_tensor * residual = dctx->octx->src[1];
+    const struct htp_tensor * post     = dctx->octx->src[2];
+    const uint32_t row_stride          = dctx->kparams->dma_row_stride;
+    const uint32_t row_size            = block->count * VLEN;
+
+    dma_queue_push(dma_q, dma_make_data(spad, residual->data + (block->it*residual->nb[2] + block->chunk*VLEN)),
+                   row_stride, residual->nb[1], row_size, 4);
+    dma_queue_push(dma_q, dma_make_data(spad + 4*row_stride, x->data + (block->it*x->nb[1] + block->chunk*VLEN)),
+                   row_stride, row_size, row_size, 1);
+    dma_queue_push(dma_q, dma_make_data(spad + 9*row_stride, post->data + (block->it*post->nb[1])),
+                   16, 16, 16, 1);
+}
+
+static inline void dsv4_hc_post_push_comb(
+        dma_queue * dma_q,
+        const struct dsv4_hc_post_context * dctx,
+        const struct dsv4_hc_block * block,
+        uint8_t * spad) {
+    const struct htp_tensor * comb = dctx->octx->src[3];
+    const uint32_t row_stride      = dctx->kparams->dma_row_stride;
+
+    dsv4_hc_post_push_identity(dma_q, dctx, block, spad);
+    dma_queue_push(dma_q, dma_make_data(spad + 9*row_stride + 64, comb->data + (block->it*comb->nb[2])),
+                   16, comb->nb[1], 16, 4);
+}
+
+static void dsv4_hc_post_thread_identity_f32(unsigned int nth, unsigned int ith, void * data) {
+    struct dsv4_hc_post_context * dctx = (struct dsv4_hc_post_context *) data;
+    struct htp_ops_context *      octx = dctx->octx;
+
+    const struct htp_tensor * dst    = octx->dst;
+    const uint64_t            d_data = dst->data;
+
+    const uint32_t block_chunks = dctx->kparams->dma_block_chunks;
+    const uint32_t row_stride   = dctx->kparams->dma_row_stride;
+    const uint32_t buffer_size  = dctx->kparams->dma_buffer_size;
+    const uint32_t start        = dctx->jobs_per_thread * ith;
+    const uint32_t end          = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    if (start >= end) {
+        return;
+    }
+
+    uint8_t * spad[2] = {
+        dctx->vtcm + 2*buffer_size*ith,
+        dctx->vtcm + 2*buffer_size*ith + buffer_size,
+    };
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    uint32_t prefetch_job = start;
+    for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
+        const struct dsv4_hc_block block = dsv4_hc_post_block_make(dctx, prefetch_job, end, block_chunks);
+        // empty dst writeback keeps pop order dst, residual, x, post in loop
+        dma_queue_push(dma_q, dma_make_data(d_data, spad[i] + 5*row_stride), dst->nb[1], row_stride, 0, 0);
+        dsv4_hc_post_push_identity(dma_q, dctx, &block, spad[i]);
+        prefetch_job += block.count;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    uint32_t job  = start;
+    uint32_t curr = 0;
+    while (job < end) {
+        const struct dsv4_hc_block block = dsv4_hc_post_block_make(dctx, job, end, block_chunks);
+
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+
+        const float * p      = (const float *) (spad[curr] + 9*row_stride);
+        uint8_t *     d_spad = spad[curr] + 5*row_stride;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+        dsv4_hc_post_identity_block_f32(spad[curr] + 4*row_stride, spad[curr], row_stride, d_spad, row_stride, block.count, p);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+
+        const uint64_t d_dst = d_data + (block.it*dst->nb[2] + block.chunk*VLEN);
+        dma_queue_push(dma_q, dma_make_data(d_dst, d_spad), dst->nb[1], row_stride, block.count*VLEN, 4);
+
+        if (prefetch_job < end) {
+            const struct dsv4_hc_block next = dsv4_hc_post_block_make(dctx, prefetch_job, end, block_chunks);
+            dsv4_hc_post_push_identity(dma_q, dctx, &next, spad[curr]);
+            prefetch_job += next.count;
+        }
+
+        job += block.count;
+        curr ^= 1;
+    }
+
+    dma_queue_flush(dma_q);
+
+    (void) nth;
+}
+
+static void dsv4_hc_post_thread_comb_f32(unsigned int nth, unsigned int ith, void * data) {
+    struct dsv4_hc_post_context * dctx = (struct dsv4_hc_post_context *) data;
+    struct htp_ops_context *      octx = dctx->octx;
+
+    const struct htp_tensor * dst    = octx->dst;
+    const uint64_t            d_data = dst->data;
+
+    const uint32_t block_chunks = dctx->kparams->dma_block_chunks;
+    const uint32_t row_stride   = dctx->kparams->dma_row_stride;
+    const uint32_t buffer_size  = dctx->kparams->dma_buffer_size;
+    const uint32_t start        = dctx->jobs_per_thread * ith;
+    const uint32_t end          = hex_smin(start + dctx->jobs_per_thread, dctx->total_jobs);
+    if (start >= end) {
+        return;
+    }
+
+    uint8_t * spad[2] = {
+        dctx->vtcm + 2*buffer_size*ith,
+        dctx->vtcm + 2*buffer_size*ith + buffer_size,
+    };
+
+    dma_queue * dma_q = octx->ctx->dma[ith];
+    uint32_t prefetch_job = start;
+    for (uint32_t i = 0; i < 2 && prefetch_job < end; ++i) {
+        const struct dsv4_hc_block block = dsv4_hc_post_block_make(dctx, prefetch_job, end, block_chunks);
+        // empty dst writeback keeps pop order dst, residual, x, post, comb in loop
+        dma_queue_push(dma_q, dma_make_data(d_data, spad[i] + 5*row_stride), dst->nb[1], row_stride, 0, 0);
+        dsv4_hc_post_push_comb(dma_q, dctx, &block, spad[i]);
+        prefetch_job += block.count;
+    }
+
+    struct htp_thread_trace * tr = &octx->ctx->trace[ith];
+
+    uint32_t job  = start;
+    uint32_t curr = 0;
+    while (job < end) {
+        const struct dsv4_hc_block block = dsv4_hc_post_block_make(dctx, job, end, block_chunks);
+
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+        dma_queue_pop(dma_q);
+
+        const float * p      = (const float *) (spad[curr] + 9*row_stride);
+        const float * c      = (const float *) (spad[curr] + 9*row_stride + 64);
+        uint8_t *     d_spad = spad[curr] + 5*row_stride;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+        dsv4_hc_post_comb_block_f32(spad[curr] + 4*row_stride, spad[curr], row_stride, d_spad, row_stride, block.count, p, c);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) block.it);
+
+        const uint64_t d_dst = d_data + (block.it*dst->nb[2] + block.chunk*VLEN);
+        dma_queue_push(dma_q, dma_make_data(d_dst, d_spad), dst->nb[1], row_stride, block.count*VLEN, 4);
+
+        if (prefetch_job < end) {
+            const struct dsv4_hc_block next = dsv4_hc_post_block_make(dctx, prefetch_job, end, block_chunks);
+            dsv4_hc_post_push_comb(dma_q, dctx, &next, spad[curr]);
+            prefetch_job += next.count;
+        }
+
+        job += block.count;
+        curr ^= 1;
+    }
+
+    dma_queue_flush(dma_q);
+
+    (void) nth;
+}
+
+int op_dsv4_hc_post(struct htp_ops_context * octx) {
+    const struct htp_tensor * x        = octx->src[0];
+    const struct htp_tensor * residual = octx->src[1];
+    const struct htp_tensor * dst      = octx->dst;
+
+    const struct htp_dsv4_hc_post_kernel_params * kparams =
+        (const struct htp_dsv4_hc_post_kernel_params *) octx->kernel_params;
+
+    const uint32_t n_threads = kparams->n_threads;
+    if (!htp_ops_context_set_n_threads_div(octx, n_threads, kparams->div_n_threads)) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    uint32_t token_start     = 0;
+    uint32_t n_tokens        = x->ne[1];
+    uint32_t total_jobs      = kparams->total_jobs;
+    uint32_t jobs_per_thread = kparams->jobs_per_thread;
+    if (octx->ctx->mdev.count > 1) {
+        const bool can_split = htp_tensor_mdev_data_aligned(dst) && ((dst->nb[2] & (HTP_TENSOR_MDEV_LINE_SIZE - 1)) == 0);
+        const struct htp_tensor_mdev_range range = htp_tensor_mdev_partition(n_tokens, can_split ? 1 : 0, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        token_start     = range.start;
+        n_tokens        = range.count;
+        total_jobs      = n_tokens * kparams->n_chunks;
+        jobs_per_thread = fastdiv(total_jobs + n_threads - 1, &kparams->div_n_threads);
+    }
+
+    if (n_tokens == 0) {
+        return HTP_STATUS_OK;
+    }
+
+    struct dsv4_hc_post_context dctx = {
+        .octx            = octx,
+        .kparams         = kparams,
+        .token_start     = token_start,
+        .total_jobs      = total_jobs,
+        .jobs_per_thread = jobs_per_thread,
+        .vtcm            = (uint8_t *) octx->ctx->vtcm_base,
+    };
+
+    work_queue_func_t worker = NULL;
+    switch (kparams->kernel_type) {
+        case HTP_DSV4_HC_POST_KERNEL_IDENTITY: worker = dsv4_hc_post_thread_identity_f32; break;
+        case HTP_DSV4_HC_POST_KERNEL_COMB:     worker = dsv4_hc_post_thread_comb_f32;     break;
+        default: return HTP_STATUS_INVAL_PARAMS;
+    }
+
+    FARF(HIGH, "dsv4-hc-post: (%ux%u) (%ux%ux%u) -> (%ux%ux%u) : ktype %u tokens %u:%u nth %u\n",
+         x->ne[0], x->ne[1], residual->ne[0], residual->ne[1], residual->ne[2],
+         dst->ne[0], dst->ne[1], dst->ne[2], kparams->kernel_type,
+         token_start, token_start + n_tokens, n_threads);
+
+    work_queue_run(octx->ctx->work_queue, worker, &dctx, n_threads);
+    return HTP_STATUS_OK;
+}

ggml/src/ggml-hexagon/htp/dsv4-hc-ops.h

@@ -0,0 +1,77 @@
+#ifndef HTP_DSV4_HC_OPS_H
+#define HTP_DSV4_HC_OPS_H
+
+#include <stdint.h>
+
+#include "hex-fastdiv.h"
+#include "htp-ops.h"
+
+#define HTP_DSV4_HC_CHUNK_ELEMS 32
+#define HTP_DSV4_HC_MAX_CHUNKS 128
+#define HTP_DSV4_HC_PRE_GATED_ROWS 9
+#define HTP_DSV4_HC_PRE_PLAIN_ROWS 6
+
+enum htp_dsv4_hc_pre_kernel_type {
+    HTP_DSV4_HC_PRE_KERNEL_PLAIN = 0,
+    HTP_DSV4_HC_PRE_KERNEL_GATED = 1,
+};
+
+struct htp_dsv4_hc_pre_kernel_params {
+    uint32_t kernel_type;
+    uint32_t n_threads;
+    uint32_t n_chunks;
+    uint32_t total_jobs;
+    uint32_t jobs_per_thread;
+
+    uint32_t dma_block_chunks;
+    uint32_t dma_row_stride;
+    uint32_t dma_buffer_size;
+    uint32_t vtcm_size;
+
+    struct fastdiv_values div_n_chunks;
+    struct fastdiv_values div_n_threads;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_dsv4_hc_pre_kernel_params) <= 128, "htp_dsv4_hc_pre_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_dsv4_hc_pre_kernel_params) <= 128, "htp_dsv4_hc_pre_kernel_params is too large for kernel_params blob");
+#endif
+
+// dma buffer holds residual rows 0..3, x at row 4, dst rows 5..8, post/comb at row 9
+#define HTP_DSV4_HC_POST_ROWS 10
+
+enum htp_dsv4_hc_post_kernel_type {
+    HTP_DSV4_HC_POST_KERNEL_IDENTITY = 0,
+    HTP_DSV4_HC_POST_KERNEL_COMB     = 1,
+};
+
+struct htp_dsv4_hc_post_kernel_params {
+    uint32_t kernel_type;
+    uint32_t n_threads;
+    uint32_t n_chunks;
+    uint32_t total_jobs;
+    uint32_t jobs_per_thread;
+    uint32_t has_comb;
+
+    uint32_t dma_block_chunks;
+    uint32_t dma_row_stride;
+    uint32_t dma_buffer_size;
+    uint32_t vtcm_size;
+
+    struct fastdiv_values div_n_chunks;
+    struct fastdiv_values div_n_threads;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_dsv4_hc_post_kernel_params) <= 128, "htp_dsv4_hc_post_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_dsv4_hc_post_kernel_params) <= 128, "htp_dsv4_hc_post_kernel_params is too large for kernel_params blob");
+#endif
+
+struct htp_ops_context;
+
+int op_dsv4_hc_pre(struct htp_ops_context * octx);
+int op_dsv4_hc_post(struct htp_ops_context * octx);
+
+#endif // HTP_DSV4_HC_OPS_H

ggml/src/ggml-hexagon/htp/glu-ops.c

@@ -186,53 +186,6 @@ static void swiglu_clamp_f32(const float * restrict src0,
 static const float GELU_COEF_A     = 0.044715f;
 static const float SQRT_2_OVER_PI  = 0.79788456080286535587989211986876f;
 
-static inline HVX_Vector hvx_vec_fast_sigmoid_f32_2it(HVX_Vector v) {
-    v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
-    v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3));
-
-    HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
-    HVX_Vector x      = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
-    HVX_Vector xx     = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x);
-
-    HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2));
-    v1            = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
-
-    HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1));
-    v2            = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx);
-    v2            = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x);
-
-    HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
-    v3            = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
-
-    HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
-    HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
-
-    // Newton-Raphson with 2 iterations
-    HVX_Vector two_sf = hvx_vec_splat_f32(2.0f);
-    HVX_Vector i_sf   = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5);
-    HVX_Vector r_qf   = Q6_Vqf32_vmpy_VsfVsf(
-        i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
-    r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
-        r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
-    HVX_Vector res = Q6_Vsf_equals_Vqf32(r_qf);
-
-    res = Q6_Vqf32_vmpy_VsfVsf(v3, res);
-
-    return Q6_Vsf_equals_Vqf32(res);
-}
-
-static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard_2it(HVX_Vector v,
-                                                            HVX_Vector one,
-                                                            HVX_Vector max_exp,
-                                                            HVX_Vector min_exp) {
-    const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v);
-    const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp);
-
-    HVX_Vector out = hvx_vec_fast_sigmoid_f32_2it(v);
-    out            = Q6_V_vmux_QVV(pred_max, out, one);
-    return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero());
-}
-
 static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) {
     assert((unsigned long) dst  % 128 == 0);
     assert((unsigned long) src0 % 128 == 0);

ggml/src/ggml-hexagon/htp/htp-ctx.h

@@ -131,21 +131,34 @@ static inline void htp_ops_context_reset(struct htp_ops_context * octx, struct h
     octx->n_threads_div = ctx->n_threads_div;
 }
 
-static inline bool htp_ops_context_set_n_threads(struct htp_ops_context * octx, uint32_t n_threads) {
+static inline bool htp_ops_context_set_n_threads_div(
+        struct htp_ops_context * octx,
+        uint32_t n_threads,
+        struct fastdiv_values n_threads_div) {
     if (n_threads == 0 || n_threads > octx->ctx->n_threads) {
         return false;
     }
 
-    if (n_threads != octx->n_threads) {
-        octx->n_threads = n_threads;
-        octx->n_threads_div = n_threads == octx->ctx->n_threads
-            ? octx->ctx->n_threads_div
-            : init_fastdiv_values(n_threads);
-    }
+    octx->n_threads     = n_threads;
+    octx->n_threads_div = n_threads_div;
 
     return true;
 }
 
+static inline bool htp_ops_context_set_n_threads(struct htp_ops_context * octx, uint32_t n_threads) {
+    if (n_threads == 0 || n_threads > octx->ctx->n_threads) {
+        return false;
+    }
+    if (n_threads == octx->n_threads) {
+        return true;
+    }
+
+    const struct fastdiv_values n_threads_div = n_threads == octx->ctx->n_threads
+        ? octx->ctx->n_threads_div
+        : init_fastdiv_values(n_threads);
+    return htp_ops_context_set_n_threads_div(octx, n_threads, n_threads_div);
+}
+
 static inline void htp_ops_context_set_status(struct htp_ops_context * octx, int status) {
     if (status > HTP_STATUS_OK && octx->status == HTP_STATUS_OK) {
         octx->status = status;
@@ -174,6 +187,7 @@ int op_argsort(struct htp_ops_context * octx);
 int op_top_k(struct htp_ops_context * octx);
 int op_ssm_conv(struct htp_ops_context * octx);
 int op_cumsum(struct htp_ops_context * octx);
+int op_top_k_mask(struct htp_ops_context * octx);
 int op_fill(struct htp_ops_context * octx);
 int op_concat(struct htp_ops_context * octx);
 int op_diag(struct htp_ops_context * octx);
@@ -189,5 +203,6 @@ int op_pool_1d(struct htp_ops_context * octx);
 int op_allreduce(struct htp_ops_context * octx);
 int op_roll(struct htp_ops_context * octx);
 int op_upscale(struct htp_ops_context * octx);
+int op_lightning_indexer(struct htp_ops_context * octx);
 
 #endif /* HTP_CTX_H */

ggml/src/ggml-hexagon/htp/htp-ops.h

@@ -117,6 +117,10 @@ enum htp_op_code {
     HTP_OP_UNARY_CEIL,
     HTP_OP_UNARY_SGN,
     HTP_OP_UPSCALE,
+    HTP_OP_DSV4_HC_PRE,
+    HTP_OP_DSV4_HC_POST,
+    HTP_OP_TOP_K_MASK,
+    HTP_OP_LIGHTNING_INDEXER,
 
     HTP_OP_INVALID
 };

ggml/src/ggml-hexagon/htp/hvx-sigmoid.h

@@ -52,6 +52,53 @@ static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard(HVX_Vector v,
     return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero());
 }
 
+static inline HVX_Vector hvx_vec_fast_sigmoid_f32_2it(HVX_Vector v) {
+    v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
+    v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3));
+
+    HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v));
+    HVX_Vector x      = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int));
+    HVX_Vector xx     = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x);
+
+    HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2));
+    v1            = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F));
+
+    HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1));
+    v2            = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx);
+    v2            = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x);
+
+    HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1));
+    v3            = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24);
+
+    HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1));
+    HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4));
+
+    // Newton-Raphson with 2 iterations
+    HVX_Vector two_sf = hvx_vec_splat_f32(2.0f);
+    HVX_Vector i_sf   = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5);
+    HVX_Vector r_qf   = Q6_Vqf32_vmpy_VsfVsf(
+        i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5)))));
+    r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32(
+        r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5))));
+    HVX_Vector res = Q6_Vsf_equals_Vqf32(r_qf);
+
+    res = Q6_Vqf32_vmpy_VsfVsf(v3, res);
+
+    return Q6_Vsf_equals_Vqf32(res);
+}
+
+static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard_2it(HVX_Vector v,
+                                                            HVX_Vector one,
+                                                            HVX_Vector max_exp,
+                                                            HVX_Vector min_exp) {
+    const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v);
+    const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp);
+
+    HVX_Vector out = hvx_vec_fast_sigmoid_f32_2it(v);
+    out            = Q6_V_vmux_QVV(pred_max, out, one);
+    return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero());
+}
+
 static inline HVX_Vector hvx_vec_tanh_f32(HVX_Vector x) {
     // tanh(x) = 2 * sigmoid(2x) - 1
     HVX_Vector two = hvx_vec_splat_f32(2.0f);

ggml/src/ggml-hexagon/htp/lightning-indexer-ops.h

@@ -0,0 +1,213 @@
+#ifndef HTP_LIGHTNING_INDEXER_OPS_H
+#define HTP_LIGHTNING_INDEXER_OPS_H
+
+#include <stdint.h>
+#include <stddef.h>
+#include <stdbool.h>
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#define HTP_LID_KV_BLOCK   64 // kv rows per block (HVX and HMX Bc)
+#define HTP_LID_TOKENS_MAX 32 // max tokens per Q tile (HVX splat path)
+#define HTP_LID_HMX_TOKENS_MAX 512 // HMX can keep more tokens; VTCM decides
+#define HTP_LID_HMX_KV_BLOCK_MAX 2048
+#define HTP_LID_HMX_MAX_QSETS    8    // streams in one HMX pipeline run
+#define HTP_LID_N_BUF      4  // K / mask / output buffers per thread (two blocks computed, two in flight)
+
+#ifndef HMX_FP16_TILE_N_COLS
+#define HMX_FP16_TILE_N_COLS 32
+#endif
+#ifndef HMX_FP16_TILE_SIZE
+#define HMX_FP16_TILE_SIZE 2048
+#endif
+
+enum htp_lid_kernel_type {
+    HTP_LID_KERNEL_HVX = 0,
+    HTP_LID_KERNEL_HMX = 1,
+};
+
+struct htp_lid_kernel_params {
+    uint8_t  kernel_type;       // enum htp_lid_kernel_type
+    uint32_t n_threads;
+    uint32_t n_tokens_tile;     // tokens per Q tile
+    uint32_t n_kv_blocks;       // HVX: per stream, HMX: per run of n_qsets sets
+    uint32_t blocks_per_thread;
+    uint32_t vtcm_size;
+    uint32_t k_row_size;        // K row bytes
+    uint32_t kv_block;          // HMX: kv rows per block (multiple of 64)
+    uint32_t n_qsets;           // HMX: Q sets kept in VTCM (streams run in one pipeline)
+    uint32_t n_qbuf;            // HMX: Q buffers, 2: the next run's Q is prepped during the current run
+    uint32_t kv_seg;            // HMX: rows per set, n_kv padded to 64
+    uint32_t n_tt;              // HMX: token tiles per stream
+
+    struct fastdiv_values div_n_head;
+    struct fastdiv_values div_m_ne3;  // mask streams (broadcast)
+    struct fastdiv_values div_kv_seg; // HMX
+    struct fastdiv_values div_n_tt;   // HMX
+    struct fastdiv_values div_n_ch;   // HMX: channel tiles per Q set
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_lid_kernel_params) <= 128, "htp_lid_kernel_params is too large");
+#else
+_Static_assert(sizeof(struct htp_lid_kernel_params) <= 128, "htp_lid_kernel_params is too large");
+#endif
+
+// VTCM layout shared by host (fit check) and device (buffer offsets).
+// Shared regions: Q splats (one f16 vector per q element) and head weight splats for the current token tile.
+// Per-thread regions: Q prep staging, K / mask / output buffers, two transposed K blocks and transpose scratch.
+struct htp_lid_vtcm_layout {
+    size_t off_q;
+    size_t off_w;
+    size_t off_qs;
+    size_t off_k;
+    size_t off_kt;
+    size_t off_m;
+    size_t off_o;
+
+    size_t size_q_row;   // f16 Q row, padded
+    size_t size_qs_row;  // f32 Q staging row, padded
+    size_t size_k_row;   // f16 K row, padded
+    size_t size_k_raw_row;
+    size_t size_m_row;   // f16 mask row of one kv block, padded
+    size_t size_o_row;   // f32 output row of one kv block, padded
+
+    size_t size_qs;      // per thread
+    size_t size_k_buf;   // one K buffer
+    size_t size_kt;      // per thread: two transposed K blocks, two 64x64 scratch tiles, Q8_0 block scales
+    size_t size_m_buf;   // one mask buffer
+    size_t size_o_buf;   // one output buffer
+
+    size_t total_bytes;
+};
+
+static inline void htp_lid_vtcm_layout_build(struct htp_lid_vtcm_layout * L, size_t D, size_t n_head, size_t n_tok, size_t n_threads, size_t k_row_size) {
+    L->size_q_row  = hex_round_up(D * sizeof(uint16_t), 128);
+    L->size_qs_row = hex_round_up(D * sizeof(float), 128);
+    L->size_k_row  = hex_round_up(D * sizeof(uint16_t), 128);
+    L->size_k_raw_row = L->size_k_row > hex_round_up(k_row_size, 128) ? L->size_k_row : hex_round_up(k_row_size, 128);
+    L->size_m_row  = hex_round_up(HTP_LID_KV_BLOCK * sizeof(uint16_t), 128);
+    L->size_o_row  = hex_round_up(HTP_LID_KV_BLOCK * sizeof(float), 128);
+
+    L->size_qs    = 2 * (n_head * L->size_qs_row + hex_round_up(n_head * sizeof(float), 128)) + L->size_q_row;
+    L->size_k_buf = HTP_LID_KV_BLOCK * L->size_k_raw_row;
+    L->size_kt    = (2 * hex_round_up(D, HTP_LID_KV_BLOCK) + 2 * HTP_LID_KV_BLOCK + D / 32) * 128;
+    L->size_m_buf = n_tok * L->size_m_row;
+    L->size_o_buf = n_tok * L->size_o_row;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_q,  n_tok * n_head * D * 128);
+    VTCM_LAYOUT_ALLOC(off, off_w,  n_tok * n_head * 128);
+    VTCM_LAYOUT_ALLOC(off, off_qs, n_threads * L->size_qs);
+    VTCM_LAYOUT_ALLOC(off, off_k,  n_threads * HTP_LID_N_BUF * L->size_k_buf);
+    VTCM_LAYOUT_ALLOC(off, off_kt, n_threads * L->size_kt);
+    VTCM_LAYOUT_ALLOC(off, off_m,  n_threads * HTP_LID_N_BUF * L->size_m_buf);
+    VTCM_LAYOUT_ALLOC(off, off_o,  n_threads * HTP_LID_N_BUF * L->size_o_buf);
+
+    L->total_bytes = off;
+}
+
+static inline size_t htp_lid_compute_vtcm_usage(size_t D, size_t n_head, size_t n_tok, size_t n_threads, size_t k_row_size) {
+    struct htp_lid_vtcm_layout L;
+    htp_lid_vtcm_layout_build(&L, D, n_head, n_tok, n_threads, k_row_size);
+    return L.total_bytes;
+}
+
+// HMX path. K is the activation (kv rows), Q the weights (one output channel per (token, head)).
+// mm1: S = relu(K Q^T) * |w| per channel (HMX convert scale + pos clamp)
+// mm2: O = S G, G[(t,h), t] = sign(w[t,h]). O is [kv, token], f16.
+// Q weights, G, scales and w are per Q set (one per stream when all tokens fit one tile).
+struct htp_lid_hmx_vtcm_layout {
+    size_t off_qw;       // Q weight tiles [n_qsets][n_ch][n_dot]
+    size_t off_g;        // G weight tiles [n_qsets][n_ch]
+    size_t off_sc;       // per channel scales, 256 bytes per channel tile, [n_qsets][n_ch]
+    size_t off_sc_id;    // identity scales
+    size_t off_s;        // S tiles of one block and token tile, [n_rt][n_sg]
+    size_t off_w;        // w of the token tile, [n_qsets][n_tok][n_head] f32
+    size_t off_qs;       // per thread: f32 and f16 staging for 32 Q rows, gather scratch
+    size_t off_k_raw[2]; // K rows, [Bc][size_k_row]
+    size_t off_k_t[2];   // K tiles [n_rt][n_dot]
+    size_t off_o[2];     // O tiles [n_rt][n_u]
+    size_t off_m[4];     // mask rows [n_tok][size_m_row]
+    size_t off_out[2];   // f32 output rows [n_tok][size_o_row]
+
+    size_t size_qs_row;  // f32 Q row, padded
+    size_t size_qh_row;  // f16 Q row, padded
+    size_t size_k_row;   // f16 K row, padded
+    size_t size_k_raw_row;
+    size_t size_m_row;
+    size_t size_o_row;
+    size_t size_qs;      // per thread
+    size_t size_w;       // per Q set
+
+    uint32_t n_ch;       // channel tiles: ceil(n_tok * n_head / 32)
+    uint32_t n_u;        // token tiles: ceil(n_tok / 32)
+    uint32_t n_dot;      // D / 32
+    uint32_t n_rt;       // kv row tiles per block: Bc / 32
+    uint32_t n_sg;       // S tiles per (row tile, token tile): min(n_head, n_ch)
+
+    size_t total_bytes;
+};
+
+static inline void htp_lid_hmx_vtcm_layout_build(struct htp_lid_hmx_vtcm_layout * L,
+                                                  size_t D, size_t n_head, size_t n_tok, size_t Bc, size_t n_qsets, size_t n_threads, size_t k_row_size) {
+    L->n_ch  = (uint32_t) ((n_tok * n_head + 31) / 32);
+    L->n_u   = (uint32_t) ((n_tok + 31) / 32);
+    L->n_dot = (uint32_t) (D / 32);
+    L->n_rt  = (uint32_t) (Bc / 32);
+    L->n_sg  = (uint32_t) (n_head < L->n_ch ? n_head : L->n_ch);
+
+    L->size_qs_row = hex_round_up(D * sizeof(float), 128);
+    L->size_qh_row = hex_round_up(D * sizeof(uint16_t), 128);
+    L->size_k_row  = hex_round_up(D * sizeof(uint16_t), 128);
+    L->size_k_raw_row = L->size_k_row > hex_round_up(k_row_size, 128) ? L->size_k_row : hex_round_up(k_row_size, 128);
+    L->size_m_row  = hex_round_up(Bc * sizeof(uint16_t), 128);
+    L->size_o_row  = hex_round_up(Bc * sizeof(float), 128);
+    L->size_qs     = 2 * 32 * L->size_qs_row + 32 * L->size_qh_row + hex_round_up(Bc / 64 * 128, 128);
+    L->size_w      = hex_round_up(n_tok * n_head * sizeof(float), 128);
+
+    const size_t tile = HMX_FP16_TILE_SIZE;
+
+    // tile regions first: they need 2 KB alignment
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC(off, off_qw,     n_qsets * L->n_ch * L->n_dot * tile);
+    VTCM_LAYOUT_ALLOC(off, off_g,      n_qsets * L->n_ch * tile);
+    VTCM_LAYOUT_ALLOC(off, off_s,      (size_t) L->n_rt * L->n_sg * tile);
+    VTCM_LAYOUT_ALLOC(off, off_k_t[0], (size_t) L->n_rt * L->n_dot * tile);
+    VTCM_LAYOUT_ALLOC(off, off_k_t[1], (size_t) L->n_rt * L->n_dot * tile);
+    VTCM_LAYOUT_ALLOC(off, off_o[0],   (size_t) L->n_rt * L->n_u * tile);
+    VTCM_LAYOUT_ALLOC(off, off_o[1],   (size_t) L->n_rt * L->n_u * tile);
+    VTCM_LAYOUT_ALLOC(off, off_sc,     n_qsets * L->n_ch * 256);
+    VTCM_LAYOUT_ALLOC(off, off_sc_id,  256);
+    VTCM_LAYOUT_ALLOC(off, off_w,      n_qsets * L->size_w);
+    VTCM_LAYOUT_ALLOC(off, off_qs,     n_threads * L->size_qs);
+    for (int i = 0; i < 2; i++) {
+        VTCM_LAYOUT_ALLOC(off, off_k_raw[i], Bc * L->size_k_raw_row);
+    }
+    for (int i = 0; i < 4; i++) {
+        VTCM_LAYOUT_ALLOC(off, off_m[i], n_tok * L->size_m_row);
+    }
+    for (int i = 0; i < 2; i++) {
+        VTCM_LAYOUT_ALLOC(off, off_out[i], n_tok * L->size_o_row);
+    }
+
+    L->total_bytes = off;
+}
+
+static inline size_t htp_lid_hmx_compute_vtcm_usage(size_t D, size_t n_head, size_t n_tok, size_t Bc, size_t n_qsets, size_t n_threads, size_t k_row_size) {
+    struct htp_lid_hmx_vtcm_layout L;
+    htp_lid_hmx_vtcm_layout_build(&L, D, n_head, n_tok, Bc, n_qsets, n_threads, k_row_size);
+    return L.total_bytes;
+}
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* HTP_LIGHTNING_INDEXER_OPS_H */

ggml/src/ggml-hexagon/htp/main.c

@@ -34,6 +34,7 @@
 #include "work-queue.h"
 #include "hex-profile.h"
 #include "allreduce-ops.h"
+#include "dsv4-hc-ops.h"
 #include "htp-fence.h"
 
 #define HMX_QUEUE_CAPACITY     128
@@ -1007,6 +1008,9 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_CUMSUM:
             return op_cumsum(octx);
 
+        case HTP_OP_TOP_K_MASK:
+            return op_top_k_mask(octx);
+
         case HTP_OP_FILL:
             return op_fill(octx);
 
@@ -1040,12 +1044,21 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_POOL_1D:
             return op_pool_1d(octx);
 
+        case HTP_OP_DSV4_HC_PRE:
+            return op_dsv4_hc_pre(octx);
+
+        case HTP_OP_DSV4_HC_POST:
+            return op_dsv4_hc_post(octx);
+
         case HTP_OP_CONCAT:
             return op_concat(octx);
 
         case HTP_OP_GATED_DELTA_NET:
             return op_gated_delta_net(octx);
 
+        case HTP_OP_LIGHTNING_INDEXER:
+            return op_lightning_indexer(octx);
+
         case HTP_OP_TRI:
             return op_unary(octx);
 

ggml/src/ggml-hexagon/htp/top-k-mask-ops.c

@@ -0,0 +1,665 @@
+#pragma clang diagnostic ignored "-Wunused-variable"
+#pragma clang diagnostic ignored "-Wunused-function"
+#pragma clang diagnostic ignored "-Wunused-but-set-variable"
+
+#include <HAP_farf.h>
+#include <HAP_perf.h>
+#include <stdatomic.h>
+
+#define GGML_COMMON_DECL_C
+#include "ggml-common.h"
+#include "hex-common.h"
+#include "hex-profile.h"
+#include "htp-ctx.h"
+#include "htp-ops.h"
+#include "htp-tensor.h"
+#include "htp-vtcm.h"
+#include "hvx-types.h"
+#include "hvx-utils.h"
+#include "dma-queue.h"
+#include "top-k-mask-ops.h"
+
+#define TOP_K_MASK_NINF_F32 0xff800000u // f32 score compare
+
+struct htp_top_k_mask_context {
+    struct htp_ops_context *                    octx;
+    const struct htp_top_k_mask_kernel_params * kparams;
+
+    uint32_t rows_per_thread;
+    uint32_t row0, row1;           // rows this device handles
+    uint32_t off_extra, off_score; // parts inside one row input slot
+    uint32_t col0, ncols_out;      // PICKS: columns this device writes back
+    atomic_uint barrier;           // PICKS: arrivals, barrier k waits for k * n_threads
+
+    uint8_t * blocks;
+    uint8_t * mask;
+    uint8_t * rowin;
+    uint8_t * tiles;
+    uint8_t * mtiles;
+    uint8_t * scratch;
+};
+
+// DMA queue pop with a trace event around the wait
+static inline void top_k_mask_pop(dma_queue * dma_q, struct htp_thread_trace * tr, uint16_t info) {
+    htp_trace_event_start(tr, HTP_TRACE_EVT_DMA, info);
+    dma_queue_pop(dma_q);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_DMA, info);
+}
+
+// Lanes 0..31 as words, in registers: inclusive scan of ones, minus one
+static inline HVX_Vector top_k_mask_lane_iota(void) {
+    const HVX_Vector zero = Q6_V_vzero();
+    const HVX_Vector one  = Q6_V_vsplat_R(1);
+    HVX_Vector       v    = one;
+    v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 4));
+    v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 8));
+    v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 16));
+    v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 32));
+    v = Q6_Vw_vadd_VwVw(v, Q6_V_vlalign_VVR(v, zero, 64));
+    return Q6_Vw_vsub_VwVw(v, one);
+}
+
+static __attribute__((noinline)) void top_k_mask_fill(uint8_t * restrict dst, uint32_t nbytes, uint32_t value) {
+    const HVX_Vector v    = Q6_V_vsplat_R(value);
+    HVX_Vector *     vdst = (HVX_Vector *) dst;
+    const uint32_t   nvec = nbytes / VLEN;
+
+    #pragma unroll(4)
+    for (uint32_t i = 0; i < nvec; i++) {
+        vdst[i] = v;
+    }
+}
+
+#define TOP_K_MASK_DEAD 0x7fffffffu // cell value of a dead lane: c - c0 >=u ncols for every tile, signed max
+
+// Byte offsets of the cells in [c0, c0 + ncols) relative to c0, 0 for the others
+static inline HVX_Vector top_k_mask_keep_off(HVX_Vector c, HVX_Vector v_c0, HVX_Vector v_ncols, const bool is_f16) {
+    const HVX_Vector cr = Q6_Vw_vsub_VwVw(c, v_c0);
+    return Q6_Vw_vasl_VwR(Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VuwVuw(v_ncols, cr), cr, Q6_V_vzero()), is_f16 ? 1 : 2);
+}
+
+// Mask value of each cell lane in [c0, c0 + ncols) into tmp (predicated gather), returns the scatter offsets
+static inline HVX_Vector top_k_mask_keep_gather(HVX_Vector c, const uint8_t * restrict mrow, HVX_Vector v_c0, HVX_Vector v_ncols,
+                                                uint32_t mu, HVX_Vector * restrict tmp, const bool is_f16) {
+    const HVX_VectorPred q   = Q6_Q_vcmp_gt_VuwVuw(v_ncols, Q6_Vw_vsub_VwVw(c, v_c0));
+    const HVX_Vector     off = top_k_mask_keep_off(c, v_c0, v_ncols, is_f16);
+
+    // Halfword data with word offsets: both halves of the pair hold the same offsets, the word predicate covers both halfwords of a lane
+    if (is_f16) {
+        Q6_vgather_AQRMWw(tmp, q, (size_t) mrow, mu, Q6_W_vcombine_VV(off, off));
+    } else {
+        Q6_vgather_AQRMVw(tmp, q, (size_t) mrow, mu, off);
+    }
+    return off;
+}
+
+static inline void top_k_mask_keep_scatter(HVX_Vector c, HVX_Vector off, uint8_t * restrict tile, HVX_Vector v_c0, HVX_Vector v_ncols,
+                                           uint32_t mu, HVX_Vector val, const bool is_f16) {
+    const HVX_VectorPred q = Q6_Q_vcmp_gt_VuwVuw(v_ncols, Q6_Vw_vsub_VwVw(c, v_c0));
+    if (is_f16) {
+        Q6_vscatter_QRMWwV(q, (size_t) tile, mu, Q6_W_vcombine_VV(off, off), val);
+    } else {
+        Q6_vscatter_QRMVwV(q, (size_t) tile, mu, off, val);
+    }
+}
+
+// Keep 4 cell vectors (dead lanes = TOP_K_MASK_DEAD): tile[c - c0] = mrow[c - c0] for c in [c0, c0 + ncols). One gather sync.
+// flag: no mask, tile[c - c0] = all ones (keep flag)
+static inline void top_k_mask_keep4(HVX_Vector c0v, HVX_Vector c1v, HVX_Vector c2v, HVX_Vector c3v, uint8_t * restrict tile,
+                                    const uint8_t * restrict mrow, HVX_Vector v_c0, HVX_Vector v_ncols, uint32_t mu,
+                                    HVX_Vector * restrict tmp, const bool is_f16, const bool flag) {
+    if (flag) {
+        const HVX_Vector ones = Q6_V_vsplat_R(0xffffffffu);
+        top_k_mask_keep_scatter(c0v, top_k_mask_keep_off(c0v, v_c0, v_ncols, is_f16), tile, v_c0, v_ncols, mu, ones, is_f16);
+        top_k_mask_keep_scatter(c1v, top_k_mask_keep_off(c1v, v_c0, v_ncols, is_f16), tile, v_c0, v_ncols, mu, ones, is_f16);
+        top_k_mask_keep_scatter(c2v, top_k_mask_keep_off(c2v, v_c0, v_ncols, is_f16), tile, v_c0, v_ncols, mu, ones, is_f16);
+        top_k_mask_keep_scatter(c3v, top_k_mask_keep_off(c3v, v_c0, v_ncols, is_f16), tile, v_c0, v_ncols, mu, ones, is_f16);
+        return;
+    }
+    const HVX_Vector o0 = top_k_mask_keep_gather(c0v, mrow, v_c0, v_ncols, mu, tmp + 0, is_f16);
+    const HVX_Vector o1 = top_k_mask_keep_gather(c1v, mrow, v_c0, v_ncols, mu, tmp + 1, is_f16);
+    const HVX_Vector o2 = top_k_mask_keep_gather(c2v, mrow, v_c0, v_ncols, mu, tmp + 2, is_f16);
+    const HVX_Vector o3 = top_k_mask_keep_gather(c3v, mrow, v_c0, v_ncols, mu, tmp + 3, is_f16);
+    hvx_gather_sync(tmp + 3);
+    top_k_mask_keep_scatter(c0v, o0, tile, v_c0, v_ncols, mu, tmp[0], is_f16);
+    top_k_mask_keep_scatter(c1v, o1, tile, v_c0, v_ncols, mu, tmp[1], is_f16);
+    top_k_mask_keep_scatter(c2v, o2, tile, v_c0, v_ncols, mu, tmp[2], is_f16);
+    top_k_mask_keep_scatter(c3v, o3, tile, v_c0, v_ncols, mu, tmp[3], is_f16);
+}
+
+// Pick vector j: picks below j1 and n_pick are valid, returns them with dead lanes marked. v_rem = j1 - lane index.
+static inline HVX_Vector top_k_mask_load_picks(const int32_t * restrict tk, uint32_t j, HVX_Vector v_rem, HVX_Vector v_npick) {
+    const HVX_Vector     p_raw = hvx_vmem(tk + j);
+    // lane + j <u j1 is the tail, p <u n_pick drops p < 0 and p >= n_pick
+    const HVX_VectorPred valid = Q6_Q_vcmp_gtand_QVuwVuw(Q6_Q_vcmp_gt_VwVw(v_rem, Q6_V_vsplat_R(j)), v_npick, p_raw);
+    return Q6_V_vmux_QVV(valid, p_raw, Q6_V_vsplat_R(TOP_K_MASK_DEAD));
+}
+
+// Scatter the kept cells of picks [j0, j1) (and the extras) into the tile of columns [c0, c0 + ncols). slot holds the row inputs, mrow the mask columns.
+// Gathers are issued in batches with one sync per batch: tmp[0..3] scores/blocks, tmp[4] score with blocks, tmp[5..8] mask values.
+static inline __attribute__((always_inline)) void top_k_mask_scatter_tile(
+        const struct htp_top_k_mask_context * tctx, const uint8_t * restrict slot, const uint8_t * restrict mrow, uint8_t * restrict tile,
+        HVX_Vector * restrict tmp, uint32_t c0, uint32_t ncols, uint32_t j0, uint32_t j1, bool do_extra,
+        const bool is_f16, const bool flag) {
+    const struct htp_top_k_mask_kernel_params * kp = tctx->kparams;
+
+    const int32_t * restrict tk  = (const int32_t *) slot;
+    const int32_t * restrict ex  = (const int32_t *) (slot + tctx->off_extra);
+    const uint8_t * restrict sc  = slot + tctx->off_score;
+    const uint8_t * restrict blk = tctx->blocks;
+
+    const bool     has_score  = kp->flags & HTP_TOP_K_MASK_HAS_SCORE;
+    const bool     has_blocks = kp->flags & HTP_TOP_K_MASK_HAS_BLOCKS;
+    const uint32_t n_blk      = kp->n_blk;
+    const uint32_t n_extra    = do_extra ? kp->n_extra : 0;
+    const uint32_t mu         = ncols * kp->elem_size - 1;
+    const uint32_t sc_mu      = kp->score_bytes - 1;
+    const uint32_t blk_mu     = kp->blocks_bytes - 1;
+    const uint32_t nb1_h      = ((uint32_t) kp->blocks_nb1 << 16) | kp->blocks_nb1;
+
+    const HVX_Vector v_lane  = top_k_mask_lane_iota(); // in the HVX_COMP scope of the caller
+    const HVX_Vector v_c0    = Q6_V_vsplat_R(c0);
+    const HVX_Vector v_ncols = Q6_V_vsplat_R(ncols);
+    const HVX_Vector v_npick = Q6_V_vsplat_R(kp->n_pick);
+    const HVX_Vector v_pmax  = Q6_V_vsplat_R(kp->n_pick - 1); // dead picks gather the last entry, every gather stays in its region
+    const HVX_Vector v_rem   = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(j1), v_lane);
+    const HVX_Vector v_ninf  = Q6_V_vsplat_R(TOP_K_MASK_NINF_F32);
+    const HVX_Vector v_dead  = Q6_V_vsplat_R(TOP_K_MASK_DEAD);
+
+    HVX_Vector * restrict mtmp = tmp + 5;
+
+    if (has_blocks) {
+        for (uint32_t j = j0; j < j1; j += VLEN_FP32) {
+            HVX_Vector       cell = top_k_mask_load_picks(tk, j, v_rem, v_npick);
+            const HVX_Vector p    = Q6_Vw_vmin_VwVw(cell, v_pmax);
+            if (has_score) {
+                Q6_vgather_ARMVw(tmp + 4, (size_t) sc, sc_mu, Q6_Vw_vasl_VwR(p, 2));
+            }
+            const HVX_Vector boff = Q6_Vw_vmpyi_VwRh(p, nb1_h);
+            for (uint32_t k = 0; k < n_blk; k += 4) {
+                // blocks k..k+3 of every pick; the missing ones (k + i >= n_blk) are dead
+                const uint32_t   nk = MIN(n_blk - k, 4);
+                const HVX_Vector bo = Q6_Vw_vadd_VwVw(boff, Q6_V_vsplat_R(k * 4));
+                for (uint32_t i = 0; i < nk; i++) {
+                    Q6_vgather_ARMVw(tmp + i, (size_t) blk, blk_mu, Q6_Vw_vadd_VwVw(bo, Q6_V_vsplat_R(i * 4)));
+                }
+                hvx_gather_sync(tmp + nk - 1);
+                if (has_score && k == 0) {
+                    // IEEE compare: NaN and -inf scores are dead
+                    cell = Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VsfVsf(tmp[4], v_ninf), cell, v_dead);
+                }
+                const HVX_VectorPred dead = Q6_Q_vcmp_eq_VwVw(cell, v_dead);
+                top_k_mask_keep4(Q6_V_vmux_QVV(dead, v_dead, tmp[0]),
+                                 nk > 1 ? Q6_V_vmux_QVV(dead, v_dead, tmp[1]) : v_dead,
+                                 nk > 2 ? Q6_V_vmux_QVV(dead, v_dead, tmp[2]) : v_dead,
+                                 nk > 3 ? Q6_V_vmux_QVV(dead, v_dead, tmp[3]) : v_dead,
+                                 tile, mrow, v_c0, v_ncols, mu, mtmp, is_f16, flag);
+            }
+        }
+    } else {
+        // 4 pick vectors per batch, the vectors at or past j1 are dead and not loaded
+        for (uint32_t j = j0; j < j1; j += 4 * VLEN_FP32) {
+            HVX_Vector c0v = top_k_mask_load_picks(tk, j, v_rem, v_npick);
+            HVX_Vector c1v = j + 1 * VLEN_FP32 < j1 ? top_k_mask_load_picks(tk, j + 1 * VLEN_FP32, v_rem, v_npick) : v_dead;
+            HVX_Vector c2v = j + 2 * VLEN_FP32 < j1 ? top_k_mask_load_picks(tk, j + 2 * VLEN_FP32, v_rem, v_npick) : v_dead;
+            HVX_Vector c3v = j + 3 * VLEN_FP32 < j1 ? top_k_mask_load_picks(tk, j + 3 * VLEN_FP32, v_rem, v_npick) : v_dead;
+            if (has_score) {
+                Q6_vgather_ARMVw(tmp + 0, (size_t) sc, sc_mu, Q6_Vw_vasl_VwR(Q6_Vw_vmin_VwVw(c0v, v_pmax), 2));
+                Q6_vgather_ARMVw(tmp + 1, (size_t) sc, sc_mu, Q6_Vw_vasl_VwR(Q6_Vw_vmin_VwVw(c1v, v_pmax), 2));
+                Q6_vgather_ARMVw(tmp + 2, (size_t) sc, sc_mu, Q6_Vw_vasl_VwR(Q6_Vw_vmin_VwVw(c2v, v_pmax), 2));
+                Q6_vgather_ARMVw(tmp + 3, (size_t) sc, sc_mu, Q6_Vw_vasl_VwR(Q6_Vw_vmin_VwVw(c3v, v_pmax), 2));
+                hvx_gather_sync(tmp + 3);
+                c0v = Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VsfVsf(tmp[0], v_ninf), c0v, v_dead);
+                c1v = Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VsfVsf(tmp[1], v_ninf), c1v, v_dead);
+                c2v = Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VsfVsf(tmp[2], v_ninf), c2v, v_dead);
+                c3v = Q6_V_vmux_QVV(Q6_Q_vcmp_gt_VsfVsf(tmp[3], v_ninf), c3v, v_dead);
+            }
+            top_k_mask_keep4(c0v, c1v, c2v, c3v, tile, mrow, v_c0, v_ncols, mu, mtmp, is_f16, flag);
+        }
+    }
+
+    const HVX_Vector v_rem_ex = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(n_extra), v_lane);
+    for (uint32_t j = 0; j < n_extra; j += VLEN_FP32) {
+        const HVX_VectorPred tail = Q6_Q_vcmp_gt_VwVw(v_rem_ex, Q6_V_vsplat_R(j));
+        top_k_mask_keep4(Q6_V_vmux_QVV(tail, hvx_vmem(ex + j), v_dead), v_dead, v_dead, v_dead, tile, mrow, v_c0, v_ncols, mu, mtmp, is_f16, flag);
+    }
+
+    // release the scatters before DMA reads the tile
+    hvx_gather_sync(tile);
+}
+
+static __attribute__((noinline)) void top_k_mask_scatter_tile_f32(const struct htp_top_k_mask_context * tctx, const uint8_t * restrict slot,
+                                        const uint8_t * restrict mrow, uint8_t * restrict tile, HVX_Vector * restrict tmp,
+                                        uint32_t c0, uint32_t ncols, uint32_t j0, uint32_t j1, bool do_extra) {
+    top_k_mask_scatter_tile(tctx, slot, mrow, tile, tmp, c0, ncols, j0, j1, do_extra, false, false);
+}
+
+static __attribute__((noinline)) void top_k_mask_scatter_tile_f16(const struct htp_top_k_mask_context * tctx, const uint8_t * restrict slot,
+                                        const uint8_t * restrict mrow, uint8_t * restrict tile, HVX_Vector * restrict tmp,
+                                        uint32_t c0, uint32_t ncols, uint32_t j0, uint32_t j1, bool do_extra) {
+    top_k_mask_scatter_tile(tctx, slot, mrow, tile, tmp, c0, ncols, j0, j1, do_extra, true, false);
+}
+
+// PICKS: keep flags (all ones) of the picks [j0, j1) into the flag tile of the whole row, the element size picks the flag width
+static __attribute__((noinline)) void top_k_mask_flag_tile_f32(const struct htp_top_k_mask_context * tctx, const uint8_t * restrict slot,
+                                        uint8_t * restrict flags, HVX_Vector * restrict tmp, uint32_t j0, uint32_t j1, bool do_extra) {
+    top_k_mask_scatter_tile(tctx, slot, NULL, flags, tmp, 0, tctx->kparams->n_kv, j0, j1, do_extra, false, true);
+}
+
+static __attribute__((noinline)) void top_k_mask_flag_tile_f16(const struct htp_top_k_mask_context * tctx, const uint8_t * restrict slot,
+                                        uint8_t * restrict flags, HVX_Vector * restrict tmp, uint32_t j0, uint32_t j1, bool do_extra) {
+    top_k_mask_scatter_tile(tctx, slot, NULL, flags, tmp, 0, tctx->kparams->n_kv, j0, j1, do_extra, true, true);
+}
+
+// out = flag ? out : -inf over nbytes (a multiple of 128) of the row tile
+static __attribute__((noinline)) void top_k_mask_select(uint8_t * restrict out, const uint8_t * restrict flags, uint32_t nbytes, uint32_t ninf) {
+    const HVX_Vector     v_ninf = Q6_V_vsplat_R(ninf);
+    HVX_Vector *         vout   = (HVX_Vector *) out;
+    const HVX_Vector *   vflag  = (const HVX_Vector *) flags;
+    const uint32_t       nvec   = nbytes / VLEN;
+
+    #pragma unroll(4)
+    for (uint32_t i = 0; i < nvec; i++) {
+        vout[i] = Q6_V_vxor_VV(v_ninf, Q6_V_vand_VV(Q6_V_vxor_VV(vout[i], v_ninf), vflag[i]));
+    }
+}
+
+// Push the row inputs of flat row r into slot: top_k, extra, score. Returns the push count.
+static inline uint32_t top_k_mask_push_rowin(const struct htp_top_k_mask_context * tctx, dma_queue * dma_q, uint32_t r, uint8_t * slot) {
+    const struct htp_ops_context *              octx  = tctx->octx;
+    const struct htp_top_k_mask_kernel_params * kp    = tctx->kparams;
+    const struct htp_tensor *                   top_k = octx->src[1];
+    const struct htp_tensor *                   score = octx->src[2];
+    const struct htp_tensor *                   extra = octx->src[4];
+
+    const uint32_t tk_bytes = kp->n_top * sizeof(int32_t);
+    dma_queue_push(dma_q, dma_make_data(slot, top_k->data + r * kp->u.rows.rs_top), tk_bytes, tk_bytes, tk_bytes, 1);
+    uint32_t n = 1;
+
+    if (kp->flags & HTP_TOP_K_MASK_HAS_EXTRA) {
+        const uint32_t ex_bytes = kp->n_extra * sizeof(int32_t);
+        dma_queue_push(dma_q, dma_make_data(slot + tctx->off_extra, extra->data + r * kp->u.rows.rs_extra), ex_bytes, ex_bytes, ex_bytes, 1);
+        n++;
+    }
+
+    if (kp->flags & HTP_TOP_K_MASK_HAS_SCORE) {
+        dma_queue_push(dma_q, dma_make_data(slot + tctx->off_score, score->data + r * kp->u.rows.rs_score), kp->score_bytes, kp->score_bytes, kp->score_bytes, 1);
+        n++;
+    }
+
+    return n;
+}
+
+static inline struct fastdiv_values top_k_mask_fd(const struct htp_top_k_mask_kernel_params * kp, uint32_t k) {
+    const struct fastdiv_values fd = { kp->u.rows.m.dec.mp[k], kp->u.rows.m.dec.l[k] };
+    return fd;
+}
+
+// Mask row of flat row r: (i1, i2, i3) of dst, each taken modulo the mask dims, rows dense above the mask ne1
+static inline uint32_t top_k_mask_mask_row(const struct htp_top_k_mask_kernel_params * kp, uint32_t r) {
+    const uint16_t *            ne  = kp->u.rows.m.dec.ne;
+    const struct fastdiv_values fd1 = top_k_mask_fd(kp, 0);
+    const struct fastdiv_values fd2 = top_k_mask_fd(kp, 1);
+    const struct fastdiv_values fm1 = top_k_mask_fd(kp, 2);
+    const struct fastdiv_values fm2 = top_k_mask_fd(kp, 3);
+    const struct fastdiv_values fm3 = top_k_mask_fd(kp, 4);
+
+    const uint32_t q1 = fastdiv(r, &fd1);
+    const uint32_t i1 = r - q1 * ne[0];
+    const uint32_t i3 = fastdiv(q1, &fd2);
+    const uint32_t i2 = q1 - i3 * ne[1];
+
+    const uint32_t m1 = fastmodulo(i1, ne[2], &fm1);
+    const uint32_t m2 = fastmodulo(i2, ne[3], &fm2);
+    const uint32_t m3 = fastmodulo(i3, ne[4], &fm3);
+    return (m3 * ne[3] + m2) * ne[2] + m1;
+}
+
+// Push the inputs of item (row r, col tile ct): row inputs on the first tile of a row, then the mask tile. Returns the push count.
+static inline uint32_t top_k_mask_push_item(const struct htp_top_k_mask_context * tctx, dma_queue * dma_q, uint32_t r, uint32_t ct,
+                                            uint8_t * slot, uint8_t * mslot) {
+    const struct htp_ops_context *              octx = tctx->octx;
+    const struct htp_top_k_mask_kernel_params * kp   = tctx->kparams;
+    const struct htp_tensor *                   mask = octx->src[0];
+
+    const uint32_t c0 = ct * kp->tile_cols;
+
+    uint32_t n = 0;
+    if (ct == 0) {
+        n += top_k_mask_push_rowin(tctx, dma_q, r, slot);
+    }
+
+    if (mslot) {
+        const uint32_t m_off = top_k_mask_mask_row(kp, r) * kp->u.rows.mask_nb1 + c0 * kp->elem_size;
+        const uint32_t nb    = MIN(kp->tile_cols, kp->n_kv - c0) * kp->elem_size;
+        dma_queue_push(dma_q, dma_make_data(mslot, mask->data + m_off), nb, nb, nb, 1);
+        n++;
+    }
+
+    return n;
+}
+
+static inline void top_k_mask_next(uint32_t * r, uint32_t * ct, uint32_t n_col_tiles) {
+    if (++(*ct) == n_col_tiles) {
+        *ct = 0;
+        ++(*r);
+    }
+}
+
+// ROWS: the thread count follows the device rows (mdev), so the slice of an n byte copy is sized here
+static inline uint32_t top_k_mask_slice(const struct htp_ops_context * octx, unsigned int nth, uint32_t n) {
+    return hex_round_up(fastdiv(n + nth - 1, &octx->n_threads_div), HTP_TOP_K_MASK_UNIT);
+}
+
+// Push this thread's slice (128B aligned) of a contiguous n byte copy. Returns the push count.
+static inline uint32_t top_k_mask_push_slice(dma_queue * dma_q, unsigned int ith, dma_addr_t dst, dma_addr_t src, uint32_t n, uint32_t slice) {
+    const uint32_t b0    = ith * slice;
+    if (b0 >= n) {
+        return 0;
+    }
+    const uint32_t nb = MIN(slice, n - b0);
+    dma_queue_push(dma_q, dma_make_data(dst + b0, src + b0), nb, nb, nb, 1);
+    return 1;
+}
+
+// Pass 1 (ROWS): load the blocks table and the resident mask row, split across threads. Fill the tiles of this thread with -inf meanwhile.
+static void top_k_mask_preload(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_top_k_mask_context *             tctx  = (struct htp_top_k_mask_context *) data;
+    const struct htp_top_k_mask_kernel_params * kp    = tctx->kparams;
+    struct htp_ops_context *                    octx  = tctx->octx;
+    dma_queue *                                 dma_q = octx->ctx->dma[ith];
+    struct htp_thread_trace *                   tr    = &octx->ctx->trace[ith];
+
+    uint32_t n_pushed = 0;
+
+    if (kp->blocks_bytes) {
+        n_pushed += top_k_mask_push_slice(dma_q, ith, (dma_addr_t) tctx->blocks, octx->src[3]->data, kp->blocks_bytes, top_k_mask_slice(octx, nth, kp->blocks_bytes));
+    }
+    if (kp->flags & HTP_TOP_K_MASK_MASK_RESIDENT) {
+        const uint32_t mask_bytes = kp->u.rows.m.mask_bytes;
+        n_pushed += top_k_mask_push_slice(dma_q, ith, (dma_addr_t) tctx->mask, octx->src[0]->data, mask_bytes, top_k_mask_slice(octx, nth, mask_bytes));
+    }
+
+    const uint32_t tile_bytes = kp->tile_units * HTP_TOP_K_MASK_UNIT;
+
+    htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ith);
+    top_k_mask_fill(tctx->tiles + ith * 2 * tile_bytes, 2 * tile_bytes, kp->ninf);
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ith);
+
+    for (uint32_t i = 0; i < n_pushed; i++) {
+        top_k_mask_pop(dma_q, tr, (uint16_t) i);
+    }
+}
+
+// Pass 2 (ROWS): depth-2 DMA pipeline over the items of rows [row_first, row_last), item = (row, col tile).
+// The row inputs alternate between 2 slots by row.
+static void top_k_mask_run_items(const struct htp_top_k_mask_context * tctx, unsigned int ith, uint32_t row_first, uint32_t row_last) {
+    struct htp_ops_context *                    octx  = tctx->octx;
+    const struct htp_top_k_mask_kernel_params * kp    = tctx->kparams;
+    const struct htp_tensor *                   dst   = octx->dst;
+    dma_queue *                                 dma_q = octx->ctx->dma[ith];
+    struct htp_thread_trace *                   tr    = &octx->ctx->trace[ith];
+
+    const bool     is_f16    = kp->flags & HTP_TOP_K_MASK_F16;
+    const uint32_t elem_size = kp->elem_size;
+
+    const uint32_t rowin_bytes = kp->rowin_units * HTP_TOP_K_MASK_UNIT;
+    const uint32_t tile_bytes  = kp->tile_units * HTP_TOP_K_MASK_UNIT;
+    const uint32_t tile_cols   = kp->tile_cols;
+    const uint32_t n_col_tiles = kp->n_col_tiles;
+    const uint32_t n_items     = (row_last - row_first) * n_col_tiles;
+
+    // 2 row input slots per thread, alternating by row
+    uint8_t * rowin  = tctx->rowin + ith * 2 * rowin_bytes;
+    uint8_t * tiles  = tctx->tiles + ith * 2 * tile_bytes;
+    uint8_t * mtiles = (kp->flags & HTP_TOP_K_MASK_MASK_RESIDENT) ? NULL : tctx->mtiles + ith * 2 * tile_bytes;
+
+    HVX_Vector * tmp = (HVX_Vector *) (tctx->scratch + ith * HTP_TOP_K_MASK_SCRATCH_BYTES);
+
+    const dma_addr_t dst_data = dst->data;
+
+    uint32_t n_in[2] = { 0, 0 };
+
+    // prefetch cursor
+    uint32_t pr = row_first, pct = 0;
+
+    for (uint32_t it = 0; it < MIN(n_items, 2); it++) {
+        // Dummy dst writeback to establish queue ordering
+        dma_queue_push(dma_q, dma_make_data(dst_data, tiles + it * tile_bytes), 0, 0, 0, 0);
+        n_in[it] = top_k_mask_push_item(tctx, dma_q, pr, pct, rowin + ((pr - row_first) & 1) * rowin_bytes, mtiles ? mtiles + it * tile_bytes : NULL);
+        top_k_mask_next(&pr, &pct, n_col_tiles);
+    }
+
+    uint32_t r = row_first, ct = 0;
+
+    for (uint32_t it = 0; it < n_items; it++) {
+        const uint32_t s    = it & 1;
+        uint8_t *      tile = tiles + s * tile_bytes;
+
+        // writeback of item it - 2 (or the dummy), then the tile is free again
+        top_k_mask_pop(dma_q, tr, (uint16_t) it);
+        if (it >= 2) {
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) it);
+            top_k_mask_fill(tile, tile_bytes, kp->ninf);
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) it);
+        }
+
+        for (uint32_t i = 0; i < n_in[s]; i++) {
+            top_k_mask_pop(dma_q, tr, (uint16_t) it);
+        }
+
+        const uint32_t c0    = ct * tile_cols;
+        const uint32_t ncols = MIN(tile_cols, kp->n_kv - c0);
+
+        const uint8_t * slot = rowin + ((r - row_first) & 1) * rowin_bytes;
+        const uint8_t * mrow = mtiles ? mtiles + s * tile_bytes : tctx->mask + c0 * elem_size;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) it);
+        if (is_f16) {
+            top_k_mask_scatter_tile_f16(tctx, slot, mrow, tile, tmp, c0, ncols, 0, kp->n_top, true);
+        } else {
+            top_k_mask_scatter_tile_f32(tctx, slot, mrow, tile, tmp, c0, ncols, 0, kp->n_top, true);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) it);
+
+        const uint32_t d_off = r * kp->u.rows.rs_dst + c0 * elem_size;
+        const uint32_t nb    = ncols * elem_size;
+        dma_queue_push(dma_q, dma_make_data(dst_data + d_off, tile), nb, nb, nb, 1);
+
+        if (it + 2 < n_items) {
+            n_in[s] = top_k_mask_push_item(tctx, dma_q, pr, pct, rowin + ((pr - row_first) & 1) * rowin_bytes, mtiles ? mtiles + s * tile_bytes : NULL);
+            top_k_mask_next(&pr, &pct, n_col_tiles);
+        }
+        top_k_mask_next(&r, &ct, n_col_tiles);
+    }
+
+    const uint32_t n_wb = MIN(n_items, 2);
+    for (uint32_t i = 0; i < n_wb; i++) {
+        top_k_mask_pop(dma_q, tr, (uint16_t) i);
+    }
+}
+
+// Pass 2 (ROWS): whole rows split across threads, the col tiles of a row in order
+static void top_k_mask_thread_rows(unsigned int nth, unsigned int ith, void * data) {
+    const struct htp_top_k_mask_context * tctx = (const struct htp_top_k_mask_context *) data;
+
+    const uint32_t ir0 = tctx->row0 + ith * tctx->rows_per_thread;
+    const uint32_t ir1 = MIN(ir0 + tctx->rows_per_thread, tctx->row1);
+    if (ir0 >= ir1) {
+        return;
+    }
+
+    top_k_mask_run_items(tctx, ith, ir0, ir1);
+}
+
+static inline void top_k_mask_barrier(atomic_uint * b, uint32_t target, struct htp_thread_trace * tr, uint16_t info) {
+    htp_trace_event_start(tr, HTP_TRACE_EVT_FENCE, info);
+    atomic_fetch_add_explicit(b, 1, memory_order_release);
+    while (atomic_load_explicit(b, memory_order_acquire) < target) {
+        hex_pause();
+    }
+    htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, info);
+}
+
+// PICKS, one pass over n_virt = kp->n_threads virtual workers (host slices); real worker ith runs the virtual workers ith, ith + nth, ...
+// A: shared preload and a slice of the flag clear per virtual worker, the mask slice of this worker into the row tile;
+// barrier; B: scatter keep flags of the picks of each virtual worker; barrier; C: select mask or -inf in this worker's slice and write it back
+static void top_k_mask_thread_picks(unsigned int nth, unsigned int ith, void * data) {
+    struct htp_top_k_mask_context *             tctx  = (struct htp_top_k_mask_context *) data;
+    const struct htp_top_k_mask_kernel_params * kp    = tctx->kparams;
+    struct htp_ops_context *                    octx  = tctx->octx;
+    dma_queue *                                 dma_q = octx->ctx->dma[ith];
+    struct htp_thread_trace *                   tr    = &octx->ctx->trace[ith];
+
+    const struct htp_tensor * top_k = octx->src[1];
+    const struct htp_tensor * score = octx->src[2];
+    const struct htp_tensor * extra = octx->src[4];
+
+    uint8_t *      slot       = tctx->rowin;
+    const uint32_t tile_bytes = kp->tile_units * HTP_TOP_K_MASK_UNIT;
+    uint8_t *      out        = tctx->tiles;
+    uint8_t *      flags      = tctx->tiles + tile_bytes;
+
+    const uint32_t n_virt = kp->n_threads;
+
+    // this worker's slice of the device columns: the mask goes straight into the row tile, it is popped after the scan.
+    // The device share of the columns depends on mdev, so this slice is sized here.
+    const uint32_t out_b0 = tctx->col0 * kp->elem_size;
+    const uint32_t out_nb = tctx->ncols_out * kp->elem_size;
+    const uint32_t wslice = hex_round_up(fastdiv(out_nb + nth - 1, &octx->n_threads_div), 128);
+    const uint32_t wb0    = out_b0 + ith * wslice;
+    const uint32_t wnb    = ith * wslice < out_nb ? MIN(wslice, out_nb - ith * wslice) : 0;
+
+    const uint32_t fslice = kp->u.picks.flags_slice_u * HTP_TOP_K_MASK_UNIT;
+
+    for (uint32_t v = ith; v < n_virt; v += nth) {
+        uint32_t n_pushed = 0;
+        if (kp->blocks_bytes) {
+            n_pushed += top_k_mask_push_slice(dma_q, v, (dma_addr_t) tctx->blocks, octx->src[3]->data, kp->blocks_bytes, kp->u.picks.blocks_slice_u * HTP_TOP_K_MASK_UNIT);
+        }
+        n_pushed += top_k_mask_push_slice(dma_q, v, (dma_addr_t) slot, top_k->data, kp->n_top * sizeof(int32_t), kp->u.picks.top_slice_u * HTP_TOP_K_MASK_UNIT);
+        if (kp->flags & HTP_TOP_K_MASK_HAS_EXTRA) {
+            n_pushed += top_k_mask_push_slice(dma_q, v, (dma_addr_t) (slot + tctx->off_extra), extra->data, kp->n_extra * sizeof(int32_t), kp->u.picks.extra_slice_u * HTP_TOP_K_MASK_UNIT);
+        }
+        if (kp->flags & HTP_TOP_K_MASK_HAS_SCORE) {
+            n_pushed += top_k_mask_push_slice(dma_q, v, (dma_addr_t) (slot + tctx->off_score), score->data, kp->score_bytes, kp->u.picks.score_slice_u * HTP_TOP_K_MASK_UNIT);
+        }
+
+        // the mask slice follows the last preload in the queue, so the pops below leave it in flight
+        if (wnb && v + nth >= n_virt) {
+            dma_queue_push(dma_q, dma_make_data(out + wb0, octx->src[0]->data + wb0), wnb, wnb, wnb, 1);
+        }
+
+        const uint32_t fb0 = v * fslice;
+        if (fb0 < tile_bytes) {
+            htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) v);
+            top_k_mask_fill(flags + fb0, MIN(fslice, tile_bytes - fb0), 0); // clear the flags
+            htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) v);
+        }
+
+        for (uint32_t i = 0; i < n_pushed; i++) {
+            top_k_mask_pop(dma_q, tr, (uint16_t) i);
+        }
+    }
+
+    top_k_mask_barrier(&tctx->barrier, nth, tr, 0);
+
+    HVX_Vector * tmp = (HVX_Vector *) (tctx->scratch + ith * HTP_TOP_K_MASK_SCRATCH_BYTES);
+
+    for (uint32_t v = ith; v < n_virt; v += nth) {
+        const uint32_t j0       = MIN(v * kp->u.picks.picks_per_thread, kp->n_top);
+        const uint32_t j1       = MIN(j0 + kp->u.picks.picks_per_thread, kp->n_top);
+        const bool     do_extra = v == n_virt - 1;
+
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) v);
+        if (kp->flags & HTP_TOP_K_MASK_F16) {
+            top_k_mask_flag_tile_f16(tctx, slot, flags, tmp, j0, j1, do_extra);
+        } else {
+            top_k_mask_flag_tile_f32(tctx, slot, flags, tmp, j0, j1, do_extra);
+        }
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) v);
+    }
+
+    top_k_mask_barrier(&tctx->barrier, 2 * nth, tr, 1);
+
+    if (wnb) {
+        top_k_mask_pop(dma_q, tr, (uint16_t) ith);
+        // wb0 is 128B aligned (col0 is a multiple of 64 columns), the select may run past the row inside the tile
+        htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ith);
+        top_k_mask_select(out + wb0, flags + wb0, hex_round_up(wnb, 128), kp->ninf);
+        htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ith);
+        dma_queue_push(dma_q, dma_make_data(octx->dst->data + wb0, out + wb0), wnb, wnb, wnb, 1);
+        top_k_mask_pop(dma_q, tr, (uint16_t) ith);
+    }
+}
+
+int op_top_k_mask(struct htp_ops_context * octx) {
+    const struct htp_top_k_mask_kernel_params * kp = (const struct htp_top_k_mask_kernel_params *) octx->kernel_params;
+
+    if (kp->vtcm_size > octx->ctx->vtcm_size) {
+        return HTP_STATUS_VTCM_TOO_SMALL;
+    }
+
+    const bool picks = kp->kernel_type == HTP_TOP_K_MASK_KERNEL_PICKS;
+
+    // mdev: ROWS gives this device a contiguous share of the rows, DMA writes are exempt from the 128B rule.
+    // PICKS: every device scans the row, the write back is split by 64 column chunks.
+    uint32_t dev_row0 = 0;
+    uint32_t dev_rows = picks ? 1 : kp->u.rows.total_rows;
+    uint32_t dev_c0   = 0;
+    uint32_t dev_cols = kp->n_kv;
+    if (octx->ctx->mdev.count > 1) {
+        const struct htp_tensor_mdev_range range = picks ?
+            htp_tensor_mdev_partition(kp->n_kv, 64, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div) :
+            htp_tensor_mdev_partition(kp->u.rows.total_rows, 1, octx->ctx->mdev.idx, octx->ctx->mdev.count, &octx->ctx->mdev.count_div);
+        if (range.count == 0) {
+            return HTP_STATUS_OK;
+        }
+        if (picks) {
+            dev_c0   = range.start;
+            dev_cols = range.count;
+        } else {
+            dev_row0 = range.start;
+            dev_rows = range.count;
+        }
+    }
+
+    // fewer workers than kparams->n_threads always fit the VTCM layout: ROWS runs one row range per worker, PICKS loops over the host slices
+    const uint32_t max_threads = MIN((uint32_t) kp->n_threads, octx->ctx->n_threads);
+    if (!htp_ops_context_set_n_threads(octx, picks ? max_threads : MIN(max_threads, dev_rows))) {
+        return HTP_STATUS_INVAL_PARAMS;
+    }
+    const uint32_t n_threads = octx->n_threads;
+
+    uint8_t * vtcm = octx->ctx->vtcm_base;
+
+    struct htp_top_k_mask_context tctx = {
+        .octx      = octx,
+        .kparams   = kp,
+        .off_extra = kp->off_extra_units * HTP_TOP_K_MASK_UNIT,
+        .off_score = kp->off_score_units * HTP_TOP_K_MASK_UNIT,
+        .row0      = dev_row0,
+        .row1      = dev_row0 + dev_rows,
+        .col0      = dev_c0,
+        .ncols_out = dev_cols,
+        .blocks    = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_blocks_u * HTP_TOP_K_MASK_UNIT),
+        .mask      = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_mask_u * HTP_TOP_K_MASK_UNIT),
+        .rowin     = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_rowin_u * HTP_TOP_K_MASK_UNIT),
+        .tiles     = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_tiles_u * HTP_TOP_K_MASK_UNIT),
+        .mtiles    = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_mtiles_u * HTP_TOP_K_MASK_UNIT),
+        .scratch   = VTCM_LAYOUT_PTR(uint8_t, vtcm, kp->off_scratch_u * HTP_TOP_K_MASK_UNIT),
+    };
+
+    if (picks) {
+        atomic_init(&tctx.barrier, 0);
+        work_queue_run(octx->ctx->work_queue, top_k_mask_thread_picks, &tctx, n_threads);
+        return HTP_STATUS_OK;
+    }
+
+    tctx.rows_per_thread = fastdiv(dev_rows + n_threads - 1, &octx->n_threads_div);
+
+    work_queue_run(octx->ctx->work_queue, top_k_mask_preload, &tctx, n_threads);
+    work_queue_run(octx->ctx->work_queue, top_k_mask_thread_rows, &tctx, n_threads);
+
+    return HTP_STATUS_OK;
+}

ggml/src/ggml-hexagon/htp/top-k-mask-ops.h

@@ -0,0 +1,131 @@
+#ifndef HTP_TOP_K_MASK_OPS_H
+#define HTP_TOP_K_MASK_OPS_H
+
+#include "hex-common.h"
+#include "hex-fastdiv.h"
+#include "htp-vtcm.h"
+#include <stdbool.h>
+#include <stdint.h>
+#include <string.h>
+
+// HVX gather results land here (gathers write to VTCM only): 10 vectors per thread
+#define HTP_TOP_K_MASK_SCRATCH_BYTES 1280
+
+enum htp_top_k_mask_kernel {
+    HTP_TOP_K_MASK_KERNEL_ROWS  = 0, // threads split rows, row inputs streamed per row
+    HTP_TOP_K_MASK_KERNEL_PICKS = 1, // one row with one resident mask row: threads split the picks, keep flags select mask or -inf
+};
+
+#define HTP_TOP_K_MASK_F16           (1u << 0)
+#define HTP_TOP_K_MASK_HAS_SCORE     (1u << 1)
+#define HTP_TOP_K_MASK_HAS_BLOCKS    (1u << 2)
+#define HTP_TOP_K_MASK_HAS_EXTRA     (1u << 3)
+#define HTP_TOP_K_MASK_MASK_RESIDENT (1u << 4) // mask has exactly 1 row
+
+#define HTP_TOP_K_MASK_UNIT 128 // unit of the compact u16 sizes in kernel_params
+
+struct htp_top_k_mask_shape {
+    uint32_t n_kv, elem_size, n_top, n_pick, n_blk, n_extra;
+    uint32_t has_score, has_blocks, mask_resident;
+};
+
+struct htp_top_k_mask_vtcm_layout {
+    size_t total_bytes;
+    size_t off_blocks, off_mask, off_rowin, off_tiles, off_mtiles; // 0 when unused
+    size_t off_scratch; // HTP_TOP_K_MASK_SCRATCH_BYTES per thread
+    size_t rowin_bytes; // one slot of top_k + extra + score row, parts at htp_top_k_mask_rowin_off_*
+    size_t tile_bytes;  // one output tile slot, 128B aligned
+    size_t blocks_bytes, mask_bytes; // exact DMA sizes, 0 when unused
+};
+
+static inline uint64_t htp_top_k_mask_align128(uint64_t n) {
+    return (n + 127) & ~(uint64_t) 127;
+}
+
+// Parts of one row input slot: top_k at 0, then extra, then score; used by the host layout
+static inline uint64_t htp_top_k_mask_rowin_off_extra(uint32_t n_top) {
+    return htp_top_k_mask_align128((uint64_t) n_top * 4);
+}
+
+static inline uint64_t htp_top_k_mask_rowin_off_score(uint32_t n_top, uint32_t n_extra) {
+    return htp_top_k_mask_rowin_off_extra(n_top) + htp_top_k_mask_align128((uint64_t) n_extra * 4);
+}
+
+// the caller compares total_bytes against the VTCM size
+static inline void htp_top_k_mask_vtcm_layout_build(struct htp_top_k_mask_vtcm_layout * L,
+        const struct htp_top_k_mask_shape * S, uint32_t kernel_type, uint32_t n_threads, uint32_t tile_cols) {
+    memset(L, 0, sizeof(*L));
+
+    // PICKS streams the mask straight into its row tile
+    const bool     picks        = kernel_type == HTP_TOP_K_MASK_KERNEL_PICKS;
+    const uint64_t blocks_exact = S->has_blocks ? (uint64_t) S->n_blk * S->n_pick * 4 : 0;
+    const uint64_t blocks_bytes = htp_top_k_mask_align128(blocks_exact);
+    const uint64_t mask_exact   = S->mask_resident && !picks ? (uint64_t) S->n_kv * S->elem_size : 0;
+    const uint64_t mask_bytes   = htp_top_k_mask_align128(mask_exact);
+    const uint64_t rowin_bytes  = htp_top_k_mask_rowin_off_score(S->n_top, S->n_extra) +
+                                  (S->has_score ? htp_top_k_mask_align128((uint64_t) S->n_pick * 4) : 0);
+    const uint64_t tile_bytes   = htp_top_k_mask_align128((uint64_t) tile_cols * S->elem_size);
+
+    // ROWS streams row inputs into 2 slots per thread. PICKS: one row, a row tile and a flag tile.
+    const uint64_t rowin_total  = (picks ? 1 : 2 * (uint64_t) n_threads) * rowin_bytes;
+    const uint64_t tiles_total  = picks ? 2 * tile_bytes : 2 * (uint64_t) n_threads * tile_bytes;
+    const uint64_t mtiles_total = S->mask_resident ? 0 : tiles_total;
+    const uint64_t scratch      = (uint64_t) n_threads * HTP_TOP_K_MASK_SCRATCH_BYTES;
+
+    size_t off = 0;
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_blocks, (size_t) blocks_bytes, blocks_bytes > 0);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_mask,   (size_t) mask_bytes,   mask_bytes > 0);
+    VTCM_LAYOUT_ALLOC(off, off_rowin, (size_t) rowin_total);
+    VTCM_LAYOUT_ALLOC(off, off_tiles, (size_t) tiles_total);
+    VTCM_LAYOUT_ALLOC_OPTIONAL(off, off_mtiles, (size_t) mtiles_total, !S->mask_resident);
+    VTCM_LAYOUT_ALLOC(off, off_scratch, (size_t) scratch);
+
+    L->rowin_bytes  = (size_t) rowin_bytes;
+    L->tile_bytes   = (size_t) tile_bytes;
+    L->blocks_bytes = (size_t) blocks_exact;
+    L->mask_bytes   = (size_t) mask_exact;
+    L->total_bytes  = off;
+}
+
+// Set by the host, used as is by the NPU: no tensor shape, stride or layout is derived on the NPU.
+// Sizes of one slot or tile, the slot parts and the VTCM offsets are in HTP_TOP_K_MASK_UNIT bytes as u16.
+// Row strides: byte offset of flat row r of top_k, extra, score and dst is r * stride (the host requires row flattenable inputs).
+// PICKS splits the work in n_threads virtual workers; fewer real workers loop over them.
+struct htp_top_k_mask_kernel_params {
+    uint8_t  n_threads, kernel_type, flags, elem_size;
+    uint32_t vtcm_size;
+    uint32_t n_kv, n_top, n_pick, n_extra;
+    uint32_t ninf;                       // -inf cell pattern of elem_size
+    uint32_t blocks_bytes, score_bytes;  // exact DMA sizes, gather regions are [0, bytes)
+    uint32_t tile_cols, n_col_tiles;
+    uint16_t n_blk, blocks_nb1;          // blocks_nb1 = n_blk * 4, fits a signed 16-bit multiply
+    uint16_t off_blocks_u, off_mask_u, off_rowin_u, off_tiles_u, off_mtiles_u, off_scratch_u;
+    uint16_t rowin_units, tile_units, off_extra_units, off_score_units;
+    union {
+        struct {
+            uint32_t total_rows;
+            uint32_t rs_top, rs_extra, rs_score, rs_dst; // flat row strides (bytes of row 1)
+            uint32_t mask_nb1;                           // mask row stride, rows of the mask are dense above it
+            union {
+                uint32_t mask_bytes;                     // resident mask (1 row)
+                struct {                                 // broadcast mask: row r -> (i1, i2, i3) -> mask row (i mod mne)
+                    uint32_t mp[5];                      // fastdiv multipliers of dst ne1, ne2 and mask ne1, ne2, ne3
+                    uint16_t ne[5];                      // the divisors
+                    uint8_t  l[5];                       // fastdiv shifts
+                } dec;
+            } m;
+        } rows;
+        struct {
+            uint32_t picks_per_thread;                   // per virtual worker, multiple of the lanes of one vector
+            uint16_t blocks_slice_u, top_slice_u, extra_slice_u, score_slice_u, flags_slice_u; // per virtual worker slices
+        } picks;
+    } u;
+};
+
+#if defined(__cplusplus)
+static_assert(sizeof(struct htp_top_k_mask_kernel_params) <= 128, "htp_top_k_mask_kernel_params is too large for kernel_params blob");
+#else
+_Static_assert(sizeof(struct htp_top_k_mask_kernel_params) <= 128, "htp_top_k_mask_kernel_params is too large for kernel_params blob");
+#endif
+
+#endif // HTP_TOP_K_MASK_OPS_H

ggml/src/ggml-metal/kernels/quantize.metal

@@ -26,11 +26,11 @@ kernel void kernel_cpy_t_t(
     const int32_t i1 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0)/args.ne0;
     const int32_t i0 = (n - i3*args.ne2*args.ne1*args.ne0 - i2*args.ne1*args.ne0 - i1*args.ne0);
 
-    device T1 * dst_data = (device T1 *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0);
+    device char * dst_data = dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0;
 
     for (int32_t i00 = iw0*ntg[0] + tpitg.x; i00 < args.ne00;) {
         device const T0 * src = (device T0 *)(src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01 + i00*args.nb00);
-        dst_data[i00] = (T1) src[0];
+        *(device T1 *)(dst_data + i00*args.nb0) = (T1) src[0];
         break;
     }
 }

ggml/src/ggml-rpc/ggml-rpc.cpp

@@ -1430,6 +1430,9 @@ public:
 
 private:
     void sync_all_backends();
+    void comm_free_device(uint32_t dev_id);
+    static bool comm_allreduce_pairwise(void * comm_ctx_v, ggml_tensor ** tensors);
+    static bool comm_bootstrap_exchange(void * user_data, const void * send_buf, void * recv_buf, size_t size);
     bool get_cached_file(uint64_t hash, std::vector<uint8_t> & data);
     ggml_tensor * deserialize_tensor(struct ggml_context * ctx, const rpc_tensor * tensor);
     ggml_tensor * create_node(uint64_t id,
@@ -1438,15 +1441,26 @@ private:
                               std::unordered_map<uint64_t, struct ggml_tensor*> & tensor_map);
 
 
-    // pairwise allreduce over a direct connection to the peer server
+    // Per-device state for one communicator.
     struct comm_state {
-        socket_ptr              peer;
+        // Sockets to the other servers. Rank 0 holds one per peer in rank order, so peers[i] is
+        // rank i + 1. Every other rank holds one to rank 0. The pairwise path reduces over peers[0].
+        std::vector<socket_ptr> peers;
         uint32_t                rank = 0;
         uint32_t                world = 0;
+        ggml_backend_t          pairwise_backend = nullptr;
         ggml_backend_buffer_ptr scratch;
         size_t                  scratch_size = 0;
         std::vector<uint8_t>    send_buf;
         std::vector<uint8_t>    recv_buf;
+
+        // The backend's communicator. Null on the pairwise path, which has none.
+        void *                              backend_comm_ctx = nullptr;
+
+        // The chosen reduction, set once a path is agreed. Receives backend_comm_ctx, or
+        // comm_state on the pairwise path. A communicator covers one device, so the caller
+        // passes a one-element tensor array.
+        ggml_backend_comm_allreduce_tensor_t allreduce_fn    = nullptr;
     };
 
     std::vector<ggml_backend_t> backends;
@@ -2225,78 +2239,188 @@ void rpc_server::sync_all_backends() {
     }
 }
 
-// The comm link between two servers uses the same caps negotiation as the client HELLO,
+void rpc_server::comm_free_device(uint32_t dev_id) {
+    comm_state & state = comm_states[dev_id];
+    if (state.backend_comm_ctx != nullptr) {
+        auto comm_free = (ggml_backend_comm_free_t) ggml_backend_reg_get_proc_address(
+            ggml_backend_dev_backend_reg(ggml_backend_get_device(backends[dev_id])), "ggml_backend_comm_free");
+        if (comm_free != nullptr) {
+            comm_free(state.backend_comm_ctx);
+        }
+    }
+    state = comm_state();
+}
+
+// Exchange RPC transport capabilities over a peer connection. recv_first selects which side
+// receives first. Both sides receiving first would deadlock.
+static bool comm_caps_handshake(const socket_ptr & peer, bool recv_first) {
+    uint8_t local_caps[RPC_CONN_CAPS_SIZE]  = {};
+    uint8_t remote_caps[RPC_CONN_CAPS_SIZE] = {};
+    peer->get_caps(local_caps);
+    if (recv_first) {
+        if (!peer->recv_data(remote_caps, sizeof(remote_caps)) ||
+            !peer->send_data(local_caps, sizeof(local_caps))) {
+            return false;
+        }
+    } else {
+        if (!peer->send_data(local_caps, sizeof(local_caps)) ||
+            !peer->recv_data(remote_caps, sizeof(remote_caps))) {
+            return false;
+        }
+    }
+    peer->update_caps(remote_caps);
+    return true;
+}
+
+// The comm link between servers uses the same caps negotiation as the client HELLO,
 // so it gets the same transport upgrades (e.g. RDMA).
 bool rpc_server::comm_init(const rpc_msg_comm_init_req & request, rpc_msg_comm_init_rsp & response) {
     response.ok = 0;
-    if (request.device >= backends.size() || request.world != 2 || request.rank >= request.world) {
+    if (request.device >= backends.size() || request.world < 2 || request.rank >= request.world) {
         return true;
     }
-    comm_state & state = comm_states[request.device];
-    if (state.peer != nullptr) {
+    const uint32_t dev_id = request.device;
+    comm_state & state = comm_states[dev_id];
+    if (state.allreduce_fn != nullptr) {
         response.ok = 1;
         return true;
     }
-    uint8_t local_caps[RPC_CONN_CAPS_SIZE] = {};
-    uint8_t remote_caps[RPC_CONN_CAPS_SIZE] = {};
+
+    // Joining a communicator blocks until every rank arrives, so drain outstanding work first.
+    sync_all_backends();
+
+    // Rank 0 listens and accepts world-1 peers; every other rank connects back to it.
     if (request.rank == 0) {
         socket_ptr srv = socket_t::create_server("0.0.0.0", request.port);
         if (srv == nullptr) {
             LOG_ERROR("[%s] failed to listen on comm port %u\n", __func__, request.port);
             return true;
         }
-        state.peer = srv->accept();
-        if (state.peer == nullptr) {
-            return true;
+        // Store peers[i] as rank i + 1, as comm_bootstrap_exchange needs.
+        // Each peer sends its rank because peers may connect in any order.
+        state.peers.resize(request.world - 1);
+        for (uint32_t i = 1; i < request.world; i++) {
+            socket_ptr peer = srv->accept();
+            uint32_t peer_rank = 0;
+            if (peer == nullptr || !comm_caps_handshake(peer, /*recv_first =*/ true) ||
+                !peer->recv_data(&peer_rank, sizeof(peer_rank)) ||
+                peer_rank == 0 || peer_rank >= request.world || state.peers[peer_rank - 1] != nullptr) {
+                LOG_ERROR("[%s] bad peer connection (rank %u)\n", __func__, peer_rank);
+                state.peers.clear();
+                return true;
+            }
+            state.peers[peer_rank - 1] = std::move(peer);
         }
-        if (!state.peer->recv_data(remote_caps, sizeof(remote_caps))) {
-            state.peer = nullptr;
-            return true;
-        }
-        state.peer->get_caps(local_caps);
-        if (!state.peer->send_data(local_caps, sizeof(local_caps))) {
-            state.peer = nullptr;
-            return true;
-        }
-        state.peer->update_caps(remote_caps);
     } else {
         const std::string host(request.host, strnlen(request.host, sizeof(request.host)));
+        socket_ptr peer;
         // rank 0 may not be listening yet, retry for a few seconds
-        for (int i = 0; i < 100 && state.peer == nullptr; i++) {
-            state.peer = socket_t::connect(host.c_str(), request.port);
-            if (state.peer == nullptr) {
+        for (int i = 0; i < 100 && peer == nullptr; i++) {
+            peer = socket_t::connect(host.c_str(), request.port);
+            if (peer == nullptr) {
                 std::this_thread::sleep_for(std::chrono::milliseconds(50));
             }
         }
-        if (state.peer == nullptr) {
+        if (peer == nullptr) {
             LOG_ERROR("[%s] failed to connect to peer %s:%u\n", __func__, host.c_str(), request.port);
             return true;
         }
-        state.peer->get_caps(local_caps);
-        if (!state.peer->send_data(local_caps, sizeof(local_caps)) ||
-            !state.peer->recv_data(remote_caps, sizeof(remote_caps))) {
-            state.peer = nullptr;
+        if (!comm_caps_handshake(peer, /*recv_first =*/ false) ||
+            !peer->send_data(&request.rank, sizeof(request.rank)) || !peer->flush()) {
             return true;
         }
-        state.peer->update_caps(remote_caps);
+        state.peers.push_back(std::move(peer));
     }
+
     state.rank  = request.rank;
     state.world = request.world;
-    LOG_INFO("[%s] device %u joined pairwise comm as rank %u\n", __func__, request.device, request.rank);
+
+    // Forces the pairwise path, for comparing the two on one build. Setting it on any one rank
+    // is enough, because the negotiation spreads the decision.
+    const bool no_backend_communicator = std::getenv("GGML_RPC_NO_BACKEND_COMMUNICATOR") != nullptr;
+
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backends[dev_id]));
+    auto init_rank    = (ggml_backend_comm_init_rank_t)
+        ggml_backend_reg_get_proc_address(reg, "ggml_backend_comm_init_rank");
+    auto allreduce_fn = (ggml_backend_comm_allreduce_tensor_t)
+        ggml_backend_reg_get_proc_address(reg, "ggml_backend_comm_allreduce_tensor");
+
+    // The ranks must agree on backend communicator or pairwise. Every rank shares whether its
+    // backend provides a communicator, and all ranks use it only if every rank can.
+    const uint8_t self_can_use_backend_comm = !no_backend_communicator && init_rank != nullptr && allreduce_fn != nullptr;
+    std::vector<uint8_t> all_can_use_backend_comm(request.world);
+    if (!comm_bootstrap_exchange(&state, &self_can_use_backend_comm, all_can_use_backend_comm.data(),
+                                 sizeof(self_can_use_backend_comm))) {
+        state.peers.clear();
+        return true;
+    }
+    const bool use_backend_communicator = std::all_of(all_can_use_backend_comm.begin(), all_can_use_backend_comm.end(),
+                                                      [](uint8_t c) { return c != 0; });
+
+    if (use_backend_communicator) {
+        void * backend_ctx = init_rank(backends[dev_id], (int) request.rank, (int) request.world, comm_bootstrap_exchange, &state);
+        if (backend_ctx == nullptr) {
+            LOG_ERROR("[%s] failed to join communicator as rank %u/%u\n", __func__, request.rank, request.world);
+            state.peers.clear();
+            return true;
+        }
+        state.backend_comm_ctx = backend_ctx;
+        state.allreduce_fn     = allreduce_fn;
+        // The sockets have done their job; the collective runs between the backends from here on.
+        state.peers.clear();
+        LOG_INFO("[%s] device %u using backend communicator, rank %u/%u\n", __func__, dev_id, request.rank, request.world);
+        response.ok = 1;
+        return true;
+    }
+
+    // No backend communicator, so fall back to the pairwise exchange, which only covers two ranks.
+    if (request.world != 2) {
+        LOG_ERROR("[%s] world size %u requires a backend communicator\n", __func__, request.world);
+        state.peers.clear();
+        return true;
+    }
+    state.pairwise_backend = backends[dev_id];
+    state.allreduce_fn     = comm_allreduce_pairwise;
+    LOG_INFO("[%s] device %u using pairwise communicator, rank %u\n", __func__, dev_id, request.rank);
     response.ok = 1;
     return true;
 }
 
-bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
-    if (request.device >= backends.size()) {
-        return false;
+// Bootstrap callback for ggml_backend_comm_init_rank_t, over the peer sockets. Rank 0 collects
+// size bytes from every rank in rank order, then sends the whole array back to every rank.
+// user_data is the comm_state.
+bool rpc_server::comm_bootstrap_exchange(void * user_data, const void * send_buf, void * recv_buf, size_t size) {
+    comm_state & state = *(comm_state *) user_data;
+    uint8_t * recv = (uint8_t *) recv_buf;
+    const size_t total = (size_t) state.world * size;
+
+    if (state.rank == 0) {
+        memcpy(recv, send_buf, size);
+        for (size_t i = 0; i < state.peers.size(); i++) {
+            if (!state.peers[i]->recv_data(recv + (i + 1) * size, size)) {
+                return false;
+            }
+        }
+        for (const socket_ptr & peer : state.peers) {
+            if (!peer->send_data(recv, total) || !peer->flush()) {
+                return false;
+            }
+        }
+        return true;
     }
-    comm_state & state = comm_states[request.device];
-    if (state.peer == nullptr) {
-        LOG_ERROR("[%s] no communicator for device %u\n", __func__, request.device);
-        return false;
-    }
-    ggml_backend_t backend = backends[request.device];
+    return state.peers[0]->send_data(send_buf, size) && state.peers[0]->flush() &&
+           state.peers[0]->recv_data(recv, total);
+}
+
+// Reduces this rank's contribution in place, exchanging partials with the peer over peers[0].
+// Used when the backend provides no communicator. comm_ctx_v is the comm_state.
+bool rpc_server::comm_allreduce_pairwise(void * comm_ctx_v, ggml_tensor ** tensors) {
+    comm_state & state = *(comm_state *) comm_ctx_v;
+    ggml_tensor * t_dst = tensors[0];
+    ggml_backend_t backend = state.pairwise_backend;
+
+    const size_t  nbytes = ggml_nbytes(t_dst);
+    const int64_t ne     = ggml_nelements(t_dst);
 
     size_t ctx_size = 16*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(8, false);
     struct ggml_init_params params = {
@@ -2307,16 +2431,7 @@ bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
     ggml_context_ptr ctx_ptr { ggml_init(params) };
     GGML_ASSERT(ctx_ptr != nullptr);
     ggml_context * ctx = ctx_ptr.get();
-    ggml_tensor * t_dst = deserialize_tensor(ctx, &request.tensor);
-    if (t_dst == nullptr || t_dst->buffer == nullptr) {
-        LOG_ERROR("[%s] error deserializing tensor\n", __func__);
-        return false;
-    }
-    const size_t  nbytes = ggml_nbytes(t_dst);
-    const int64_t ne     = ggml_nelements(t_dst);
-    if (nbytes == 0) {
-        return true;
-    }
+
     // reduce large partials in bf16 to halve the wire bytes; small (decode-sized) ones
     // stay f32 since the extra casts and sync cost more than the bytes saved
     const bool   wire_bf16  = t_dst->type == GGML_TYPE_F32 && ne >= 32768;
@@ -2372,13 +2487,13 @@ bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
 
     // rank 0 sends first, rank 1 receives first, so large payloads cannot deadlock
     if (state.rank == 0) {
-        if (!state.peer->send_data(state.send_buf.data(), wire_bytes) || !state.peer->flush() ||
-            !state.peer->recv_data(state.recv_buf.data(), wire_bytes)) {
+        if (!state.peers[0]->send_data(state.send_buf.data(), wire_bytes) || !state.peers[0]->flush() ||
+            !state.peers[0]->recv_data(state.recv_buf.data(), wire_bytes)) {
             return false;
         }
     } else {
-        if (!state.peer->recv_data(state.recv_buf.data(), wire_bytes) ||
-            !state.peer->send_data(state.send_buf.data(), wire_bytes) || !state.peer->flush()) {
+        if (!state.peers[0]->recv_data(state.recv_buf.data(), wire_bytes) ||
+            !state.peers[0]->send_data(state.send_buf.data(), wire_bytes) || !state.peers[0]->flush()) {
             return false;
         }
     }
@@ -2408,11 +2523,50 @@ bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
     return true;
 }
 
+bool rpc_server::comm_allreduce(const rpc_msg_comm_allreduce_req & request) {
+    if (request.device >= backends.size()) {
+        return false;
+    }
+    comm_state & state = comm_states[request.device];
+    if (state.allreduce_fn == nullptr) {
+        LOG_ERROR("[%s] no communicator for device %u\n", __func__, request.device);
+        return false;
+    }
+
+    struct ggml_init_params params = {
+        /*.mem_size   =*/ ggml_tensor_overhead(),
+        /*.mem_buffer =*/ NULL,
+        /*.no_alloc   =*/ true,
+    };
+    ggml_context_ptr ctx_ptr { ggml_init(params) };
+    GGML_ASSERT(ctx_ptr != nullptr);
+    ggml_tensor * t_dst = deserialize_tensor(ctx_ptr.get(), &request.tensor);
+    if (t_dst == nullptr || t_dst->buffer == nullptr) {
+        LOG_ERROR("[%s] error deserializing tensor\n", __func__);
+        return false;
+    }
+    if (ggml_nbytes(t_dst) == 0) {
+        return true;
+    }
+
+    // Neither path needs a wait here. The pairwise path waits internally, and the backend orders
+    // its reduction after the work that produced the partial sum.
+    void * comm_ctx = state.backend_comm_ctx != nullptr ? state.backend_comm_ctx : &state;
+    if (!state.allreduce_fn(comm_ctx, &t_dst)) {
+        LOG_ERROR("[%s] allreduce failed for tensor %s\n", __func__, t_dst->name);
+        return false;
+    }
+    return true;
+}
+
 bool rpc_server::comm_free(const rpc_msg_comm_free_req & request) {
     if (request.device >= backends.size()) {
         return false;
     }
-    comm_states[request.device] = comm_state();
+    // Wait for outstanding work. Destroying a communicator while its kernels are still running
+    // is unsafe.
+    sync_all_backends();
+    comm_free_device(request.device);
     return true;
 }
 
@@ -2432,6 +2586,14 @@ bool rpc_server::get_device_memory(const rpc_msg_get_device_memory_req & request
 }
 
 rpc_server::~rpc_server() {
+    // Wait for outstanding work. Destroying a communicator while its kernels are still running
+    // is unsafe.
+    sync_all_backends();
+
+    // Free any communicator still open. A client can disconnect without sending COMM_FREE.
+    for (uint32_t dev_id = 0; dev_id < comm_states.size(); dev_id++) {
+        comm_free_device(dev_id);
+    }
     for (auto buffer : buffers) {
         ggml_backend_buffer_free(buffer);
     }
@@ -2997,7 +3159,7 @@ static void ggml_backend_rpc_comm_free(void * comm_ctx_v) {
 }
 
 static void * ggml_backend_rpc_comm_init(ggml_backend_t * backends, size_t n_backends) {
-    if (n_backends != 2 || std::getenv("GGML_RPC_NO_COMM") != nullptr) {
+    if (n_backends < 2 || std::getenv("GGML_RPC_NO_COMM") != nullptr) {
         return nullptr;
     }
     std::vector<ggml_backend_rpc_comm_context::rank_info> ranks;

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -15550,13 +15550,14 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         }
         return true;
     };
-    // reject any tensors larger than the max buffer size
+    // non-compute ops are never executed here; their size is bounded by the buffer allocation limit
+    const bool non_compute = ggml_op_is_empty(op->op);
     for (int i = 0; i < GGML_MAX_SRC; i++) {
-        if (op->src[i] && !tensor_size_supported(ggml_nbytes(op->src[i]))) {
+        if (op->src[i] && !non_compute && !tensor_size_supported(ggml_nbytes(op->src[i]))) {
             return false;
         }
     }
-    if (!tensor_size_supported(ggml_nbytes(op))) {
+    if (!non_compute && !tensor_size_supported(ggml_nbytes(op))) {
         return false;
     }
 

ggml/src/ggml-zdnn/fattn.cpp

@@ -20,8 +20,8 @@ bool ggml_zdnn_supports_flash_attn_ext(const ggml_backend_zdnn_device_context *
     memcpy(&max_bias,      (const float *) op->op_params + 1, sizeof(float));
     memcpy(&logit_softcap, (const float *) op->op_params + 2, sizeof(float));
 
-    // sinks, alibi, softcap and sparse k/v sets are not implemented
-    if (op->src[4] != nullptr || max_bias != 0.0f || logit_softcap != 0.0f || ggml_get_op_params_i32(op, 4) != 0) {
+    // sinks, alibi, softcap, sparse k/v sets and k/v row lists are not implemented
+    if (op->src[4] != nullptr || op->src[5] != nullptr || max_bias != 0.0f || logit_softcap != 0.0f || ggml_get_op_params_i32(op, 4) != 0) {
         return false;
     }
 

ggml/src/ggml.c

@@ -5390,10 +5390,10 @@ struct ggml_tensor * ggml_roll(
         int                   shift2,
         int                   shift3) {
     GGML_ASSERT(a->nb[0] == ggml_type_size(a->type));
-    GGML_ASSERT(abs(shift0) < a->ne[0]);
-    GGML_ASSERT(abs(shift1) < a->ne[1]);
-    GGML_ASSERT(abs(shift2) < a->ne[2]);
-    GGML_ASSERT(abs(shift3) < a->ne[3]);
+    GGML_ASSERT(shift0 > -a->ne[0] && shift0 < a->ne[0]);
+    GGML_ASSERT(shift1 > -a->ne[1] && shift1 < a->ne[1]);
+    GGML_ASSERT(shift2 > -a->ne[2] && shift2 < a->ne[2]);
+    GGML_ASSERT(shift3 > -a->ne[3] && shift3 < a->ne[3]);
 
     struct ggml_tensor * result = ggml_dup_tensor(ctx, a);
 

ggml/src/gguf.cpp

@@ -1463,6 +1463,7 @@ struct gguf_writer_base {
     // we bet on devirtualization
     virtual void write(int8_t val) = 0;
     virtual void write(const std::vector<int8_t> & val) = 0;
+    virtual void write(const void * data, size_t size) = 0;
     virtual void write_tensor_data(const struct gguf_tensor_info & info, size_t offset_data, size_t alignment) = 0;
 
     template <typename T>
@@ -1579,21 +1580,31 @@ struct gguf_writer_buf final : public gguf_writer_base {
         written_bytes += val.size();
     }
 
+    void write(const void * data, const size_t size) override {
+        buf.insert(buf.end(), (const int8_t *) data, (const int8_t *) data + size);
+        written_bytes += size;
+    }
+
     void write_tensor_data(const struct gguf_tensor_info & info, const size_t offset_data, const size_t alignment) override {
         GGML_ASSERT(buf.size() - offset_data == info.offset);
 
         GGML_ASSERT(ggml_is_contiguous(&info.t));
-        const size_t offset = buf.size();
         const size_t nbytes = ggml_nbytes(&info.t);
 
-        buf.resize(offset + nbytes);
-        if (info.t.buffer) {
-            ggml_backend_tensor_get(&info.t, buf.data() + offset, 0, nbytes);
+        if (info.t.buffer && ggml_backend_buffer_is_host(info.t.buffer)) {
+            GGML_ASSERT(info.t.data != NULL && "tensor not allocated");
+            write(info.t.data, nbytes);
         } else {
-            GGML_ASSERT(info.t.data);
-            memcpy(buf.data() + offset, info.t.data, nbytes);
+            const size_t offset = buf.size();
+            buf.resize(offset + nbytes);
+            if (info.t.buffer) {
+                ggml_backend_tensor_get(&info.t, buf.data() + offset, 0, nbytes);
+            } else {
+                GGML_ASSERT(info.t.data);
+                memcpy(buf.data() + offset, info.t.data, nbytes);
+            }
+            written_bytes += nbytes;
         }
-        written_bytes += nbytes;
 
         pad(alignment);
     }
@@ -1624,20 +1635,33 @@ struct gguf_writer_file final : public gguf_writer_base {
         }
     }
 
+    void write(const void * data, const size_t size) override {
+        const auto ret = fwrite(data, 1, size, file);
+        written_bytes += size;
+        if (ret != size) {
+            throw std::runtime_error("unexpected fwrite number of bytes written, '" + std::to_string(ret) + "' instead of '" + std::to_string(size) + "'");
+        }
+    }
+
     void write_tensor_data(const struct gguf_tensor_info & info, const size_t offset_data, const size_t alignment) override {
         GGML_ASSERT(written_bytes - offset_data == info.offset);
 
         GGML_ASSERT(ggml_is_contiguous(&info.t));
         const size_t nbytes = ggml_nbytes(&info.t);
 
-        std::vector<int8_t> buf(nbytes);
-        if (info.t.buffer) {
-            ggml_backend_tensor_get(&info.t, buf.data(), 0, nbytes);
+        if (info.t.buffer && ggml_backend_buffer_is_host(info.t.buffer)) {
+            GGML_ASSERT(info.t.data != NULL && "tensor not allocated");
+            write(info.t.data, nbytes);
         } else {
-            GGML_ASSERT(info.t.data);
-            memcpy(buf.data(), info.t.data, nbytes);
+            std::vector<int8_t> buf(nbytes);
+            if (info.t.buffer) {
+                ggml_backend_tensor_get(&info.t, buf.data(), 0, nbytes);
+            } else {
+                GGML_ASSERT(info.t.data);
+                memcpy(buf.data(), info.t.data, nbytes);
+            }
+            write(buf);
         }
-        write(buf);
 
         pad(alignment);
     }

gguf-py/gguf/constants.py

@@ -203,6 +203,7 @@ class Keys:
         ICLR_LORA_RANK               = "{arch}.attention.iclr_lora_rank"
         VALUE_RESIDUAL_MIX_LORA_RANK = "{arch}.attention.value_residual_mix_lora_rank"
         GATE_LORA_RANK               = "{arch}.attention.gate_lora_rank"
+        PASS_LORA_RANK               = "{arch}.attention.pass_lora_rank"
         REL_BUCKETS_COUNT            = "{arch}.attention.relative_buckets_count"
         SLIDING_WINDOW               = "{arch}.attention.sliding_window"
         SCALE                        = "{arch}.attention.scale"
@@ -315,6 +316,12 @@ class Keys:
         L_CYCLES         = "{arch}.hrm.l_cycles"
         PREFIX_LM        = "{arch}.hrm.prefix_lm"
 
+    class Metis:
+        PASS_COUNT               = "{arch}.metis.pass_count"
+        MEMORY_LENGTH            = "{arch}.metis.memory_length"
+        ROUTE_LENGTH             = "{arch}.metis.route_length"
+        CONTINUATION_LOGIT_LIMIT = "{arch}.metis.continuation_logit_limit"
+
     class SSM:
         CONV_KERNEL    = "{arch}.ssm.conv_kernel"
         INNER_SIZE     = "{arch}.ssm.inner_size"
@@ -322,6 +329,7 @@ class Keys:
         TIME_STEP_RANK = "{arch}.ssm.time_step_rank"
         GROUP_COUNT    = "{arch}.ssm.group_count"
         DT_B_C_RMS     = "{arch}.ssm.dt_b_c_rms"
+        NORM_EPSILON   = "{arch}.ssm.norm_epsilon"
 
     class KDA:
         HEAD_DIM         = "{arch}.kda.head_dim"
@@ -506,6 +514,11 @@ class Keys:
         WINDOW_SIZE         = "clip.audio.window_size"
         LOCAL_BLOCK_COUNT   = "clip.audio.local_block_count" # mimo-v2.5: input_local_transformer layer count
         LOCAL_GROUP_SIZE    = "clip.audio.local_group_size"  # mimo-v2.5: input_local_transformer grouping size
+        ROPE_DIMS           = "clip.audio.rope_dims"         # nemotron-transcribe
+
+        class Speaker: # nemotron-transcribe
+            BLOCK_COUNT         = "clip.audio.spk.block_count"
+            HEAD_COUNT          = "clip.audio.spk.head_count"
 
         class Attention:
             HEAD_COUNT      = "clip.audio.attention.head_count"
@@ -676,6 +689,7 @@ class MODEL_ARCH(IntEnum):
     NEMOTRON         = auto()
     NEMOTRON_H       = auto()
     NEMOTRON_H_MOE   = auto()
+    METIS16          = auto()
     EXAONE           = auto()
     EXAONE4          = auto()
     EXAONE_MOE       = auto()
@@ -825,6 +839,36 @@ class MODEL_TENSOR(IntEnum):
     FFN_GATE_TID2EID     = auto()
     MOE_LATENT_DOWN      = auto() # nemotron 3 super
     MOE_LATENT_UP        = auto() # nemotron 3 super
+    STREAM_EMBD          = auto()
+    DEPTH_MEM_STATE_WRITE = auto()
+    DEPTH_MEM_META_WRITE = auto()
+    DEPTH_MEM_ENTRY_BIAS = auto()
+    DEPTH_MEM_QUERY      = auto()
+    DEPTH_MEM_KEY        = auto()
+    DEPTH_MEM_VALUE      = auto()
+    DEPTH_MEM_OUT        = auto()
+    DEPTH_MEM_ROUTE_STATE = auto()
+    DEPTH_MEM_ROUTE_RETR = auto()
+    DEPTH_MEM_ROUTE_DIFF = auto()
+    DEPTH_MEM_ROUTE_HIST = auto()
+    DEPTH_MEM_ROUTE      = auto()
+    DEPTH_MEM_STREAM_GATE = auto()
+    CONT_HIDDEN_STATE    = auto()
+    CONT_HIDDEN_RETR     = auto()
+    CONT_HIDDEN_DIFF     = auto()
+    CONT_HIDDEN_ROUTE    = auto()
+    CONT_HIDDEN          = auto()
+    CONT_OUTPUT          = auto()
+    HC_ATTN_MIX          = auto()
+    HC_ATTN_READ         = auto()
+    HC_ATTN_WRITE        = auto()
+    HC_FFN_MIX           = auto()
+    HC_FFN_READ          = auto()
+    HC_FFN_WRITE         = auto()
+    ATTN_PASS_LORA_DOWN  = auto()
+    ATTN_PASS_LORA_UP    = auto()
+    ATTN_PASS_GATE       = auto()
+    FFN_EXP_EMBD         = auto()
     FFN_RESID_ALPHA      = auto()
     FFN_RESID_BETA       = auto()
     FFN_PRE_EXP_NORM     = auto()
@@ -1411,6 +1455,24 @@ class MODEL_TENSOR(IntEnum):
     A_ENC_CONV_NORM_VAR    = auto() # parakeet
     A_ENC_MEL_FILTERS      = auto() # parakeet
     A_ENC_WINDOW           = auto() # parakeet
+    A_ENC_ATTN_QKV         = auto() # nemotron-transcribe
+    A_SPK_INP_PROJ         = auto() # nemotron-transcribe
+    A_SPK_PRE_NORM         = auto() # nemotron-transcribe
+    A_SPK_POST_NORM        = auto() # nemotron-transcribe
+    A_SPK_ATTN_NORM        = auto() # nemotron-transcribe
+    A_SPK_ATTN_QKV         = auto() # nemotron-transcribe
+    A_SPK_ATTN_OUT         = auto() # nemotron-transcribe
+    A_SPK_FFN_NORM         = auto() # nemotron-transcribe
+    A_SPK_FFN_UP           = auto() # nemotron-transcribe
+    A_SPK_FFN_DOWN         = auto() # nemotron-transcribe
+    A_SPK_HEAD_PROJ        = auto() # nemotron-transcribe
+    A_SPK_HEAD_UPSAMPLE    = auto() # nemotron-transcribe
+    A_SPK_HEAD_FC1         = auto() # nemotron-transcribe
+    A_SPK_HEAD_FC2         = auto() # nemotron-transcribe
+    A_SPK_SIL_EMB          = auto() # nemotron-transcribe
+    A_FUSION_ASR_NORM      = auto() # nemotron-transcribe
+    A_FUSION_SPK_NORM      = auto() # nemotron-transcribe
+    A_FUSION_SPK_KERNEL    = auto() # nemotron-transcribe
     A_CTC_OUT              = auto()
     A_CTC_OUT_MID          = auto()
     A_ENC_ATTN_REL_POS_EMB = auto()
@@ -1535,6 +1597,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.NEMOTRON:         "nemotron",
     MODEL_ARCH.NEMOTRON_H:       "nemotron_h",
     MODEL_ARCH.NEMOTRON_H_MOE:   "nemotron_h_moe",
+    MODEL_ARCH.METIS16:          "metis16",
     MODEL_ARCH.EXAONE:           "exaone",
     MODEL_ARCH.EXAONE4:          "exaone4",
     MODEL_ARCH.EXAONE_MOE:       "exaone-moe",
@@ -1685,6 +1748,36 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     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.STREAM_EMBD:               "stream_embd",
+    MODEL_TENSOR.DEPTH_MEM_STATE_WRITE:     "depth_mem_state_write",
+    MODEL_TENSOR.DEPTH_MEM_META_WRITE:      "depth_mem_meta_write",
+    MODEL_TENSOR.DEPTH_MEM_ENTRY_BIAS:      "depth_mem_entry_bias",
+    MODEL_TENSOR.DEPTH_MEM_QUERY:           "depth_mem_query",
+    MODEL_TENSOR.DEPTH_MEM_KEY:             "depth_mem_key",
+    MODEL_TENSOR.DEPTH_MEM_VALUE:           "depth_mem_value",
+    MODEL_TENSOR.DEPTH_MEM_OUT:             "depth_mem_out",
+    MODEL_TENSOR.DEPTH_MEM_ROUTE_STATE:     "depth_mem_route_state",
+    MODEL_TENSOR.DEPTH_MEM_ROUTE_RETR:      "depth_mem_route_retr",
+    MODEL_TENSOR.DEPTH_MEM_ROUTE_DIFF:      "depth_mem_route_diff",
+    MODEL_TENSOR.DEPTH_MEM_ROUTE_HIST:      "depth_mem_route_hist",
+    MODEL_TENSOR.DEPTH_MEM_ROUTE:           "depth_mem_route",
+    MODEL_TENSOR.DEPTH_MEM_STREAM_GATE:     "depth_mem_stream_gate",
+    MODEL_TENSOR.CONT_HIDDEN_STATE:         "cont_hidden_state",
+    MODEL_TENSOR.CONT_HIDDEN_RETR:          "cont_hidden_retr",
+    MODEL_TENSOR.CONT_HIDDEN_DIFF:          "cont_hidden_diff",
+    MODEL_TENSOR.CONT_HIDDEN_ROUTE:         "cont_hidden_route",
+    MODEL_TENSOR.CONT_HIDDEN:               "cont_hidden",
+    MODEL_TENSOR.CONT_OUTPUT:               "cont_output",
+    MODEL_TENSOR.HC_ATTN_MIX:               "blk.{bid}.hc_attn_mix",
+    MODEL_TENSOR.HC_ATTN_READ:              "blk.{bid}.hc_attn_read",
+    MODEL_TENSOR.HC_ATTN_WRITE:             "blk.{bid}.hc_attn_write",
+    MODEL_TENSOR.HC_FFN_MIX:                "blk.{bid}.hc_ffn_mix",
+    MODEL_TENSOR.HC_FFN_READ:               "blk.{bid}.hc_ffn_read",
+    MODEL_TENSOR.HC_FFN_WRITE:              "blk.{bid}.hc_ffn_write",
+    MODEL_TENSOR.ATTN_PASS_LORA_DOWN:       "blk.{bid}.attn_pass_lora_down",
+    MODEL_TENSOR.ATTN_PASS_LORA_UP:         "blk.{bid}.attn_pass_lora_up",
+    MODEL_TENSOR.ATTN_PASS_GATE:            "blk.{bid}.attn_pass_gate",
+    MODEL_TENSOR.FFN_EXP_EMBD:              "blk.{bid}.ffn_exp_embd",
     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",
@@ -2219,6 +2312,24 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.A_ENC_CONV_NORM_VAR:       "a.blk.{bid}.conv_norm_var",
     MODEL_TENSOR.A_ENC_MEL_FILTERS:         "a.mel_filters",
     MODEL_TENSOR.A_ENC_WINDOW:              "a.window",
+    MODEL_TENSOR.A_ENC_ATTN_QKV:            "a.blk.{bid}.attn_qkv",
+    MODEL_TENSOR.A_SPK_INP_PROJ:            "a.spk.input_projection",
+    MODEL_TENSOR.A_SPK_PRE_NORM:            "a.spk.pre_ln",
+    MODEL_TENSOR.A_SPK_POST_NORM:           "a.spk.post_ln",
+    MODEL_TENSOR.A_SPK_ATTN_NORM:           "a.spk.blk.{bid}.ln1",
+    MODEL_TENSOR.A_SPK_ATTN_QKV:            "a.spk.blk.{bid}.attn_qkv",
+    MODEL_TENSOR.A_SPK_ATTN_OUT:            "a.spk.blk.{bid}.attn_out",
+    MODEL_TENSOR.A_SPK_FFN_NORM:            "a.spk.blk.{bid}.ln2",
+    MODEL_TENSOR.A_SPK_FFN_UP:              "a.spk.blk.{bid}.ffn_up",
+    MODEL_TENSOR.A_SPK_FFN_DOWN:            "a.spk.blk.{bid}.ffn_down",
+    MODEL_TENSOR.A_SPK_HEAD_PROJ:           "a.spk.head.proj",
+    MODEL_TENSOR.A_SPK_HEAD_UPSAMPLE:       "a.spk.head.upsample",
+    MODEL_TENSOR.A_SPK_HEAD_FC1:            "a.spk.head.fc1",
+    MODEL_TENSOR.A_SPK_HEAD_FC2:            "a.spk.head.fc2",
+    MODEL_TENSOR.A_SPK_SIL_EMB:             "a.spk.head.sil_emb",
+    MODEL_TENSOR.A_FUSION_ASR_NORM:         "a.fusion.asr_norm",
+    MODEL_TENSOR.A_FUSION_SPK_NORM:         "a.fusion.spk_norm",
+    MODEL_TENSOR.A_FUSION_SPK_KERNEL:       "a.fusion.spk_kernel",
     MODEL_TENSOR.A_CTC_OUT:                 "a.enc_ctc_out",
     MODEL_TENSOR.A_CTC_OUT_MID:             "a.enc_ctc_out_mid",
     MODEL_TENSOR.A_ENC_ATTN_REL_POS_EMB:    "a.blk.{bid}.attn_rel_pos_emb",
@@ -2587,6 +2698,24 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.A_ENC_CONV_NORM_VAR,
         MODEL_TENSOR.A_ENC_MEL_FILTERS,
         MODEL_TENSOR.A_ENC_WINDOW,
+        MODEL_TENSOR.A_ENC_ATTN_QKV,
+        MODEL_TENSOR.A_SPK_INP_PROJ,
+        MODEL_TENSOR.A_SPK_PRE_NORM,
+        MODEL_TENSOR.A_SPK_POST_NORM,
+        MODEL_TENSOR.A_SPK_ATTN_NORM,
+        MODEL_TENSOR.A_SPK_ATTN_QKV,
+        MODEL_TENSOR.A_SPK_ATTN_OUT,
+        MODEL_TENSOR.A_SPK_FFN_NORM,
+        MODEL_TENSOR.A_SPK_FFN_UP,
+        MODEL_TENSOR.A_SPK_FFN_DOWN,
+        MODEL_TENSOR.A_SPK_HEAD_PROJ,
+        MODEL_TENSOR.A_SPK_HEAD_UPSAMPLE,
+        MODEL_TENSOR.A_SPK_HEAD_FC1,
+        MODEL_TENSOR.A_SPK_HEAD_FC2,
+        MODEL_TENSOR.A_SPK_SIL_EMB,
+        MODEL_TENSOR.A_FUSION_ASR_NORM,
+        MODEL_TENSOR.A_FUSION_SPK_NORM,
+        MODEL_TENSOR.A_FUSION_SPK_KERNEL,
         MODEL_TENSOR.A_MM_INP_PROJ,
         MODEL_TENSOR.A_MM_SOFT_EMB_NORM,
         MODEL_TENSOR.A_MM_EMBEDDING,
@@ -4858,6 +4987,61 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_DOWN,
         MODEL_TENSOR.FFN_UP,
     ],
+    MODEL_ARCH.METIS16: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.STREAM_EMBD,
+        MODEL_TENSOR.DEPTH_MEM_STATE_WRITE,
+        MODEL_TENSOR.DEPTH_MEM_META_WRITE,
+        MODEL_TENSOR.DEPTH_MEM_ENTRY_BIAS,
+        MODEL_TENSOR.DEPTH_MEM_QUERY,
+        MODEL_TENSOR.DEPTH_MEM_KEY,
+        MODEL_TENSOR.DEPTH_MEM_VALUE,
+        MODEL_TENSOR.DEPTH_MEM_OUT,
+        MODEL_TENSOR.DEPTH_MEM_ROUTE_STATE,
+        MODEL_TENSOR.DEPTH_MEM_ROUTE_RETR,
+        MODEL_TENSOR.DEPTH_MEM_ROUTE_DIFF,
+        MODEL_TENSOR.DEPTH_MEM_ROUTE_HIST,
+        MODEL_TENSOR.DEPTH_MEM_ROUTE,
+        MODEL_TENSOR.DEPTH_MEM_STREAM_GATE,
+        MODEL_TENSOR.CONT_HIDDEN_STATE,
+        MODEL_TENSOR.CONT_HIDDEN_RETR,
+        MODEL_TENSOR.CONT_HIDDEN_DIFF,
+        MODEL_TENSOR.CONT_HIDDEN_ROUTE,
+        MODEL_TENSOR.CONT_HIDDEN,
+        MODEL_TENSOR.CONT_OUTPUT,
+        MODEL_TENSOR.HC_ATTN_MIX,
+        MODEL_TENSOR.HC_ATTN_READ,
+        MODEL_TENSOR.HC_ATTN_WRITE,
+        MODEL_TENSOR.HC_FFN_MIX,
+        MODEL_TENSOR.HC_FFN_READ,
+        MODEL_TENSOR.HC_FFN_WRITE,
+        MODEL_TENSOR.ATTN_PASS_LORA_DOWN,
+        MODEL_TENSOR.ATTN_PASS_LORA_UP,
+        MODEL_TENSOR.ATTN_PASS_GATE,
+        MODEL_TENSOR.FFN_EXP_EMBD,
+        MODEL_TENSOR.SSM_IN,
+        MODEL_TENSOR.SSM_CONV1D,
+        MODEL_TENSOR.SSM_DT,
+        MODEL_TENSOR.SSM_A,
+        MODEL_TENSOR.SSM_D,
+        MODEL_TENSOR.SSM_NORM,
+        MODEL_TENSOR.SSM_OUT,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.MOE_LATENT_DOWN,
+        MODEL_TENSOR.MOE_LATENT_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_UP_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+    ],
     MODEL_ARCH.NEMOTRON_H: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,
@@ -6732,6 +6916,8 @@ class VisionProjectorType:
     QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor
     POCKETTTS_SPKENC = "pockettts_spkenc" # audio: mimi encoder as voice-prompt encoder
     POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder
+    CHATTERBOX_SPKENC = "chatterbox_spkenc" # audio: CAMPPlus speaker encoder
+    CHATTERBOX = "chatterbox" # audio generation: s3gen flow matching + HiFT vocoder
     SOPRANO = "soprano" # audio generation: Vocos decoder
     HUNYUANVL      = "hunyuanvl"
     PARAKEET       = "parakeet"  # audio
@@ -6740,6 +6926,7 @@ class VisionProjectorType:
     MINICPMV4_7    = "minicpmv4_7"
     MINICPM_OCR_1_0 = "minicpm_ocr_1_0"
     GRANITE_SPEECH = "granite_speech"  # audio
+    NEMOTRON_TRANSCRIBE = "nemotron_transcribe"  # audio
     MIMOVL         = "mimovl"
     MIMO_AUDIO     = "mimo_audio"
     GRANITE4_VISION = "granite4_vision"

gguf-py/gguf/gguf_writer.py

@@ -975,6 +975,18 @@ class GGUFWriter:
     def add_hrm_prefix_lm(self, value: bool) -> None:
         self.add_bool(Keys.HRM.PREFIX_LM.format(arch=self.arch), value)
 
+    def add_metis_pass_count(self, value: int) -> None:
+        self.add_uint32(Keys.Metis.PASS_COUNT.format(arch=self.arch), value)
+
+    def add_metis_memory_length(self, value: int) -> None:
+        self.add_uint32(Keys.Metis.MEMORY_LENGTH.format(arch=self.arch), value)
+
+    def add_metis_route_length(self, value: int) -> None:
+        self.add_uint32(Keys.Metis.ROUTE_LENGTH.format(arch=self.arch), value)
+
+    def add_metis_continuation_logit_limit(self, value: float) -> None:
+        self.add_float32(Keys.Metis.CONTINUATION_LOGIT_LIMIT.format(arch=self.arch), value)
+
     def add_adapter_count(self, count: int) -> None:
         self.add_uint32(Keys.Adapters.COUNT.format(arch=self.arch), count)
 
@@ -1086,6 +1098,9 @@ class GGUFWriter:
     def add_attention_output_lora_rank(self, length: int) -> None:
         self.add_uint32(Keys.Attention.OUTPUT_LORA_RANK.format(arch=self.arch), length)
 
+    def add_attention_pass_lora_rank(self, length: int) -> None:
+        self.add_uint32(Keys.Attention.PASS_LORA_RANK.format(arch=self.arch), length)
+
     def add_attention_compress_ratios(self, values: Sequence[int]) -> None:
         self.add_array(Keys.Attention.COMPRESS_RATIOS.format(arch=self.arch), values)
 
@@ -1267,6 +1282,9 @@ class GGUFWriter:
     def add_ssm_dt_b_c_rms(self, value: bool) -> None:
         self.add_bool(Keys.SSM.DT_B_C_RMS.format(arch=self.arch), value)
 
+    def add_ssm_norm_epsilon(self, value: float) -> None:
+        self.add_float32(Keys.SSM.NORM_EPSILON.format(arch=self.arch), value)
+
     def add_expert_latent_length(self, value: int) -> None:
         self.add_uint32(Keys.LLM.EXPERT_LATENT_LENGTH.format(arch=self.arch), value)
 
@@ -1712,6 +1730,15 @@ class GGUFWriter:
     def add_audio_projector_head_count(self, value: int) -> None:
         self.add_uint32(Keys.ClipAudio.Projector.HEAD_COUNT, value)
 
+    def add_audio_rope_dims(self, value: int) -> None:
+        self.add_uint32(Keys.ClipAudio.ROPE_DIMS, value)
+
+    def add_audio_spk_block_count(self, value: int) -> None:
+        self.add_uint32(Keys.ClipAudio.Speaker.BLOCK_COUNT, value)
+
+    def add_audio_spk_head_count(self, value: int) -> None:
+        self.add_uint32(Keys.ClipAudio.Speaker.HEAD_COUNT, value)
+
     # audio generation (mmproj)
 
     def add_clip_gen_audio_projector_type(self, value: str) -> None:

gguf-py/gguf/tensor_mapping.py

@@ -35,6 +35,7 @@ class TensorNameMap:
             "model.transformer.wte",                     # llada
             "embed_tokens",                              # qwen3-embedding
             "model.embed",                               # talkie
+            "embedding",                                 # metis16
         ),
 
         # Masked embeddings
@@ -224,6 +225,7 @@ class TensorNameMap:
     block_mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = {
         # Attention norm
         MODEL_TENSOR.ATTN_NORM: (
+            "layers.{bid}.mixer_connection.norm",                   # metis16
             "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
@@ -294,6 +296,7 @@ class TensorNameMap:
             "model.layers.{bid}.linear_attn.in_proj_qkv",                          # qwen3.5
             "head.layers.{bid}.self_attn.in_proj",  # laya
             "layers.{bid}.linear_attn.in_proj_qkv",                                # qwen3.5 text
+            "layers.{bid}.mixer.qkv",                                              # metis16
         ),
 
         # Attention query
@@ -401,6 +404,7 @@ class TensorNameMap:
             "model.layers.{bid}.self_attn.language_expert_dense",           # cogvlm
             "model.blocks.{bid}.attn.attn_resid",                           # talkie
             "head.layers.{bid}.self_attn.out_proj",  # laya
+            "layers.{bid}.mixer.out",                                       # metis16
         ),
 
         # Attention output norm
@@ -449,6 +453,7 @@ class TensorNameMap:
 
         # Feed-forward norm
         MODEL_TENSOR.FFN_NORM: (
+            "layers.{bid}.moe_connection.norm",                              # metis16
             "model.lilicorr.layers.{bid}.mlp_norm",                          # lilicorr
             "model.lilicorr.feature_mlp.{bid}.norm",                         # lilicorr
             "gpt_neox.layers.{bid}.post_attention_layernorm",                # gptneox
@@ -536,6 +541,7 @@ class TensorNameMap:
             "backbone.layers.{bid}.mixer.gate",                 # nemotron-h-moe
             "model.layers.{bid}.moe.gate",                      # step3.5
             "model.layers.{bid}.router.proj",                   # gemma4
+            "layers.{bid}.moe.expert_router",                   # metis16
         ),
 
         MODEL_TENSOR.FFN_GATE_INP_SHEXP: (
@@ -543,6 +549,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.FFN_EXP_PROBS_B: (
+            "layers.{bid}.moe.selection_bias",                              # metis16
             "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
@@ -617,6 +624,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.FFN_UP_SHEXP: (
+            "layers.{bid}.moe.shared_expert.up",                     # metis16
             "model.layers.{bid}.mlp.shared_expert.up_proj",          # qwen2moe
             "model.layers.{bid}.mlp.shared_experts.up_proj",         # deepseek deepseek2
             "model.layers.{bid}.feed_forward.shared_expert.up_proj", # llama4
@@ -670,6 +678,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.FFN_GATE_SHEXP: (
+            "layers.{bid}.moe.shared_expert.gate",                     # metis16
             "model.layers.{bid}.mlp.shared_expert.gate_proj",          # qwen2moe
             "model.layers.{bid}.mlp.shared_experts.gate_proj",         # deepseek deepseek2
             "model.layers.{bid}.feed_forward.shared_expert.gate_proj", # llama4
@@ -686,18 +695,94 @@ class TensorNameMap:
         MODEL_TENSOR.FFN_GATE_UP_EXP: (
             "model.layers.{bid}.mlp.experts.gate_up_proj",
             "model.layers.{bid}.experts.gate_up_proj", # gemma4
+            "layers.{bid}.moe.local_experts.gate_up",  # metis16
         ),
 
         MODEL_TENSOR.MOE_LATENT_DOWN: (
             "backbone.layers.{bid}.mixer.fc1_latent_proj",                 # nemotron 3 super
             "model.layers.{bid}.mlp.pre_expert_proj",                      # muse-spark
+            "layers.{bid}.moe.latent_down",                                # metis16
         ),
 
         MODEL_TENSOR.MOE_LATENT_UP: (
             "backbone.layers.{bid}.mixer.fc2_latent_proj",                 # nemotron 3 super
             "model.layers.{bid}.mlp.post_expert_proj",                     # muse-spark
+            "layers.{bid}.moe.latent_up",                                  # metis16
         ),
 
+        MODEL_TENSOR.STREAM_EMBD: (
+            "stream_embeddings",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_STATE_WRITE: (
+            "depth_memory.state_write",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_META_WRITE: (
+            "depth_memory.metadata_write",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_QUERY: (
+            "depth_memory.query",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_KEY: (
+            "depth_memory.key",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_VALUE: (
+            "depth_memory.value",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_OUT: (
+            "depth_memory.output",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_ROUTE_STATE: (
+            "depth_memory.route_state",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_ROUTE_RETR: (
+            "depth_memory.route_retr",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_ROUTE_DIFF: (
+            "depth_memory.route_diff",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_ROUTE_HIST: (
+            "depth_memory.route_hist",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_ROUTE: (
+            "depth_memory.route_projection",  # metis16
+        ),
+        MODEL_TENSOR.DEPTH_MEM_STREAM_GATE: (
+            "depth_memory.stream_gate",  # metis16
+        ),
+        MODEL_TENSOR.CONT_HIDDEN_STATE: (
+            "continuation.hidden_state",  # metis16
+        ),
+        MODEL_TENSOR.CONT_HIDDEN_RETR: (
+            "continuation.hidden_retr",  # metis16
+        ),
+        MODEL_TENSOR.CONT_HIDDEN_DIFF: (
+            "continuation.hidden_diff",  # metis16
+        ),
+        MODEL_TENSOR.CONT_HIDDEN_ROUTE: (
+            "continuation.hidden_route",  # metis16
+        ),
+        MODEL_TENSOR.CONT_HIDDEN: (
+            "continuation.hidden",  # metis16
+        ),
+        MODEL_TENSOR.CONT_OUTPUT: (
+            "continuation.output",  # metis16
+        ),
+        MODEL_TENSOR.ATTN_PASS_LORA_DOWN: (
+            "layers.{bid}.mixer.pass_lora_down",  # metis16
+        ),
+        MODEL_TENSOR.ATTN_PASS_LORA_UP: (
+            "layers.{bid}.mixer.pass_lora_up",  # metis16
+        ),
+        MODEL_TENSOR.ATTN_PASS_GATE: (
+            "layers.{bid}.mixer.pass_gates",  # metis16
+        ),
+        MODEL_TENSOR.FFN_EXP_EMBD: (
+            "layers.{bid}.moe.expert_embeddings",  # metis16
+        ),
+
+
+
+
         MODEL_TENSOR.FFN_PRE_EXP_NORM: (
             "model.layers.{bid}.mlp.pre_expert_norm",                      # muse-spark
         ),
@@ -759,6 +844,7 @@ class TensorNameMap:
             "model.layers.{bid}.block_sparse_moe.experts.down",     # smallthinker
             "model.layers.{bid}.moe.down_proj",                     # step3.5
             "model.layers.{bid}.experts.down_proj",                 # gemma4
+            "layers.{bid}.moe.local_experts.down",                  # metis16
         ),
 
         MODEL_TENSOR.FFN_DOWN_SHEXP: (
@@ -771,6 +857,7 @@ class TensorNameMap:
             "backbone.layers.{bid}.mixer.shared_experts.down_proj",    # nemotron-h-moe
             "model.layers.{bid}.block_sparse_moe.shared_experts.down_proj", # kimi
             "model.layers.{bid}.share_expert.down_proj",               # step3.5
+            "layers.{bid}.moe.shared_expert.down",                     # metis16
         ),
 
         MODEL_TENSOR.FFN_DOWN_CHEXP: (
@@ -906,6 +993,7 @@ class TensorNameMap:
             "model.layers.{bid}.mamba.in_proj",             # jamba falcon-h1 granite-hybrid
             "model.layers.layers.{bid}.mixer.in_proj",      # plamo2
             "model.layers.{bid}.linear_attn.in_proj_qkvz",  # qwen3next
+            "layers.{bid}.mixer.impl.mixer.in_proj",        # metis16
         ),
 
         MODEL_TENSOR.SSM_CONV1D: (
@@ -915,6 +1003,7 @@ class TensorNameMap:
             "model.layers.layers.{bid}.mixer.conv1d",  # plamo2
             "model.layers.{bid}.linear_attn.conv1d",   # qwen3next
             "layers.{bid}.linear_attn.conv1d",         # qwen3.5 text
+            "layers.{bid}.mixer.impl.mixer.conv1d",    # metis16
         ),
 
         MODEL_TENSOR.SSM_X: (
@@ -925,6 +1014,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.SSM_DT: (
+            "layers.{bid}.mixer.impl.mixer.dt_bias",    # metis16
             "model.layers.{bid}.dt_proj",               # mamba-hf
             "backbone.layers.{bid}.mixer.dt_proj",      # mamba
             "model.layers.{bid}.mamba.dt_proj",         # jamba falcon-h1 granite-hybrid
@@ -942,6 +1032,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.SSM_A: (
+            "layers.{bid}.mixer.impl.mixer.A_log",    # metis16
             "model.layers.{bid}.A_log",               # mamba-hf
             "backbone.layers.{bid}.mixer.A_log",      # mamba
             "model.layers.{bid}.mamba.A_log",         # jamba falcon-h1 granite-hybrid
@@ -965,6 +1056,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.SSM_D: (
+            "layers.{bid}.mixer.impl.mixer.D",    # metis16
             "model.layers.{bid}.D",               # mamba-hf
             "backbone.layers.{bid}.mixer.D",      # mamba
             "model.layers.{bid}.mamba.D",         # jamba falcon-h1 granite-hybrid
@@ -972,6 +1064,7 @@ class TensorNameMap:
         ),
 
         MODEL_TENSOR.SSM_NORM: (
+            "layers.{bid}.mixer.impl.mixer.norm",   # metis16
             "model.layers.{bid}.mamba.norm",        # falcon-h1 granite-hybrid
             "model.layers.{bid}.linear_attn.norm",  # qwen3next
             "layers.{bid}.linear_attn.norm",        # qwen3.5 text
@@ -987,6 +1080,7 @@ class TensorNameMap:
             "model.layers.{bid}.linear_attn.out_proj",   # qwen3next
             "layers.{bid}.linear_attn.out_proj",         # qwen3.5 text
             "model.layers.layers.{bid}.mixer.out_proj",  # plamo2
+            "layers.{bid}.mixer.impl.mixer.out_proj",    # metis16
         ),
 
         MODEL_TENSOR.SSM_ALPHA: (
@@ -2532,6 +2626,7 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_INP_PROJ: (
             "conformer.subsample_conv_projection.input_proj_linear", # gemma4
+            "perception.encoder.asr_encoder.pre_encode.proj", # nemotron-transcribe (feature stacking)
             "encoder.input_linear",
         ),
 
@@ -2546,13 +2641,16 @@ class TensorNameMap:
             "speaker_encoder.mfa.conv", # qwen3tts speaker encoder: multi-layer feature aggregation
         ),
 
-        MODEL_TENSOR.A_PRE_NORM: (),
+        MODEL_TENSOR.A_PRE_NORM: (
+            "perception.encoder.asr_encoder.embed_norm", # nemotron-transcribe
+        ),
 
         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
+            "perception.encoder.asr_encoder.final_norm", # nemotron-transcribe
             "encoder.layer_norm", # mimo-audio-tokenizer
             "encoder.lfm.lfm.embedding_norm", # lfm3a
         ),
@@ -2595,10 +2693,12 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_ATTN_Q_NORM: (
             "encoder.lfm.lfm.layers.{bid}.self_attn.q_layernorm", # lfm3a
+            "perception.encoder.asr_encoder.layers.{bid}.attn.q_norm", # nemotron-transcribe
         ),
 
         MODEL_TENSOR.A_ENC_ATTN_K_NORM: (
             "encoder.lfm.lfm.layers.{bid}.self_attn.k_layernorm", # lfm3a
+            "perception.encoder.asr_encoder.layers.{bid}.attn.k_norm", # nemotron-transcribe
         ),
 
         MODEL_TENSOR.A_ENC_ATTN_K_REL: (
@@ -2629,6 +2729,7 @@ class TensorNameMap:
             "encoder.lfm.lfm.layers.{bid}.operator_norm", # lfm3a
             "conformer.layers.{bid}.attention.pre_attn_norm", # gemma3n
             "sound_encoder.encoder.layers.{bid}.norm_self_att", # parakeet
+            "perception.encoder.asr_encoder.layers.{bid}.norm1", # nemotron-transcribe
             "encoder.layers.{bid}.attn.pre_norm", # granite_speech
             "encoder.layers.{bid}.self_attn_layer_norm", # mimo-audio-tokenizer
         ),
@@ -2642,6 +2743,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.attention.post", # gemma3n
             "conformer.layers.{bid}.self_attn.post", # gemma4
             "sound_encoder.encoder.layers.{bid}.self_attn.o_proj", # parakeet
+            "perception.encoder.asr_encoder.layers.{bid}.attn.out_proj", # nemotron-transcribe
             "encoder.layers.{bid}.attn.to_out", # granite_speech
             "encoder.layers.{bid}.self_attn.out_proj", # mimo-audio-tokenizer
         ),
@@ -2653,6 +2755,7 @@ class TensorNameMap:
             "encoder.lfm.lfm.layers.{bid}.ffn_norm", # lfm3a
             "conformer.layers.{bid}.attention.post_norm", # gemma3n
             "sound_encoder.encoder.layers.{bid}.norm_out", # parakeet
+            "perception.encoder.asr_encoder.layers.{bid}.norm2", # nemotron-transcribe
             "encoder.layers.{bid}.post_norm", # granite_speech
             "encoder.layers.{bid}.final_layer_norm", # mimo-audio-tokenizer
         ),
@@ -2683,6 +2786,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n
             "conformer.layers.{bid}.feed_forward1.ffw_layer_1", # gemma4
             "sound_encoder.encoder.layers.{bid}.feed_forward1.linear1", # parakeet
+            "perception.encoder.asr_encoder.layers.{bid}.ffn.net.0", # nemotron-transcribe
             "encoder.layers.{bid}.ff1.up_proj", # granite_speech
             "encoder.layers.{bid}.fc1", # mimo-audio-tokenizer
         ),
@@ -2702,6 +2806,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n
             "conformer.layers.{bid}.feed_forward1.ffw_layer_2", # gemma4
             "sound_encoder.encoder.layers.{bid}.feed_forward1.linear2", # parakeet
+            "perception.encoder.asr_encoder.layers.{bid}.ffn.net.3", # nemotron-transcribe
             "encoder.layers.{bid}.ff1.down_proj", # granite_speech
             "encoder.layers.{bid}.fc2", # mimo-audio-tokenizer
         ),
@@ -2792,6 +2897,7 @@ class TensorNameMap:
             "audio_tower.proj", # qwen2omni
             "model.audio_tower.output_proj", # gemma4
             "speaker_encoder.fc", # qwen3tts speaker encoder: final speaker embedding projection
+            "perception.proj", # nemotron-transcribe
         ),
 
         MODEL_TENSOR.A_MM_NORM_PRE: (
@@ -2914,12 +3020,85 @@ class TensorNameMap:
 
         MODEL_TENSOR.A_ENC_MEL_FILTERS: (
             "sound_encoder.encoder.feature_extractor.featurizer.fb", # parakeet
+            "perception.preprocessor.featurizer.fb", # nemotron-transcribe
         ),
 
         MODEL_TENSOR.A_ENC_WINDOW: (
             "sound_encoder.encoder.feature_extractor.featurizer.window", # parakeet
+            "perception.preprocessor.featurizer.window", # nemotron-transcribe
         ),
 
+        MODEL_TENSOR.A_ENC_ATTN_QKV: (
+            "perception.encoder.asr_encoder.layers.{bid}.attn.w_qkv", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_INP_PROJ: (
+            "perception.encoder.diarization_model.encoder.pre_encode.proj", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_PRE_NORM: (
+            "perception.encoder.diarization_model.encoder.embed_norm", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_POST_NORM: (
+            "perception.encoder.diarization_model.encoder.final_norm", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_ATTN_NORM: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.norm1", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_ATTN_QKV: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.attn.w_qkv", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_ATTN_OUT: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.attn.out_proj", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_FFN_NORM: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.norm2", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_FFN_UP: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.ffn.net.0", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_FFN_DOWN: (
+            "perception.encoder.diarization_model.encoder.layers.{bid}.ffn.net.3", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_HEAD_PROJ: (
+            "perception.encoder.diarization_model.sortformer_modules.encoder_proj", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_HEAD_UPSAMPLE: (
+            "perception.encoder.diarization_model.sortformer_modules.subpixel_upsample", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_HEAD_FC1: (
+            "perception.encoder.diarization_model.sortformer_modules.first_hidden_to_hidden", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_HEAD_FC2: (
+            "perception.encoder.diarization_model.sortformer_modules.single_hidden_to_spks", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_SPK_SIL_EMB: (
+            "perception.encoder.diarization_model.sortformer_modules.learnable_sil_emb", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_FUSION_ASR_NORM: (
+            "perception.encoder.asr_norm", # nemotron-transcribe
+        ),
+
+        MODEL_TENSOR.A_FUSION_SPK_NORM: (
+            "perception.encoder.diar_norm", # nemotron-transcribe
+        ),
+
+        # generated in conversion code (nemotron-transcribe)
+        MODEL_TENSOR.A_FUSION_SPK_KERNEL: (),
+
         MODEL_TENSOR.A_MM_EMBEDDING: (
             "model.embed_audio.embedding", # gemma3n
         ),

scripts/snapdragon/run.py

@@ -33,6 +33,7 @@ MANAGED_ENV_NAMES = (
     "GGML_HEXAGON_MM_SELECT",
     "GGML_HEXAGON_FA_SELECT",
     "GGML_HEXAGON_FA_HEAD_SPLIT",
+    "GGML_HEXAGON_LI_SELECT",
     "GGML_HEXAGON_GDN_SELECT",
     "GGML_HEXAGON_AR_SELECT",
     "GGML_HEXAGON_AR_SCATTER",
@@ -177,11 +178,17 @@ def device_arg_from_devices(devices):
 
 
 def normalize_cmd_device_args(cmd_args):
+    exec_name = os.path.basename(cmd_args[0]) if cmd_args else ""
+    if exec_name.endswith(".exe"):
+        exec_name = exec_name[:-4]
+
     for i, arg in enumerate(cmd_args):
         if arg == "--device" and i + 1 < len(cmd_args):
             cmd_args[i + 1] = device_arg_from_devices(cmd_args[i + 1])
         elif arg.startswith("--device="):
             cmd_args[i] = "--device=" + device_arg_from_devices(arg.split("=", 1)[1])
+        elif exec_name in ("test-backend-ops", "test-llama-archs") and arg in ("-b", "--backend") and i + 1 < len(cmd_args):
+            cmd_args[i + 1] = device_arg_from_devices(cmd_args[i + 1])
 
 
 def main():
@@ -229,6 +236,7 @@ def main():
     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")
     parser.add_argument("--hex-fa-select", help="Select Flash Attention kernel (GGML_HEXAGON_FA_SELECT) 2:HMX,1:HVX,0:disable")
     parser.add_argument("--hex-fa-head-split", help="Enable (1) or disable (0) head-parallel flash_attn partitioning (GGML_HEXAGON_FA_HEAD_SPLIT)")
+    parser.add_argument("--hex-li-select", help="Select Lightning Indexer kernel (GGML_HEXAGON_LI_SELECT) 2:HMX,1:HVX,0:disable")
     parser.add_argument("--hex-gdn-select", help="Select Gated Delta Net kernel (GGML_HEXAGON_GDN_SELECT) 2:HMX,1:HVX,0:disable")
     parser.add_argument("--hex-ar-select", help="Select All-Reduce kernel (GGML_HEXAGON_AR_SELECT) 1:enable,0:disable")
     parser.add_argument("--hex-ar-scatter", help="Enable (1) or disable (0) reduce-scatter for fused ALLREDUCE+ADD (GGML_HEXAGON_AR_SCATTER)")
@@ -299,6 +307,10 @@ def main():
             env_vars[env_name] = str(opt_val)
 
     # Resolve and filter devices (HTP vs OpenCL)
+    exec_name = os.path.basename(cmd_args[0]) if cmd_args else ""
+    if exec_name.endswith(".exe"):
+        exec_name = exec_name[:-4]
+
     device_in_cmd = None
     for i, arg in enumerate(cmd_args):
         if arg == "--device" and i + 1 < len(cmd_args):
@@ -307,6 +319,9 @@ def main():
         elif arg.startswith("--device="):
             device_in_cmd = arg.split("=", 1)[1]
             break
+        elif exec_name in ("test-backend-ops", "test-llama-archs") and arg in ("-b", "--backend") and i + 1 < len(cmd_args):
+            device_in_cmd = cmd_args[i + 1]
+            break
 
     if args.devices is not None:
         devices_val = args.devices
@@ -375,6 +390,7 @@ def main():
     set_env("GGML_HEXAGON_MM_SELECT", args.hex_mm_select)
     set_env("GGML_HEXAGON_FA_SELECT", args.hex_fa_select)
     set_env("GGML_HEXAGON_FA_HEAD_SPLIT", args.hex_fa_head_split)
+    set_env("GGML_HEXAGON_LI_SELECT", args.hex_li_select)
     set_env("GGML_HEXAGON_GDN_SELECT", args.hex_gdn_select)
     set_env("GGML_HEXAGON_AR_SELECT", args.hex_ar_select)
     set_env("GGML_HEXAGON_AR_SCATTER", args.hex_ar_scatter)
@@ -402,7 +418,7 @@ def main():
 
     # Resolve executable path
     executable = cmd_args[0]
-    known_binaries = ["llama-cli", "llama-bench", "llama-completion", "llama-mtmd-cli", "test-backend-ops"]
+    known_binaries = ["llama-cli", "llama-bench", "llama-completion", "llama-mtmd-cli", "test-backend-ops", "test-llama-archs"]
     if executable in known_binaries:
         if target_type in ("android", "linux"):
             resolved_exec = f"./bin/{executable}"
@@ -419,22 +435,23 @@ def main():
         basename = basename[:-4]
 
     device_val = None
-    if basename == "test-backend-ops":
-        for i in range(len(cmd_args)):
-            if cmd_args[i] in ("-p", "--params") and i + 1 < len(cmd_args):
-                val = cmd_args[i + 1]
-                new_val = ""
-                for j, char in enumerate(val):
-                    if char in ('[', ']'):
-                        if j > 0 and val[j - 1] == '\\':
-                            new_val += char
+    if basename in ("test-backend-ops", "test-llama-archs"):
+        if basename == "test-backend-ops":
+            for i in range(len(cmd_args)):
+                if cmd_args[i] in ("-p", "--params") and i + 1 < len(cmd_args):
+                    val = cmd_args[i + 1]
+                    new_val = ""
+                    for j, char in enumerate(val):
+                        if char in ('[', ']'):
+                            if j > 0 and val[j - 1] == '\\':
+                                new_val += char
+                            else:
+                                new_val += '\\' + char
                         else:
-                            new_val += '\\' + char
-                    else:
-                        new_val += char
-                cmd_args[i + 1] = new_val
+                            new_val += char
+                    cmd_args[i + 1] = new_val
 
-        has_b = any(arg == "-b" for arg in cmd_args)
+        has_b = any(arg in ("-b", "--backend") for arg in cmd_args)
         if not has_b:
             if args.devices:
                 device_val = device_arg_from_devices(args.devices)
@@ -462,9 +479,13 @@ def main():
         or args.hex_profile is not None
         or args.hex_optrace is not None
     )
-    if verbose_trigger and basename in ("llama-cli", "llama-completion", "llama-bench", "llama-server", "llama-mtmd-cli"):
-        if "-v" not in cmd_args and "--verbose" not in cmd_args:
-            cmd_args.append("-v")
+    if verbose_trigger:
+        if basename == "test-llama-archs":
+            if "-v" not in cmd_args:
+                cmd_args += ["-v", "5"]
+        elif basename in ("llama-cli", "llama-completion", "llama-bench", "llama-server", "llama-mtmd-cli"):
+            if "-v" not in cmd_args and "--verbose" not in cmd_args:
+                cmd_args.append("-v")
 
     # Inject defaults for llama-cli, llama-completion, and llama-server if not overridden by the user
     if basename in ("llama-cli", "llama-completion", "llama-server"):

src/llama-arch.cpp

@@ -102,6 +102,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
     { LLM_ARCH_NEMOTRON,         "nemotron"         },
     { LLM_ARCH_NEMOTRON_H,       "nemotron_h"       },
     { LLM_ARCH_NEMOTRON_H_MOE,   "nemotron_h_moe"   },
+    { LLM_ARCH_METIS16,          "metis16"          },
     { LLM_ARCH_EXAONE,           "exaone"           },
     { LLM_ARCH_EXAONE4,          "exaone4"          },
     { LLM_ARCH_EXAONE_MOE,       "exaone-moe"       },
@@ -272,6 +273,10 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_HRM_H_CYCLES,                      "%s.hrm.h_cycles"                      },
     { LLM_KV_HRM_L_CYCLES,                      "%s.hrm.l_cycles"                      },
     { LLM_KV_HRM_PREFIX_LM,                     "%s.hrm.prefix_lm"                     },
+    { LLM_KV_METIS_PASS_COUNT,                  "%s.metis.pass_count"                  },
+    { LLM_KV_METIS_MEMORY_LENGTH,               "%s.metis.memory_length"               },
+    { LLM_KV_METIS_ROUTE_LENGTH,                "%s.metis.route_length"                },
+    { LLM_KV_METIS_CONTINUATION_LOGIT_LIMIT,    "%s.metis.continuation_logit_limit"    },
     { LLM_KV_MUDD_FEED_FORWARD_LENGTH,          "%s.mudd.feed_forward_length"          },
     { LLM_KV_MUDD_TAP_COUNT,                    "%s.mudd.tap_count"                    },
     { LLM_KV_MUDD_TAP_INDICES,                  "%s.mudd.tap_indices"                  },
@@ -329,6 +334,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_ATTENTION_KV_SOURCE_LAYERS,               "%s.attention.kv_source_layers"               },
     { LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,           "%s.attention.output_group_count"           },
     { LLM_KV_ATTENTION_OUTPUT_LORA_RANK,             "%s.attention.output_lora_rank"             },
+    { LLM_KV_ATTENTION_PASS_LORA_RANK,               "%s.attention.pass_lora_rank"               },
     { LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,      "%s.attention.compress_rope_freq_base"      },
     { LLM_KV_ATTENTION_COMPRESS_RATIOS,              "%s.attention.compress_ratios"              },
     { LLM_KV_ATTENTION_SHARED_KV_LAYERS,             "%s.attention.shared_kv_layers"             },
@@ -387,6 +393,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_SSM_TIME_STEP_RANK, "%s.ssm.time_step_rank" },
     { LLM_KV_SSM_GROUP_COUNT,    "%s.ssm.group_count"    },
     { LLM_KV_SSM_DT_B_C_RMS,     "%s.ssm.dt_b_c_rms"     },
+    { LLM_KV_SSM_NORM_EPS,       "%s.ssm.norm_epsilon"   },
 
     { LLM_KV_KDA_HEAD_DIM,         "%s.kda.head_dim"         },
     { LLM_KV_KDA_SAFE_GATE,        "%s.kda.safe_gate"        },
@@ -548,6 +555,36 @@ 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_STREAM_EMBD,                            "stream_embd" },
+    { LLM_TENSOR_DEPTH_MEM_STATE_WRITE,                  "depth_mem_state_write" },
+    { LLM_TENSOR_DEPTH_MEM_META_WRITE,                   "depth_mem_meta_write" },
+    { LLM_TENSOR_DEPTH_MEM_ENTRY_BIAS,                   "depth_mem_entry_bias" },
+    { LLM_TENSOR_DEPTH_MEM_QUERY,                        "depth_mem_query" },
+    { LLM_TENSOR_DEPTH_MEM_KEY,                          "depth_mem_key" },
+    { LLM_TENSOR_DEPTH_MEM_VALUE,                        "depth_mem_value" },
+    { LLM_TENSOR_DEPTH_MEM_OUT,                          "depth_mem_out" },
+    { LLM_TENSOR_DEPTH_MEM_ROUTE_STATE,                  "depth_mem_route_state" },
+    { LLM_TENSOR_DEPTH_MEM_ROUTE_RETR,                   "depth_mem_route_retr" },
+    { LLM_TENSOR_DEPTH_MEM_ROUTE_DIFF,                   "depth_mem_route_diff" },
+    { LLM_TENSOR_DEPTH_MEM_ROUTE_HIST,                   "depth_mem_route_hist" },
+    { LLM_TENSOR_DEPTH_MEM_ROUTE,                        "depth_mem_route" },
+    { LLM_TENSOR_DEPTH_MEM_STREAM_GATE,                  "depth_mem_stream_gate" },
+    { LLM_TENSOR_CONT_HIDDEN_STATE,                      "cont_hidden_state" },
+    { LLM_TENSOR_CONT_HIDDEN_RETR,                       "cont_hidden_retr" },
+    { LLM_TENSOR_CONT_HIDDEN_DIFF,                       "cont_hidden_diff" },
+    { LLM_TENSOR_CONT_HIDDEN_ROUTE,                      "cont_hidden_route" },
+    { LLM_TENSOR_CONT_HIDDEN,                            "cont_hidden" },
+    { LLM_TENSOR_CONT_OUTPUT,                            "cont_output" },
+    { LLM_TENSOR_HC_ATTN_MIX,                            "blk.%d.hc_attn_mix" },
+    { LLM_TENSOR_HC_ATTN_READ,                           "blk.%d.hc_attn_read" },
+    { LLM_TENSOR_HC_ATTN_WRITE,                          "blk.%d.hc_attn_write" },
+    { LLM_TENSOR_HC_FFN_MIX,                             "blk.%d.hc_ffn_mix" },
+    { LLM_TENSOR_HC_FFN_READ,                            "blk.%d.hc_ffn_read" },
+    { LLM_TENSOR_HC_FFN_WRITE,                           "blk.%d.hc_ffn_write" },
+    { LLM_TENSOR_ATTN_PASS_LORA_DOWN,                    "blk.%d.attn_pass_lora_down" },
+    { LLM_TENSOR_ATTN_PASS_LORA_UP,                      "blk.%d.attn_pass_lora_up" },
+    { LLM_TENSOR_ATTN_PASS_GATE,                         "blk.%d.attn_pass_gate" },
+    { LLM_TENSOR_FFN_EXP_EMBD,                           "blk.%d.ffn_exp_embd" },
     { 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" },
@@ -1158,6 +1195,36 @@ 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_STREAM_EMBD,                {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_ADD}},
+    {LLM_TENSOR_DEPTH_MEM_STATE_WRITE,      {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_META_WRITE,       {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_ENTRY_BIAS,       {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_QUERY,            {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_KEY,              {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_VALUE,            {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_OUT,              {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_ROUTE_STATE,      {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_ROUTE_RETR,       {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_ROUTE_DIFF,       {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_ROUTE_HIST,       {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_ROUTE,            {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_DEPTH_MEM_STREAM_GATE,      {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_CONT_HIDDEN_STATE,          {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_CONT_HIDDEN_RETR,           {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_CONT_HIDDEN_DIFF,           {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_CONT_HIDDEN_ROUTE,          {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_CONT_HIDDEN,                {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_CONT_OUTPUT,                {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_HC_ATTN_MIX,                {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_HC_ATTN_READ,               {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_HC_ATTN_WRITE,              {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_HC_FFN_MIX,                 {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_HC_FFN_READ,                {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_HC_FFN_WRITE,               {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_ATTN_PASS_LORA_DOWN,        {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_ATTN_PASS_LORA_UP,          {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+    {LLM_TENSOR_ATTN_PASS_GATE,             {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+    {LLM_TENSOR_FFN_EXP_EMBD,               {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
     {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}},
@@ -1286,6 +1353,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) {
         case LLM_ARCH_LFM3:
         case LLM_ARCH_NEMOTRON_H:
         case LLM_ARCH_NEMOTRON_H_MOE:
+        case LLM_ARCH_METIS16:
         case LLM_ARCH_QWEN3NEXT:
         case LLM_ARCH_KIMI_LINEAR:
         case LLM_ARCH_BAILINGMOE3:
@@ -1390,7 +1458,6 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
         case LLM_ARCH_KIMI_LINEAR:
         case LLM_ARCH_BAILINGMOE3:
         case LLM_ARCH_KIMI_K3:
-        case LLM_ARCH_GLM5_NEXT:
         case LLM_ARCH_QWEN3TTS:
         case LLM_ARCH_LIMITE:
             return false;

src/llama-arch.h

@@ -106,6 +106,7 @@ enum llm_arch {
     LLM_ARCH_NEMOTRON,
     LLM_ARCH_NEMOTRON_H,
     LLM_ARCH_NEMOTRON_H_MOE,
+    LLM_ARCH_METIS16,
     LLM_ARCH_EXAONE,
     LLM_ARCH_EXAONE4,
     LLM_ARCH_EXAONE_MOE,
@@ -276,6 +277,10 @@ enum llm_kv {
     LLM_KV_HRM_H_CYCLES,
     LLM_KV_HRM_L_CYCLES,
     LLM_KV_HRM_PREFIX_LM,
+    LLM_KV_METIS_PASS_COUNT,
+    LLM_KV_METIS_MEMORY_LENGTH,
+    LLM_KV_METIS_ROUTE_LENGTH,
+    LLM_KV_METIS_CONTINUATION_LOGIT_LIMIT,
     LLM_KV_REALTIME_DELAY_TOKENS,
     LLM_KV_MUDD_FEED_FORWARD_LENGTH,
     LLM_KV_MUDD_TAP_COUNT,
@@ -333,6 +338,7 @@ enum llm_kv {
     LLM_KV_ATTENTION_KV_SOURCE_LAYERS,
     LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,
     LLM_KV_ATTENTION_OUTPUT_LORA_RANK,
+    LLM_KV_ATTENTION_PASS_LORA_RANK,
     LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,
     LLM_KV_ATTENTION_COMPRESS_RATIOS,
     LLM_KV_ATTENTION_SHARED_KV_LAYERS,
@@ -391,6 +397,7 @@ enum llm_kv {
     LLM_KV_SSM_TIME_STEP_RANK,
     LLM_KV_SSM_GROUP_COUNT,
     LLM_KV_SSM_DT_B_C_RMS,
+    LLM_KV_SSM_NORM_EPS,
 
     LLM_KV_KDA_HEAD_DIM,
     LLM_KV_KDA_SAFE_GATE,
@@ -566,6 +573,36 @@ enum llm_tensor {
     LLM_TENSOR_FFN_EXP_PROBS_B_VL,
     LLM_TENSOR_FFN_LATENT_DOWN,
     LLM_TENSOR_FFN_LATENT_UP,
+    LLM_TENSOR_STREAM_EMBD,
+    LLM_TENSOR_DEPTH_MEM_STATE_WRITE,
+    LLM_TENSOR_DEPTH_MEM_META_WRITE,
+    LLM_TENSOR_DEPTH_MEM_ENTRY_BIAS,
+    LLM_TENSOR_DEPTH_MEM_QUERY,
+    LLM_TENSOR_DEPTH_MEM_KEY,
+    LLM_TENSOR_DEPTH_MEM_VALUE,
+    LLM_TENSOR_DEPTH_MEM_OUT,
+    LLM_TENSOR_DEPTH_MEM_ROUTE_STATE,
+    LLM_TENSOR_DEPTH_MEM_ROUTE_RETR,
+    LLM_TENSOR_DEPTH_MEM_ROUTE_DIFF,
+    LLM_TENSOR_DEPTH_MEM_ROUTE_HIST,
+    LLM_TENSOR_DEPTH_MEM_ROUTE,
+    LLM_TENSOR_DEPTH_MEM_STREAM_GATE,
+    LLM_TENSOR_CONT_HIDDEN_STATE,
+    LLM_TENSOR_CONT_HIDDEN_RETR,
+    LLM_TENSOR_CONT_HIDDEN_DIFF,
+    LLM_TENSOR_CONT_HIDDEN_ROUTE,
+    LLM_TENSOR_CONT_HIDDEN,
+    LLM_TENSOR_CONT_OUTPUT,
+    LLM_TENSOR_HC_ATTN_MIX,
+    LLM_TENSOR_HC_ATTN_READ,
+    LLM_TENSOR_HC_ATTN_WRITE,
+    LLM_TENSOR_HC_FFN_MIX,
+    LLM_TENSOR_HC_FFN_READ,
+    LLM_TENSOR_HC_FFN_WRITE,
+    LLM_TENSOR_ATTN_PASS_LORA_DOWN,
+    LLM_TENSOR_ATTN_PASS_LORA_UP,
+    LLM_TENSOR_ATTN_PASS_GATE,
+    LLM_TENSOR_FFN_EXP_EMBD,
     LLM_TENSOR_FFN_PRE_EXP_NORM,
     LLM_TENSOR_FFN_POST_EXP_NORM,
     LLM_TENSOR_FFN_RESID_ALPHA,

src/llama-context.cpp

@@ -2703,6 +2703,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
         res = std::max<uint32_t>(n_tokens * 80, 64u * model.n_tensors());
     } else if (model.arch == LLM_ARCH_QWEN3NEXT ||
         model.arch == LLM_ARCH_KIMI_LINEAR ||
+        model.arch == LLM_ARCH_METIS16 ||
         model.arch == LLM_ARCH_BAILINGMOE3 ||
         model.arch == LLM_ARCH_QWEN35 ||
         model.arch == LLM_ARCH_QWEN35MOE ||

src/llama-hparams.h

@@ -209,6 +209,7 @@ struct llama_hparams {
     uint32_t ssm_d_state = 0;
     uint32_t ssm_dt_rank = 0;
     uint32_t ssm_n_group = 0;
+    float    ssm_norm_eps = 0.0f; // 0 = use f_norm_rms_eps
 
     // for MiniMax-Text-01 linear attention
     uint32_t n_embd_head_la = 0;
@@ -230,6 +231,13 @@ struct llama_hparams {
     uint32_t n_hrm_l_cycles = 0;
     bool     hrm_prefix_lm = false;
 
+    // metis16 (passes over shared layers, depth memory)
+    uint32_t metis_n_pass         = 0;
+    uint32_t metis_n_mem          = 0;
+    uint32_t metis_n_route        = 0;
+    uint32_t metis_pass_lora_rank = 0;
+    float    metis_cont_limit     = 0.0f;
+
     // limite
     float    f_xsa_eps = 0.0f;
     std::array<float, 3> final_logit_sigmoid_capping = {}; // scale, shift, temperature

src/llama-model-saver.cpp

@@ -247,6 +247,10 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_HRM_H_CYCLES,                      hparams.n_hrm_h_cycles);
     add_kv(LLM_KV_HRM_L_CYCLES,                      hparams.n_hrm_l_cycles);
     add_kv(LLM_KV_HRM_PREFIX_LM,                     hparams.hrm_prefix_lm);
+    add_kv(LLM_KV_METIS_PASS_COUNT,                  hparams.metis_n_pass);
+    add_kv(LLM_KV_METIS_MEMORY_LENGTH,               hparams.metis_n_mem);
+    add_kv(LLM_KV_METIS_ROUTE_LENGTH,                hparams.metis_n_route);
+    add_kv(LLM_KV_METIS_CONTINUATION_LOGIT_LIMIT,    hparams.metis_cont_limit);
     if (hparams.final_logit_sigmoid_capping[0] != 0.0f) {
         add_kv(LLM_KV_FINAL_LOGIT_SIGMOID_CAPPING, std::vector<float>(
                 hparams.final_logit_sigmoid_capping.begin(), hparams.final_logit_sigmoid_capping.end()));
@@ -323,6 +327,7 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS,        hparams.is_recr_impl, true);
     add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,      hparams.dsv4_o_group_count);
     add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK,        hparams.dsv4_o_lora_rank);
+    add_kv(LLM_KV_ATTENTION_PASS_LORA_RANK,          hparams.metis_pass_lora_rank);
     add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, hparams.dsv4_compress_rope_base);
     if (model->arch == LLM_ARCH_DEEPSEEK4 || hparams.dsv4_hc_mult > 0) {
         // the loader requires one compress ratio per layer, including nextn layers
@@ -419,6 +424,7 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_SSM_TIME_STEP_RANK,                hparams.ssm_dt_rank);
     add_kv(LLM_KV_SSM_GROUP_COUNT,                   hparams.ssm_n_group);
     add_kv(LLM_KV_SSM_DT_B_C_RMS,                    hparams.ssm_dt_b_c_rms);
+    add_kv(LLM_KV_SSM_NORM_EPS,                      hparams.ssm_norm_eps);
 
     add_kv(LLM_KV_KDA_HEAD_DIM,                      hparams.n_embd_head_kda);
     add_kv(LLM_KV_KDA_SAFE_GATE,                     hparams.kda_safe_gate);
@@ -521,6 +527,32 @@ void llama_model_saver::add_tensors_from_model() {
     add_tensor(model->cls_out_b);
     add_tensor(model->cls_norm);
     add_tensor(model->hrm_z_l_init);
+    add_tensor(model->stream_embd);
+    add_tensor(model->depth_mem_state_write);
+    add_tensor(model->depth_mem_meta_write);
+    add_tensor(model->depth_mem_entry_bias);
+    add_tensor(model->depth_mem_query);
+    add_tensor(model->depth_mem_query_b);
+    add_tensor(model->depth_mem_key);
+    add_tensor(model->depth_mem_key_b);
+    add_tensor(model->depth_mem_value);
+    add_tensor(model->depth_mem_value_b);
+    add_tensor(model->depth_mem_out);
+    add_tensor(model->depth_mem_out_b);
+    add_tensor(model->depth_mem_route_state);
+    add_tensor(model->depth_mem_route_retr);
+    add_tensor(model->depth_mem_route_diff);
+    add_tensor(model->depth_mem_route_hist);
+    add_tensor(model->depth_mem_route_b);
+    add_tensor(model->depth_mem_stream_gate);
+    add_tensor(model->depth_mem_stream_gate_b);
+    add_tensor(model->cont_hidden_state);
+    add_tensor(model->cont_hidden_retr);
+    add_tensor(model->cont_hidden_diff);
+    add_tensor(model->cont_hidden_route);
+    add_tensor(model->cont_hidden_b);
+    add_tensor(model->cont_output);
+    add_tensor(model->cont_output_b);
     add_tensor(model->hc_head_fn);
     add_tensor(model->hc_head_base);
     add_tensor(model->hc_head_scale);

src/llama-model.cpp

@@ -242,6 +242,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
             return new llama_model_nemotron_h(params);
         case LLM_ARCH_NEMOTRON_H_MOE:
             return new llama_model_nemotron_h_moe(params);
+        case LLM_ARCH_METIS16:
+            return new llama_model_metis16(params);
         case LLM_ARCH_EXAONE:
             return new llama_model_exaone(params);
         case LLM_ARCH_EXAONE4:
@@ -410,6 +412,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     const std::string tensor_name = tensor->name;
     const bool is_dsv4 = ud->model->arch == LLM_ARCH_DEEPSEEK4 || ud->model->arch == LLM_ARCH_DEEPSEEK41 ||
         (ud->model->arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0);
+    const bool is_glm5_next = ud->model->arch == LLM_ARCH_GLM5_NEXT;
 
     static const std::regex pattern_q_weight        ("blk\\.\\d*\\.attn_q.weight");
     static const std::regex pattern_kv_weight       ("blk\\.\\d*\\.attn_(k|v).weight");
@@ -429,6 +432,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     static const std::regex pattern_attn_out_a_weight("blk\\.\\d*\\.attn_output_a\\.weight");
     static const std::regex pattern_attn_out_b_weight("blk\\.\\d*\\.attn_output_b\\.weight");
     static const std::regex pattern_attn_q_b_weight ("blk\\.\\d*\\.attn_q_b\\.weight");
+    static const std::regex pattern_attn_kv_b_weight("blk\\.\\d*\\.attn_(k|v)_b\\.weight");
     static const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight");
 
     static const std::regex pattern_ssm_dt          ("blk\\.\\d*\\.ssm_dt.bias");
@@ -441,6 +445,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     static const std::regex pattern_s_cache         ("cache_s_l\\d*");
     static const std::regex pattern_ssm_conv1d      ("blk\\.\\d*\\.ssm_conv1d.weight");
     static const std::regex pattern_ssm_out_weight  ("blk\\.\\d*\\.ssm_out.weight");
+    static const std::regex pattern_kda_conv1d      ("blk\\.\\d*\\.ssm_conv1d_(q|k|v)\\.weight");
+    static const std::regex pattern_kda_gate_b      ("blk\\.\\d*\\.ssm_(f|g)_b\\.weight");
 
     static const std::regex pattern_ffn_up_weight     ("blk\\.\\d*\\.ffn_up(_exps)?.weight");
     static const std::regex pattern_ffn_up_bias       ("blk\\.\\d*\\.ffn_up(_exps)?.bias");
@@ -511,7 +517,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
     };
 
     auto get_tensor_config = [&]() -> tensor_config {
-        if (ud->model->arch == LLM_ARCH_HRM_TEXT) {
+        if (ud->model->arch == LLM_ARCH_HRM_TEXT || ud->model->arch == LLM_ARCH_METIS16) {
             // aliased cache slots cannot satisfy the meta-split invariants, so replicate all tensors
             return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, tensor, 0, 0};
         }
@@ -532,6 +538,25 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
             if (std::regex_match(tensor_name, pattern_attn_out_b_weight)) {
                 return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0);
             }
+        }
+        if (is_glm5_next) {
+            // split the KDA and MLA heads, the MLA latent cache and the indexer stay mirrored
+            if (std::regex_match(tensor_name, pattern_kv_cache)) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED);
+            }
+            if (std::regex_match(tensor_name, pattern_attn_q_b_weight) || std::regex_match(tensor_name, pattern_kda_gate_b) ||
+                    std::regex_match(tensor_name, pattern_ssm_beta)) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight");
+            }
+            if (std::regex_match(tensor_name, pattern_attn_kv_b_weight) || std::regex_match(tensor_name, pattern_kda_conv1d)) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_2, "attn_output.weight");
+            }
+            if (std::regex_match(tensor_name, pattern_ssm_dt) || std::regex_match(tensor_name, pattern_ssm_a) ||
+                    std::regex_match(tensor_name, pattern_r_cache) || std::regex_match(tensor_name, pattern_s_cache)) {
+                return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "attn_output.weight");
+            }
+        }
+        if (is_dsv4 || is_glm5_next) {
             if (std::regex_match(tensor_name, pattern_ffn_up_shexp_weight) ||
                     std::regex_match(tensor_name, pattern_ffn_gate_shexp_weight)) {
                 return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ffn_down_shexp.weight");
@@ -645,6 +670,12 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
 
     // if a model has fused tensors they need to be separated into "segments", see the comments on ggml_backend_meta_split_state struct
     auto get_split_segments = [&](int axis, uint32_t il) -> std::vector<std::pair<int64_t, uint32_t>> {
+        if (is_glm5_next && std::regex_match(tensor_name, pattern_r_cache)) {
+            // the conv state holds the q, k and v windows one after the other
+            const int64_t conv_state_size = (hparams.ssm_d_conv - 1) * hparams.n_head(il) * hparams.n_embd_head_kda;
+            GGML_ASSERT(tensor->ne[axis] == 3*conv_state_size);
+            return {{conv_state_size, 3}};
+        }
         // TODO: deduplicate condition [TAG_SPLIT_QGATE_QWEN]
         if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || ud->model->arch == LLM_ARCH_CLEF || ud->model->arch == LLM_ARCH_CLOVER_MOE ||
                 ud->model->arch == LLM_ARCH_QWEN4EXP) {
@@ -736,7 +767,39 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str
 
     auto get_split_granularity = [&](int64_t blck_size, uint32_t il, const std::vector<std::pair<int64_t, uint32_t>> & segments) -> std::vector<int64_t> {
         // for better performance it may make sense to round up blck_size to a higher power of 2 so that more efficient kernels can be used
-        if (hparams.is_recr(il)) {
+        if (is_glm5_next && hparams.is_recr(il)) {
+            // KDA, each device holds whole heads
+            const int64_t head_dim        = hparams.n_embd_head_kda;
+            const int64_t blck_size_perf  = std::lcm(blck_size, 128);
+            const int64_t granularity_qkv = std::lcm(blck_size_perf, head_dim);
+            if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_kv_weight) ||
+                    std::regex_match(tensor_name, pattern_kda_conv1d) || std::regex_match(tensor_name, pattern_kda_gate_b) ||
+                    std::regex_match(tensor_name, pattern_ssm_dt) || std::regex_match(tensor_name, pattern_attn_out_weight)) {
+                return {granularity_qkv};
+            }
+            if (std::regex_match(tensor_name, pattern_ssm_a) || std::regex_match(tensor_name, pattern_ssm_beta)) {
+                return {granularity_qkv / head_dim};
+            }
+            if (std::regex_match(tensor_name, pattern_r_cache)) {
+                return std::vector<int64_t>(segments.size(), granularity_qkv * (hparams.ssm_d_conv - 1));
+            }
+            if (std::regex_match(tensor_name, pattern_s_cache)) {
+                return {granularity_qkv * head_dim};
+            }
+        } else if (is_glm5_next) {
+            // MLA, each device holds whole heads
+            // an even number of heads per device, the HIP tile FA kernel for D=512 needs an even GQA ratio
+            const int64_t granularity_head = std::lcm(std::lcm(hparams.n_embd_head_v_mla(), blck_size) / hparams.n_embd_head_v_mla(), int64_t(2));
+            if (std::regex_match(tensor_name, pattern_attn_q_b_weight)) {
+                return {granularity_head * hparams.n_embd_head_k_mla()};
+            }
+            if (std::regex_match(tensor_name, pattern_attn_kv_b_weight)) {
+                return {granularity_head};
+            }
+            if (std::regex_match(tensor_name, pattern_attn_out_weight)) {
+                return {granularity_head * hparams.n_embd_head_v_mla()};
+            }
+        } else if (hparams.is_recr(il)) {
             // linear attention
             const int64_t head_dim        = hparams.ssm_d_state;
             const int64_t blck_size_perf  = std::lcm(blck_size, 128);
@@ -988,6 +1051,7 @@ const char * llm_type_name(llm_type type) {
         case LLM_TYPE_7B:            return "7B";
         case LLM_TYPE_8B:            return "8B";
         case LLM_TYPE_9B:            return "9B";
+        case LLM_TYPE_10B:           return "10B";
         case LLM_TYPE_11B:           return "11B";
         case LLM_TYPE_12B:           return "12B";
         case LLM_TYPE_13B:           return "13B";
@@ -3557,6 +3621,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
         case LLM_ARCH_WAVTOKENIZER_DEC:
         case LLM_ARCH_NEMOTRON_H:
         case LLM_ARCH_NEMOTRON_H_MOE:
+        case LLM_ARCH_METIS16:
         case LLM_ARCH_KIMI_LINEAR:
         case LLM_ARCH_KIMI_K3:
         case LLM_ARCH_GLM5_NEXT:

src/llama-model.h

@@ -76,6 +76,7 @@ enum llm_type {
     LLM_TYPE_7B,
     LLM_TYPE_8B,
     LLM_TYPE_9B,
+    LLM_TYPE_10B,
     LLM_TYPE_11B,
     LLM_TYPE_12B,
     LLM_TYPE_13B,
@@ -646,6 +647,18 @@ struct llama_layer {
     // gemma4 layer output scale, reused for talkie embedding skip scale
     struct ggml_tensor * out_scale = nullptr;
 
+    // metis16 per-pass hyper-connection mixes, attention pass LoRA, expert route embeddings
+    struct ggml_tensor * hc_attn_mix         = nullptr;
+    struct ggml_tensor * hc_attn_read        = nullptr;
+    struct ggml_tensor * hc_attn_write       = nullptr;
+    struct ggml_tensor * hc_ffn_mix          = nullptr;
+    struct ggml_tensor * hc_ffn_read         = nullptr;
+    struct ggml_tensor * hc_ffn_write        = nullptr;
+    struct ggml_tensor * attn_pass_lora_down = nullptr;
+    struct ggml_tensor * attn_pass_lora_up   = nullptr;
+    struct ggml_tensor * attn_pass_gate      = nullptr;
+    struct ggml_tensor * ffn_exp_embd        = nullptr;
+
     // limite
     struct ggml_tensor * attn_ve_gate     = nullptr;
     struct ggml_tensor * attn_xsa_alpha   = nullptr;
@@ -779,6 +792,34 @@ struct llama_model {
     // hrm-text initial low-cycle state
     struct ggml_tensor * hrm_z_l_init = nullptr;
 
+    // metis16 hyper-connection stream embeddings, depth memory and continuation head
+    struct ggml_tensor * stream_embd             = nullptr;
+    struct ggml_tensor * depth_mem_state_write   = nullptr;
+    struct ggml_tensor * depth_mem_meta_write    = nullptr;
+    struct ggml_tensor * depth_mem_entry_bias    = nullptr;
+    struct ggml_tensor * depth_mem_query         = nullptr;
+    struct ggml_tensor * depth_mem_query_b       = nullptr;
+    struct ggml_tensor * depth_mem_key           = nullptr;
+    struct ggml_tensor * depth_mem_key_b         = nullptr;
+    struct ggml_tensor * depth_mem_value         = nullptr;
+    struct ggml_tensor * depth_mem_value_b       = nullptr;
+    struct ggml_tensor * depth_mem_out           = nullptr;
+    struct ggml_tensor * depth_mem_out_b         = nullptr;
+    struct ggml_tensor * depth_mem_route_state   = nullptr;
+    struct ggml_tensor * depth_mem_route_retr    = nullptr;
+    struct ggml_tensor * depth_mem_route_diff    = nullptr;
+    struct ggml_tensor * depth_mem_route_hist    = nullptr;
+    struct ggml_tensor * depth_mem_route_b       = nullptr;
+    struct ggml_tensor * depth_mem_stream_gate   = nullptr;
+    struct ggml_tensor * depth_mem_stream_gate_b = nullptr;
+    struct ggml_tensor * cont_hidden_state       = nullptr;
+    struct ggml_tensor * cont_hidden_retr        = nullptr;
+    struct ggml_tensor * cont_hidden_diff        = nullptr;
+    struct ggml_tensor * cont_hidden_route       = nullptr;
+    struct ggml_tensor * cont_hidden_b           = nullptr;
+    struct ggml_tensor * cont_output             = nullptr;
+    struct ggml_tensor * cont_output_b           = nullptr;
+
     // DeepSeek-V4
     struct ggml_tensor * hc_head_fn    = nullptr;
     struct ggml_tensor * hc_head_base  = nullptr;

src/llama-quant.cpp

@@ -361,6 +361,16 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param
     quantize &= name.find("indexer.k_proj.weight") == std::string::npos;
     quantize &= name.find("indexer.q_proj.weight") == std::string::npos;
 
+    // metis16: small control tensors that steer every layer; the stream and hc_ tables (ne[0] = 4) would fall back to F16
+    if (arch == LLM_ARCH_METIS16) {
+        quantize &= name.rfind("depth_mem_", 0) == std::string::npos;
+        quantize &= name.rfind("cont_", 0)      == std::string::npos;
+        quantize &= name.find("stream_embd")    == std::string::npos;
+        quantize &= name.find("attn_pass_lora") == std::string::npos;
+        quantize &= name.find("ffn_exp_embd")   == std::string::npos;
+        quantize &= name.find("hc_")            == std::string::npos;
+    }
+
     // glm5-next
     if (arch == LLM_ARCH_GLM5_NEXT) {
         quantize &= name.find("hc_")                     == std::string::npos;

src/models/gemma4.cpp

@@ -154,8 +154,7 @@ std::unique_ptr<llm_graph_context> llama_model_gemma4::build_arch_graph(const ll
 // get 2D slice view from a 3D tensor, the idx corresponds to the 3rd dim
 static ggml_tensor * gemma4_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {
     GGML_ASSERT(idx < (int) x->ne[2]);
-    return ggml_view_2d(ctx0, x, x->ne[0], x->ne[1], ggml_row_size(x->type, x->ne[0]),
-                        idx * x->ne[0] * x->ne[1] * ggml_element_size(x));
+    return ggml_view_2d(ctx0, x, x->ne[0], x->ne[1], x->nb[1], idx * x->nb[2]);
 }
 
 llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_params & params) :
@@ -183,10 +182,14 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
     // TODO: is causal == true correct? might need some changes
     auto * inp_attn = build_attn_inp_kv_iswa();
 
+    auto inp_attn_full  = *inp_attn;
+    auto inp_attn_prune = *inp_attn;
+
     ggml_tensor * inp_out_ids = build_inp_out_ids();
 
     if (model.per_layer_tok_embd) {
-        const float tok_embd_scale = sqrtf((float) n_embd_per_layer);
+        // fold per_layer_input_scale (1/sqrt(2)) from project_per_layer_inputs() into tok_embd_scale
+        const float tok_embd_scale = sqrtf((float) n_embd_per_layer) / sqrtf(2.0f);
 
         inp_per_layer = ggml_scale     (ctx0, inp_per_layer, tok_embd_scale);
         inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, n_embd_per_layer, n_layer, inp_per_layer->ne[1]);
@@ -195,7 +198,36 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
         inp_per_layer = project_per_layer_inputs(inpL, inp_per_layer);
     }
 
+    // KV-shared layers do not write to the KV cache, so prune non-output token rows before the first KV-shared layer
+    bool prune_shared = crop_before_nextn(inp_out_ids) &&
+                        hparams.n_layer_kv_from_start > 0 && hparams.n_layer_kv_from_start < n_layer &&
+                        !cparams.embeddings && cparams.n_seq_max == 1;
+    for (int il = (int) hparams.n_layer_kv_from_start; prune_shared && il <= n_layer; ++il) {
+        prune_shared = !cparams.embeddings_layer_inp[il];
+    }
+
+    int64_t n_rows = n_tokens;
+
     for (int il = 0; il < n_layer; ++il) {
+        if (prune_shared && il == (int) hparams.n_layer_kv_from_start) {
+            inpL    = ggml_get_rows(ctx0, inpL, inp_out_ids);
+            inp_pos = ggml_reshape_1d(ctx0, ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, inp_pos, 1, n_tokens), inp_out_ids), n_outputs);
+
+            // TODO: instead of doing get_rows here, extend llm_attn_mask to have a `out` mask that
+            //       is construted during set_input to contain just the pruned output cells
+            //       this can be used to enable pruning with more than 1 sequence, both with unified and split KV cache
+            //       ref: https://github.com/ggml-org/llama.cpp/pull/31356#discussion_r4455286628
+            inp_attn_prune.self_attn_mask.raw     = ggml_get_rows(ctx0, inp_attn_full.self_attn_mask.raw,     inp_out_ids);
+            inp_attn_prune.self_attn_mask_swa.raw = ggml_get_rows(ctx0, inp_attn_full.self_attn_mask_swa.raw, inp_out_ids);
+
+            inp_attn_prune.self_attn_mask.raw     = ggml_cast(ctx0, inp_attn_prune.self_attn_mask.raw,     GGML_TYPE_F16);
+            inp_attn_prune.self_attn_mask_swa.raw = ggml_cast(ctx0, inp_attn_prune.self_attn_mask_swa.raw, GGML_TYPE_F16);
+
+            n_rows = n_outputs;
+
+            inp_attn = &inp_attn_prune;
+        }
+
         const int64_t n_embd_head = hparams.n_embd_head_k(il);
         GGML_ASSERT(n_embd_head == hparams.n_embd_head_v(il));
 
@@ -226,14 +258,14 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
             qkv_fused = build_lora_mm(model.layers[il].wqkv, cur, model.layers[il].wqkv_s);
             cb(qkv_fused, "wqkv", il);
             const int64_t q_dim = n_embd_head * n_head;
-            Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, q_dim, n_tokens, qkv_fused->nb[1], 0));
+            Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, q_dim, n_rows, qkv_fused->nb[1], 0));
         } else {
             Qcur = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s);
         }
         {
             cb(Qcur, "Qcur", il);
 
-            Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
+            Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_rows);
 
             Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);
             cb(Qcur, "Qcur_normed", il);
@@ -287,9 +319,8 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
                     Qcur, nullptr, nullptr, nullptr, nullptr, nullptr, hparams.f_attention_scale, il);
         }
 
-        // TODO @ngxson : strip unused token right after the last KV layer to speed up prompt processing
         // keep all rows when extracting unmasked nextn embeddings (MTP target needs the hidden state for every token)
-        if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
+        if (!prune_shared && il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
             cur  = ggml_get_rows(ctx0,  cur, inp_out_ids);
             inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
         }
@@ -331,7 +362,7 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
             ggml_tensor * tmp = ggml_rms_norm(ctx0, attn_out, hparams.f_norm_rms_eps);
             tmp = ggml_scale(ctx0, tmp, 1.0f / sqrtf((float) n_embd));
             tmp = ggml_mul(ctx0, tmp, model.layers[il].ffn_gate_inp_s);
-            ggml_tensor * logits = build_lora_mm(model.layers[il].ffn_gate_inp, tmp); // [n_expert, n_tokens]
+            ggml_tensor * logits = build_lora_mm(model.layers[il].ffn_gate_inp, tmp); // [n_expert, n_rows]
             cb(logits, "ffn_moe_logits", il);
 
             cur_moe = build_moe_ffn(cur_moe,
@@ -383,18 +414,17 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
             ggml_tensor * pe_in = cur;
             cb(cur, "pe_in", il);
 
-            cur = build_lora_mm(model.layers[il].per_layer_inp_gate, cur); // [n_embd_per_layer, n_tokens]
+            cur = build_lora_mm(model.layers[il].per_layer_inp_gate, cur); // [n_embd_per_layer, n_rows]
             cur = ggml_gelu(ctx0, cur);
 
             ggml_tensor * inp_this_layer = gemma4_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens]
 
-            // TODO @ngxson : improve this
-            if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
+            if ((prune_shared && !hparams.has_kv(il)) || (il == n_layer - 1 && crop_before_nextn(inp_out_ids))) {
                 inp_this_layer = ggml_get_rows(ctx0, inp_this_layer, inp_out_ids);
             }
 
             cur = ggml_mul(ctx0, cur, inp_this_layer);
-            cur = build_lora_mm(model.layers[il].per_layer_proj, cur); // [n_embd, n_tokens]
+            cur = build_lora_mm(model.layers[il].per_layer_proj, cur); // [n_embd, n_rows]
             cur = build_norm(cur, model.layers[il].per_layer_post_norm, nullptr, LLM_NORM_RMS, il);
             cb(cur, "per_layer_embd_out", il);
 
@@ -522,19 +552,22 @@ ggml_tensor * llama_model_gemma4::graph::project_per_layer_inputs(ggml_tensor *
     const float per_layer_input_scale      = 1.0f / sqrtf(2.0f);
 
     // note: this matrix multiplication will be performed in the input layer (i.e. on the CPU)
+    // fold per_layer_input_scale into proj_norm_w and per_layer_projection_scale into proj_eps
+    ggml_tensor * proj_norm_w = ggml_scale(ctx0, model.per_layer_proj_norm, per_layer_input_scale);
+    ggml_build_forward_expand(gf, proj_norm_w);
+
     ggml_tensor * per_layer_proj;
     per_layer_proj = ggml_mul_mat   (ctx0, model.per_layer_model_proj, inp_batch);
-    per_layer_proj = ggml_scale     (ctx0, per_layer_proj, per_layer_projection_scale);
     per_layer_proj = ggml_reshape_3d(ctx0, per_layer_proj, n_embd_per_layer, n_layer, n_tokens);
 
-    per_layer_proj = build_norm(per_layer_proj, model.per_layer_proj_norm, nullptr, LLM_NORM_RMS, -1);
+    const float proj_eps = hparams.f_norm_rms_eps / (per_layer_projection_scale * per_layer_projection_scale);
+    per_layer_proj = ggml_rms_norm(ctx0, per_layer_proj, proj_eps);
+    per_layer_proj = ggml_mul(ctx0, per_layer_proj, proj_norm_w);
     cb(per_layer_proj, "per_layer_proj", -1);
 
-    inp_per_layer = ggml_add  (ctx0, per_layer_proj, inp_per_layer);
-    inp_per_layer = ggml_scale(ctx0, inp_per_layer, per_layer_input_scale);
+    inp_per_layer = ggml_add(ctx0, per_layer_proj, inp_per_layer);
     cb(inp_per_layer, "inp_per_layer", -1);
 
     // permute to shape: [n_embd_per_layer, n_tokens, n_layer]
-    inp_per_layer = ggml_cont(ctx0, ggml_permute(ctx0, inp_per_layer, 0, 2, 1, 3));
-    return inp_per_layer;
+    return ggml_permute(ctx0, inp_per_layer, 0, 2, 1, 3);
 }

src/models/gpt2.cpp

@@ -22,6 +22,7 @@ void llama_model_gpt2::load_arch_tensors(llama_model_loader &) {
     output_norm   = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
     output_norm_b = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "bias"),   {n_embd}, 0);
     output        = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
+    output_b      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "bias"),   {n_vocab}, TENSOR_NOT_REQUIRED);
 
     // if output is NULL, init from the input tok embed
     if (output == NULL) {
@@ -140,6 +141,9 @@ llama_model_gpt2::graph::graph(const llama_model & model, const llm_graph_params
     res->t_embd = cur;
 
     cur = build_lora_mm(model.output, cur, model.output_s);
+    if (model.output_b) {
+        cur = ggml_add(ctx0, cur, model.output_b);
+    }
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;

src/models/lfm-base.cpp

@@ -117,9 +117,9 @@ ggml_tensor * llm_build_lfm_base::build_lfm_dense_ffn(const llama_model & model,
     GGML_ASSERT(!model.layers[il].ffn_gate_b);
     GGML_ASSERT(!model.layers[il].ffn_down_b);
     return build_ffn(cur,
-        model.layers[il].ffn_up,   NULL, NULL,
-        model.layers[il].ffn_gate, NULL, NULL,
-        model.layers[il].ffn_down, NULL, NULL,
+        model.layers[il].ffn_up,   NULL, model.layers[il].ffn_up_s,
+        model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,
+        model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s,
         NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
 }
 
@@ -136,5 +136,10 @@ ggml_tensor * llm_build_lfm_base::build_lfm_moe_ffn(const llama_model & model,
             LLM_FFN_SILU, true,
             hparams.expert_weights_scale,
             static_cast<llama_expert_gating_func_type>(hparams.expert_gating_func),
-            il);
+            il,
+            nullptr,
+            nullptr,
+            model.layers[il].ffn_up_exps_s,
+            model.layers[il].ffn_gate_exps_s,
+            model.layers[il].ffn_down_exps_s);
 }

src/models/lfm2.cpp

@@ -399,9 +399,9 @@ llama_model_lfm2::graph_decision::graph_decision(const llama_model & model, cons
 
         ggml_tensor * ffn_out = build_norm(cur, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
         ffn_out = build_ffn(ffn_out,
-                layer.ffn_up,   NULL, NULL,
-                layer.ffn_gate, NULL, NULL,
-                layer.ffn_down, NULL, NULL,
+                layer.ffn_up,   NULL, layer.ffn_up_s,
+                layer.ffn_gate, NULL, layer.ffn_gate_s,
+                layer.ffn_down, NULL, layer.ffn_down_s,
                 NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
 
         cur = ggml_add(ctx0, cur, ffn_out);

src/models/lfm3.cpp

@@ -129,7 +129,7 @@ llama_model_lfm3::graph<decision>::graph(const llama_model & model, const llm_gr
                 q, k, v, nullptr, nullptr, nullptr, 1.0f / sqrtf(float(n_embd_head)), il);
         cb(attn_out, "attn_out", il);
 
-        ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, cur_in);
+        ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, cur_in, model.layers[il].wqkv_gate_s);
         gate = ggml_sigmoid(ctx0, gate);
         cb(gate, "model.layers.{}.self_attn.gate", il);
 
@@ -157,6 +157,8 @@ llama_model_lfm3::graph<decision>::graph(const llama_model & model, const llm_gr
     ggml_tensor * inp_out_ids = build_inp_out_ids();
 
     for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = cur;
+
         auto * prev_cur = cur;
 
         cur = build_norm(cur, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);

src/models/mamba-base.cpp

@@ -290,7 +290,11 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
         // grouped RMS norm
         if (model.layers[il].ssm_norm) {
             y = ggml_reshape_4d(ctx0, y, d_inner / n_group, n_group, n_seq_tokens, n_seqs);
-            y = build_norm(y, model.layers[il].ssm_norm, NULL, LLM_NORM_RMS, il);
+            if (hparams.ssm_norm_eps > 0.0f) {
+                y = ggml_mul(ctx0, ggml_rms_norm(ctx0, y, hparams.ssm_norm_eps), model.layers[il].ssm_norm);
+            } else {
+                y = build_norm(y, model.layers[il].ssm_norm, NULL, LLM_NORM_RMS, il);
+            }
         }
 
         y = ggml_reshape_2d(ctx0, y, d_inner, n_seq_tokens * n_seqs);

src/models/metis16.cpp

@@ -0,0 +1,401 @@
+#include "models.h"
+
+#include <cmath>
+
+void llama_model_metis16::load_arch_hparams(llama_model_loader & ml) {
+    ml.get_key(LLM_KV_SSM_CONV_KERNEL,    hparams.ssm_d_conv);
+    ml.get_key(LLM_KV_SSM_INNER_SIZE,     hparams.ssm_d_inner);
+    ml.get_key(LLM_KV_SSM_STATE_SIZE,     hparams.ssm_d_state);
+    ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
+    ml.get_key(LLM_KV_SSM_GROUP_COUNT,    hparams.ssm_n_group);
+    ml.get_key(LLM_KV_SSM_NORM_EPS,       hparams.ssm_norm_eps);
+
+    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_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);
+    ml.get_key(LLM_KV_MOE_LATENT_SIZE,                   hparams.moe_latent_size);
+
+    ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT,            hparams.dsv4_hc_mult);
+    ml.get_key(LLM_KV_METIS_PASS_COUNT,                  hparams.metis_n_pass);
+    ml.get_key(LLM_KV_METIS_MEMORY_LENGTH,               hparams.metis_n_mem);
+    ml.get_key(LLM_KV_METIS_ROUTE_LENGTH,                hparams.metis_n_route);
+    ml.get_key(LLM_KV_ATTENTION_PASS_LORA_RANK,          hparams.metis_pass_lora_rank);
+    ml.get_key(LLM_KV_METIS_CONTINUATION_LOGIT_LIMIT,    hparams.metis_cont_limit);
+
+    n_pass     = hparams.metis_n_pass;
+    n_stream   = hparams.dsv4_hc_mult;
+    n_mem      = hparams.metis_n_mem;
+    n_route    = hparams.metis_n_route;
+    cont_limit = hparams.metis_cont_limit;
+
+    if (n_pass < 1 || hparams.n_layer() % n_pass != 0) {
+        throw std::runtime_error(format("metis16: %u layers do not split into %u passes", hparams.n_layer(), n_pass));
+    }
+    n_layer_phys = hparams.n_layer() / n_pass;
+
+    bool has_attn = false;
+    for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
+        if (hparams.n_head_kv(il) != hparams.n_head_kv(il % n_layer_phys)) {
+            throw std::runtime_error(format("metis16: layer %u does not repeat the layout of layer %u", il, il % n_layer_phys));
+        }
+        hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0;
+        has_attn = has_attn || !hparams.is_recr_impl[il];
+    }
+    if (!has_attn) {
+        throw std::runtime_error("metis16: at least one attention layer is required");
+    }
+
+    switch (n_layer_phys) {
+        case 12: type = LLM_TYPE_3B;  break; // Praxis
+        case 20: type = LLM_TYPE_10B; break; // Logos
+        default: type = LLM_TYPE_UNKNOWN;
+    }
+}
+
+void llama_model_metis16::load_arch_tensors(llama_model_loader &) {
+    LLAMA_LOAD_LOCALS;
+
+    const int64_t d_conv     = hparams.ssm_d_conv;
+    const int64_t d_inner    = hparams.ssm_d_inner;
+    const int64_t d_state    = hparams.ssm_d_state;
+    const int64_t n_ssm_head = hparams.ssm_dt_rank;
+    const int64_t n_group    = hparams.ssm_n_group;
+    const int64_t d_in_proj  = 2*d_inner + 2*n_group*d_state + n_ssm_head;
+    const int64_t n_latent   = hparams.moe_latent_size;
+    const int64_t n_ff_shexp = hparams.n_ff_shexp;
+    const int64_t n_s = n_stream;
+    const int64_t n_p = n_pass;
+    const int64_t n_m = n_mem;
+    const int64_t n_r = n_route;
+
+    int64_t n_attn = 0;
+    for (uint32_t il = 0; il < n_layer_phys; ++il) {
+        n_attn += !hparams.is_recr(il);
+    }
+
+    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);
+    }
+
+    stream_embd = create_tensor(tn(LLM_TENSOR_STREAM_EMBD, "weight"), {n_s, n_embd}, 0);
+
+    depth_mem_state_write   = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_STATE_WRITE, "weight"), {n_embd, n_m}, 0);
+    depth_mem_meta_write    = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_META_WRITE,  "weight"), {n_r + 3, n_m}, 0);
+    depth_mem_entry_bias    = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_ENTRY_BIAS,  "weight"), {n_m, n_p*(n_attn + 1)}, 0);
+    depth_mem_query         = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_QUERY,       "weight"), {n_embd, n_m}, 0);
+    depth_mem_query_b       = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_QUERY,       "bias"),   {n_m}, 0);
+    depth_mem_key           = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_KEY,         "weight"), {n_m, n_m}, 0);
+    depth_mem_key_b         = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_KEY,         "bias"),   {n_m}, 0);
+    depth_mem_value         = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_VALUE,       "weight"), {n_m, n_m}, 0);
+    depth_mem_value_b       = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_VALUE,       "bias"),   {n_m}, 0);
+    depth_mem_out           = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_OUT,         "weight"), {n_m, n_embd}, 0);
+    depth_mem_out_b         = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_OUT,         "bias"),   {n_embd}, 0);
+    depth_mem_route_state   = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_ROUTE_STATE, "weight"), {n_embd, n_r}, 0);
+    depth_mem_route_retr    = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_ROUTE_RETR,  "weight"), {n_embd, n_r}, 0);
+    depth_mem_route_diff    = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_ROUTE_DIFF,  "weight"), {n_embd, n_r}, 0);
+    depth_mem_route_hist    = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_ROUTE_HIST,  "weight"), {n_r, n_r}, 0);
+    depth_mem_route_b       = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_ROUTE,       "bias"),   {n_r}, 0);
+    depth_mem_stream_gate   = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_STREAM_GATE, "weight"), {n_embd, n_s}, 0);
+    depth_mem_stream_gate_b = create_tensor(tn(LLM_TENSOR_DEPTH_MEM_STREAM_GATE, "bias"),   {n_s}, 0);
+
+    cont_hidden_state = create_tensor(tn(LLM_TENSOR_CONT_HIDDEN_STATE, "weight"), {n_embd, n_r}, 0);
+    cont_hidden_retr  = create_tensor(tn(LLM_TENSOR_CONT_HIDDEN_RETR,  "weight"), {n_embd, n_r}, 0);
+    cont_hidden_diff  = create_tensor(tn(LLM_TENSOR_CONT_HIDDEN_DIFF,  "weight"), {n_embd, n_r}, 0);
+    cont_hidden_route = create_tensor(tn(LLM_TENSOR_CONT_HIDDEN_ROUTE, "weight"), {n_r, n_r}, 0);
+    cont_hidden_b     = create_tensor(tn(LLM_TENSOR_CONT_HIDDEN,       "bias"),   {n_r}, 0);
+    cont_output       = create_tensor(tn(LLM_TENSOR_CONT_OUTPUT,       "weight"), {n_r, 1}, 0);
+    cont_output_b     = create_tensor(tn(LLM_TENSOR_CONT_OUTPUT,       "bias"),   {1}, 0);
+
+    for (int il = 0; il < (int) n_layer_phys; ++il) {
+        auto & layer = layers[il];
+
+        layer.attn_norm     = create_tensor(tn(LLM_TENSOR_ATTN_NORM,     "weight", il), {n_embd}, 0);
+        layer.hc_attn_mix   = create_tensor(tn(LLM_TENSOR_HC_ATTN_MIX,   "weight", il), {n_s, n_s, n_p}, 0);
+        layer.hc_attn_read  = create_tensor(tn(LLM_TENSOR_HC_ATTN_READ,  "weight", il), {n_s, n_p}, 0);
+        layer.hc_attn_write = create_tensor(tn(LLM_TENSOR_HC_ATTN_WRITE, "weight", il), {n_s, n_p}, 0);
+        layer.hc_ffn_mix    = create_tensor(tn(LLM_TENSOR_HC_FFN_MIX,    "weight", il), {n_s, n_s, n_p}, 0);
+        layer.hc_ffn_read   = create_tensor(tn(LLM_TENSOR_HC_FFN_READ,   "weight", il), {n_s, n_p}, 0);
+        layer.hc_ffn_write  = create_tensor(tn(LLM_TENSOR_HC_FFN_WRITE,  "weight", il), {n_s, n_p}, 0);
+        layer.ffn_norm      = create_tensor(tn(LLM_TENSOR_FFN_NORM,      "weight", il), {n_embd}, 0);
+
+        if (hparams.is_recr(il)) {
+            layer.ssm_in       = create_tensor(tn(LLM_TENSOR_SSM_IN,     "weight", il), {n_embd, d_in_proj}, 0);
+            layer.ssm_conv1d   = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), {d_conv, d_inner + 2*n_group*d_state}, 0);
+            layer.ssm_conv1d_b = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "bias",   il), {d_inner + 2*n_group*d_state}, 0);
+            layer.ssm_dt_b     = create_tensor(tn(LLM_TENSOR_SSM_DT,     "bias",   il), {n_ssm_head}, 0);
+            layer.ssm_norm     = create_tensor(tn(LLM_TENSOR_SSM_NORM,   "weight", il), {d_inner / n_group, n_group}, 0);
+            layer.ssm_out      = create_tensor(tn(LLM_TENSOR_SSM_OUT,    "weight", il), {d_inner, n_embd}, 0);
+
+            // no "weight" suffix for these
+            layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, il), {1, n_ssm_head}, 0);
+            layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, il), {1, n_ssm_head}, 0);
+        } else {
+            const int64_t n_qkv = n_embd_head_k * (hparams.n_head(il) + 2*hparams.n_head_kv(il));
+            const int64_t rank  = hparams.metis_pass_lora_rank;
+
+            layer.wqkv                = create_tensor(tn(LLM_TENSOR_ATTN_QKV,            "weight", il), {n_embd, n_qkv}, 0);
+            layer.wo                  = create_tensor(tn(LLM_TENSOR_ATTN_OUT,            "weight", il), {n_embd_head_k * hparams.n_head(il), n_embd}, 0);
+            layer.attn_pass_lora_down = create_tensor(tn(LLM_TENSOR_ATTN_PASS_LORA_DOWN, "weight", il), {n_embd, rank}, 0);
+            layer.attn_pass_lora_up   = create_tensor(tn(LLM_TENSOR_ATTN_PASS_LORA_UP,   "weight", il), {rank, n_qkv}, 0);
+            layer.attn_pass_gate      = create_tensor(tn(LLM_TENSOR_ATTN_PASS_GATE,      "weight", il), {n_p}, 0);
+        }
+
+        const int64_t n_ff_exp = hparams.n_ff_exp(il);
+
+        layer.ffn_latent_down  = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN,  "weight", il), {n_embd, n_latent}, 0);
+        layer.ffn_latent_up    = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP,    "weight", il), {n_latent, n_embd}, 0);
+        layer.ffn_gate_inp     = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,     "weight", il), {n_latent + n_r, n_expert}, 0);
+        layer.ffn_gate_inp_b   = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,     "bias",   il), {n_expert}, 0);
+        layer.ffn_exp_probs_b  = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B,  "bias",   il), {n_expert}, 0);
+        layer.ffn_exp_embd     = create_tensor(tn(LLM_TENSOR_FFN_EXP_EMBD,     "weight", il), {n_r, n_expert}, 0);
+        layer.ffn_gate_up_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_UP_EXPS, "weight", il), {n_latent, 2*n_ff_exp, n_expert}, 0);
+        layer.ffn_down_exps    = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS,    "weight", il), {n_ff_exp, n_latent, n_expert}, 0);
+        layer.ffn_gate_shexp   = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP,   "weight", il), {n_latent, n_ff_shexp}, 0);
+        layer.ffn_up_shexp     = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,     "weight", il), {n_latent, n_ff_shexp}, 0);
+        layer.ffn_down_shexp   = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP,   "weight", il), {n_ff_shexp, n_latent}, 0);
+    }
+
+    // the virtual layers of later passes share the physical layer's weights
+    for (int il = n_layer_phys; il < n_layer; ++il) {
+        layers[il] = layers[il % n_layer_phys];
+    }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_metis16::build_arch_graph(const llm_graph_params & params) const {
+    return std::make_unique<graph>(*this, params);
+}
+
+// streams are kept stream-innermost: [n_stream, n_embd, n_tokens]
+ggml_tensor * llama_model_metis16::graph::stream_mean(ggml_tensor * x) const {
+    ggml_tensor * mean = ggml_mean(ctx0, ggml_reshape_2d(ctx0, x, x->ne[0], x->ne[1]*x->ne[2]));
+    return ggml_reshape_2d(ctx0, mean, x->ne[1], x->ne[2]);
+}
+
+// returns the normalized sublayer input
+ggml_tensor * llama_model_metis16::graph::hc_read(ggml_tensor * x, ggml_tensor * mix, ggml_tensor * read, ggml_tensor * norm,
+        int p, ggml_tensor ** mixed, int il) const {
+    const int64_t n_s = x->ne[0];
+    ggml_tensor * x2 = ggml_reshape_2d(ctx0, x, n_s, x->ne[1]*x->ne[2]);
+
+    ggml_tensor * m = ggml_view_2d(ctx0, mix, n_s, n_s, mix->nb[1], p*mix->nb[2]);
+    *mixed = ggml_mul_mat(ctx0, m, x2);
+
+    ggml_tensor * r   = ggml_reshape_2d(ctx0, ggml_view_1d(ctx0, read, n_s, p*read->nb[1]), n_s, 1);
+    ggml_tensor * src = ggml_reshape_2d(ctx0, ggml_mul_mat(ctx0, r, x2), x->ne[1], x->ne[2]);
+    return build_norm(src, norm, NULL, LLM_NORM_RMS, il);
+}
+
+ggml_tensor * llama_model_metis16::graph::hc_write(ggml_tensor * mixed, ggml_tensor * write, ggml_tensor * update, int p) const {
+    const int64_t n_s = mixed->ne[0];
+    ggml_tensor * w   = ggml_reshape_2d(ctx0, ggml_view_1d(ctx0, write, n_s, p*write->nb[1]), 1, n_s);
+    ggml_tensor * u   = ggml_reshape_2d(ctx0, update, 1, update->ne[0]*update->ne[1]);
+    ggml_tensor * out = ggml_add(ctx0, mixed, ggml_mul_mat(ctx0, w, u));
+    return ggml_reshape_3d(ctx0, out, n_s, update->ne[0], update->ne[1]);
+}
+
+// blends the retrieved entries into the streams; returns their summary
+ggml_tensor * llama_model_metis16::graph::mem_read(const llama_model_metis16 & model, ggml_tensor ** x) const {
+    const int64_t n_s = (*x)->ne[0];
+    const int64_t n_m = model.n_mem;
+
+    ggml_tensor * x_e = ggml_cont(ctx0, ggml_permute(ctx0, *x, 1, 0, 2, 3));
+
+    ggml_tensor * q = ggml_mul_mat(ctx0, model.depth_mem_query, ggml_reshape_2d(ctx0, x_e, n_embd, n_s*n_tokens));
+    q = ggml_reshape_3d(ctx0, ggml_add(ctx0, q, model.depth_mem_query_b), n_m, n_s, n_tokens);
+
+    ggml_tensor * kq = ggml_mul_mat(ctx0, bank_k, q);
+    kq = ggml_soft_max_ext(ctx0, kq, nullptr, 1.0f/sqrtf(float(n_m)), 0.0f);
+
+    ggml_tensor * ret  = ggml_mul_mat(ctx0, bank_vt, kq);
+    ggml_tensor * proj = ggml_mul_mat(ctx0, model.depth_mem_out, ggml_reshape_2d(ctx0, ret, n_m, n_s*n_tokens));
+    proj = ggml_reshape_3d(ctx0, ggml_add(ctx0, proj, model.depth_mem_out_b), n_embd, n_s, n_tokens);
+
+    ggml_tensor * gate = ggml_sum_rows(ctx0, ggml_mul(ctx0, x_e, model.depth_mem_stream_gate));
+    gate = ggml_sigmoid(ctx0, ggml_add(ctx0, gate, ggml_reshape_2d(ctx0, model.depth_mem_stream_gate_b, 1, n_s)));
+
+    ggml_tensor * proj_s = ggml_cont(ctx0, ggml_permute(ctx0, proj, 1, 0, 2, 3));
+    *x = ggml_add(ctx0, *x, ggml_mul(ctx0, proj_s, ggml_reshape_3d(ctx0, gate, n_s, 1, n_tokens)));
+
+    return stream_mean(proj_s);
+}
+
+void llama_model_metis16::graph::mem_write(const llama_model_metis16 & model, ggml_tensor * x, ggml_tensor * summary,
+        ggml_tensor * entropy, ggml_tensor * confidence, ggml_tensor * cont, int entry) {
+    const int64_t n_m = model.n_mem;
+
+    if (cont == nullptr) {
+        cont = ggml_scale(ctx0, confidence, 0.0f);
+    }
+    ggml_tensor * meta = ggml_concat(ctx0, ggml_concat(ctx0, summary, entropy, 0), ggml_concat(ctx0, confidence, cont, 0), 0);
+
+    ggml_tensor * e = ggml_add(ctx0, ggml_mul_mat(ctx0, model.depth_mem_state_write, stream_mean(x)), ggml_mul_mat(ctx0, model.depth_mem_meta_write, meta));
+    e = ggml_add(ctx0, e, ggml_view_1d(ctx0, model.depth_mem_entry_bias, n_m, entry*model.depth_mem_entry_bias->nb[1]));
+
+    ggml_tensor * k = ggml_add(ctx0, ggml_mul_mat(ctx0, model.depth_mem_key, e), model.depth_mem_key_b);
+    ggml_tensor * v = ggml_add(ctx0, ggml_mul_mat(ctx0, model.depth_mem_value, e), model.depth_mem_value_b);
+    k = ggml_reshape_3d(ctx0, k, n_m, 1, n_tokens);
+    v = ggml_reshape_3d(ctx0, v, 1, n_m, n_tokens);
+    bank_k  = bank_k  ? ggml_concat(ctx0, bank_k,  k, 1) : k;
+    bank_vt = bank_vt ? ggml_concat(ctx0, bank_vt, v, 0) : v;
+}
+
+ggml_tensor * llama_model_metis16::graph::route_features(const llama_model_metis16 & model, ggml_tensor * state,
+        ggml_tensor * summary, ggml_tensor * diff, ggml_tensor * history) const {
+    ggml_tensor * f = ggml_mul_mat(ctx0, model.depth_mem_route_state, state);
+    if (summary) {
+        f = ggml_add(ctx0, f, ggml_mul_mat(ctx0, model.depth_mem_route_retr, summary));
+    }
+    if (diff) {
+        f = ggml_add(ctx0, f, ggml_mul_mat(ctx0, model.depth_mem_route_diff, diff));
+    }
+    if (history) {
+        f = ggml_add(ctx0, f, ggml_mul_mat(ctx0, model.depth_mem_route_hist, history));
+    }
+    return ggml_tanh(ctx0, ggml_add(ctx0, f, model.depth_mem_route_b));
+}
+
+llama_model_metis16::graph::graph(const llama_model_metis16 & model, const llm_graph_params & params) :
+    llm_build_mamba_base(params) {
+    const int64_t n_s      = model.n_stream;
+    const int     n_p      = model.n_pass;
+    const int     n_phys   = model.n_layer_phys;
+    const int64_t n_r      = model.n_route;
+    const int64_t n_hd     = hparams.n_embd_head_k();
+    const int64_t n_exp    = hparams.n_expert;
+    const int64_t n_used   = hparams.n_expert_used();
+
+    ggml_tensor * inpL = build_inp_embd(model.tok_embd);
+    auto * inp = build_inp_mem_hybrid();
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    ggml_tensor * x = ggml_repeat_4d(ctx0, ggml_reshape_3d(ctx0, inpL, 1, n_embd, n_tokens), n_s, n_embd, n_tokens, 1);
+    x = ggml_add(ctx0, x, model.stream_embd);
+    cb(x, "streams_in", -1);
+
+    ggml_tensor * prev_state = stream_mean(x);
+    ggml_tensor * history    = nullptr; // routing history, [n_route, n_tokens]
+    ggml_tensor * cont       = nullptr; // depth head reading of the last pass, [1, n_tokens]
+    int n_entry = 0;
+
+    for (int p = 0; p < n_p; ++p) {
+        ggml_tensor * summ = nullptr;
+        ggml_tensor * ent  = nullptr;
+        ggml_tensor * conf = nullptr;
+
+        for (int il_p = 0; il_p < n_phys; ++il_p) {
+            const int il = p*n_phys + il_p;
+            const auto & layer = model.layers[il];
+            const bool is_attn  = !hparams.is_recr(il);
+            const bool writes   = is_attn || il_p == n_phys - 1;
+
+            ggml_tensor * retrieved = bank_k ? mem_read(model, &x) : nullptr;
+            ggml_tensor * feats = route_features(model, stream_mean(x), retrieved, nullptr, history);
+
+            ggml_tensor * mixed = nullptr;
+            ggml_tensor * cur = hc_read(x, layer.hc_attn_mix, layer.hc_attn_read, layer.attn_norm, p, &mixed, il);
+            if (is_attn) {
+                ggml_tensor * qkv  = build_lora_mm(layer.wqkv, cur);
+                ggml_tensor * lora = ggml_mul_mat(ctx0, layer.attn_pass_lora_up, ggml_mul_mat(ctx0, layer.attn_pass_lora_down, cur));
+                qkv = ggml_add(ctx0, qkv, ggml_mul(ctx0, lora, ggml_view_1d(ctx0, layer.attn_pass_gate, 1, p*layer.attn_pass_gate->nb[0])));
+
+                const int64_t n_head    = hparams.n_head(il);
+                const int64_t n_head_kv = hparams.n_head_kv(il);
+                const size_t  es        = ggml_element_size(qkv);
+                ggml_tensor * Q = ggml_cont(ctx0, ggml_view_3d(ctx0, qkv, n_hd, n_head,    n_tokens, n_hd*es, qkv->nb[1], 0));
+                ggml_tensor * K = ggml_cont(ctx0, ggml_view_3d(ctx0, qkv, n_hd, n_head_kv, n_tokens, n_hd*es, qkv->nb[1], n_hd*n_head*es));
+                ggml_tensor * V = ggml_cont(ctx0, ggml_view_3d(ctx0, qkv, n_hd, n_head_kv, n_tokens, n_hd*es, qkv->nb[1], n_hd*(n_head + n_head_kv)*es));
+
+                cur = build_attn(inp->get_attn(), layer.wo, nullptr, nullptr, Q, K, V, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_hd)), il);
+            } else {
+                cur = build_mamba2_layer(inp->get_recr(), cur, model, ubatch, il);
+            }
+            cb(cur, "mixer_out", il);
+            x = hc_write(mixed, layer.hc_attn_write, cur, p);
+
+            cur = hc_read(x, layer.hc_ffn_mix, layer.hc_ffn_read, layer.ffn_norm, p, &mixed, il);
+            {
+                ggml_tensor * latent = ggml_mul_mat(ctx0, layer.ffn_latent_down, cur);
+                ggml_tensor * logits = ggml_mul_mat(ctx0, layer.ffn_gate_inp, ggml_concat(ctx0, latent, feats, 0));
+                logits = ggml_add(ctx0, logits, layer.ffn_gate_inp_b);
+
+                // the balance bias picks the experts; the weights come from the unbiased logits
+                ggml_tensor * selected = ggml_argsort_top_k(ctx0, ggml_add(ctx0, logits, layer.ffn_exp_probs_b), n_used);
+
+                ggml_tensor * moe_out = build_moe_ffn(latent, nullptr, nullptr, nullptr, layer.ffn_down_exps, nullptr,
+                        n_exp, n_used, LLM_FFN_SILU, false, 0.0f, LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX_WEIGHT, il_p,
+                        logits, layer.ffn_gate_up_exps, nullptr, nullptr, nullptr, selected);
+                ggml_tensor * shexp = build_ffn(latent,
+                        layer.ffn_up_shexp,   NULL, NULL,
+                        layer.ffn_gate_shexp, NULL, NULL,
+                        layer.ffn_down_shexp, NULL, NULL,
+                        NULL, LLM_FFN_SILU, LLM_FFN_PAR, il_p);
+                cur = ggml_mul_mat(ctx0, layer.ffn_latent_up, ggml_add(ctx0, moe_out, shexp));
+                cb(cur, "ffn_out", il);
+
+                // routing summary: the chosen experts' identity vectors blended by their weights
+                ggml_tensor * logits3 = ggml_reshape_3d(ctx0, logits, 1, n_exp, n_tokens);
+                ggml_tensor * weights = ggml_soft_max(ctx0, ggml_reshape_2d(ctx0, ggml_get_rows(ctx0, logits3, selected), n_used, n_tokens));
+                ggml_tensor * rows    = ggml_get_rows(ctx0, layer.ffn_exp_embd, ggml_reshape_1d(ctx0, ggml_cont(ctx0, selected), n_used*n_tokens));
+                rows = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, rows, n_r, n_used, n_tokens), 1, 0, 2, 3));
+                summ = ggml_reshape_2d(ctx0, ggml_mul_mat(ctx0, rows, ggml_reshape_3d(ctx0, weights, n_used, 1, n_tokens)), n_r, n_tokens);
+
+                if (writes) {
+                    ggml_tensor * probs = ggml_soft_max(ctx0, logits);
+                    ent  = ggml_sum_rows(ctx0, ggml_mul(ctx0, probs, ggml_log(ctx0, ggml_clamp(ctx0, probs, 1.0e-9f, 1.0f))));
+                    ent  = ggml_scale(ctx0, ent, -1.0f);
+                    conf = ggml_get_rows(ctx0, ggml_reshape_3d(ctx0, probs, 1, n_exp, n_tokens),
+                            ggml_reshape_2d(ctx0, ggml_argmax(ctx0, probs), 1, n_tokens));
+                    conf = ggml_reshape_2d(ctx0, conf, 1, n_tokens);
+                }
+            }
+            x = hc_write(mixed, layer.hc_ffn_write, cur, p);
+            cb(x, "streams", il);
+
+            history = history ? ggml_add(ctx0, ggml_scale(ctx0, history, 0.5f), ggml_scale(ctx0, summ, 0.5f)) : ggml_scale(ctx0, summ, 0.5f);
+
+            if (is_attn) {
+                mem_write(model, x, summ, ent, conf, cont, n_entry++);
+            }
+        }
+
+        // end of the pass: one more memory read, the depth head, and an entry recording where the pass left the token
+        ggml_tensor * retrieved = mem_read(model, &x);
+        ggml_tensor * state = stream_mean(x);
+        ggml_tensor * diff  = ggml_sub(ctx0, state, prev_state);
+        ggml_tensor * feats = route_features(model, state, retrieved, diff, history);
+
+        ggml_tensor * h = ggml_mul_mat(ctx0, model.cont_hidden_state, state);
+        h = ggml_add(ctx0, h, ggml_mul_mat(ctx0, model.cont_hidden_retr,  retrieved));
+        h = ggml_add(ctx0, h, ggml_mul_mat(ctx0, model.cont_hidden_diff,  diff));
+        h = ggml_add(ctx0, h, ggml_mul_mat(ctx0, model.cont_hidden_route, feats));
+        h = ggml_silu(ctx0, ggml_add(ctx0, h, model.cont_hidden_b));
+
+        ggml_tensor * logit = ggml_add(ctx0, ggml_mul_mat(ctx0, model.cont_output, h), model.cont_output_b);
+        ggml_tensor * scale = ggml_clamp(ctx0, ggml_sqrt(ctx0, ggml_mean(ctx0, ggml_sqr(ctx0, state))), 1.0e-6f, INFINITY);
+        ggml_tensor * innov = ggml_sqrt(ctx0, ggml_mean(ctx0, ggml_sqr(ctx0, diff)));
+        logit = ggml_add(ctx0, logit, ggml_scale_bias(ctx0, ggml_tanh(ctx0, ggml_div(ctx0, innov, scale)), 1.0f, -0.5f));
+        cont  = ggml_sigmoid(ctx0, ggml_clamp(ctx0, logit, -model.cont_limit, model.cont_limit));
+        cb(cont, "continuation", p);
+
+        mem_write(model, x, summ, ent, conf, cont, n_entry++);
+        prev_state = state;
+    }
+
+    ggml_tensor * cur = stream_mean(x);
+    if (inp_out_ids) {
+        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+    }
+
+    cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+    cb(cur, "result_norm", -1);
+    res->t_embd = cur;
+
+    cur = build_lora_mm(model.output, cur);
+    cb(cur, "result_output", -1);
+    res->t_logits = cur;
+
+}

src/models/models.h

@@ -1799,6 +1799,39 @@ struct llama_model_nemotron_h_moe : public llama_model_nemotron_h {
 };
 
 
+struct llama_model_metis16 : public llama_model_base {
+    llama_model_metis16(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;
+
+    uint32_t n_pass       = 1;
+    uint32_t n_stream     = 1;
+    uint32_t n_mem        = 0;
+    uint32_t n_route      = 0;
+    uint32_t n_layer_phys = 0;
+    float    cont_limit   = 0.0f;
+
+    struct graph : public llm_build_mamba_base {
+        graph(const llama_model_metis16 & model, const llm_graph_params & params);
+
+        // memory bank of the current forward: keys [n_mem, n_entry, n_tokens], values [n_entry, n_mem, n_tokens]
+        ggml_tensor * bank_k  = nullptr;
+        ggml_tensor * bank_vt = nullptr;
+
+        ggml_tensor * stream_mean(ggml_tensor * x) const;
+        ggml_tensor * hc_read(ggml_tensor * x, ggml_tensor * mix, ggml_tensor * read, ggml_tensor * norm, int p, ggml_tensor ** mixed, int il) const;
+        ggml_tensor * hc_write(ggml_tensor * mixed, ggml_tensor * write, ggml_tensor * update, int p) const;
+        ggml_tensor * mem_read(const llama_model_metis16 & model, ggml_tensor ** x) const;
+        void          mem_write(const llama_model_metis16 & model, ggml_tensor * x, ggml_tensor * summary,
+                                ggml_tensor * entropy, ggml_tensor * confidence, ggml_tensor * cont, int entry);
+        ggml_tensor * route_features(const llama_model_metis16 & model, ggml_tensor * state, ggml_tensor * summary,
+                                     ggml_tensor * diff, ggml_tensor * history) const;
+    };
+
+    std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
+
 struct llama_model_exaone : public llama_model_base {
     llama_model_exaone(const struct llama_model_params & params) : llama_model_base(params) {}
     void load_arch_hparams(llama_model_loader & ml) override;

src/models/qwen4exp.cpp

@@ -1443,7 +1443,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(
                 conv_states_all->nb[1],
                 (slot * mem_size + kv_head) * row_size);
 
-        ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst));
+        ggml_build_forward_expand(gf, ggml_cpy(ctx0, tail, dst));
     }
 
     return conv_input;

tests/fusion/MTL.csv

@@ -188,6 +188,8 @@ maple               ,0   ,any     ,SOFT_MAX+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV,
 mellum              ,0   ,any     ,MUL+ADD                     ,      4
 mellum              ,0   ,any     ,RMS_NORM+MUL                ,      9
 mellum              ,0   ,any     ,SOFT_MAX+ARGSORT+GET_ROWS+SUM_ROWS+CLAMP+DIV,      2
+metis16             ,0   ,any     ,MUL+ADD                     ,      8
+metis16             ,0   ,any     ,RMS_NORM+MUL                ,     11
 mimo2               ,0   ,any     ,MUL+ADD                     ,      4
 mimo2               ,0   ,any     ,RMS_NORM+MUL                ,      5
 minicpm             ,0   ,any     ,RMS_NORM+MUL                ,      5

tests/test-backend-ops.cpp

@@ -10234,6 +10234,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_dsv4_hc_post(31, 17, true, true));
     test_cases.emplace_back(new test_dsv4_hc_post(2560, 21, true, true));
     test_cases.emplace_back(new test_dsv4_hc_post(31, 17, false, true));
+    test_cases.emplace_back(new test_dsv4_hc_post(4096, 1));             // DeepSeek-V4-Flash decode
+    test_cases.emplace_back(new test_dsv4_hc_post(2560, 1, true, true)); // Qwen4-Exp decode
 
     for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
         // blocks with extra cells and scores (k-pool)
@@ -10247,6 +10249,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         test_cases.emplace_back(new test_top_k_mask(type, 256, 0, 0, 16, 0, false, 32));
         test_cases.emplace_back(new test_top_k_mask(type, 1001, 0, 0, 64, 0, false, 7, 1, 3));
         test_cases.emplace_back(new test_top_k_mask(type, 4096, 0, 0, 2048, 0, false, 5));
+        // one row, no blocks, fewer than 128 picks: no read past the pick list
+        test_cases.emplace_back(new test_top_k_mask(type, 1001, 0, 0, 16, 3, false, 1));
+        test_cases.emplace_back(new test_top_k_mask(type, 256, 0, 0, 40, 0, true, 1));
+        // one row (single-row kernel) with each optional input combination: no score and no extra, blocks only, blocks and extra without score
+        test_cases.emplace_back(new test_top_k_mask(type, 512, 0, 0, 16, 0, false, 1));
+        test_cases.emplace_back(new test_top_k_mask(type, 1024, 4, 200, 32, 0, false, 1));
+        test_cases.emplace_back(new test_top_k_mask(type, 1024, 4, 200, 32, 5, false, 1));
         // cells with scores and extra cells
         test_cases.emplace_back(new test_top_k_mask(type, 512, 0, 0, 32, 5, true, 9));
         // blocks picked per head, the mask broadcast over the heads
@@ -12771,6 +12780,24 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
     // delta-net conv input: conv state followed by the transposed qkv of the ubatch
     test_cases.emplace_back(new test_concat(GGML_TYPE_F32, {3, 8192, 1, 1}, 512, 0, 16));
 
+    // DSV4_HC_PRE: n_tokens 1 is decode, 128 and 512 are prompt batches
+    for (int64_t n_tokens : {1, 128, 512}) {
+        test_cases.emplace_back(new test_dsv4_hc_pre(4096, 4, n_tokens));       // DeepSeek-V4-Flash
+        test_cases.emplace_back(new test_dsv4_hc_pre(4096, 4, n_tokens, true)); // synthetic, no gated model uses this width
+        test_cases.emplace_back(new test_dsv4_hc_pre(2560, 4, n_tokens, true)); // Qwen4-Exp (Qwen3.8-Flash-Next)
+    }
+
+    // DSV4_HC_POST: n_tokens 1 is decode, 128 and 512 are prompt batches
+    for (int64_t n_tokens : {1, 128, 512}) {
+        test_cases.emplace_back(new test_dsv4_hc_post(4096, n_tokens));             // DeepSeek-V4-Flash
+        test_cases.emplace_back(new test_dsv4_hc_post(7168, n_tokens));             // DeepSeek-V4-Pro
+        test_cases.emplace_back(new test_dsv4_hc_post(2560, n_tokens, true, true)); // Qwen4-Exp, with the post gate chain
+        test_cases.emplace_back(new test_dsv4_hc_post(2560, n_tokens, true));       // synthetic, Qwen4-Exp kernel without the gate chain
+    }
+
+    // Qwen3.8-Flash-Next pooled indexer prefill.
+    test_cases.emplace_back(new test_lightning_indexer(128, 4, 1088, 1024, 1, 1, GGML_TYPE_F32));
+    test_cases.emplace_back(new test_lightning_indexer(128, 4, 1088, 1024, 1, 1, GGML_TYPE_Q4_0));
     // SWIGLU at a 27B-class FFN width, fused [gate|up] vs split operands
     // note: same bytes either way, so a backend that indexes them differently shows it here
     for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_BF16}) {

tests/test-backend-tensor.cpp

@@ -148,6 +148,28 @@ static bool test_tensor_n(ggml_backend_buffer_type_t buft) {
         }
     }
 
+    {
+        uint8_t input[24], expected[128], out[128] = {};
+        for (size_t i = 0; i < sizeof(input); i++) {
+            input[i] = (uint8_t) (i + 1);
+        }
+        ggml_backend_tensor_set(t_big, input, 0, sizeof(input));
+
+        // Verify scattered byte reads, including repeated source offsets.
+        std::vector<ggml_tensor *> tensors(128, t_big);
+        std::vector<size_t> offsets(128), sizes(128, 1), data_offsets(128);
+        for (size_t i = 0; i < tensors.size(); i++) {
+            offsets[i] = (i * 5) % sizeof(input);
+            data_offsets[i] = i;
+            expected[i] = input[offsets[i]];
+        }
+
+        ggml_backend_tensor_get_n(tensors.data(), out, offsets.data(), sizes.data(), data_offsets.data(), tensors.size());
+        if (memcmp(out, expected, sizeof(out)) != 0) {
+            return false;
+        }
+    }
+
     return true;
 }
 

tests/test-llama-archs.cpp

@@ -279,7 +279,6 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
             || arch == LLM_ARCH_KIMI_LINEAR
             || arch == LLM_ARCH_BAILINGMOE3
             || arch == LLM_ARCH_KIMI_K3
-            || arch == LLM_ARCH_GLM5_NEXT
             || arch == LLM_ARCH_MISTRAL4
             || arch == LLM_ARCH_MISTRAL_LARGE4
             || arch == LLM_ARCH_HY_V4) {
@@ -289,8 +288,15 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         if (arch == LLM_ARCH_MISTRAL_LARGE4) {
             n_layer = 4; // cover NoPE, shared indexer and KV sharing layers
         }
+    } else if (arch == LLM_ARCH_GLM5_NEXT) {
+        // several heads so that -sm tensor can split them and the MLA flash attention has GQA > 1
+        n_embd = 128;
+        n_head = 16;
+        n_ff   = 192;
     } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) {
         n_layer = 3;
+    } else if (arch == LLM_ARCH_METIS16) {
+        n_layer = 4; // 2 passes over 2 physical layers, one attention and one Mamba-2
     } else if (arch == LLM_ARCH_CHAMELEON) {
         n_vocab = 10240;
     } else if (arch == LLM_ARCH_QWEN3TTS) {
@@ -357,10 +363,23 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         // GLM5 next KDA heads come from the uniform head count, only head_count_kv is per layer.
         if (arch == LLM_ARCH_GLM5_NEXT) {
             ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head);
+            // the MLA cache holds one latent head
+            for (uint32_t & n : n_head_per_layer) {
+                n = std::min(n, uint32_t(1));
+            }
         } else {
             ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head_per_layer);
         }
         ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer);
+    } else if (arch == LLM_ARCH_METIS16) {
+        // the layer types repeat in every pass
+        std::vector<uint32_t> n_head_kv_per_layer;
+        n_head_kv_per_layer.reserve(n_layer);
+        for (uint32_t il = 0; il < n_layer; il++) {
+            n_head_kv_per_layer.push_back(il % 2 == 0 ? n_head_kv : 0);
+        }
+        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head);
+        ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_kv_per_layer);
     } else {
         ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head);
         ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_DEEPSEEK41) ? uint32_t(1) : n_head_kv);
@@ -608,6 +627,17 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
         ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 7.0f);
     }
 
+    if (arch == LLM_ARCH_METIS16) {
+        ms.add_kv(LLM_KV_METIS_PASS_COUNT,               uint32_t(2));
+        ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT,         uint32_t(4));
+        ms.add_kv(LLM_KV_METIS_MEMORY_LENGTH,            uint32_t(64));
+        ms.add_kv(LLM_KV_METIS_ROUTE_LENGTH,             uint32_t(32));
+        ms.add_kv(LLM_KV_ATTENTION_PASS_LORA_RANK,       uint32_t(8));
+        ms.add_kv(LLM_KV_METIS_CONTINUATION_LOGIT_LIMIT, 4.0f);
+        ms.add_kv(LLM_KV_SSM_NORM_EPS,                   1e-5f);
+        ms.add_kv(LLM_KV_MOE_LATENT_SIZE,                n_embd/2);
+    }
+
     if (arch == LLM_ARCH_LIMITE) {
         ms.add_kv(LLM_KV_ATTENTION_SCALE,             0.1f);
         ms.add_kv(LLM_KV_ATTENTION_XSA_EPS,           1e-4f);
@@ -964,6 +994,7 @@ static bool moe_mandatory(const llm_arch arch) {
         case LLM_ARCH_MAPLE:
         case LLM_ARCH_MUSE_SPARK:
         case LLM_ARCH_KOLIBRI1:
+        case LLM_ARCH_METIS16:
             return true;
         default:
             return false;

tests/test-model-resolution.cpp

@@ -153,6 +153,20 @@ static const std::vector<std::string> vendors = {
     "BlackSheep-24B.Q8_0.gguf",
 };
 
+// a tag must match the full token: F16 ends inside BF16, Q4_XS inside IQ4_XS
+static const std::vector<std::string> twins = {
+    "model-BF16.gguf",
+    "model-F16.gguf",
+    "model-IQ4_XS.gguf",
+    "model-Q4_XS.gguf",
+};
+
+// an underscore separates the quant like a dash, as in step3p5_flash_Q4_K_S.gguf
+static const std::vector<std::string> underscore = {
+    "model_Q4_K_S.gguf",
+    "model_Q8_0.gguf",
+};
+
 // every speculative sidecar type at the same quant
 static const std::vector<std::string> quad = {
     "model-Q8_0.gguf",
@@ -270,6 +284,32 @@ static const plan_case plan_cases[] = {
      "TheDrummer_Model-24B-v4.1-Q8_0.gguf", {"TheDrummer_Model-24B-v4.1-Q8_0.gguf"},
      "", "", "", "", ""},
 
+    // the tag picks the file carrying it as a full token, never a longer one
+    {"twins f16 tag", twins, "test/repo:F16", "", false, false,
+     "model-F16.gguf", {"model-F16.gguf"},
+     "", "", "", "", ""},
+
+    {"twins bf16 tag", twins, "test/repo:BF16", "", false, false,
+     "model-BF16.gguf", {"model-BF16.gguf"},
+     "", "", "", "", ""},
+
+    {"twins q4_xs tag", twins, "test/repo:Q4_XS", "", false, false,
+     "model-Q4_XS.gguf", {"model-Q4_XS.gguf"},
+     "", "", "", "", ""},
+
+    {"twins iq4_xs tag", twins, "test/repo:IQ4_XS", "", false, false,
+     "model-IQ4_XS.gguf", {"model-IQ4_XS.gguf"},
+     "", "", "", "", ""},
+
+    // the token boundary accepts an underscore, for the tag and the default lookup
+    {"underscore tag", underscore, "test/repo:Q4_K_S", "", false, false,
+     "model_Q4_K_S.gguf", {"model_Q4_K_S.gguf"},
+     "", "", "", "", ""},
+
+    {"underscore default", underscore, "test/repo", "", false, false,
+     "model_Q8_0.gguf", {"model_Q8_0.gguf"},
+     "", "", "", "", ""},
+
     // every sidecar type resolves at the tag
     {"quad exact tag", quad, "test/repo:Q8_0", "", true, false,
      "model-Q8_0.gguf", {"model-Q8_0.gguf"},
@@ -868,7 +908,7 @@ static const real_repo real_repos[] = {
       {"Q8_0", "MTP/Qwen3.8-27B-Uncensored-OrcaRouter-MTP-Q8_0.gguf", 1, "", "", "", "", ""},
       {"BF16", "BF16/Qwen3.8-27B-Uncensored-OrcaRouter-FP8-derived-BF16-00001-of-00002.gguf", 2, "", "", "", "", ""},
       {"Q5_K_M", "Qwen3.8-27B-Uncensored-OrcaRouter-Q5_K_M.gguf", 1, "", "", "", "", ""},
-      {"F16", "BF16/Qwen3.8-27B-Uncensored-OrcaRouter-FP8-derived-BF16-00001-of-00002.gguf", 2, "", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_0", "Q5_K_M", "Q4_K_M", "BF16"}},
 
     {"ChrisColeTech/qwen-image-edit-uncensored-GGUF",
@@ -980,7 +1020,7 @@ static const real_repo real_repos[] = {
       {"Q8_0", "Qwen3.5-9B-The-Defiant-Fable-Uncnr-Heretic-NEO-MAX-MTP-Q8_0.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
       {"IQ2_M", "Qwen3.5-9B-The-Defiant-Fable-Uncnr-Heretic-NEO-MAX-IQ2_M.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
       {"BF16", "Qwen3.5-9B-The-Defiant-Fable-Uncnr-Heretic-plusIQ-NEO-MAX-MTP-bf16.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
-      {"F16", "Qwen3.5-9B-The-Defiant-Fable-Uncnr-Heretic-plusIQ-NEO-MAX-MTP-bf16.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_0", "BF16", "Q4_K_S", "Q4_K_M", "IQ2_M"}},
 
     {"DavidAU/Qwen3.6-27B-Heretic-Uncensored-FINETUNE-NEO-CODE-Di-IMatrix-MAX-GGUF",
@@ -1044,7 +1084,7 @@ static const real_repo real_repos[] = {
       {"Q4_K_M", "Qwythos-27B-MTP-Q4_K_M.gguf", 1, "mmproj-Qwythos-27B-F16.gguf", "", "", "", ""},
       {"Q8_0", "Qwythos-27B-MTP-Q8_0.gguf", 1, "mmproj-Qwythos-27B-F16.gguf", "", "", "", ""},
       {"BF16", "Qwythos-27B-BF16.gguf", 1, "mmproj-Qwythos-27B-F16.gguf", "", "", "", ""},
-      {"F16", "Qwythos-27B-BF16.gguf", 1, "mmproj-Qwythos-27B-F16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_0", "Q5_K_M", "Q4_K_M", "BF16"}},
 
     {"empero-ai/Qwythos-9B-v2-GGUF",
@@ -1199,6 +1239,20 @@ static const real_repo real_repos[] = {
       {"BF16", "", 0, "", "", "", "eagle3-gpt-oss-120b-BF16.gguf", ""}},
      {"MXFP4"}},
 
+    // F16 and BF16 in one repo, the F16 tag must not pick the BF16 file
+    {"ggml-org/Llama-2-7B-GGUF",
+     {{"Llama-2-7B-BF16.gguf"},
+      {"Llama-2-7B-F16.gguf"},
+      {"Llama-2-7B-Q4_0.gguf"},
+      {"Llama-2-7B-Q8_0.gguf"}},
+     {
+      {"", "Llama-2-7B-Q8_0.gguf", 1, "", "", "", "", ""},
+      {"Q8_0", "Llama-2-7B-Q8_0.gguf", 1, "", "", "", "", ""},
+      {"Q4_0", "Llama-2-7B-Q4_0.gguf", 1, "", "", "", "", ""},
+      {"F16", "Llama-2-7B-F16.gguf", 1, "", "", "", "", ""},
+      {"BF16", "Llama-2-7B-BF16.gguf", 1, "", "", "", "", ""}},
+     {"F16", "Q4_0", "Q8_0", "BF16"}},
+
     {"ggml-org/MiMo-V2.6-Distill-Qwen-9B-GGUF",
      {{"MiMo-V2.6-Distill-Qwen-9B-Q8_0.gguf"},
       {"mmproj-MiMo-V2.6-Distill-Qwen-9B-Q8_0.gguf"}},
@@ -1412,7 +1466,7 @@ static const real_repo real_repos[] = {
       {"Q4_K_M", "Qwen3.5-27B-Uncensored-HauhauCS-Aggressive-Q4_K_M.gguf", 1, "mmproj-Qwen3.5-27B-Uncensored-HauhauCS-Aggressive-f16.gguf", "", "", "", ""},
       {"Q8_0", "Qwen3.5-27B-Uncensored-HauhauCS-Aggressive-Q8_0.gguf", 1, "mmproj-Qwen3.5-27B-Uncensored-HauhauCS-Aggressive-f16.gguf", "", "", "", ""},
       {"BF16", "Qwen3.5-27B-Uncensored-HauhauCS-Aggressive-BF16.gguf", 1, "mmproj-Qwen3.5-27B-Uncensored-HauhauCS-Aggressive-f16.gguf", "", "", "", ""},
-      {"F16", "Qwen3.5-27B-Uncensored-HauhauCS-Aggressive-BF16.gguf", 1, "mmproj-Qwen3.5-27B-Uncensored-HauhauCS-Aggressive-f16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_0", "Q4_K_M", "Q3_K_M", "BF16"}},
 
     {"HauhauCS/Qwen3.5-2B-Uncensored-HauhauCS-Aggressive",
@@ -1426,7 +1480,7 @@ static const real_repo real_repos[] = {
       {"Q4_K_M", "Qwen3.5-2B-Uncensored-HauhauCS-Aggressive-Q4_K_M.gguf", 1, "mmproj-Qwen3.5-2B-Uncensored-HauhauCS-Aggressive-f16.gguf", "", "", "", ""},
       {"Q8_0", "Qwen3.5-2B-Uncensored-HauhauCS-Aggressive-Q8_0.gguf", 1, "mmproj-Qwen3.5-2B-Uncensored-HauhauCS-Aggressive-f16.gguf", "", "", "", ""},
       {"BF16", "Qwen3.5-2B-Uncensored-HauhauCS-Aggressive-BF16.gguf", 1, "mmproj-Qwen3.5-2B-Uncensored-HauhauCS-Aggressive-f16.gguf", "", "", "", ""},
-      {"F16", "Qwen3.5-2B-Uncensored-HauhauCS-Aggressive-BF16.gguf", 1, "mmproj-Qwen3.5-2B-Uncensored-HauhauCS-Aggressive-f16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_0", "Q6_K", "Q4_K_M", "BF16"}},
 
     {"HauhauCS/Qwen3.5-4B-Uncensored-HauhauCS-Aggressive",
@@ -1622,7 +1676,7 @@ static const real_repo real_repos[] = {
       {"IQ2_M", "Qwen3.8-27B-Uncensored-IQ2_M.gguf", 1, "mmproj-Qwen3.8-27B-Uncensored-F16.gguf", "", "", "", ""},
       {"Q4_0", "Qwen3.8-27B-Uncensored-draft-Q4_0.gguf", 1, "mmproj-Qwen3.8-27B-Uncensored-F16.gguf", "", "", "", ""},
       {"F16", "", 0, "", "", "", "", ""},
-      {"mtp", "Qwen3.8-27B-Uncensored-noMTP-IQ2_M.gguf", 1, "mmproj-Qwen3.8-27B-Uncensored-F16.gguf", "", "", "", ""}},
+      {"mtp", "", 0, "", "", "", "", ""}},
      {"Q8_0", "Q5_K_M", "Q4_K_M", "IQ2_M", "Q4_0"}},
 
     {"kai-os/Carnice-V2-27b-GGUF",
@@ -1903,7 +1957,7 @@ static const real_repo real_repos[] = {
       {"Q4_K_M", "Huihui-Qwen3-VL-4B-Instruct-abliterated-Q4_K_M.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
       {"Q8_0", "Huihui-Qwen3-VL-4B-Instruct-abliterated-Q8_0.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
       {"BF16", "Huihui-Qwen3-VL-4B-Instruct-abliterated-BF16.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
-      {"F16", "Huihui-Qwen3-VL-4B-Instruct-abliterated-BF16.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_0", "Q4_K_M", "IQ4_XS", "BF16"}},
 
     {"OBLITERATUS/Gemma-4-12B-OBLITERATED",
@@ -2291,7 +2345,7 @@ static const real_repo real_repos[] = {
       {"Q8_0", "embeddinggemma-2-Q8_0.gguf", 1, "mmproj-Q8_0.gguf", "", "", "", ""},
       {"BF16", "embeddinggemma-2-BF16.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
       {"UD-Q6_K_XL", "embeddinggemma-2-UD-Q6_K_XL.gguf", 1, "mmproj-Q8_0.gguf", "", "", "", ""},
-      {"F16", "embeddinggemma-2-BF16.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q6_K_XL", "Q4_K_XL", "Q8_0", "BF16"}},
 
     {"unsloth/GLM-4.7-Flash-GGUF",
@@ -2320,7 +2374,7 @@ static const real_repo real_repos[] = {
       {"Q8_0", "Q8_0/GLM-5.3-Flash-Q8_0-00001-of-00008.gguf", 8, "mmproj-BF16.gguf", "", "", "", ""},
       {"BF16", "BF16/GLM-5.3-Flash-BF16-00001-of-00014.gguf", 14, "mmproj-BF16.gguf", "", "", "", ""},
       {"UD-Q6_K_XL", "UD-Q6_K_XL/GLM-5.3-Flash-UD-Q6_K_XL-00001-of-00007.gguf", 7, "mmproj-BF16.gguf", "", "", "", ""},
-      {"F16", "BF16/GLM-5.3-Flash-BF16-00001-of-00014.gguf", 14, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q6_K_XL", "IQ4_XS", "Q8_0", "BF16"}},
 
     {"unsloth/GLM-5.3-GGUF",
@@ -2359,7 +2413,7 @@ static const real_repo real_repos[] = {
       {"", "BF16/Kimi-K2.6-BF16-00001-of-00046.gguf", 46, "mmproj-BF16.gguf", "", "", "", ""},
       {"BF16", "BF16/Kimi-K2.6-BF16-00001-of-00046.gguf", 46, "mmproj-BF16.gguf", "", "", "", ""},
       {"UD-Q8_K_XL", "UD-Q8_K_XL/Kimi-K2.6-UD-Q8_K_XL-00001-of-00014.gguf", 14, "mmproj-BF16.gguf", "", "", "", ""},
-      {"F16", "BF16/Kimi-K2.6-BF16-00001-of-00046.gguf", 46, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_K_XL", "Q4_K_XL", "BF16"}},
 
     {"unsloth/Kimi-K2.7-Code-GGUF",
@@ -2405,7 +2459,7 @@ static const real_repo real_repos[] = {
       {"Q8_0", "LFM2.5-VL-3B-Q8_0.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
       {"BF16", "LFM2.5-VL-3B-BF16.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
       {"UD-Q8_K_XL", "LFM2.5-VL-3B-UD-Q8_K_XL.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
-      {"F16", "LFM2.5-VL-3B-BF16.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_K_XL", "Q8_0", "Q5_K_S", "Q4_K_M", "BF16"}},
 
     {"unsloth/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-GGUF",
@@ -2449,7 +2503,7 @@ static const real_repo real_repos[] = {
       {"Q8_0", "Qwen3-VL-30B-A3B-Instruct-Q8_0.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
       {"BF16", "BF16/Qwen3-VL-30B-A3B-Instruct-BF16-00001-of-00002.gguf", 2, "mmproj-BF16.gguf", "", "", "", ""},
       {"UD-TQ1_0", "Qwen3-VL-30B-A3B-Instruct-UD-TQ1_0.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
-      {"F16", "BF16/Qwen3-VL-30B-A3B-Instruct-BF16-00001-of-00002.gguf", 2, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"TQ1_0", "Q8_0", "Q6_K", "Q4_K_M", "BF16"}},
 
     {"unsloth/Qwen3.5-27B-GGUF",
@@ -2467,7 +2521,7 @@ static const real_repo real_repos[] = {
       {"Q8_0", "Qwen3.5-27B-Q8_0.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
       {"BF16", "BF16/Qwen3.5-27B-BF16-00001-of-00002.gguf", 2, "mmproj-BF16.gguf", "", "", "", ""},
       {"UD-Q8_K_XL", "Qwen3.5-27B-UD-Q8_K_XL.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
-      {"F16", "BF16/Qwen3.5-27B-BF16-00001-of-00002.gguf", 2, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_K_XL", "Q8_0", "Q6_K", "Q4_K_M", "BF16"}},
 
     {"unsloth/Qwen3.5-35B-A3B-GGUF",
@@ -2485,7 +2539,7 @@ static const real_repo real_repos[] = {
       {"Q8_0", "Qwen3.5-35B-A3B-Q8_0.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
       {"BF16", "BF16/Qwen3.5-35B-A3B-BF16-00001-of-00002.gguf", 2, "mmproj-BF16.gguf", "", "", "", ""},
       {"UD-Q8_K_XL", "Qwen3.5-35B-A3B-UD-Q8_K_XL.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
-      {"F16", "BF16/Qwen3.5-35B-A3B-BF16-00001-of-00002.gguf", 2, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_K_XL", "IQ3_S", "Q8_0", "Q4_K_M", "BF16"}},
 
     {"unsloth/Qwen3.5-397B-A17B-GGUF",
@@ -2503,7 +2557,7 @@ static const real_repo real_repos[] = {
       {"Q8_0", "Q8_0/Qwen3.5-397B-A17B-Q8_0-00001-of-00010.gguf", 10, "mmproj-BF16.gguf", "", "", "", ""},
       {"BF16", "BF16/Qwen3.5-397B-A17B-BF16-00001-of-00017.gguf", 17, "mmproj-BF16.gguf", "", "", "", ""},
       {"UD-Q8_K_XL", "UD-Q8_K_XL/Qwen3.5-397B-A17B-UD-Q8_K_XL-00001-of-00010.gguf", 10, "mmproj-BF16.gguf", "", "", "", ""},
-      {"F16", "BF16/Qwen3.5-397B-A17B-BF16-00001-of-00017.gguf", 17, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_K_XL", "IQ3_XXS", "Q8_0", "Q4_K_M", "BF16"}},
 
     {"unsloth/Qwen3.6-35B-A3B-MTP-GGUF",
@@ -2519,7 +2573,7 @@ static const real_repo real_repos[] = {
       {"Q8_0", "Qwen3.6-35B-A3B-Q8_0.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
       {"BF16", "BF16/Qwen3.6-35B-A3B-BF16-00001-of-00002.gguf", 2, "mmproj-BF16.gguf", "", "", "", ""},
       {"UD-Q8_K_XL", "Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf", 1, "mmproj-BF16.gguf", "", "", "", ""},
-      {"F16", "BF16/Qwen3.6-35B-A3B-BF16-00001-of-00002.gguf", 2, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_K_XL", "Q3_K_M", "Q8_0", "BF16"}},
 
     {"unsloth/Qwen3.8-Flash-Next-GGUF",
@@ -2557,7 +2611,7 @@ static const real_repo real_repos[] = {
       {"Q8_0", "Q8_0/Step-3.7-Flash-Q8_0-00001-of-00006.gguf", 6, "mmproj-BF16.gguf", "", "", "", ""},
       {"BF16", "BF16/Step-3.7-Flash-BF16-00001-of-00008.gguf", 8, "mmproj-BF16.gguf", "", "", "", ""},
       {"UD-Q8_K_XL", "UD-Q8_K_XL/Step-3.7-Flash-UD-Q8_K_XL-00001-of-00006.gguf", 6, "mmproj-BF16.gguf", "", "", "", ""},
-      {"F16", "BF16/Step-3.7-Flash-BF16-00001-of-00008.gguf", 8, "mmproj-BF16.gguf", "", "", "", ""}},
+      {"F16", "", 0, "", "", "", "", ""}},
      {"Q8_K_XL", "Q3_K_M", "Q8_0", "BF16"}},
 
     {"vmlinux/Qwen3.8-Flash-Next-Uncensored-Gufo-Q4Mix-GGUF",

tests/test-mtmd-impl.cpp

@@ -1,11 +1,13 @@
 #include "testing.h"
 
 #include "mtmd-audio.h"
+#include "mtmd-helper.h"
 #include "mtmd-image.h"
 #include "mtmd-internal.h"
 
 #include <cmath>
 #include <cstdint>
+#include <cstdlib>
 #include <iostream>
 #include <stdexcept>
 #include <string>
@@ -92,6 +94,33 @@ MAKE_TEST(test_image_preprocessor_dyn_size) {
     }
 }
 
+MAKE_TEST(test_image_preprocessor_qwen3vl_video) {
+    clip_hparams hparams;
+    hparams.patch_size = 16;
+    hparams.n_merge = 2;
+    const mtmd_image_preprocessor_qwen3vl preproc(hparams);
+
+    // { n_frames, frame size, expected size from Qwen3VLVideoProcessor smart_resize }
+    const std::vector<std::tuple<int, clip_image_size, clip_image_size>> cases = {
+        {  20, {  720,  358 }, {  704,  352 } }, // 720 / 32 = 22.5 rounds half to even
+        { 768, { 1920, 1080 }, {  224,  128 } }, // over the longest_edge budget
+        {   4, {   16,   40 }, {   32,   64 } }, // edge below the factor is upscaled first
+        {   2, {   32,   32 }, {   64,   64 } }, // under the shortest_edge budget
+        {   5, { 2760, 1552 }, { 2752, 1536 } }, // 5 / 2 = 2.5 rounds half to even
+        {   4, {  250, 6000 }, {  256, 6016 } },
+    };
+
+    for (const auto & [n_frames, input, expected] : cases) {
+        clip_image_size actual;
+        const std::string name = std::to_string(n_frames) + " frames of " + std::to_string(input.width) + "x" + std::to_string(input.height);
+        if (!t.assert_true(name + " has a frame size", preproc.get_video_frame_size(n_frames, input.width, input.height, actual))) {
+            continue;
+        }
+        t.assert_equal(name + " width", expected.width, actual.width);
+        t.assert_equal(name + " height", expected.height, actual.height);
+    }
+}
+
 static std::string image_size_to_str(const clip_image_size & size) {
     return std::to_string(size.width) + "x" + std::to_string(size.height);
 }
@@ -301,6 +330,130 @@ MAKE_TEST(test_temporal_merge_grouping) {
     }
 }
 
+//
+// mtmd_helper video
+//
+
+MAKE_TEST(test_video_config) {
+    clip_hparams hparams;
+    hparams.patch_size = 16;
+    hparams.n_merge = 2;
+    const mtmd_image_preprocessor_qwen3vl preproc(hparams);
+
+    // { template, n_frames and fps of the source, expected fps, n_frames and frame size }
+    // the 1920x1080 sizes are from Qwen3VLVideoProcessor smart_resize
+    const std::vector<std::tuple<mtmd_video_template, int32_t, float, float, int32_t, clip_image_size>> cases = {
+        { MTMD_VIDEO_TEMPLATE_QWEN3VL,   36000, 30.0f, 2.0f, 768, { 224, 128 } },
+        { MTMD_VIDEO_TEMPLATE_LFM3,      36000, 30.0f, 2.0f, 768, { 224, 128 } },
+        { MTMD_VIDEO_TEMPLATE_LFM3_OMNI, 36000, 30.0f, 2.0f, 384, { 320, 192 } },
+        { MTMD_VIDEO_TEMPLATE_QWEN3VL,      -1, 30.0f, 2.0f,   0, {   0,   0 } }, // frame count unknown
+        { MTMD_VIDEO_TEMPLATE_DEFAULT,   36000, 30.0f, 0.0f,   0, {   0,   0 } },
+    };
+
+    for (const auto & [tmpl, n_frames, fps, expected_fps, expected_n_frames, expected_size] : cases) {
+        const mtmd_video_config cfg = mtmd_video_config_init(tmpl, &preproc, 1920, 1080, n_frames, fps);
+        const std::string name = "template " + std::to_string(tmpl) + " with " + std::to_string(n_frames) + " frames";
+        t.assert_equal(name + " template", (int) tmpl, (int) cfg.tmpl);
+        t.assert_equal(name + " fps", expected_fps, cfg.fps);
+        t.assert_equal(name + " n_frames", expected_n_frames, cfg.n_frames);
+        t.assert_equal(name + " frame size", image_size_to_str(expected_size), image_size_to_str({ cfg.frame_width, cfg.frame_height }));
+    }
+}
+
+// n_frames source frames, each one starts with its index
+struct test_video_source : mtmd_helper_video_seq {
+    int32_t n_frames;
+    int32_t n_read = 0;
+
+    test_video_source(int32_t n_frames) : n_frames(n_frames) {}
+
+    bool read_raw_frame() override {
+        if (n_read >= n_frames) {
+            return false;
+        }
+        frame_buf[0] = static_cast<uint8_t>(n_read & 255);
+        frame_buf[1] = static_cast<uint8_t>(n_read >> 8);
+        n_read++;
+        return true;
+    }
+};
+
+// text as is, a frame as [source index], an audio bitmap as {n_samples}
+static std::string video_sequence_to_str(testing & t, mtmd_helper_video_seq & video) {
+    std::string out;
+    mtmd_bitmap * bitmap;
+    char * text;
+    int32_t ret;
+    while ((ret = video.read_next(&bitmap, &text)) == 0) {
+        if (text) {
+            out += text;
+            free(text);
+        } else if (mtmd_bitmap_is_audio(bitmap)) {
+            out += "{" + std::to_string(mtmd_bitmap_get_n_bytes(bitmap) / sizeof(float)) + "}";
+            mtmd_bitmap_free(bitmap);
+        } else {
+            const unsigned char * data = mtmd_bitmap_get_data(bitmap);
+            out += "[" + std::to_string(data[0] | data[1] << 8) + "]";
+            mtmd_bitmap_free(bitmap);
+        }
+    }
+    t.assert_equal("end of stream", -1, ret);
+    return out;
+}
+
+MAKE_TEST(test_video_sequence_qwen3vl) {
+    clip_hparams hparams;
+    hparams.patch_size = 16;
+    hparams.n_merge = 2;
+    const mtmd_image_preprocessor_qwen3vl preproc(hparams);
+
+    // { source video, expected frame size and sequence from Qwen3VLProcessor (transformers 4.57.1) }
+    const std::vector<std::tuple<mtmd_helper_video_src, clip_image_size, std::string>> cases = {
+        { {  720, 358, 24.0f, 240 }, { 704, 352 },
+            "<0.3 seconds>[0][13]<1.3 seconds>[25][38]<2.4 seconds>[50][63]<3.4 seconds>[75][88]<4.5 seconds>[101][113]"
+            "<5.5 seconds>[126][138]<6.6 seconds>[151][164]<7.6 seconds>[176][189]<8.6 seconds>[201][214]<9.7 seconds>[226][239]" },
+        // odd frame count: the last group repeats its last frame
+        { {  640, 480, 10.0f,  35 }, { 640, 480 },
+            "<0.3 seconds>[0][6]<1.4 seconds>[11][17]<2.5 seconds>[23][28]<3.4 seconds>[34]" },
+        // fewer frames than the minimum of 4
+        { {  320, 240, 30.0f,   3 }, { 320, 256 },
+            "<0.0 seconds>[0][1]<0.1 seconds>[2]" },
+        { { 1280, 720, 30000.0f / 1001.0f, 300 }, { 1280, 704 },
+            "<0.3 seconds>[0][16]<1.3 seconds>[31][47]<2.4 seconds>[63][79]<3.4 seconds>[94][110]<4.5 seconds>[126][142]"
+            "<5.5 seconds>[157][173]<6.6 seconds>[189][205]<7.6 seconds>[220][236]<8.7 seconds>[252][268]<9.7 seconds>[283][299]" },
+    };
+
+    for (const auto & [src, expected_size, expected] : cases) {
+        const std::string name = std::to_string(src.n_frames) + " frames of " + std::to_string(src.width) + "x" + std::to_string(src.height);
+        const mtmd_video_config cfg = mtmd_video_config_init(MTMD_VIDEO_TEMPLATE_QWEN3VL, &preproc, src.width, src.height, src.n_frames, src.fps);
+        test_video_source video(src.n_frames);
+        video.configure(cfg, src, 0.0f);
+        video.src_alive = true;
+        t.assert_equal(name + " frame size", image_size_to_str(expected_size), image_size_to_str({ (int) video.info.width, (int) video.info.height }));
+        t.assert_equal(name + " sequence", expected, video_sequence_to_str(t, video));
+    }
+}
+
+// the soundtrack of each group comes after its frames
+MAKE_TEST(test_video_sequence_lfm3_omni) {
+    clip_hparams hparams;
+    hparams.patch_size = 16;
+    hparams.n_merge = 2;
+    const mtmd_image_preprocessor_lfm3vl preproc(hparams);
+
+    const mtmd_helper_video_src src = { 64, 64, 2.0f, 9 };
+    const mtmd_video_config cfg = mtmd_video_config_init(MTMD_VIDEO_TEMPLATE_LFM3_OMNI, &preproc, src.width, src.height, src.n_frames, src.fps);
+    test_video_source video(src.n_frames);
+    video.configure(cfg, src, 0.0f);
+    video.src_alive  = true;
+    video.audio_rate = 100;
+    video.soundtrack.assign(450, 0.0f);
+
+    t.assert_equal("sequence",
+        std::string("<0.2 seconds>[0][1]{100}<1.2 seconds>[2][3]{100}<2.2 seconds>[4][5]{100}<3.2 seconds>[6][7]{100}<4.0 seconds>[8]{50}"),
+        video_sequence_to_str(t, video));
+}
+
 //
 // mtmd_audio
 //
@@ -357,6 +510,13 @@ int main(int argc, char ** argv) {
     testing t(std::cout);
     t.verbose = true;
 
+    // the video helper logs every frame at debug level
+    mtmd_helper_log_set([](ggml_log_level level, const char * text, void *) {
+        if (level != GGML_LOG_LEVEL_DEBUG) {
+            fputs(text, stderr);
+        }
+    }, nullptr);
+
     // usage: test-mtmd-impl [filter_regex]
     for (int i = 1; i < argc; i++) {
         t.set_filter(argv[i]);

tools/gguf-split/gguf-split.cpp

@@ -454,6 +454,17 @@ static void gguf_merge(const split_params & split_params) {
         ctx_ggufs.push_back(ctx_gguf);
         ctx_metas.push_back(ctx_meta);
 
+        // Verify that this split belongs to the same set as the first one
+        if (i_split > 0) {
+            auto key_split_no    = gguf_find_key(ctx_gguf, LLM_KV_SPLIT_NO);
+            auto key_split_count = gguf_find_key(ctx_gguf, LLM_KV_SPLIT_COUNT);
+            if (key_split_no < 0 || gguf_get_val_u16(ctx_gguf, key_split_no) != i_split ||
+                key_split_count < 0 || gguf_get_val_u16(ctx_gguf, key_split_count) != n_split) {
+                fprintf(stderr, "\n%s: %s is not split %d of %d\n", __func__, split_path, i_split + 1, n_split);
+                exit(EXIT_FAILURE);
+            }
+        }
+
         if (i_split == 0) {
             auto key_n_split = gguf_find_key(ctx_gguf, LLM_KV_SPLIT_COUNT);
             if (key_n_split < 0) {
@@ -508,6 +519,14 @@ static void gguf_merge(const split_params & split_params) {
 
         fprintf(stderr, "\033[3Ddone\n");
     }
+
+    // Verify that no tensor is missing
+    auto key_n_tensors = gguf_find_key(ctx_ggufs[0], LLM_KV_SPLIT_TENSORS_COUNT);
+    if (key_n_tensors < 0 || gguf_get_val_i32(ctx_ggufs[0], key_n_tensors) != total_tensors) {
+        fprintf(stderr, "%s: splits contain %d tensors, which does not match %s\n", __func__, total_tensors, LLM_KV_SPLIT_TENSORS_COUNT);
+        exit(EXIT_FAILURE);
+    }
+
     std::ofstream fout;
     if (!split_params.dry_run) {
         fout.open(split_params.output.c_str(), std::ios::binary);

tools/mtmd/CMakeLists.txt

@@ -28,6 +28,8 @@ add_library(mtmd
             clip-model.h
             clip-graph.h
             models/models.h
+            models/chatterbox-gen.cpp
+            models/chatterbox-spkenc.cpp
             models/cogvlm.cpp
             models/conformer.cpp
             models/deepseek4v.cpp
@@ -47,6 +49,7 @@ add_library(mtmd
             models/kimik25.cpp
             models/nemotron-v2-vl.cpp
             models/nemotron-omni.cpp
+            models/nemotron-transcribe.cpp
             models/muse-glimmer.cpp
             models/muse-spark-audio.cpp
             models/llama4.cpp

tools/mtmd/clip-impl.h

@@ -113,6 +113,11 @@
 #define KEY_GEN_AUDIO_PENALTY_LAST_N "clip.gen.audio.sampling.penalty_last_n"
 #define KEY_AUDIO_SUBSMPL_FACTOR   "clip.audio.subsampling_factor"
 
+// nemotron transcribe
+#define KEY_A_ROPE_DIMS            "clip.audio.rope_dims"
+#define KEY_A_SPK_N_BLOCK          "clip.audio.spk.block_count"
+#define KEY_A_SPK_N_HEAD           "clip.audio.spk.head_count"
+
 //
 // tensor name constants
 //
@@ -469,6 +474,13 @@
 #define TN_CONV_NORM_MEAN        "%s.blk.%d.conv_norm_mean"
 #define TN_CONV_NORM_VAR         "%s.blk.%d.conv_norm_var"
 
+// nemotron transcribe
+#define TN_A_SPK_INP_PROJ        "a.spk.input_projection.%s"
+#define TN_A_SPK_HEAD            "a.spk.head.%s.%s" // proj, upsample, fc1, fc2
+#define TN_A_SPK_SIL_EMB         "a.spk.head.sil_emb"
+#define TN_A_FUSION_NORM         "a.fusion.%s.%s"   // asr_norm, spk_norm
+#define TN_A_FUSION_KERNEL       "a.fusion.spk_kernel"
+
 // align x to upper multiple of n
 #define CLIP_ALIGN(x, n) ((((x) + (n) - 1) / (n)) * (n))
 
@@ -487,6 +499,11 @@ struct clip_enc_state {
 
     // lfm3a: same order as list_lfm3a_state_slots(), F32 [ne0, ne1], updated in place
     std::vector<ggml_tensor *> slots;
+
+    // nemotron transcribe: the window's normalized log-mel, cache and FIFO rows in; its speaker rows and predictions out
+    std::vector<float> mel_norm;
+    std::vector<float> spk_prev;
+    std::vector<float> spk_chunk, spk_preds;
 };
 
 // per-sequence state of audio generation, see clip_gen_state_init()
@@ -567,6 +584,7 @@ enum projector_type {
     PROJECTOR_TYPE_MINICPMV4_7,
     PROJECTOR_TYPE_MINICPM_OCR_1_0,
     PROJECTOR_TYPE_GRANITE_SPEECH,
+    PROJECTOR_TYPE_NEMOTRON_TRANSCRIBE,
     PROJECTOR_TYPE_MIMOVL,
     PROJECTOR_TYPE_MINIMAX_M3,
     PROJECTOR_TYPE_GRANITE4_VISION,
@@ -581,6 +599,8 @@ enum projector_type {
     PROJECTOR_TYPE_SENSENOVA_U1,
     PROJECTOR_TYPE_COHERE2V,
     PROJECTOR_TYPE_MUSE_SPARK_AUDIO,
+    PROJECTOR_TYPE_CHATTERBOX,
+    PROJECTOR_TYPE_CHATTERBOX_SPKENC,
     PROJECTOR_TYPE_UNKNOWN,
 };
 
@@ -647,6 +667,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
     { PROJECTOR_TYPE_MINICPMV4_7,       "minicpmv4_7"},
     { PROJECTOR_TYPE_MINICPM_OCR_1_0,   "minicpm_ocr_1_0"},
     { PROJECTOR_TYPE_GRANITE_SPEECH,    "granite_speech"},
+    { PROJECTOR_TYPE_NEMOTRON_TRANSCRIBE, "nemotron_transcribe"},
     { PROJECTOR_TYPE_MIMOVL,            "mimovl"},
     { PROJECTOR_TYPE_MINIMAX_M3,        "minimax_m3"},
     { PROJECTOR_TYPE_GRANITE4_VISION,   "granite4_vision"},
@@ -662,6 +683,8 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
     { PROJECTOR_TYPE_SENSENOVA_U1,      "sensenova_u1"},
     { PROJECTOR_TYPE_COHERE2V,          "cohere2v"},
     { PROJECTOR_TYPE_MUSE_SPARK_AUDIO,  "muse-spark-audio"},
+    { PROJECTOR_TYPE_CHATTERBOX,        "chatterbox"},
+    { PROJECTOR_TYPE_CHATTERBOX_SPKENC, "chatterbox_spkenc"},
 };
 
 static projector_type clip_projector_type_from_string(const std::string & str) {

tools/mtmd/clip-model.h

@@ -6,6 +6,7 @@
 
 #include <algorithm>
 #include <array>
+#include <map>
 #include <vector>
 #include <unordered_set>
 #include <cstdint>
@@ -165,6 +166,12 @@ struct clip_hparams {
     std::vector<float> mel_filters;
     std::vector<float> window;
 
+    // nemotron transcribe
+    int32_t audio_rope_dims = 0;
+    int32_t spk_n_embd      = 0;
+    int32_t spk_n_layer     = 0;
+    int32_t spk_n_head      = 0;
+
     // mimo-audio-tokenizer: residual vector quantizer
     int32_t rvq_num_quantizers = 0;
     std::vector<int32_t> rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17)
@@ -292,6 +299,8 @@ struct clip_layer {
 
     ggml_tensor * k_norm = nullptr;
     ggml_tensor * q_norm = nullptr;
+    ggml_tensor * k_norm_b = nullptr;
+    ggml_tensor * q_norm_b = nullptr;
 
     // lfm3: per-head sigmoid gate on the attention output
     ggml_tensor * attn_gate_w = nullptr;
@@ -603,6 +612,235 @@ struct clip_code2wav {
     ggml_tensor * dac_post_conv_b      = nullptr;
 };
 
+struct clip_chatterbox {
+    // espnet rel-pos conformer layer of the flow encoder
+    struct enc_layer {
+        ggml_tensor * norm_mha_w = nullptr;
+        ggml_tensor * norm_mha_b = nullptr;
+        ggml_tensor * attn_q_w = nullptr;
+        ggml_tensor * attn_q_b = nullptr;
+        ggml_tensor * attn_k_w = nullptr;
+        ggml_tensor * attn_k_b = nullptr;
+        ggml_tensor * attn_v_w = nullptr;
+        ggml_tensor * attn_v_b = nullptr;
+        ggml_tensor * attn_pos_w = nullptr;      // linear_pos, no bias
+        ggml_tensor * attn_pos_bias_u = nullptr;
+        ggml_tensor * attn_pos_bias_v = nullptr;
+        ggml_tensor * attn_out_w = nullptr;
+        ggml_tensor * attn_out_b = nullptr;
+        ggml_tensor * norm_ff_w = nullptr;
+        ggml_tensor * norm_ff_b = nullptr;
+        ggml_tensor * ffn_1_w = nullptr;
+        ggml_tensor * ffn_1_b = nullptr;
+        ggml_tensor * ffn_2_w = nullptr;
+        ggml_tensor * ffn_2_b = nullptr;
+    };
+
+    // causal conv block of the estimator (conv k3 left padded + layer norm)
+    struct causal_block {
+        ggml_tensor * conv_w = nullptr;
+        ggml_tensor * conv_b = nullptr;
+        ggml_tensor * norm_w = nullptr;
+        ggml_tensor * norm_b = nullptr;
+    };
+
+    // time conditioned resnet of the estimator
+    struct resnet {
+        causal_block block1;
+        causal_block block2;
+        ggml_tensor * mlp_w = nullptr;      // time projection (mlp.1)
+        ggml_tensor * mlp_b = nullptr;
+        ggml_tensor * res_conv_w = nullptr;
+        ggml_tensor * res_conv_b = nullptr;
+    };
+
+    // diffusers style transformer block of the estimator
+    struct tfm_block {
+        ggml_tensor * norm1_w = nullptr;
+        ggml_tensor * norm1_b = nullptr;
+        ggml_tensor * attn_q_w = nullptr;   // no bias on qkv
+        ggml_tensor * attn_k_w = nullptr;
+        ggml_tensor * attn_v_w = nullptr;
+        ggml_tensor * attn_out_w = nullptr; // to_out.0
+        ggml_tensor * attn_out_b = nullptr;
+        ggml_tensor * norm3_w = nullptr;
+        ggml_tensor * norm3_b = nullptr;
+        ggml_tensor * ff_in_w = nullptr;    // ff.net.0.proj
+        ggml_tensor * ff_in_b = nullptr;
+        ggml_tensor * ff_out_w = nullptr;   // ff.net.2
+        ggml_tensor * ff_out_b = nullptr;
+    };
+
+    // one estimator stage: resnet, transformer stack, boundary causal conv (null on mid stages)
+    struct est_stage {
+        resnet res;
+        std::vector<tfm_block> tfm;
+        ggml_tensor * conv_w = nullptr;
+        ggml_tensor * conv_b = nullptr;
+    };
+
+    // snake resblock unit of the hift vocoder
+    struct hift_res_unit {
+        ggml_tensor * act1_alpha = nullptr;
+        ggml_tensor * act2_alpha = nullptr;
+        ggml_tensor * conv1_w = nullptr;
+        ggml_tensor * conv1_b = nullptr;
+        ggml_tensor * conv2_w = nullptr;
+        ggml_tensor * conv2_b = nullptr;
+    };
+    struct hift_res {
+        std::vector<hift_res_unit> units;   // dilations 1/3/5 on conv1
+    };
+
+    // one hift upsample stage: conv transpose, source injection, 3 resblocks
+    struct hift_up {
+        ggml_tensor * up_w = nullptr;
+        ggml_tensor * up_b = nullptr;
+        ggml_tensor * source_down_w = nullptr;
+        ggml_tensor * source_down_b = nullptr;
+        hift_res source_res;
+        hift_res res_0;
+        hift_res res_1;
+        hift_res res_2;
+    };
+
+    // f0 predictor conv of the hift vocoder
+    struct f0_conv {
+        ggml_tensor * w = nullptr;
+        ggml_tensor * b = nullptr;
+    };
+
+    // s3 tokenizer attention block (neox rope on q/k, fsmn memory on v)
+    struct s3tok_block {
+        ggml_tensor * attn_ln_w = nullptr;
+        ggml_tensor * attn_ln_b = nullptr;
+        ggml_tensor * attn_q_w = nullptr;
+        ggml_tensor * attn_q_b = nullptr;
+        ggml_tensor * attn_k_w = nullptr;   // no bias
+        ggml_tensor * attn_v_w = nullptr;
+        ggml_tensor * attn_v_b = nullptr;
+        ggml_tensor * fsmn_w = nullptr;
+        ggml_tensor * attn_out_w = nullptr;
+        ggml_tensor * attn_out_b = nullptr;
+        ggml_tensor * mlp_ln_w = nullptr;
+        ggml_tensor * mlp_ln_b = nullptr;
+        ggml_tensor * mlp_in_w = nullptr;
+        ggml_tensor * mlp_in_b = nullptr;
+        ggml_tensor * mlp_out_w = nullptr;
+        ggml_tensor * mlp_out_b = nullptr;
+    };
+
+    // folded batchnorm (w/b stay null on the affine free variant)
+    struct bn {
+        ggml_tensor * mean = nullptr;
+        ggml_tensor * var  = nullptr;
+        ggml_tensor * w    = nullptr;
+        ggml_tensor * b    = nullptr;
+    };
+
+    // fcm residual 2d block of the speaker encoder, the shortcut is null on stride 1 blocks
+    struct spk_res2d {
+        ggml_tensor * conv1_w = nullptr;
+        bn bn1;
+        ggml_tensor * conv2_w = nullptr;
+        bn bn2;
+        ggml_tensor * shortcut_w = nullptr;
+        bn shortcut_bn;
+    };
+
+    // cam dense tdnn layer of the speaker encoder
+    struct spk_cam_layer {
+        bn nl1_bn;
+        ggml_tensor * linear1_w = nullptr;
+        bn nl2_bn;
+        ggml_tensor * local_w = nullptr;    // cam_layer.linear_local
+        ggml_tensor * ctx1_w = nullptr;     // cam_layer.linear1
+        ggml_tensor * ctx1_b = nullptr;
+        ggml_tensor * ctx2_w = nullptr;     // cam_layer.linear2
+        ggml_tensor * ctx2_b = nullptr;
+    };
+    struct spk_cam_block {
+        std::vector<spk_cam_layer> layers;
+        bn transit_bn;
+        ggml_tensor * transit_w = nullptr;
+    };
+
+    // flow encoder
+    ggml_tensor * input_embedding_w = nullptr;  // flow.input_embedding
+    ggml_tensor * embed_linear_w = nullptr;     // fenc.embed.out.0
+    ggml_tensor * embed_linear_b = nullptr;
+    ggml_tensor * embed_norm_w = nullptr;       // fenc.embed.out.1
+    ggml_tensor * embed_norm_b = nullptr;
+    ggml_tensor * pre_conv1_w = nullptr;        // pre_lookahead_layer
+    ggml_tensor * pre_conv1_b = nullptr;
+    ggml_tensor * pre_conv2_w = nullptr;
+    ggml_tensor * pre_conv2_b = nullptr;
+    std::vector<enc_layer> enc;
+    ggml_tensor * up_conv_w = nullptr;          // fenc.up_layer.conv
+    ggml_tensor * up_conv_b = nullptr;
+    ggml_tensor * up_embed_linear_w = nullptr;  // fenc.up_embed.out.0
+    ggml_tensor * up_embed_linear_b = nullptr;
+    ggml_tensor * up_embed_norm_w = nullptr;    // fenc.up_embed.out.1
+    ggml_tensor * up_embed_norm_b = nullptr;
+    std::vector<enc_layer> up_enc;
+    ggml_tensor * after_norm_w = nullptr;
+    ggml_tensor * after_norm_b = nullptr;
+    ggml_tensor * encoder_proj_w = nullptr;     // flow.encoder_proj
+    ggml_tensor * encoder_proj_b = nullptr;
+
+    // cfm estimator
+    ggml_tensor * time_mlp_1_w = nullptr;
+    ggml_tensor * time_mlp_1_b = nullptr;
+    ggml_tensor * time_mlp_2_w = nullptr;
+    ggml_tensor * time_mlp_2_b = nullptr;
+    ggml_tensor * time_embed_mixer_w = nullptr; // meanflow variant only
+    est_stage est_down;
+    std::vector<est_stage> est_mid;
+    est_stage est_up;
+    causal_block est_final_block;
+    ggml_tensor * est_final_proj_w = nullptr;
+    ggml_tensor * est_final_proj_b = nullptr;
+
+    // hift vocoder
+    ggml_tensor * hift_pre_w = nullptr;
+    ggml_tensor * hift_pre_b = nullptr;
+    std::vector<hift_up> hift_ups;
+    ggml_tensor * hift_post_w = nullptr;
+    ggml_tensor * hift_post_b = nullptr;
+    std::vector<f0_conv> f0_condnet;
+    ggml_tensor * f0_classifier_w = nullptr;
+    ggml_tensor * f0_classifier_b = nullptr;
+
+    // s3 tokenizer
+    ggml_tensor * s3tok_conv1_w = nullptr;
+    ggml_tensor * s3tok_conv1_b = nullptr;
+    ggml_tensor * s3tok_conv2_w = nullptr;
+    ggml_tensor * s3tok_conv2_b = nullptr;
+    std::vector<s3tok_block> s3tok_blocks;
+    ggml_tensor * s3tok_down_w = nullptr;       // quantizer._codebook.project_down
+    ggml_tensor * s3tok_down_b = nullptr;
+
+    // CAMPPlus speaker encoder
+    ggml_tensor * spk_conv1_w = nullptr;
+    bn spk_bn1;
+    spk_res2d spk_layer1_0;
+    spk_res2d spk_layer1_1;
+    spk_res2d spk_layer2_0;
+    spk_res2d spk_layer2_1;
+    ggml_tensor * spk_conv2_w = nullptr;
+    bn spk_bn2;
+    ggml_tensor * spk_tdnn_w = nullptr;
+    bn spk_tdnn_bn;
+    spk_cam_block spk_block1;
+    spk_cam_block spk_block2;
+    spk_cam_block spk_block3;
+    bn spk_out_bn;
+    ggml_tensor * spk_dense_w = nullptr;
+    bn spk_dense_bn;
+    ggml_tensor * spk_affine_w = nullptr;
+    ggml_tensor * spk_affine_b = nullptr;
+};
+
 struct clip_vocos {
     ggml_tensor * input_w        = nullptr;
     ggml_tensor * input_b        = nullptr;
@@ -888,8 +1126,13 @@ struct clip_model {
     ggml_tensor * gen_upsample_w  = nullptr; // depthwise convtr, frame rate -> encoder frame rate
     std::vector<clip_layer> gen_tfm_layers; // mimi decoder_transformer
 
+    // chatterbox flow encoder, cfm estimator, hift vocoder, s3 tokenizer and CAMPPlus speaker encoder
+    clip_chatterbox cbx;
+
+    // chatterbox tensors read on the host, see clip_cbx_read_tensor()
+    std::map<std::string, ggml_tensor *> cbx_tensors;
     // predefined voices, no rows if the voice is named in the prompt (lfm3a)
-    std::map<std::string, std::vector<float>> gen_voices;
+    std::map<std::string, clip_gen_voice> gen_voices;
 
     // soprano: Vocos decoder, hidden states -> STFT coefficients
     clip_vocos vocos;
@@ -926,6 +1169,28 @@ struct clip_model {
 
     std::vector<clip_layer> sam_layers;
 
+    // nemotron transcribe: speaker expert, Sortformer head, fusion
+    std::vector<clip_layer> spk_layers;
+    ggml_tensor * spk_inp_proj_w     = nullptr;
+    ggml_tensor * spk_pre_ln_w       = nullptr;
+    ggml_tensor * spk_pre_ln_b       = nullptr;
+    ggml_tensor * spk_post_ln_w      = nullptr;
+    ggml_tensor * spk_post_ln_b      = nullptr;
+    ggml_tensor * spk_head_proj_w    = nullptr;
+    ggml_tensor * spk_head_proj_b    = nullptr;
+    ggml_tensor * spk_head_up_w      = nullptr;
+    ggml_tensor * spk_head_up_b      = nullptr;
+    ggml_tensor * spk_head_fc1_w     = nullptr;
+    ggml_tensor * spk_head_fc1_b     = nullptr;
+    ggml_tensor * spk_head_fc2_w     = nullptr;
+    ggml_tensor * spk_head_fc2_b     = nullptr;
+    ggml_tensor * fusion_asr_norm_w  = nullptr;
+    ggml_tensor * fusion_asr_norm_b  = nullptr;
+    ggml_tensor * fusion_spk_norm_w  = nullptr;
+    ggml_tensor * fusion_spk_norm_b  = nullptr;
+    ggml_tensor * fusion_spk_kernel  = nullptr;
+    std::vector<float> spk_sil_emb;
+
     // deepseek-ocr-2
     ggml_tensor * resample_query_768 = nullptr;
     ggml_tensor * resample_query_1024 = nullptr;

tools/mtmd/clip.h

@@ -96,7 +96,18 @@ enum clip_gen_process_type {
     CLIP_GEN_PROCESS_GEN_UNKNOWN,
     CLIP_GEN_PROCESS_GEN_CODE, // h_state to codes
     CLIP_GEN_PROCESS_GEN_WAV,  // codes to raw PCM audio
+    CLIP_GEN_PROCESS_TTS,        // full utterance of codes to raw PCM audio
+    CLIP_GEN_PROCESS_TTS_VOCODE, // internal: mel + source stft to istft input
+    CLIP_GEN_PROCESS_TOKENIZE,   // raw PCM audio to semantic speech tokens
 };
+
+// gen-side voice reference (chatterbox): speech tokens, mel-rate features and the speaker vector
+struct clip_gen_ref {
+    std::vector<int32_t> tokens;
+    std::vector<float>   feat;
+    std::vector<float>   spk;
+};
+
 struct clip_encode_params {
     int n_threads = 1;
     const clip_image_f32_batch * imgs = nullptr;
@@ -114,10 +125,20 @@ struct clip_encode_params {
     float   temp = 0.0f;                      // sampling temperature, noise scale for flow-matching decoders
     bool * out_is_eos = nullptr;
 
+    // TOKENIZE: speech embedding rows of the produced codes
+    std::vector<float> * out_code_embd = nullptr;
+
     // GEN_WAV
     const std::vector<int32_t> * codes = nullptr;     // this frame's 16 RVQ codes
     const std::vector<float> *   feats = nullptr;     // continuous counterpart of codes
     std::vector<float> * out_audio = nullptr;         // decoded PCM samples, F32
+    const clip_gen_ref * ref = nullptr; // TTS reference conditioning of the voice
+    // TTS_VOCODE internal stage inputs
+    const std::vector<float> * mel_in   = nullptr;
+    const std::vector<float> * sstft_in = nullptr;
+    int vocode_n_mel  = 0;
+    int vocode_n_stft = 0;
+    std::vector<float> * out_spec = nullptr;
 
     // audio gen: state of one sequence, clip_encode() updates it in place
     struct clip_gen_state * gen_state = nullptr;
@@ -139,7 +160,16 @@ struct clip_gen_state * clip_gen_state_init(struct clip_ctx * ctx, uint32_t seed
 void                    clip_gen_state_free(struct clip_gen_state * state);
 void                    clip_gen_state_reset(struct clip_gen_state * state);
 
-const std::map<std::string, std::vector<float>> & clip_gen_get_voices(const struct clip_ctx * ctx);
+// read a chatterbox tensor by name as F32, returns the element count (0 if not found), out may be null to query the size
+size_t clip_cbx_read_tensor(struct clip_ctx * ctx, const char * name, float * out, size_t n_max);
+
+// a predefined voice: backbone rows, plus the gen-side reference when the gen stack needs one
+struct clip_gen_voice {
+    std::vector<float> rows;
+    clip_gen_ref       ref;
+};
+
+const std::map<std::string, clip_gen_voice> & clip_gen_get_voices(const struct clip_ctx * ctx);
 
 bool clip_is_llava(const struct clip_ctx * ctx);
 // note for contributor: this clip_is_(model) pattern is deprecated

tools/mtmd/models/chatterbox-gen.cpp

@@ -0,0 +1,542 @@
+#define _USE_MATH_DEFINES // for M_PI
+#include "models.h"
+
+// Chatterbox generation graphs: flow encoder, cfm estimator, s3 tokenizer and hift vocoder
+
+// x [C, T]: y = W x + b with torch Linear weights stored as [in, out]
+ggml_tensor * clip_graph_chatterbox_base::cbx_linear(ggml_tensor * w, ggml_tensor * b, ggml_tensor * x) const {
+    ggml_tensor * y = ggml_mul_mat(ctx0, w, x);
+    if (b) {
+        y = ggml_add(ctx0, y, b);
+    }
+    return y;
+}
+
+// conv1d over time with asymmetric zero padding: x [C, T] -> [OC, T_out], kernel [K, IC, OC]
+ggml_tensor * clip_graph_chatterbox_base::cbx_conv1d(ggml_tensor * k, ggml_tensor * b, ggml_tensor * x,
+                                                     int stride, int pad_l, int pad_r) const {
+    ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x)); // [T, C]
+    if (pad_l > 0) {
+        ggml_tensor * z = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, pad_l, xt->ne[1]);
+        z = ggml_scale(ctx0, z, 0.0f);
+        xt = ggml_concat(ctx0, z, xt, 0);
+    }
+    if (pad_r > 0) {
+        ggml_tensor * z = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, pad_r, xt->ne[1]);
+        z = ggml_scale(ctx0, z, 0.0f);
+        xt = ggml_concat(ctx0, xt, z, 0);
+    }
+    ggml_tensor * y = ggml_conv_1d(ctx0, k, xt, stride, 0, 1); // [T_out, OC]
+    y = ggml_cont(ctx0, ggml_transpose(ctx0, y));              // [OC, T_out]
+    if (b) {
+        y = ggml_add(ctx0, y, b);
+    }
+    return y;
+}
+
+// x [C, T] -> symmetric-padded dilated conv -> [OC, T]
+ggml_tensor * clip_graph_chatterbox_base::cbx_conv1d_dil(ggml_tensor * k, ggml_tensor * b, ggml_tensor * x,
+                                                         int pad, int dil) const {
+    ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x));
+    ggml_tensor * y = ggml_conv_1d(ctx0, k, xt, 1, pad, dil);
+    y = ggml_cont(ctx0, ggml_transpose(ctx0, y));
+    if (b) {
+        y = ggml_add(ctx0, y, b);
+    }
+    return y;
+}
+
+// Transformer-XL relative shift: bd [2T-1, T, H] -> [T, T, H] with out[j, i, h] = bd[(T-1) - i + j, i, h]
+// same buffer walk as the espnet rel_shift: pad one column, drop the first row, keep the first T columns
+static ggml_tensor * cbx_rel_shift(ggml_context * ctx0, ggml_tensor * bd, int T) {
+    const int H = (int) bd->ne[2];
+    ggml_tensor * z = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, T, H);
+    z = ggml_scale(ctx0, z, 0.0f);
+    ggml_tensor * p = ggml_cont(ctx0, ggml_concat(ctx0, z, bd, 0));           // [2T, T, H]
+    p = ggml_reshape_3d(ctx0, p, T, 2 * T, H);                                // [T, 2T, H]
+    p = ggml_view_3d(ctx0, p, T, 2 * T - 1, H, p->nb[1], p->nb[2], p->nb[1]); // drop first row
+    p = ggml_cont(ctx0, p);
+    p = ggml_reshape_3d(ctx0, p, 2 * T - 1, T, H);                            // [2T-1, T, H]
+    p = ggml_view_3d(ctx0, p, T, T, H, p->nb[1], p->nb[2], 0);                // first T columns
+    return ggml_cont(ctx0, p);
+}
+
+// espnet rel-pos self attention block, pre-norm, x [512, T], pos [512, 2T-1]
+ggml_tensor * clip_graph_chatterbox::enc_layer(const clip_chatterbox::enc_layer & l, ggml_tensor * x,
+                                               ggml_tensor * pos, int T) {
+    const int n_head = 8;
+    const int d_head = 64;
+    const float scale = 1.0f / sqrtf((float) d_head);
+
+    ggml_tensor * res = x;
+    ggml_tensor * cur = build_norm(x, l.norm_mha_w, l.norm_mha_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+
+    ggml_tensor * q = cbx_linear(l.attn_q_w, l.attn_q_b, cur);
+    ggml_tensor * k = cbx_linear(l.attn_k_w, l.attn_k_b, cur);
+    ggml_tensor * v = cbx_linear(l.attn_v_w, l.attn_v_b, cur);
+    ggml_tensor * pe = cbx_linear(l.attn_pos_w, nullptr, pos); // [512, 2T-1]
+
+    q = ggml_reshape_3d(ctx0, q, d_head, n_head, T);
+    k = ggml_reshape_3d(ctx0, k, d_head, n_head, T);
+    v = ggml_reshape_3d(ctx0, v, d_head, n_head, T);
+    pe = ggml_reshape_3d(ctx0, pe, d_head, n_head, 2 * T - 1);
+
+    ggml_tensor * u = l.attn_pos_bias_u; // [64, 8]
+    ggml_tensor * w = l.attn_pos_bias_v;
+
+    ggml_tensor * qu = ggml_add(ctx0, q, ggml_reshape_3d(ctx0, u, d_head, n_head, 1));
+    ggml_tensor * qv = ggml_add(ctx0, q, ggml_reshape_3d(ctx0, w, d_head, n_head, 1));
+
+    // per head: [64, T] tensors, scores [T(k), T(q)]
+    qu = ggml_cont(ctx0, ggml_permute(ctx0, qu, 0, 2, 1, 3)); // [64, T, 8]
+    qv = ggml_cont(ctx0, ggml_permute(ctx0, qv, 0, 2, 1, 3));
+    k  = ggml_cont(ctx0, ggml_permute(ctx0, k,  0, 2, 1, 3));
+    v  = ggml_cont(ctx0, ggml_permute(ctx0, v,  0, 2, 1, 3));
+    pe = ggml_cont(ctx0, ggml_permute(ctx0, pe, 0, 2, 1, 3)); // [64, 2T-1, 8]
+
+    ggml_tensor * ac = ggml_mul_mat(ctx0, k, qu);   // [T(k), T(q), 8]
+    ggml_tensor * bd = ggml_mul_mat(ctx0, pe, qv);  // [2T-1, T(q), 8]
+    bd = cbx_rel_shift(ctx0, bd, T);                // [T(k), T(q), 8]
+
+    ggml_tensor * scores = ggml_scale(ctx0, ggml_add(ctx0, ac, bd), scale);
+    ggml_tensor * probs = ggml_soft_max(ctx0, scores);
+
+    ggml_tensor * o = ggml_mul_mat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, v)), probs); // [64, T, 8]
+    o = ggml_cont(ctx0, ggml_permute(ctx0, o, 0, 2, 1, 3));                                 // [64, 8, T]
+    o = ggml_reshape_2d(ctx0, o, n_head * d_head, T);
+    o = cbx_linear(l.attn_out_w, l.attn_out_b, o);
+    x = ggml_add(ctx0, res, o);
+
+    res = x;
+    cur = build_norm(x, l.norm_ff_w, l.norm_ff_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+    cur = cbx_linear(l.ffn_1_w, l.ffn_1_b, cur);
+    cur = ggml_silu(ctx0, cur); // swish
+    cur = cbx_linear(l.ffn_2_w, l.ffn_2_b, cur);
+    x = ggml_add(ctx0, res, cur);
+    return x;
+}
+
+// mish = x * tanh(softplus(x))
+static ggml_tensor * cbx_mish(ggml_context * ctx0, ggml_tensor * x) {
+    return ggml_mul(ctx0, x, ggml_tanh(ctx0, ggml_softplus(ctx0, x)));
+}
+
+// causal block: conv k3 left-padded, layer norm over channels, mish; x [C, T]
+ggml_tensor * clip_graph_chatterbox::causal_block(const clip_chatterbox::causal_block & b, ggml_tensor * x) {
+    x = cbx_conv1d(b.conv_w, b.conv_b, x, 1, (int) b.conv_w->ne[0] - 1, 0);
+    x = build_norm(x, b.norm_w, b.norm_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+    return cbx_mish(ctx0, x);
+}
+
+// resnet block with time conditioning; x [C, T], temb [1024]
+ggml_tensor * clip_graph_chatterbox::resnet(const clip_chatterbox::resnet & r, ggml_tensor * x, ggml_tensor * temb) {
+    ggml_tensor * h = causal_block(r.block1, x);
+    ggml_tensor * tproj = cbx_linear(r.mlp_w, r.mlp_b, cbx_mish(ctx0, temb));
+    h = ggml_add(ctx0, h, tproj); // broadcast [256, 1] over T
+    h = causal_block(r.block2, h);
+    ggml_tensor * res = cbx_conv1d(r.res_conv_w, r.res_conv_b, x, 1, 0, 0);
+    return ggml_add(ctx0, h, res);
+}
+
+// diffusers-style transformer block, full attention; x [256, T]
+ggml_tensor * clip_graph_chatterbox::tfm_block(const clip_chatterbox::tfm_block & b, ggml_tensor * x) {
+    const int n_head = 8;
+    const int d_head = 64;
+    const int T = (int) x->ne[1];
+
+    ggml_tensor * cur = build_norm(x, b.norm1_w, b.norm1_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+    ggml_tensor * q = ggml_reshape_3d(ctx0, ggml_mul_mat(ctx0, b.attn_q_w, cur), d_head, n_head, T);
+    ggml_tensor * k = ggml_reshape_3d(ctx0, ggml_mul_mat(ctx0, b.attn_k_w, cur), d_head, n_head, T);
+    ggml_tensor * v = ggml_reshape_3d(ctx0, ggml_mul_mat(ctx0, b.attn_v_w, cur), d_head, n_head, T);
+    x = ggml_add(ctx0, x, build_attn(b.attn_out_w, b.attn_out_b, q, k, v, nullptr, 1.0f / sqrtf((float) d_head), -1));
+
+    cur = build_norm(x, b.norm3_w, b.norm3_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+    cur = cbx_linear(b.ff_in_w, b.ff_in_b, cur);
+    cur = ggml_gelu_erf(ctx0, cur);
+    cur = cbx_linear(b.ff_out_w, b.ff_out_b, cur);
+    return ggml_add(ctx0, x, cur);
+}
+
+// one estimator evaluation; x_noise, mu, cond [80, T], spks [80], temb [1024]
+ggml_tensor * clip_graph_chatterbox::estimator(ggml_tensor * x_noise, ggml_tensor * mu, ggml_tensor * spks,
+                                               ggml_tensor * cond, ggml_tensor * temb, int T) {
+    const auto & c = model.cbx;
+
+    // channels live on ne0, time on ne1: pack along ne0
+    ggml_tensor * x = ggml_concat(ctx0, x_noise, mu, 0);                   // [160, T]
+    ggml_tensor * spks_b = ggml_repeat(ctx0, ggml_reshape_2d(ctx0, spks, 80, 1), ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 80, T));
+    x = ggml_concat(ctx0, x, spks_b, 0);                                   // [240, T]
+    x = ggml_concat(ctx0, x, cond, 0);                                     // [320, T]
+    x = ggml_cont(ctx0, x);
+
+    // down
+    ggml_tensor * skip;
+    x = resnet(c.est_down.res, x, temb);
+    for (const auto & b : c.est_down.tfm) {
+        x = tfm_block(b, x);
+    }
+    skip = x;
+    x = cbx_conv1d(c.est_down.conv_w, c.est_down.conv_b, x, 1, (int) c.est_down.conv_w->ne[0] - 1, 0);
+
+    // mid
+    for (const auto & m : c.est_mid) {
+        x = resnet(m.res, x, temb);
+        for (const auto & b : m.tfm) {
+            x = tfm_block(b, x);
+        }
+    }
+
+    // up with skip
+    x = ggml_concat(ctx0, x, skip, 0);                                     // [512, T]
+    x = resnet(c.est_up.res, x, temb);
+    for (const auto & b : c.est_up.tfm) {
+        x = tfm_block(b, x);
+    }
+    x = cbx_conv1d(c.est_up.conv_w, c.est_up.conv_b, x, 1, (int) c.est_up.conv_w->ne[0] - 1, 0);
+
+    x = causal_block(c.est_final_block, x);
+    x = cbx_conv1d(c.est_final_proj_w, c.est_final_proj_b, x, 1, 0, 0);    // [80, T]
+    return x;
+}
+
+// s3 tokenizer encoder: log-mel [T, 128] -> fsq code [8, T / 4], rounded to base 3 tokens on the host
+ggml_cgraph * clip_graph_chatterbox::build_s3tok(int T) {
+    const auto & c = model.cbx;
+    const int n_head = 20;
+    const int d_head = 64;
+    const int T1 = (T - 1) / 2 + 1;
+    const int T2 = (T1 - 1) / 2 + 1;
+
+    ggml_tensor * inp = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, T, 128);
+    ggml_set_name(inp, "inp_raw");
+    ggml_set_input(inp);
+
+    ggml_tensor * pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, T2);
+    ggml_set_name(pos, "inp_pos");
+    ggml_set_input(pos);
+
+    ggml_tensor * x = ggml_cont(ctx0, ggml_transpose(ctx0, inp)); // [128, T]
+    x = cbx_conv1d(c.s3tok_conv1_w, c.s3tok_conv1_b, x, 2, 1, 1);
+    x = ggml_gelu_erf(ctx0, x);
+    x = cbx_conv1d(c.s3tok_conv2_w, c.s3tok_conv2_b, x, 2, 1, 1);
+    x = ggml_gelu_erf(ctx0, x); // [1280, T2]
+
+    for (const auto & blk : c.s3tok_blocks) {
+        ggml_tensor * res = x;
+        ggml_tensor * cur = build_norm(x, blk.attn_ln_w, blk.attn_ln_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+        ggml_tensor * q = cbx_linear(blk.attn_q_w, blk.attn_q_b, cur);
+        ggml_tensor * k = ggml_mul_mat(ctx0, blk.attn_k_w, cur);
+        ggml_tensor * v = cbx_linear(blk.attn_v_w, blk.attn_v_b, cur);
+
+        // fsmn memory: depthwise conv over time on v plus residual, added to the attention output
+        ggml_tensor * fsm = ggml_cont(ctx0, ggml_transpose(ctx0, v)); // [T2, 1280]
+        {
+            ggml_tensor * w = blk.fsmn_w;
+            ggml_tensor * m = ggml_conv_1d_dw(ctx0, w, fsm, 1, ((int) w->ne[0] - 1) / 2, 1);
+            fsm = ggml_add(ctx0, ggml_reshape_2d(ctx0, m, fsm->ne[0], fsm->ne[1]), fsm);
+        }
+        fsm = ggml_cont(ctx0, ggml_transpose(ctx0, fsm)); // [1280, T2]
+
+        q = ggml_reshape_3d(ctx0, q, d_head, n_head, T2);
+        k = ggml_reshape_3d(ctx0, k, d_head, n_head, T2);
+        v = ggml_reshape_3d(ctx0, v, d_head, n_head, T2);
+        q = ggml_rope_ext(ctx0, q, pos, nullptr, d_head, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+        k = ggml_rope_ext(ctx0, k, pos, nullptr, d_head, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+        ggml_tensor * o = build_attn(blk.attn_out_w, blk.attn_out_b, q, k, v, nullptr, 1.0f / sqrtf((float) d_head), -1);
+        x = ggml_add(ctx0, res, ggml_add(ctx0, o, fsm));
+
+        res = x;
+        cur = build_norm(x, blk.mlp_ln_w, blk.mlp_ln_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+        cur = cbx_linear(blk.mlp_in_w, blk.mlp_in_b, cur);
+        cur = ggml_gelu_erf(ctx0, cur);
+        cur = cbx_linear(blk.mlp_out_w, blk.mlp_out_b, cur);
+        x = ggml_add(ctx0, res, cur);
+    }
+
+    x = cbx_linear(c.s3tok_down_w, c.s3tok_down_b, x); // [8, T2]
+    x = ggml_tanh(ctx0, x);
+
+    ggml_set_name(x, "out_fsq");
+    ggml_set_output(x);
+    ggml_build_forward_expand(gf, x);
+    return gf;
+}
+
+ggml_cgraph * clip_graph_chatterbox::build() {
+    if (gen_process == CLIP_GEN_PROCESS_TTS_VOCODE) {
+        return build_vocoder(vocode_n_mel, vocode_n_stft);
+    }
+    if (gen_process == CLIP_GEN_PROCESS_TOKENIZE) {
+        return build_s3tok(img.nx());
+    }
+    if (gen_process != CLIP_GEN_PROCESS_TTS) {
+        // load-time buffer sizing path
+        ggml_tensor * inp = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
+        ggml_set_name(inp, "inp_stub");
+        ggml_set_input(inp);
+        ggml_tensor * cur = ggml_dup(ctx0, inp);
+        ggml_set_name(cur, "out_stub");
+        ggml_set_output(cur);
+        ggml_build_forward_expand(gf, cur);
+        return gf;
+    }
+
+    const auto & c = model.cbx;
+
+    const int T1 = n_tokens;      // token-rate length (prompt + generated)
+    const int T2 = 2 * n_tokens;  // mel-rate length after the x2 upsample
+
+    ggml_tensor * inp_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, T1);
+    ggml_set_name(inp_tokens, "inp_tokens");
+    ggml_set_input(inp_tokens);
+
+    ggml_tensor * pos1 = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 512, 2 * T1 - 1);
+    ggml_set_name(pos1, "inp_pos1");
+    ggml_set_input(pos1);
+
+    ggml_tensor * pos2 = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 512, 2 * T2 - 1);
+    ggml_set_name(pos2, "inp_pos2");
+    ggml_set_input(pos2);
+
+    // token embedding
+    ggml_tensor * x = ggml_get_rows(ctx0, c.input_embedding_w, inp_tokens); // [512, T1]
+
+    // embed: linear + layer norm, then the espnet xscale
+    x = cbx_linear(c.embed_linear_w, c.embed_linear_b, x);
+    x = build_norm(x, c.embed_norm_w, c.embed_norm_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+    x = ggml_scale(ctx0, x, sqrtf(512.0f));
+    cb(x, "fenc_embd", -1);
+
+    // pre-lookahead: conv k=4 right-padded 3, leaky 0.01, conv k=3 left-padded 2, residual
+    {
+        ggml_tensor * res = x;
+        ggml_tensor * cur = cbx_conv1d(c.pre_conv1_w, c.pre_conv1_b, x, 1, 0, 3);
+        cur = ggml_leaky_relu(ctx0, cur, 0.01f, false);
+        cur = cbx_conv1d(c.pre_conv2_w, c.pre_conv2_b, cur, 1, 2, 0);
+        x = ggml_add(ctx0, res, cur);
+        cb(x, "fenc_pre_lookahead", -1);
+    }
+
+    for (size_t i = 0; i < c.enc.size(); i++) {
+        x = enc_layer(c.enc[i], x, pos1, T1);
+        cb(x, "fenc_enc", (int) i);
+    }
+
+    // upsample x2: nearest repeat, left pad 4, conv k=5
+    {
+        ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x));            // [T1, 512]
+        xt = ggml_interpolate(ctx0, xt, 2 * T1, 512, 1, 1, GGML_SCALE_MODE_NEAREST);
+        ggml_tensor * z = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 4, 512);
+        z = ggml_scale(ctx0, z, 0.0f);
+        xt = ggml_concat(ctx0, z, xt, 0);
+        ggml_tensor * y = ggml_conv_1d(ctx0, c.up_conv_w, xt, 1, 0, 1);
+        x = ggml_cont(ctx0, ggml_transpose(ctx0, y));                            // [512, T2]
+        x = ggml_add(ctx0, x, c.up_conv_b);
+        cb(x, "fenc_upsample", -1);
+    }
+
+    // up embed: linear + layer norm + xscale
+    x = cbx_linear(c.up_embed_linear_w, c.up_embed_linear_b, x);
+    x = build_norm(x, c.up_embed_norm_w, c.up_embed_norm_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+    x = ggml_scale(ctx0, x, sqrtf(512.0f));
+
+    for (size_t i = 0; i < c.up_enc.size(); i++) {
+        x = enc_layer(c.up_enc[i], x, pos2, T2);
+        cb(x, "fenc_up_enc", (int) i);
+    }
+
+    x = build_norm(x, c.after_norm_w, c.after_norm_b, NORM_TYPE_NORMAL, 1e-5f, -1);
+
+    // encoder projection to the mel channel count
+    ggml_tensor * mu = cbx_linear(c.encoder_proj_w, c.encoder_proj_b, x); // [80, T2]
+    cb(mu, "flow_mu", -1);
+
+    // cfm solver unrolled in the graph: meanflow takes 2 euler steps without cfg, its time embeds mix t and r
+    // the classic solver takes 10 euler steps on the cosine schedule with cfg 0.7
+    const bool meanflow = c.time_embed_mixer_w != nullptr;
+    const int n_steps = meanflow ? 2 : 10;
+
+    // span points, same schedule as the host side sinusoid fill in clip.cpp
+    float span[11];
+    for (int i = 0; i <= n_steps; i++) {
+        const float u = (float) i / n_steps;
+        span[i] = meanflow ? u : 1.0f - cosf(u * (float) M_PI / 2.0f);
+    }
+
+    ggml_tensor * noise = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 80, T2);
+    ggml_set_name(noise, "inp_noise");
+    ggml_set_input(noise);
+    // sinusoidal time embeddings, one row per span point
+    ggml_tensor * temb_sin = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 320, n_steps + 1);
+    ggml_set_name(temb_sin, "inp_temb");
+    ggml_set_input(temb_sin);
+
+    auto time_mlp = [&](ggml_tensor * e) {
+        e = cbx_linear(c.time_mlp_1_w, c.time_mlp_1_b, e);
+        e = ggml_silu(ctx0, e);
+        e = cbx_linear(c.time_mlp_2_w, c.time_mlp_2_b, e);
+        return e;
+    };
+    auto span_emb = [&](int i) {
+        return ggml_view_2d(ctx0, temb_sin, 320, 1, temb_sin->nb[1], (size_t) i * temb_sin->nb[1]);
+    };
+    auto step_temb = [&](int i) {
+        if (!meanflow) {
+            return time_mlp(span_emb(i));
+        }
+        ggml_tensor * e = ggml_concat(ctx0, time_mlp(span_emb(i)), time_mlp(span_emb(i + 1)), 0); // [2048, 1]
+        return ggml_mul_mat(ctx0, c.time_embed_mixer_w, e);                                       // [1024, 1]
+    };
+
+    // mel-rate conditions of the voice: prompt features then zeros, and the speaker vector
+    ggml_tensor * pf = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 80, n_prompt_mel);
+    ggml_set_name(pf, "inp_prompt_feat");
+    ggml_set_input(pf);
+    ggml_tensor * zc = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 80, T2 - n_prompt_mel);
+    zc = ggml_scale(ctx0, zc, 0.0f);
+    ggml_tensor * cond = ggml_concat(ctx0, pf, zc, 1); // [80, T2]
+    // raw 192-dim campplus x-vector, normalized then through the flow speaker affine
+    ggml_tensor * xvec = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 192);
+    ggml_set_name(xvec, "inp_xvec");
+    ggml_set_input(xvec);
+    ggml_tensor * n2 = ggml_sqrt(ctx0, ggml_sum(ctx0, ggml_mul(ctx0, xvec, xvec)));
+    ggml_tensor * unit = ggml_div(ctx0, xvec, n2);
+    ggml_tensor * spks = cbx_linear(c.spk_affine_w, c.spk_affine_b,
+                                    ggml_reshape_2d(ctx0, unit, 192, 1));
+    spks = ggml_cont(ctx0, ggml_reshape_1d(ctx0, spks, 80));
+
+    const float cfg = 0.7f;
+    ggml_tensor * mu_zero   = meanflow ? nullptr : ggml_scale(ctx0, mu, 0.0f);
+    ggml_tensor * cond_zero = meanflow ? nullptr : ggml_scale(ctx0, cond, 0.0f);
+    ggml_tensor * spks_zero = meanflow ? nullptr : ggml_scale(ctx0, spks, 0.0f);
+
+    ggml_tensor * mx = noise;
+    for (int i = 0; i < n_steps; i++) {
+        ggml_tensor * temb = step_temb(i);
+        ggml_tensor * d = estimator(mx, mu, spks, cond, temb, T2);
+        if (!meanflow) {
+            ggml_tensor * du = estimator(mx, mu_zero, spks_zero, cond_zero, temb, T2);
+            d = ggml_add(ctx0, ggml_scale(ctx0, d, 1.0f + cfg), ggml_scale(ctx0, du, -cfg));
+        }
+        mx = ggml_add(ctx0, mx, ggml_scale(ctx0, d, span[i + 1] - span[i]));
+        cb(mx, "cfm_step", i);
+    }
+    ggml_tensor * mel = mx;
+
+    // trim the prompt frames at mel rate
+    mel = ggml_view_2d(ctx0, mel, 80, T2 - n_prompt_mel, mel->nb[1], (size_t) n_prompt_mel * mel->nb[1]);
+    mel = ggml_cont(ctx0, mel);
+    ggml_set_name(mel, "out_mel");
+    ggml_set_output(mel);
+    ggml_build_forward_expand(gf, mel);
+
+    // f0 predictor on the trimmed mel: 5x (conv k3 same-pad + elu), abs(linear)
+    {
+        ggml_tensor * fx = mel;
+        for (const auto & fc : c.f0_condnet) {
+            fx = cbx_conv1d(fc.w, fc.b, fx, 1, 1, 1);
+            fx = ggml_elu(ctx0, fx);
+        }
+        fx = cbx_linear(c.f0_classifier_w, c.f0_classifier_b, fx);
+        fx = ggml_abs(ctx0, fx); // [1, T]
+        ggml_set_name(fx, "out_f0");
+        ggml_set_output(fx);
+        ggml_build_forward_expand(gf, fx);
+    }
+    return gf;
+}
+
+// snake activation with per-channel alpha: x + sin^2(alpha x) / alpha
+static ggml_tensor * cbx_snake(ggml_context * ctx0, ggml_tensor * x, ggml_tensor * alpha) {
+    ggml_tensor * sx = ggml_sin(ctx0, ggml_mul(ctx0, x, alpha));
+    sx = ggml_mul(ctx0, sx, sx);
+    sx = ggml_div(ctx0, sx, alpha);
+    return ggml_add(ctx0, x, sx);
+}
+
+// hifigan-snake resblock; kernels with dilations 1/3/5 on convs1, 1 on convs2
+ggml_tensor * clip_graph_chatterbox::hift_resblock(const clip_chatterbox::hift_res & r, ggml_tensor * x) {
+    static const int dil[3] = {1, 3, 5};
+    for (size_t j = 0; j < r.units.size(); j++) {
+        const auto & u = r.units[j];
+        const int d = dil[j % 3];
+        const int p1 = (int) (u.conv1_w->ne[0] - 1) / 2 * d;
+        const int p2 = (int) (u.conv2_w->ne[0] - 1) / 2;
+        ggml_tensor * xt = cbx_snake(ctx0, x, u.act1_alpha);
+        xt = cbx_conv1d_dil(u.conv1_w, u.conv1_b, xt, p1, d);
+        xt = cbx_snake(ctx0, xt, u.act2_alpha);
+        xt = cbx_conv1d_dil(u.conv2_w, u.conv2_b, xt, p2, 1);
+        x = ggml_add(ctx0, x, xt);
+    }
+    return x;
+}
+
+// mel [80, T] + source stft [18, T_stft] -> conv_post output [18, T_stft2]
+ggml_cgraph * clip_graph_chatterbox::build_vocoder(int n_mel, int n_stft) {
+    const auto & c = model.cbx;
+
+    ggml_tensor * mel = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 80, n_mel);
+    ggml_set_name(mel, "inp_mel");
+    ggml_set_input(mel);
+    ggml_tensor * sstft = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 18, n_stft);
+    ggml_set_name(sstft, "inp_sstft");
+    ggml_set_input(sstft);
+
+    ggml_tensor * x = cbx_conv1d_dil(c.hift_pre_w, c.hift_pre_b, mel, 3, 1);
+
+    for (size_t i = 0; i < c.hift_ups.size(); i++) {
+        const auto & up = c.hift_ups[i];
+        ggml_tensor * uk = up.up_w;
+        const int K = (int) uk->ne[0];
+        const int S = K / 2;
+        const int P = (K - S) / 2;
+
+        x = ggml_leaky_relu(ctx0, x, 0.1f, false);
+        // conv transpose then trim the torch padding P on both sides
+        ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x));
+        xt = ggml_conv_transpose_1d(ctx0, uk, xt, S, 0, 1);
+        xt = ggml_cont(ctx0, ggml_view_2d(ctx0, xt, xt->ne[0] - 2 * P, xt->ne[1], xt->nb[1], (size_t) P * ggml_element_size(xt)));
+        x = ggml_cont(ctx0, ggml_transpose(ctx0, xt));
+        x = ggml_add(ctx0, x, up.up_b);
+
+        const bool is_last = i + 1 == c.hift_ups.size();
+        if (is_last) {
+            ggml_tensor * xr = ggml_cont(ctx0, ggml_transpose(ctx0, x));
+            xr = ggml_pad_reflect_1d(ctx0, xr, 1, 0);
+            x = ggml_cont(ctx0, ggml_transpose(ctx0, xr));
+        }
+
+        // source injection: strided conv on the source stft, one resblock
+        ggml_tensor * sk = up.source_down_w;
+        const int SK = (int) sk->ne[0];
+        const int SS = SK > 1 ? SK / 2 : 1;
+        const int SP = SK > 1 ? SS / 2 : 0;
+        ggml_tensor * si;
+        {
+            ggml_tensor * st = ggml_cont(ctx0, ggml_transpose(ctx0, sstft));
+            st = ggml_conv_1d(ctx0, sk, st, SS, SP, 1);
+            si = ggml_cont(ctx0, ggml_transpose(ctx0, st));
+            si = ggml_add(ctx0, si, up.source_down_b);
+        }
+        si = hift_resblock(up.source_res, si);
+        // align lengths: the reflection pad on the last stage adds one step
+        if ((int) si->ne[1] != (int) x->ne[1]) {
+            const int n = (int) std::min(si->ne[1], x->ne[1]);
+            si = ggml_cont(ctx0, ggml_view_2d(ctx0, si, si->ne[0], n, si->nb[1], 0));
+            x  = ggml_cont(ctx0, ggml_view_2d(ctx0, x,  x->ne[0],  n, x->nb[1],  0));
+        }
+        x = ggml_add(ctx0, x, si);
+
+        ggml_tensor * acc = hift_resblock(up.res_0, x);
+        acc = ggml_add(ctx0, acc, hift_resblock(up.res_1, x));
+        acc = ggml_add(ctx0, acc, hift_resblock(up.res_2, x));
+        x = ggml_scale(ctx0, acc, 1.0f / 3.0f);
+    }
+
+    x = ggml_leaky_relu(ctx0, x, 0.01f, false);
+    x = cbx_conv1d_dil(c.hift_post_w, c.hift_post_b, x, 3, 1);
+    ggml_set_name(x, "out_spec");
+    ggml_set_output(x);
+    ggml_build_forward_expand(gf, x);
+    return gf;
+}

tools/mtmd/models/chatterbox-spkenc.cpp

@@ -0,0 +1,156 @@
+#include "models.h"
+
+// Chatterbox speaker encoder: CAMPPlus x-vector on kaldi fbank features, see s3gen/xvector.py
+
+static constexpr float bn_eps = 1e-5f;
+
+// per-channel batchnorm on x [C, T], w/b are null on the affine=False variant
+ggml_tensor * clip_graph_chatterbox_spkenc::bn1d(const clip_chatterbox::bn & n, ggml_tensor * x) {
+    ggml_tensor * sd = ggml_sqrt(ctx0, ggml_scale_bias(ctx0, n.var, 1.0f, bn_eps));
+    if (!n.w) {
+        return ggml_div(ctx0, ggml_sub(ctx0, x, n.mean), sd);
+    }
+    ggml_tensor * a     = ggml_div(ctx0, n.w, sd);
+    ggml_tensor * shift = ggml_sub(ctx0, n.b, ggml_mul(ctx0, n.mean, a));
+    return ggml_add(ctx0, ggml_mul(ctx0, x, a), shift);
+}
+
+// batchnorm on a conv2d activation [W=T, H=F, C, 1], stats on ne2
+static ggml_tensor * cbx_bn2d(ggml_context * ctx0, const clip_chatterbox::bn & n, ggml_tensor * x) {
+    const int C = (int) x->ne[2];
+    ggml_tensor * mean = ggml_reshape_4d(ctx0, n.mean, 1, 1, C, 1);
+    ggml_tensor * var  = ggml_reshape_4d(ctx0, n.var,  1, 1, C, 1);
+    ggml_tensor * w    = ggml_reshape_4d(ctx0, n.w, 1, 1, C, 1);
+    ggml_tensor * b    = ggml_reshape_4d(ctx0, n.b, 1, 1, C, 1);
+    ggml_tensor * a     = ggml_div(ctx0, w, ggml_sqrt(ctx0, ggml_scale_bias(ctx0, var, 1.0f, bn_eps)));
+    ggml_tensor * shift = ggml_sub(ctx0, b, ggml_mul(ctx0, mean, a));
+    return ggml_add(ctx0, ggml_mul(ctx0, x, a), shift);
+}
+
+ggml_tensor * clip_graph_chatterbox_spkenc::bn2d_relu(const clip_chatterbox::bn & n, ggml_tensor * x) {
+    return ggml_relu(ctx0, cbx_bn2d(ctx0, n, x));
+}
+
+// fcm residual 2d block, stride on the frequency axis only
+ggml_tensor * clip_graph_chatterbox_spkenc::res2d(const clip_chatterbox::spk_res2d & r, ggml_tensor * x,
+                                                  int stride) {
+    ggml_tensor * cur = ggml_conv_2d(ctx0, r.conv1_w, x, 1, stride, 1, 1, 1, 1);
+    cur = bn2d_relu(r.bn1, cur);
+    cur = ggml_conv_2d(ctx0, r.conv2_w, cur, 1, 1, 1, 1, 1, 1);
+    // bn2 without the relu, applied before the residual add
+    cur = cbx_bn2d(ctx0, r.bn2, cur);
+    ggml_tensor * res = x;
+    if (r.shortcut_w) {
+        res = ggml_conv_2d(ctx0, r.shortcut_w, x, 1, stride, 0, 0, 1, 1);
+        res = cbx_bn2d(ctx0, r.shortcut_bn, res);
+    }
+    return ggml_relu(ctx0, ggml_add(ctx0, cur, res));
+}
+
+// cam dense tdnn layer: bottleneck then context-gated conv; x [C_in, T] -> [growth, T]
+ggml_tensor * clip_graph_chatterbox_spkenc::cam_layer(const clip_chatterbox::spk_cam_layer & l, ggml_tensor * x,
+                                                      int dil, ggml_tensor * segfix) {
+    ggml_tensor * h = ggml_relu(ctx0, bn1d(l.nl1_bn, x));
+    h = cbx_conv1d(l.linear1_w, nullptr, h, 1, 0, 0);
+    h = ggml_relu(ctx0, bn1d(l.nl2_bn, h));
+
+    ggml_tensor * k = l.local_w;
+    const int pad = ((int) k->ne[0] - 1) / 2 * dil;
+    ggml_tensor * y = cbx_conv1d_dil(k, nullptr, h, pad, dil);
+
+    // context: global mean plus ceil-mode segment means of length 100
+    const int T = (int) h->ne[1];
+    const int C = (int) h->ne[0];
+    const int S = (T + 99) / 100;
+    ggml_tensor * ht = ggml_cont(ctx0, ggml_transpose(ctx0, h)); // [T, C]
+    ggml_tensor * gmean = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_mean(ctx0, ht))); // [C, 1]
+    ggml_tensor * seg;
+    {
+        ggml_tensor * padded = ht;
+        if (S * 100 != T) {
+            ggml_tensor * z = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, S * 100 - T, C);
+            z = ggml_scale(ctx0, z, 0.0f);
+            padded = ggml_concat(ctx0, ht, z, 0);
+        }
+        ggml_tensor * pooled = ggml_pool_1d(ctx0, padded, GGML_OP_POOL_AVG, 100, 100, 0); // [S, C]
+        pooled = ggml_mul(ctx0, pooled, segfix);
+        ggml_tensor * exp = ggml_interpolate(ctx0, pooled, S * 100, C, 1, 1, GGML_SCALE_MODE_NEAREST);
+        exp = ggml_cont(ctx0, ggml_view_2d(ctx0, exp, T, C, exp->nb[1], 0));
+        seg = ggml_cont(ctx0, ggml_transpose(ctx0, exp));        // [C, T]
+    }
+    ggml_tensor * context = ggml_add(ctx0, seg, gmean);
+    context = cbx_conv1d(l.ctx1_w, l.ctx1_b, context, 1, 0, 0);
+    context = ggml_relu(ctx0, context);
+    context = cbx_conv1d(l.ctx2_w, l.ctx2_b, context, 1, 0, 0);
+    ggml_tensor * m = ggml_sigmoid(ctx0, context);
+    return ggml_mul(ctx0, y, m);
+}
+
+ggml_cgraph * clip_graph_chatterbox_spkenc::build() {
+    const auto & c = model.cbx;
+
+    const int T  = img.nx();
+    const int T1 = (T - 1) / 2 + 1;
+    const int S  = (T1 + 99) / 100;
+    ggml_tensor * segfix = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, S);
+    ggml_set_name(segfix, "inp_segfix");
+    ggml_set_input(segfix);
+
+    // fbank features [T, 80] from the preprocessor
+    ggml_tensor * inp = build_inp_raw(1);
+
+    // fcm 2d front: [W=T, H=F=80, C=1] -> [T, 10, 32] -> [320, T]
+    ggml_tensor * x = ggml_reshape_4d(ctx0, inp, T, 80, 1, 1);
+    x = ggml_conv_2d(ctx0, c.spk_conv1_w, x, 1, 1, 1, 1, 1, 1);
+    x = bn2d_relu(c.spk_bn1, x);
+    x = res2d(c.spk_layer1_0, x, 2);
+    x = res2d(c.spk_layer1_1, x, 1);
+    x = res2d(c.spk_layer2_0, x, 2);
+    x = res2d(c.spk_layer2_1, x, 1);
+    x = ggml_conv_2d(ctx0, c.spk_conv2_w, x, 1, 2, 1, 1, 1, 1);
+    x = bn2d_relu(c.spk_bn2, x);
+    x = ggml_reshape_2d(ctx0, x, T, 320);
+    x = ggml_cont(ctx0, ggml_transpose(ctx0, x));               // [320, T]
+    cb(x, "spk_fcm", -1);
+
+    // tdnn k5 stride 2 over time, then the three cam dense blocks
+    x = cbx_conv1d(c.spk_tdnn_w, nullptr, x, 2, 2, 2); // [128, T1]
+    x = ggml_relu(ctx0, bn1d(c.spk_tdnn_bn, x));
+    cb(x, "spk_tdnn", -1);
+
+    static const int block_dil[3] = {1, 2, 2};
+    const clip_chatterbox::spk_cam_block * blocks[3] = { &c.spk_block1, &c.spk_block2, &c.spk_block3 };
+    for (int bi = 1; bi <= 3; bi++) {
+        const auto & blk = *blocks[bi - 1];
+        for (const auto & l : blk.layers) {
+            ggml_tensor * out = cam_layer(l, x, block_dil[bi - 1], segfix);
+            x = ggml_concat(ctx0, x, out, 0);
+        }
+        x = ggml_relu(ctx0, bn1d(blk.transit_bn, x));
+        x = cbx_conv1d(blk.transit_w, nullptr, x, 1, 0, 0);
+        cb(x, "spk_block", bi);
+    }
+    x = ggml_relu(ctx0, bn1d(c.spk_out_bn, x)); // [512, T1]
+
+    // statistics pooling: mean and unbiased std over time -> [1024, 1]
+    ggml_tensor * xt = ggml_cont(ctx0, ggml_transpose(ctx0, x));                 // [T1, 512]
+    ggml_tensor * mean = ggml_mean(ctx0, xt);                                    // [1, 512]
+    ggml_tensor * m2   = ggml_mean(ctx0, ggml_mul(ctx0, xt, xt));
+    ggml_tensor * var  = ggml_sub(ctx0, m2, ggml_mul(ctx0, mean, mean));
+    var = ggml_scale(ctx0, var, (float) T1 / (float) (T1 - 1));
+    ggml_tensor * sd = ggml_sqrt(ctx0, ggml_relu(ctx0, var));
+    ggml_tensor * stats = ggml_concat(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, mean)),
+                                            ggml_cont(ctx0, ggml_transpose(ctx0, sd)), 0); // [1024, 1]
+    cb(stats, "spk_stats_pool", -1);
+
+    // the x-vector stays raw, the gen graph normalizes it
+    ggml_tensor * dw = ggml_reshape_2d(ctx0, c.spk_dense_w, 1024, 192);
+    ggml_tensor * emb = ggml_mul_mat(ctx0, dw, stats);                           // [192, 1]
+    emb = bn1d(c.spk_dense_bn, emb);
+    emb = ggml_cont(ctx0, ggml_reshape_1d(ctx0, emb, 192));
+    cb(emb, "spk_embd", -1);
+    ggml_set_name(emb, "out_xvec");
+    ggml_set_output(emb);
+    ggml_build_forward_expand(gf, emb);
+    return gf;
+}

tools/mtmd/models/models.h

@@ -30,6 +30,22 @@ struct clip_graph_gemma4uv : clip_graph {
     ggml_cgraph * build() override;
 };
 
+struct clip_graph_nemotron_transcribe : clip_graph {
+    const clip_enc_state * state; // nullptr during warmup
+    clip_graph_nemotron_transcribe(clip_ctx * ctx, const clip_image_f32 & img, const clip_enc_state * state)
+        : clip_graph(ctx, img), state(state) {}
+    ggml_cgraph * build() override;
+    ggml_tensor * build_spk_encoder(ggml_tensor * cur, ggml_tensor * pos);
+
+    // one recording: encodes its windows in turn, carrying the speaker state between them
+    static bool encode_windows(clip_ctx * ctx, const clip_model & model, clip_encode_params * params);
+
+    // NeMo ParallelExpertEncoder online defaults, in output frames: 40 s windows with 4 s of context on each side
+    static constexpr int window_core  = 500;
+    static constexpr int window_left  = 50;
+    static constexpr int window_right = 50;
+};
+
 struct clip_graph_pixtral : clip_graph {
     clip_graph_pixtral(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
     ggml_cgraph * build() override;
@@ -190,6 +206,53 @@ struct clip_graph_conformer : clip_graph {
     ggml_cgraph * build() override;
 };
 
+// builders shared by the chatterbox gen and spkenc graphs, defined in chatterbox-gen.cpp
+struct clip_graph_chatterbox_base : clip_graph {
+    using clip_graph::clip_graph;
+
+  protected:
+    ggml_tensor * cbx_linear(ggml_tensor * w, ggml_tensor * b, ggml_tensor * x) const;
+    ggml_tensor * cbx_conv1d(ggml_tensor * k, ggml_tensor * b, ggml_tensor * x,
+                             int stride, int pad_l, int pad_r) const;
+    ggml_tensor * cbx_conv1d_dil(ggml_tensor * k, ggml_tensor * b, ggml_tensor * x,
+                                 int pad, int dil) const;
+};
+
+struct clip_graph_chatterbox_spkenc : clip_graph_chatterbox_base {
+    clip_graph_chatterbox_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph_chatterbox_base(ctx, img) {}
+    ggml_cgraph * build() override;
+
+  private:
+    ggml_tensor * bn1d(const clip_chatterbox::bn & n, ggml_tensor * x);
+    ggml_tensor * bn2d_relu(const clip_chatterbox::bn & n, ggml_tensor * x);
+    ggml_tensor * res2d(const clip_chatterbox::spk_res2d & r, ggml_tensor * x, int stride);
+    ggml_tensor * cam_layer(const clip_chatterbox::spk_cam_layer & l, ggml_tensor * x, int dil, ggml_tensor * segfix);
+};
+
+struct clip_graph_chatterbox : clip_graph_chatterbox_base {
+    clip_gen_process_type gen_process = CLIP_GEN_PROCESS_GEN_CODE;
+    int n_tokens = 0;
+    int n_prompt_mel  = 0;
+    int vocode_n_mel  = 0;
+    int vocode_n_stft = 0;
+    clip_graph_chatterbox(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, int n_tokens,
+                          int n_prompt_mel, int vocode_n_mel, int vocode_n_stft)
+        : clip_graph_chatterbox_base(ctx, img), gen_process(gen_process), n_tokens(n_tokens), n_prompt_mel(n_prompt_mel),
+          vocode_n_mel(vocode_n_mel), vocode_n_stft(vocode_n_stft) {}
+    ggml_cgraph * build() override;
+
+  private:
+    ggml_tensor * enc_layer(const clip_chatterbox::enc_layer & l, ggml_tensor * x, ggml_tensor * pos, int T);
+    ggml_tensor * causal_block(const clip_chatterbox::causal_block & b, ggml_tensor * x);
+    ggml_tensor * resnet(const clip_chatterbox::resnet & r, ggml_tensor * x, ggml_tensor * temb);
+    ggml_tensor * tfm_block(const clip_chatterbox::tfm_block & b, ggml_tensor * x);
+    ggml_tensor * estimator(ggml_tensor * x_noise, ggml_tensor * mu, ggml_tensor * spks,
+                            ggml_tensor * cond, ggml_tensor * temb, int T);
+    ggml_tensor * hift_resblock(const clip_chatterbox::hift_res & r, ggml_tensor * x);
+    ggml_cgraph * build_s3tok(int T);
+    ggml_cgraph * build_vocoder(int n_mel, int n_stft);
+};
+
 struct clip_graph_lfm3a : clip_graph {
     clip_graph_lfm3a(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
     ggml_cgraph * build() override;

tools/mtmd/models/nemotron-transcribe.cpp

@@ -0,0 +1,343 @@
+#include "models.h"
+
+#include <algorithm>
+#include <cmath>
+#include <limits>
+#include <numeric>
+
+// NeMo ParallelExpertEncoder: ASR and speaker (Sortformer) transformers, speaker activity fused into the ASR states
+
+// a speaker is active in an 80 ms frame where its prediction exceeds this (speaker_activity_threshold)
+static constexpr float SPK_ACTIVITY_THR = 0.5f;
+
+// NeMo streaming Sortformer: PEE online defaults and the checkpoint's sortformer_modules config
+static constexpr int   SPK_FIFO_LEN          = 40;
+static constexpr int   SPK_CACHE_LEN         = 188;
+static constexpr int   SPK_SIL_FRAMES        = 1;
+static constexpr float SPK_PRED_SCORE_THR    = 0.25f;
+static constexpr float SPK_BOOST_LATEST      = 0.05f;
+static constexpr float SPK_STRONG_BOOST_RATE = 0.75f;
+static constexpr float SPK_WEAK_BOOST_RATE   = 1.5f;
+static constexpr float SPK_MIN_POS_RATE      = 0.5f;
+
+// _compress_spkcache runs here on exp(score): its additive log terms become factors and -inf becomes 0
+static constexpr float SPK_SCORE_OFFSET = 2.0f; // -log(0.5) in _get_log_pred_scores
+static constexpr float SPK_STRONG_BOOST = 4.0f; // _boost_topk_scores, scale_factor 2
+static constexpr float SPK_WEAK_BOOST   = 2.0f; // _boost_topk_scores, scale_factor 1
+static constexpr float SPK_SPEECH_THR   = 0.5f; // _disable_low_scores: a speaker's frames at or below it are not cached
+
+// the speaker expert has its own layers, width and head count, which build_vit does not take; full rotary
+ggml_tensor * clip_graph_nemotron_transcribe::build_spk_encoder(ggml_tensor * cur, ggml_tensor * pos) {
+    const int     n_head = hparams.spk_n_head;
+    const int64_t n_embd = cur->ne[0];
+    const int64_t n_pos  = cur->ne[1];
+    const int64_t d_head = n_embd / n_head;
+
+    cur = build_norm(cur, model.spk_pre_ln_w, model.spk_pre_ln_b, NORM_TYPE_NORMAL, eps, -1);
+    for (int il = 0; il < (int) model.spk_layers.size(); il++) {
+        const auto & layer = model.spk_layers[il];
+        ggml_tensor * x = build_norm(cur, layer.ln_1_w, layer.ln_1_b, NORM_TYPE_NORMAL, eps, il);
+
+        ggml_tensor * qkv = build_mm(layer.qkv_w, x);
+        const size_t es = ggml_element_size(qkv);
+        ggml_tensor * q = ggml_view_3d(ctx0, qkv, d_head, n_head, n_pos, d_head * es, qkv->nb[1], 0);
+        ggml_tensor * k = ggml_view_3d(ctx0, qkv, d_head, n_head, n_pos, d_head * es, qkv->nb[1], n_embd * es);
+        ggml_tensor * v = ggml_view_3d(ctx0, qkv, d_head, n_head, n_pos, d_head * es, qkv->nb[1], 2 * n_embd * es);
+        q = ggml_rope_ext(ctx0, q, pos, nullptr, d_head, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+        k = ggml_rope_ext(ctx0, k, pos, nullptr, d_head, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+
+        x = build_attn(layer.o_w, layer.o_b, q, k, v, nullptr, 1.0f / sqrtf((float) d_head), il);
+        cur = ggml_add(ctx0, cur, x);
+
+        x = build_norm(cur, layer.ln_2_w, layer.ln_2_b, NORM_TYPE_NORMAL, eps, il);
+        x = build_ffn(x, layer.ff_up_w, layer.ff_up_b, nullptr, nullptr, layer.ff_down_w, layer.ff_down_b, FFN_GELU_ERF, il);
+        cur = ggml_add(ctx0, cur, x);
+    }
+
+    return build_norm(cur, model.spk_post_ln_w, model.spk_post_ln_b, NORM_TYPE_NORMAL, eps, -1);
+}
+
+ggml_cgraph * clip_graph_nemotron_transcribe::build() {
+    const int     sub          = hparams.subsampling_factor;
+    const int64_t n_mel        = hparams.n_mel_bins;
+    const int64_t n_mel_frames = img.nx();
+    const int64_t n_pos        = (n_mel_frames + sub - 1) / sub;
+    GGML_ASSERT(img.ny() == n_mel);
+
+    // warmup: the largest graph, with a full speaker cache and FIFO
+    const int64_t n_prev = state == nullptr ? SPK_CACHE_LEN + SPK_FIFO_LEN
+        : (int64_t) state->spk_prev.size() / hparams.spk_n_embd;
+
+    // [n_frames, n_mel] log-mel: raw for the speaker expert, normalized over the recording (encode_windows) for the ASR one
+    ggml_tensor * inp = build_inp_raw(1);
+    ggml_tensor * inp_norm = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_mel_frames, n_mel);
+    ggml_set_name(inp_norm, "inp_norm");
+    ggml_set_input(inp_norm);
+    auto stack_mel = [&](ggml_tensor * mel) {
+        return build_stack(ggml_cont(ctx0, ggml_transpose(ctx0, mel)), sub, n_mel);
+    };
+
+    // positions over [cache | FIFO | window]; the ASR expert takes the first n_pos
+    ggml_tensor * pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_prev + n_pos);
+    ggml_set_name(pos, "positions");
+    ggml_set_input(pos);
+
+    ggml_tensor * spk_prev = nullptr;
+    if (n_prev > 0) {
+        spk_prev = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.spk_n_embd, n_prev);
+        ggml_set_name(spk_prev, "spk_prev");
+        ggml_set_input(spk_prev);
+    }
+
+    ggml_tensor * asr_pos = ggml_view_1d(ctx0, pos, n_pos, 0);
+    auto add_pos = [&](ggml_tensor * cur, const clip_layer &) {
+        return ggml_rope_ext(ctx0, cur, asr_pos, nullptr, hparams.audio_rope_dims, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f,
+                             0.0f, 1.0f, 0.0f, 0.0f);
+    };
+    ggml_tensor * asr = build_vit(build_mm(model.inp_proj_w, stack_mel(inp_norm)), n_pos, NORM_TYPE_NORMAL, FFN_GELU_ERF, nullptr,
+                                  add_pos);
+    cb(asr, "asr_out", -1);
+
+    // speaker expert over [cache | FIFO | window]; the window's rows and predictions feed the next window's state
+    ggml_tensor * chunk = build_mm(model.spk_inp_proj_w, stack_mel(inp));
+    ggml_set_name(chunk, "spk_chunk");
+    ggml_set_output(chunk);
+    ggml_tensor * spk = spk_prev ? ggml_concat(ctx0, spk_prev, chunk, 1) : chunk;
+    const int64_t n_spk_pos = n_prev + n_pos;
+    spk = build_spk_encoder(spk, pos);
+    cb(spk, "spk_out", -1);
+
+    // Sortformer head, 10 ms frames
+    ggml_tensor * cur = ggml_add(ctx0, build_mm(model.spk_head_proj_w, spk), model.spk_head_proj_b);
+    const int64_t d_tf = cur->ne[0];
+    cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
+    cur = ggml_conv_1d(ctx0, model.spk_head_up_w, cur, 1, 1, 1);                            // [n_spk_pos, sub*d_tf]
+    cur = ggml_add(ctx0, cur, ggml_reshape_2d(ctx0, model.spk_head_up_b, 1, sub * d_tf));
+    cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));
+    cur = ggml_reshape_2d(ctx0, cur, d_tf, sub * n_spk_pos);
+    cur = ggml_add(ctx0, build_mm(model.spk_head_fc1_w, ggml_relu(ctx0, cur)), model.spk_head_fc1_b);
+    cur = ggml_add(ctx0, build_mm(model.spk_head_fc2_w, ggml_relu(ctx0, cur)), model.spk_head_fc2_b);
+    ggml_tensor * probs = ggml_sigmoid(ctx0, cur);                                          // [n_spk, sub*n_spk_pos]
+    cb(probs, "spk_probs", -1);
+
+    // 80 ms predictions
+    ggml_tensor * preds = ggml_pool_2d(ctx0, probs, GGML_OP_POOL_AVG, 1, sub, 1, sub, 0, 0);  // [n_spk, n_spk_pos]
+    ggml_set_name(preds, "spk_preds");
+    ggml_set_output(preds);
+
+    // fusion, on the window's rows
+    preds = ggml_view_2d(ctx0, preds, preds->ne[0], n_pos, preds->nb[1], n_prev * preds->nb[1]);
+    ggml_tensor * act = ggml_step(ctx0, ggml_scale_bias(ctx0, preds, 1.0f, -SPK_ACTIVITY_THR));
+    act = build_norm(act, model.fusion_spk_norm_w, model.fusion_spk_norm_b, NORM_TYPE_NORMAL, eps, -1);
+    cur = build_norm(asr, model.fusion_asr_norm_w, model.fusion_asr_norm_b, NORM_TYPE_NORMAL, eps, -1);
+    cur = ggml_add(ctx0, cur, build_mm(model.fusion_spk_kernel, act));
+    cb(cur, "fused", -1);
+
+    cur = ggml_add(ctx0, build_mm(model.mm_fc_w, cur), model.mm_fc_b);
+    cb(cur, "projected", -1);
+
+    ggml_build_forward_expand(gf, cur);
+
+    return gf;
+}
+
+// indices of the k largest of v[0, n), ties to the lower index
+static std::vector<int> spk_topk(const float * v, int n, int k) {
+    std::vector<int> idx(n);
+    std::iota(idx.begin(), idx.end(), 0);
+    k = std::clamp(k, 0, n);
+    std::partial_sort(idx.begin(), idx.begin() + k, idx.end(), [&](int a, int b) {
+        return v[a] > v[b] || (v[a] == v[b] && a < b);
+    });
+    idx.resize(k);
+    return idx;
+}
+
+// NeMo SortformerModules._compress_spkcache: keep the SPK_CACHE_LEN highest-scoring rows, speaker by speaker
+static void spk_compress_cache(std::vector<float> & emb, std::vector<float> & preds, int n, int d, int n_spk,
+                               const std::vector<float> & sil_emb) {
+    const int   per_spk  = std::max(0, SPK_CACHE_LEN / n_spk - SPK_SIL_FRAMES);
+    const int   strong   = (int) std::floor(per_spk * SPK_STRONG_BOOST_RATE);
+    const int   weak     = (int) std::floor(per_spk * SPK_WEAK_BOOST_RATE);
+    const int   min_pos  = (int) std::floor(per_spk * SPK_MIN_POS_RATE);
+    const float disabled = 0.0f;
+    const float latest   = std::exp(SPK_BOOST_LATEST);
+
+    // speaker-major: each speaker's n rows, then its silence rows (+inf)
+    const int nt = n + SPK_SIL_FRAMES;
+    std::vector<float> scores((size_t) n_spk * nt, std::numeric_limits<float>::infinity());
+    for (int s = 0; s < n_spk; s++) {
+        float * sc = scores.data() + (size_t) s * nt;
+        int n_pos = 0;
+        for (int t = 0; t < n; t++) {
+            const float * pt = preds.data() + (size_t) t * n_spk;
+            float prod = 1.0f;
+            for (int j = 0; j < n_spk; j++) {
+                if (j != s) {
+                    prod *= std::max(1.0f - pt[j], SPK_PRED_SCORE_THR);
+                }
+            }
+            prod *= std::max(pt[s], SPK_PRED_SCORE_THR);
+            sc[t] = pt[s] > SPK_SPEECH_THR ? prod * SPK_SCORE_OFFSET : disabled;
+            n_pos += sc[t] > 1.0f;
+        }
+        if (n_pos >= min_pos) {
+            for (int t = 0; t < n; t++) {
+                if (!(sc[t] > 1.0f)) {
+                    sc[t] = disabled;
+                }
+            }
+        }
+        for (int t = SPK_CACHE_LEN; t < n; t++) {
+            sc[t] *= latest;
+        }
+        for (int t : spk_topk(sc, n, strong)) {
+            sc[t] *= SPK_STRONG_BOOST;
+        }
+        for (int t : spk_topk(sc, n, weak)) {
+            sc[t] *= SPK_WEAK_BOOST;
+        }
+    }
+
+    std::vector<int> pick = spk_topk(scores.data(), (int) scores.size(), SPK_CACHE_LEN);
+    const int placeholder = (int) scores.size();
+    for (int & i : pick) {
+        if (scores[i] == disabled) {
+            i = placeholder;
+        }
+    }
+    std::sort(pick.begin(), pick.end());
+
+    std::vector<float> out_emb((size_t) SPK_CACHE_LEN * d);
+    std::vector<float> out_preds((size_t) SPK_CACHE_LEN * n_spk, 0.0f);
+    for (int i = 0; i < SPK_CACHE_LEN; i++) {
+        const int t = pick[i] == placeholder ? -1 : pick[i] % nt;
+        float * e = out_emb.data() + (size_t) i * d;
+        if (t < 0 || t >= n) {
+            std::copy(sil_emb.begin(), sil_emb.end(), e);
+        } else {
+            std::copy_n(emb.data() + (size_t) t * d, d, e);
+            std::copy_n(preds.data() + (size_t) t * n_spk, n_spk, out_preds.data() + (size_t) i * n_spk);
+        }
+    }
+    emb.swap(out_emb);
+    preds.swap(out_preds);
+}
+
+// speaker cache, with its predictions, and FIFO of one recording
+struct spk_stream {
+    std::vector<float> cache, cache_preds;
+    std::vector<float> fifo;
+};
+
+// every update has a full window, which overflows the FIFO and the cache (NeMo's update period of 300 never applies)
+static_assert(clip_graph_nemotron_transcribe::window_core - SPK_FIFO_LEN > SPK_CACHE_LEN);
+
+// NeMo SortformerModules.streaming_update: core rows join the FIFO, its overflow moves to the cache, which is compressed
+static void update_spk_state(spk_stream & ss, const clip_enc_state & st, const clip_model & model, int64_t core_begin,
+                             int64_t n_core) {
+    const int64_t d       = model.hparams.spk_n_embd;
+    const int64_t n_spk   = model.spk_head_fc2_w->ne[1];
+    const int64_t n_cache = (int64_t) ss.cache.size() / d;
+    const int64_t n_fifo  = (int64_t) ss.fifo.size()  / d;
+    const int64_t n_rows  = (int64_t) st.spk_chunk.size() / d;
+    GGML_ASSERT((n_cache + n_fifo + n_rows) * n_spk == (int64_t) st.spk_preds.size() && core_begin + n_core <= n_rows);
+    const float * preds      = st.spk_preds.data();
+    const float * core_preds = preds + (n_cache + n_fifo + core_begin) * n_spk;
+
+    std::vector<float> fifo_preds(preds + n_cache * n_spk, preds + (n_cache + n_fifo) * n_spk);
+    ss.fifo.insert(ss.fifo.end(), st.spk_chunk.begin() + core_begin * d, st.spk_chunk.begin() + (core_begin + n_core) * d);
+    fifo_preds.insert(fifo_preds.end(), core_preds, core_preds + n_core * n_spk);
+
+    const int64_t n_pop = n_fifo + n_core - SPK_FIFO_LEN;
+    ss.cache.insert(ss.cache.end(), ss.fifo.begin(), ss.fifo.begin() + n_pop * d);
+    ss.cache_preds.insert(ss.cache_preds.end(), fifo_preds.begin(), fifo_preds.begin() + n_pop * n_spk);
+    ss.fifo.erase(ss.fifo.begin(), ss.fifo.begin() + n_pop * d);
+    spk_compress_cache(ss.cache, ss.cache_preds, (int) (ss.cache.size() / d), (int) d, (int) n_spk, model.spk_sil_emb);
+}
+
+// NeMo ParallelExpertEncoder._forward_online: windows with left/right context, one graph each, the speaker state carried
+bool clip_graph_nemotron_transcribe::encode_windows(clip_ctx * ctx, const clip_model & model, clip_encode_params * params) {
+    GGML_ASSERT(params->imgs->entries.size() == 1 && params->out_embd != nullptr);
+    const clip_image_f32 & rec = params->imgs->entries[0];
+    const auto &  buf      = rec.get_ro_buf();
+    const int64_t n_frames = rec.nx();
+    const int64_t n_rows   = rec.ny();
+    const int64_t n_embd   = clip_n_mmproj_embd(ctx);
+    const int     sub      = model.hparams.subsampling_factor;
+    const int64_t core     = (int64_t) window_core  * sub;
+    const int64_t left     = (int64_t) window_left  * sub;
+    const int64_t right    = (int64_t) window_right * sub;
+    const int64_t n_win    = std::max<int64_t>(1, (n_frames + core - 1) / core);
+    GGML_ASSERT((int64_t) buf.size() == n_frames * n_rows);
+
+    auto & out = *params->out_embd;
+    GGML_ASSERT(out.size() == (size_t) n_embd * clip_n_output_tokens(ctx, &rec));
+
+    // the ASR expert's input: per-feature normalization over the whole recording (NeMo normalize_batch, per_feature)
+    std::vector<double> mean(n_rows), denom(n_rows);
+    for (int64_t r = 0; r < n_rows; r++) {
+        const float * row = buf.data() + r * n_frames;
+        double sum = 0.0;
+        double sq_diff_sum = 0.0;
+        for (int64_t t = 0; t < n_frames; t++) {
+            sum += (double) row[t];
+        }
+        mean[r] = sum / n_frames;
+        for (int64_t t = 0; t < n_frames; t++) {
+            const double diff = (double) row[t] - mean[r];
+            sq_diff_sum += diff * diff;
+        }
+        denom[r] = std::sqrt(sq_diff_sum / (n_frames - 1.0)) + 1e-5;
+    }
+
+    spk_stream     ss;
+    clip_enc_state st;
+    size_t n_out = 0;
+    for (int64_t w = 0; w < n_win; w++) {
+        const int64_t core_begin = w * core;
+        const int64_t core_end   = std::min(core_begin + core, n_frames);
+        const int64_t win_begin  = std::max<int64_t>(core_begin - left, 0);
+        const int64_t win_end    = std::min(core_end + right, n_frames);
+        const int64_t win_len    = win_end - win_begin;
+
+        std::vector<float> data((size_t) win_len * n_rows);
+        st.mel_norm.resize(data.size());
+        for (int64_t r = 0; r < n_rows; r++) {
+            const float * src = buf.data() + r * n_frames + win_begin;
+            std::copy_n(src, win_len, data.data() + r * win_len);
+            for (int64_t t = 0; t < win_len; t++) {
+                st.mel_norm[r * win_len + t] = (float) ((src[t] - mean[r]) / denom[r]);
+            }
+        }
+        clip_image_f32_batch batch;
+        batch.is_audio = true;
+        batch.entries.emplace_back();
+        batch.entries[0].set_size({(int) win_len, (int) n_rows}, false, true);
+        batch.entries[0].cpy_buf(data);
+
+        std::vector<float> embd((size_t) n_embd * clip_n_output_tokens(ctx, &batch.entries[0]));
+        clip_encode_params wp = *params;
+        wp.imgs      = &batch;
+        wp.out_embd  = &embd;
+        wp.enc_state = &st;
+        if (!clip_encode(ctx, &wp)) {
+            return false;
+        }
+
+        // output frames of the core, rounded up as NeMo's _asr_output_frame_boundary
+        const int64_t n_left = (core_begin - win_begin) / sub;
+        const int64_t n_core = (core_end - core_begin + sub - 1) / sub;
+        GGML_ASSERT((n_out + n_core) * n_embd <= out.size());
+        std::copy_n(embd.data() + n_left * n_embd, n_core * n_embd, out.data() + n_out * n_embd);
+        n_out += n_core;
+
+        if (w + 1 < n_win) {
+            update_spk_state(ss, st, model, n_left, n_core);
+            st.spk_prev = ss.cache;
+            st.spk_prev.insert(st.spk_prev.end(), ss.fifo.begin(), ss.fifo.end());
+        }
+    }
+    GGML_ASSERT(n_out * n_embd == out.size());
+    return true;
+}

tools/mtmd/mtmd-audio.cpp

@@ -1,6 +1,7 @@
 #include "mtmd-audio.h"
 
 #define _USE_MATH_DEFINES // for M_PI
+#include <cfloat>
 #include <cmath>
 #include <cstdint>
 #include <cstring>
@@ -970,6 +971,480 @@ bool mtmd_audio_preprocessor_qwen3tts_spk::preprocess(const float *
     return true;
 }
 
+// windowed frame of a centered, reflect padded signal: dst[i] = x[start + i] * win[i], zeros past the window
+static void mtmd_audio_reflect_frame(const float * x, int n, int start, const std::vector<float> & win, std::vector<float> & dst) {
+    std::fill(dst.begin(), dst.end(), 0.0f);
+    for (int i = 0; i < (int) win.size(); i++) {
+        int idx = start + i;
+        if (idx < 0) {
+            idx = -idx;
+        }
+        if (idx >= n) {
+            idx = 2 * n - 2 - idx;
+        }
+        dst[(size_t) i] = x[idx] * win[(size_t) i];
+    }
+}
+
+// mel energies of the frame in fft_in: power (or magnitude) spectrum through cache.filters
+static void mtmd_audio_mel_energies(const mtmd_audio_cache & cache, std::vector<float> & fft_in, std::vector<float> & fft_out,
+                                    bool magnitude, std::vector<double> & mel) {
+    const int n_fft  = (int) cache.filters.n_fft;
+    const int n_bins = n_fft / 2 + 1;
+    fft(cache, fft_in.data(), n_fft, fft_out.data());
+    for (int k = 0; k < n_bins; k++) {
+        const float p = fft_out[2 * k + 0] * fft_out[2 * k + 0] + fft_out[2 * k + 1] * fft_out[2 * k + 1];
+        fft_out[k] = magnitude ? sqrtf(p + 1e-9f) : p;
+    }
+    mel.assign(cache.filters.n_mel, 0.0);
+    for (int64_t m = 0; m < cache.filters.n_mel; m++) {
+        const float * w = cache.filters.data.data() + (size_t) m * n_bins;
+        for (int k = 0; k < n_bins; k++) {
+            mel[(size_t) m] += w[k] * fft_out[k];
+        }
+    }
+}
+
+// whisper style log-mel of the chatterbox s3 tokenizer (s3tokenizer.py), mel major
+static bool mtmd_audio_s3tok_log_mel(const mtmd_audio_cache & cache, const float * samples, size_t n_samples,
+                                     std::vector<float> & out, int & n_frames) {
+    const int n_fft = 400;
+    const int hop   = 160;
+    const int n_mel = (int) cache.filters.n_mel;
+    const int n     = (int) n_samples;
+
+    n_frames = n / hop;
+    if (n_frames <= 0) {
+        return false;
+    }
+
+    std::vector<float>  fft_in((size_t) n_fft * 2), fft_out((size_t) n_fft * 8);
+    std::vector<double> e;
+    std::vector<double> mel((size_t) n_mel * n_frames);
+    for (int fr = 0; fr < n_frames; fr++) {
+        mtmd_audio_reflect_frame(samples, n, fr * hop - n_fft / 2, cache.hann_window, fft_in);
+        mtmd_audio_mel_energies(cache, fft_in, fft_out, false, e);
+        for (int m = 0; m < n_mel; m++) {
+            mel[(size_t) m * n_frames + fr] = log10(std::max(e[(size_t) m], 1e-10));
+        }
+    }
+
+    const double mx = *std::max_element(mel.begin(), mel.end());
+    out.resize(mel.size());
+    for (size_t i = 0; i < mel.size(); i++) {
+        out[i] = (float) ((std::max(mel[i], mx - 8.0) + 4.0) / 4.0);
+    }
+    return true;
+}
+
+// 16 kHz -> 24 kHz upsampler: hann windowed sinc cut just under the source nyquist
+static void mtmd_audio_upsample_3_2(const float * samples, size_t n_samples, std::vector<float> & out) {
+    const int    W  = 16;    // sinc half width in source samples
+    const double fc = 0.495; // cutoff, normalized to the source rate
+
+    const size_t n_out = n_samples * 3 / 2;
+    out.assign(n_out, 0.0f);
+    for (size_t n = 0; n < n_out; n++) {
+        const double t  = (double) (2 * n) / 3.0;
+        const int    k0 = (int) floor(t) - W + 1;
+        double acc = 0.0;
+        for (int k = k0; k < k0 + 2 * W; k++) {
+            if (k < 0 || k >= (int) n_samples) {
+                continue;
+            }
+            const double x = t - k;
+            const double s = x == 0.0 ? 1.0 : sin(2.0 * M_PI * fc * x) / (M_PI * x);
+            acc += samples[k] * s * 0.5 * (1.0 + cos(M_PI * x / (W + 1)));
+        }
+        out[n] = (float) acc;
+    }
+}
+
+// matcha style log-mel of the s3gen prompt features (s3gen/utils/mel.py), frame major
+static bool mtmd_audio_matcha_log_mel(const mtmd_audio_cache & cache, const float * samples, size_t n_samples,
+                                      std::vector<float> & out, int & n_frames) {
+    const int n_fft = 1920;
+    const int hop   = 480;
+    const int pad   = (n_fft - hop) / 2;
+    const int n_mel = 80;
+    const int n     = (int) n_samples;
+
+    n_frames = 1 + (n + 2 * pad - n_fft) / hop;
+    if (n <= pad || n_frames <= 0) {
+        return false;
+    }
+
+    std::vector<float>  fft_in((size_t) n_fft * 2), fft_out((size_t) n_fft * 8);
+    std::vector<double> e;
+    out.resize((size_t) n_mel * n_frames);
+    for (int fr = 0; fr < n_frames; fr++) {
+        mtmd_audio_reflect_frame(samples, n, fr * hop - pad, cache.hann_window, fft_in);
+        mtmd_audio_mel_energies(cache, fft_in, fft_out, true, e);
+        for (int m = 0; m < n_mel; m++) {
+            out[(size_t) fr * n_mel + m] = logf(std::max((float) e[(size_t) m], 1e-5f));
+        }
+    }
+    return true;
+}
+
+// librosa.effects.trim: non-silent where the rms window is within top_db of the peak
+static void mtmd_audio_trim_silence(const float * samples, size_t n_samples, float top_db,
+                             size_t & start, size_t & end) {
+    const int win = 2048;
+    const int hop = 512;
+    const int n   = (int) n_samples;
+
+    const int n_fr = 1 + n / hop;
+    std::vector<double> rms((size_t) n_fr);
+    double mx = 0.0;
+    for (int fr = 0; fr < n_fr; fr++) {
+        double acc = 0.0;
+        for (int i = 0; i < win; i++) {
+            const int idx = fr * hop - win / 2 + i;
+            if (idx >= 0 && idx < n) {
+                acc += (double) samples[idx] * samples[idx];
+            }
+        }
+        rms[(size_t) fr] = sqrt(acc / win);
+        mx = std::max(mx, rms[(size_t) fr]);
+    }
+
+    const double thr = mx * pow(10.0, -top_db / 20.0);
+    int first = -1, last = -1;
+    for (int fr = 0; fr < n_fr; fr++) {
+        if (rms[(size_t) fr] > thr) {
+            if (first < 0) {
+                first = fr;
+            }
+            last = fr;
+        }
+    }
+    if (first < 0) {
+        start = end = 0;
+        return;
+    }
+    start = (size_t) first * hop;
+    end   = std::min((size_t) (last + 1) * hop, n_samples);
+}
+
+// power mel of the voice encoder front-end, frame major, no log
+static bool mtmd_audio_ve_mel(const mtmd_audio_cache & cache, const float * samples, size_t n_samples,
+                              std::vector<float> & out, int & n_frames) {
+    const int n_fft = 400;
+    const int hop   = 160;
+    const int n_mel = 40;
+    const int n     = (int) n_samples;
+
+    n_frames = 1 + n / hop;
+    if (n < 2) {
+        return false;
+    }
+
+    std::vector<float>  fft_in((size_t) n_fft * 2), fft_out((size_t) n_fft * 8);
+    std::vector<double> e;
+    out.resize((size_t) n_mel * n_frames);
+    for (int fr = 0; fr < n_frames; fr++) {
+        mtmd_audio_reflect_frame(samples, n, fr * hop - n_fft / 2, cache.hann_window, fft_in);
+        mtmd_audio_mel_energies(cache, fft_in, fft_out, false, e);
+        for (int m = 0; m < n_mel; m++) {
+            out[(size_t) fr * n_mel + m] = (float) e[(size_t) m];
+        }
+    }
+    return true;
+}
+
+// ITU-R BS.1770 integrated loudness of a mono signal, matching pyloudnorm
+static float mtmd_audio_lufs(const float * samples, size_t n_samples, int sample_rate) {
+    std::vector<double> y(samples, samples + n_samples);
+
+    auto biquad = [&](double b0, double b1, double b2, double a1, double a2) {
+        double x1 = 0.0, x2 = 0.0, y1 = 0.0, y2 = 0.0;
+        for (double & v : y) {
+            const double x0 = v;
+            v = b0 * x0 + b1 * x1 + b2 * x2 - a1 * y1 - a2 * y2;
+            x2 = x1; x1 = x0;
+            y2 = y1; y1 = v;
+        }
+    };
+
+    // stage 1: high shelf, f0 1681.9744509555319 Hz, +3.99984385397 dB, Q 0.7071752369554196
+    {
+        const double A     = pow(10.0, 3.99984385397 / 40.0);
+        const double w0    = 2.0 * M_PI * 1681.9744509555319 / sample_rate;
+        const double alpha = sin(w0) / (2.0 * 0.7071752369554196);
+        const double c     = cos(w0);
+        const double sq    = 2.0 * sqrt(A) * alpha;
+        const double b0 =      A * ((A + 1.0) + (A - 1.0) * c + sq);
+        const double b1 = -2.0 * A * ((A - 1.0) + (A + 1.0) * c);
+        const double b2 =      A * ((A + 1.0) + (A - 1.0) * c - sq);
+        const double a0 =           (A + 1.0) - (A - 1.0) * c + sq;
+        const double a1 =  2.0 *    ((A - 1.0) - (A + 1.0) * c);
+        const double a2 =           (A + 1.0) - (A - 1.0) * c - sq;
+        biquad(b0 / a0, b1 / a0, b2 / a0, a1 / a0, a2 / a0);
+    }
+    // stage 2: high pass, f0 38.13547087602444 Hz, Q 0.5003270373238773
+    {
+        const double w0    = 2.0 * M_PI * 38.13547087602444 / sample_rate;
+        const double alpha = sin(w0) / (2.0 * 0.5003270373238773);
+        const double c     = cos(w0);
+        const double b0 =  (1.0 + c) / 2.0;
+        const double b1 = -(1.0 + c);
+        const double b2 =  (1.0 + c) / 2.0;
+        const double a0 =   1.0 + alpha;
+        const double a1 =  -2.0 * c;
+        const double a2 =   1.0 - alpha;
+        biquad(b0 / a0, b1 / a0, b2 / a0, a1 / a0, a2 / a0);
+    }
+
+    const int win = (int) (0.4 * sample_rate);
+    const int hop = win / 4;
+    if ((int) n_samples < win) {
+        return -HUGE_VALF;
+    }
+    const int n_blocks = 1 + ((int) n_samples - win) / hop;
+    std::vector<double> z((size_t) n_blocks);
+    for (int b = 0; b < n_blocks; b++) {
+        double acc = 0.0;
+        for (int i = 0; i < win; i++) {
+            acc += y[(size_t) b * hop + i] * y[(size_t) b * hop + i];
+        }
+        z[(size_t) b] = acc / win;
+    }
+
+    auto gated_mean = [&](double thr_lufs) {
+        double acc = 0.0;
+        int    n   = 0;
+        for (const double v : z) {
+            if (-0.691 + 10.0 * log10(std::max(v, 1e-30)) > thr_lufs) {
+                acc += v;
+                n++;
+            }
+        }
+        return n > 0 ? acc / n : 0.0;
+    };
+
+    const double z_abs = gated_mean(-70.0);
+    if (z_abs <= 0.0) {
+        return -HUGE_VALF;
+    }
+    const double thr_rel = -0.691 + 10.0 * log10(z_abs) - 10.0;
+    const double z_rel   = gated_mean(thr_rel);
+    if (z_rel <= 0.0) {
+        return -HUGE_VALF;
+    }
+    return (float) (-0.691 + 10.0 * log10(z_rel));
+}
+
+//
+// mtmd_audio_preprocessor_chatterbox_ref
+//
+
+void mtmd_audio_preprocessor_chatterbox_ref::initialize() {
+    // kaldi fbank: povey window, kaldi mel banks over the 512-point spectrum, the nyquist bin left out
+    {
+        const int frame_len = 400;
+        const int n_fft     = 512;
+        const int n_bins    = n_fft / 2 + 1;
+        const int n_mel     = hparams.n_mel_bins;
+        const double sr     = (double) hparams.audio_sample_rate;
+
+        fbank_cache.fill_sin_cos_table(n_fft);
+        fbank_cache.hann_window.resize(frame_len);
+        for (int i = 0; i < frame_len; i++) {
+            fbank_cache.hann_window[(size_t) i] = (float) pow(0.5 - 0.5 * cos(2.0 * M_PI * i / (frame_len - 1)), 0.85);
+        }
+
+        auto mel = [](double f) { return 1127.0 * log(1.0 + f / 700.0); };
+        const double mel_lo = mel(20.0);
+        const double mel_hi = mel(sr / 2.0);
+        const double delta  = (mel_hi - mel_lo) / (n_mel + 1);
+        const double bin_hz = sr / n_fft;
+
+        auto & filters = fbank_cache.filters;
+        filters.n_mel = n_mel;
+        filters.n_fft = n_fft;
+        filters.data.assign((size_t) n_mel * n_bins, 0.0f);
+        for (int m = 0; m < n_mel; m++) {
+            const double left   = mel_lo + m * delta;
+            const double center = left + delta;
+            const double right  = center + delta;
+            for (int i = 0; i < n_bins - 1; i++) {
+                const double f = mel(bin_hz * i);
+                if (f > left && f < right) {
+                    const double w = f <= center ? (f - left) / (center - left)
+                                                 : (right - f) / (right - center);
+                    filters.data[(size_t) m * n_bins + i] = (float) w;
+                }
+            }
+        }
+    }
+
+    // s3 tokenizer: librosa slaney mel, 128 bands
+    s3tok_cache.fill_sin_cos_table(400);
+    s3tok_cache.fill_hann_window(400, true);
+    s3tok_cache.fill_mel_filterbank_matrix(128, 400, 16000);
+
+    // voice encoder: slaney mel, 40 bands up to 8 kHz
+    ve_cache.fill_sin_cos_table(400);
+    ve_cache.fill_hann_window(400, true);
+    ve_cache.fill_mel_filterbank_matrix(40, 400, 16000, 0.0f, 8000.0f);
+
+    // s3gen prompt features at the 24 kHz decoder rate
+    matcha_cache.fill_sin_cos_table(1920);
+    matcha_cache.fill_hann_window(1920, true);
+    matcha_cache.fill_mel_filterbank_matrix(80, 1920, 24000, 0.0f, 8000.0f);
+}
+
+bool mtmd_audio_preprocessor_chatterbox_ref::preprocess(const float *                 samples,
+                                                        size_t                        n_samples,
+                                                        std::vector<mtmd_audio_mel> & output) const {
+    output.clear();
+    const int sr = (int) hparams.audio_sample_rate;
+
+    // the turbo variant loudness-normalizes the whole clip before any feature
+    std::vector<float> pcm(samples, samples + n_samples);
+    if (!is_mtl) {
+        const float lufs = mtmd_audio_lufs(pcm.data(), pcm.size(), sr);
+        if (lufs != -HUGE_VALF) {
+            const float gain = powf(10.0f, (-27.0f - lufs) / 20.0f);
+            if (std::isfinite(gain) && gain > 0.0f) {
+                for (float & v : pcm) {
+                    v *= gain;
+                }
+            }
+        }
+    }
+
+    // entry 0: torchaudio kaldi.fbank of CAMPPlus (snip edges, dc removal, preemphasis 0.97), mean subtracted over time
+    {
+        const int frame_len = 400;
+        const int hop       = 160;
+        const int n_fft     = 512;
+        const int n_mel     = hparams.n_mel_bins;
+
+        if ((int) pcm.size() < frame_len) {
+            return false;
+        }
+        const int n_frames = 1 + ((int) pcm.size() - frame_len) / hop;
+
+        mtmd_audio_mel out;
+        out.n_len     = n_frames;
+        out.n_len_org = n_frames;
+        out.n_mel     = n_mel;
+        out.data.assign((size_t) n_mel * n_frames, 0.0f);
+
+        const std::vector<float> & win = fbank_cache.hann_window;
+        std::vector<float>  fft_in((size_t) n_fft * 2), fft_out((size_t) n_fft * 8);
+        std::vector<double> e;
+        for (int fr = 0; fr < n_frames; fr++) {
+            const float * x = pcm.data() + (size_t) fr * hop;
+
+            double mean = 0.0;
+            for (int i = 0; i < frame_len; i++) {
+                mean += x[i];
+            }
+            mean /= frame_len;
+
+            std::fill(fft_in.begin(), fft_in.end(), 0.0f);
+            fft_in[0] = (float) ((x[0] - mean) * (1.0 - 0.97) * win[0]);
+            for (int i = 1; i < frame_len; i++) {
+                fft_in[(size_t) i] = (float) (((x[i] - mean) - 0.97 * (x[i - 1] - mean)) * win[(size_t) i]);
+            }
+            mtmd_audio_mel_energies(fbank_cache, fft_in, fft_out, false, e);
+            for (int m = 0; m < n_mel; m++) {
+                out.data[(size_t) m * n_frames + fr] = (float) log(std::max(e[(size_t) m], (double) FLT_EPSILON));
+            }
+        }
+
+        for (int m = 0; m < n_mel; m++) {
+            float * row = out.data.data() + (size_t) m * n_frames;
+            double mean = 0.0;
+            for (int fr = 0; fr < n_frames; fr++) {
+                mean += row[fr];
+            }
+            mean /= n_frames;
+            for (int fr = 0; fr < n_frames; fr++) {
+                row[fr] -= (float) mean;
+            }
+        }
+
+        output.push_back(std::move(out));
+    }
+
+    // entries 1 and 2: s3 tokenizer log-mels at the flow and t3 caps, padded to whole 40 ms tokens
+    auto s3tok_entry = [&](size_t cap) -> bool {
+        std::vector<float> clip(pcm.begin(), pcm.begin() + std::min(pcm.size(), cap));
+        clip.resize((clip.size() + 639) / 640 * 640, 0.0f);
+        std::vector<float> mel;
+        int n_frames = 0;
+        if (!mtmd_audio_s3tok_log_mel(s3tok_cache, clip.data(), clip.size(), mel, n_frames)) {
+            return false;
+        }
+        mtmd_audio_mel out;
+        out.n_len     = n_frames;
+        out.n_len_org = n_frames;
+        out.n_mel     = s3tok_cache.filters.n_mel;
+        out.data      = std::move(mel);
+        output.push_back(std::move(out));
+        return true;
+    };
+    const size_t gen_cap = (size_t) 10 * sr;
+    const size_t t3_cap  = (size_t) (is_mtl ? 6 : 15) * sr;
+    if (!s3tok_entry(gen_cap) || !s3tok_entry(t3_cap)) {
+        return false;
+    }
+
+    // entry 3: s3gen prompt features of the flow-capped clip at 24 kHz
+    {
+        std::vector<float> clip(pcm.begin(), pcm.begin() + std::min(pcm.size(), gen_cap));
+        clip.resize((clip.size() + 639) / 640 * 640, 0.0f);
+        std::vector<float> pcm24;
+        mtmd_audio_upsample_3_2(clip.data(), clip.size(), pcm24);
+        std::vector<float> feat;
+        int n_frames = 0;
+        if (!mtmd_audio_matcha_log_mel(matcha_cache, pcm24.data(), pcm24.size(), feat, n_frames)) {
+            return false;
+        }
+        mtmd_audio_mel out;
+        out.n_len     = n_frames;
+        out.n_len_org = n_frames;
+        out.n_mel     = 80;
+        out.data      = std::move(feat);
+        output.push_back(std::move(out));
+    }
+
+    // entry 4: voice encoder mel of the silence-trimmed clip, padded to the 160-frame partial grid (rate 1.3)
+    {
+        size_t t0 = 0, t1 = 0;
+        mtmd_audio_trim_silence(pcm.data(), pcm.size(), 20.0f, t0, t1);
+        if (t1 <= t0) {
+            return false;
+        }
+        std::vector<float> mel;
+        int n_frames = 0;
+        if (!mtmd_audio_ve_mel(ve_cache, pcm.data() + t0, t1 - t0, mel, n_frames)) {
+            return false;
+        }
+        const int n_mel     = 40;
+        const int n_partial = 160;
+        const int step      = (int) lround((16000.0 / 1.3) / n_partial);
+        int n_wins = std::max(n_frames - n_partial + step, 0) / step;
+        const int rem = std::max(n_frames - n_partial + step, 0) % step;
+        if (n_wins == 0 || (double) (rem + n_partial - step) / n_partial >= 0.8) {
+            n_wins++;
+        }
+        const int target = n_partial + step * (n_wins - 1);
+        mel.resize((size_t) target * n_mel, 0.0f);
+        mtmd_audio_mel out;
+        out.n_len     = target;
+        out.n_len_org = target;
+        out.n_mel     = n_mel;
+        out.data      = std::move(mel);
+        output.push_back(std::move(out));
+    }
+    return true;
+}
+
 //
 // mtmd_audio_preprocessor_conformer
 //
@@ -1385,13 +1860,8 @@ void mtmd_audio_preprocessor_parakeet::initialize() {
     cache.hann_window = hparams.window;
 }
 
-bool mtmd_audio_preprocessor_parakeet::preprocess(const float * samples,
-                                                       size_t   n_samples_in,
-                                  std::vector<mtmd_audio_mel> & output) const {
-    if (n_samples_in == 0) {
-        return false;
-    }
-
+// log-mel without normalization
+void mtmd_audio_preprocessor_parakeet::log_mel(const float * samples, size_t n_samples_in, mtmd_audio_mel & out_full) const {
     filter_params params;
     params.n_mel            = hparams.n_mel_bins;
     params.n_fft_bins       = 1 + (hparams.audio_n_fft / 2);
@@ -1422,7 +1892,6 @@ bool mtmd_audio_preprocessor_parakeet::preprocess(const float * samples,
     std::vector<float> samples_padded(n_samples_in + 2 * pad, 0.0f);
     std::copy(samples_preprocessed.begin(), samples_preprocessed.end(), samples_padded.begin() + pad);
 
-    mtmd_audio_mel out_full;
     out_full.n_mel = params.n_mel;
     out_full.n_len = (samples_padded.size() - frame_size) / frame_step + 1;
     out_full.n_len_org = out_full.n_len;
@@ -1461,11 +1930,22 @@ bool mtmd_audio_preprocessor_parakeet::preprocess(const float * samples,
     for (int iw = 0; iw < n_threads - 1; ++iw) {
         workers[iw].join();
     }
+}
+
+bool mtmd_audio_preprocessor_parakeet::preprocess(const float * samples,
+                                                       size_t   n_samples_in,
+                                  std::vector<mtmd_audio_mel> & output) const {
+    if (n_samples_in == 0) {
+        return false;
+    }
+
+    mtmd_audio_mel out_full;
+    log_mel(samples, n_samples_in, out_full);
 
     // Per-feature normalization (only on valid frames)
     {
         const double eps = 1e-5;
-        int valid_frames = n_samples_in / frame_step;
+        int valid_frames = n_samples_in / hparams.audio_hop_len;
 
         for (int j = 0; j < out_full.n_mel; j++) {
             double sum = 0.0;
@@ -1497,6 +1977,32 @@ bool mtmd_audio_preprocessor_parakeet::preprocess(const float * samples,
     return true;
 }
 
+//
+// mtmd_audio_preprocessor_nemotron_transcribe implementation
+//
+
+bool mtmd_audio_preprocessor_nemotron_transcribe::preprocess(const float *                 samples,
+                                                             size_t                        n_samples,
+                                                             std::vector<mtmd_audio_mel> & output) const {
+    if (n_samples < (size_t) hparams.audio_n_fft) {
+        return false;
+    }
+
+    mtmd_audio_mel mel;
+    log_mel(samples, n_samples, mel);
+
+    // frames past n_samples / hop are padding (NeMo get_seq_len); kept raw, encode_windows normalizes them for ASR
+    const int64_t n_frames = n_samples / hparams.audio_hop_len;
+    for (int64_t j = 1; j < mel.n_mel; j++) {
+        std::copy_n(mel.data.begin() + j * mel.n_len, n_frames, mel.data.begin() + j * n_frames);
+    }
+    mel.n_len     = n_frames;
+    mel.n_len_org = n_frames;
+    mel.data.resize(mel.n_mel * n_frames);
+    output.push_back(std::move(mel));
+    return true;
+}
+
 
 // mtmd_audio_preprocessor_gemma4ua
 //

tools/mtmd/mtmd-audio.h

@@ -171,11 +171,34 @@ struct mtmd_audio_preprocessor_pockettts : mtmd_audio_preprocessor {
     bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) const override;
 };
 
+// chatterbox reference front-end: one clip gives five entries, used by position in mtmd.cpp
+//   0: CAMPPlus kaldi fbank            [80 x n_frames]  (mel major)
+//   1: s3 tokenizer log-mel, flow cap  [n_mel x n_frames] (mel major)
+//   2: s3 tokenizer log-mel, t3 cap    [n_mel x n_frames] (mel major)
+//   3: s3gen prompt features           [n_frames x 80]  (frame major, 24 kHz mel rate)
+//   4: voice encoder power mel         [n_frames x 40]  (frame major, trimmed, partial grid)
+// turbo loudness-normalizes the clip to -27 LUFS and caps the t3 clip at 15 s instead of 6 s
+struct mtmd_audio_preprocessor_chatterbox_ref : mtmd_audio_preprocessor {
+    mtmd_audio_preprocessor_chatterbox_ref(const clip_ctx * ctx, bool is_mtl) : mtmd_audio_preprocessor(ctx), is_mtl(is_mtl) {}
+    void initialize() override;
+    bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) const override;
+
+  private:
+    bool             is_mtl;
+    mtmd_audio_cache fbank_cache;  // entry 0, the window slot holds the povey window
+    mtmd_audio_cache s3tok_cache;  // entries 1 and 2
+    mtmd_audio_cache matcha_cache; // entry 3
+    mtmd_audio_cache ve_cache;     // entry 4
+};
+
 struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor {
     mtmd_audio_preprocessor_parakeet(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;
 
+  protected:
+    void log_mel(const float * samples, size_t n_samples, mtmd_audio_mel & out) const;
+
   private:
     mtmd_audio_cache cache;
 
@@ -186,6 +209,12 @@ struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor {
                               const mtmd_audio_cache & cache, mtmd_audio_mel & mel);
 };
 
+// one chunk per recording: the raw log-mel of its valid frames
+struct mtmd_audio_preprocessor_nemotron_transcribe : mtmd_audio_preprocessor_parakeet {
+    mtmd_audio_preprocessor_nemotron_transcribe(clip_ctx * ctx) : mtmd_audio_preprocessor_parakeet(ctx) { }
+    bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) const override;
+};
+
 //
 // streaming ISTFT - converts spectrogram frames back to audio one frame at a time
 //

tools/mtmd/mtmd-helper-gen.cpp

@@ -124,6 +124,39 @@ protected:
         return true;
     }
 
+    // encodes a speaker clip or a voice, out receives the encoder embeddings
+    bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
+        if (!mtmd_bitmap_is_voice(bitmap) && !mtmd_support_audio(mctx)) {
+            LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n");
+            return false;
+        }
+        const std::string  marker = mtmd_get_marker(mctx);
+        mtmd_input_text     text{ marker.c_str(), marker.size(), false, true };
+        mtmd_input_chunks * chunks = mtmd_input_chunks_init();
+        const mtmd_bitmap * bptr = bitmap;
+        bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
+        if (ok) {
+            ok = false;
+            for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
+                const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
+                if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
+                    continue;
+                }
+                if (mtmd_encode_chunk(mctx, chunk) != 0) {
+                    LOG_ERR("mtmd_helper_gen_audio: speaker encode failed\n");
+                    break;
+                }
+                const float * embd = mtmd_get_output_embd(mctx);
+                const size_t  n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
+                out.assign(embd, embd + n);
+                ok = true;
+                break;
+            }
+        }
+        mtmd_input_chunks_free(chunks);
+        return ok;
+    }
+
     llama_context * lctx;
     mtmd_context  * mctx;
     const llama_model * model;
@@ -428,39 +461,6 @@ private:
         return true;
     }
 
-    // runs the reference wav through the speaker encoder, returns one x-vector embedding row
-    bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
-        if (!mtmd_bitmap_is_voice(bitmap) && !mtmd_support_audio(mctx)) {
-            LOG_ERR("mtmd_helper_gen_audio: mmproj has no speaker/audio encoder\n");
-            return false;
-        }
-        const std::string  marker = mtmd_get_marker(mctx);
-        mtmd_input_text     text{ marker.c_str(), marker.size(), false, true };
-        mtmd_input_chunks * chunks = mtmd_input_chunks_init();
-        const mtmd_bitmap * bptr = bitmap;
-        bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
-        if (ok) {
-            ok = false;
-            for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
-                const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
-                if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
-                    continue;
-                }
-                if (mtmd_encode_chunk(mctx, chunk) != 0) {
-                    LOG_ERR("mtmd_helper_gen_audio: speaker encode failed\n");
-                    break;
-                }
-                const float * embd = mtmd_get_output_embd(mctx);
-                const size_t  n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
-                out.assign(embd, embd + n);
-                ok = true;
-                break;
-            }
-        }
-        mtmd_input_chunks_free(chunks);
-        return ok;
-    }
-
     // one GEN_WAV process() call over the buffered codes, state is carried across batches
     bool flush_gen_wav() {
         if (codes_buf.empty()) {
@@ -996,39 +996,6 @@ private:
         return n;
     }
 
-    // runs the reference wav through the mimi encoder, returns one row per 12.5Hz frame
-    bool encode_speaker(mtmd_bitmap * bitmap, std::vector<float> & out) {
-        if (!mtmd_bitmap_is_voice(bitmap) && !mtmd_support_audio(mctx)) {
-            LOG_ERR("mtmd_helper_gen_audio: mmproj has no voice encoder\n");
-            return false;
-        }
-        const std::string  marker = mtmd_get_marker(mctx);
-        mtmd_input_text     text{ marker.c_str(), marker.size(), false, true };
-        mtmd_input_chunks * chunks = mtmd_input_chunks_init();
-        const mtmd_bitmap * bptr = bitmap;
-        bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0;
-        if (ok) {
-            ok = false;
-            for (size_t i = 0; i < mtmd_input_chunks_size(chunks); i++) {
-                const mtmd_input_chunk * chunk = mtmd_input_chunks_get(chunks, i);
-                if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) {
-                    continue;
-                }
-                if (mtmd_encode_chunk(mctx, chunk) != 0) {
-                    LOG_ERR("mtmd_helper_gen_audio: voice encode failed\n");
-                    break;
-                }
-                const float * embd = mtmd_get_output_embd(mctx);
-                const size_t  n = (size_t) llama_model_n_embd_inp(model) * mtmd_input_chunk_get_n_tokens(chunk);
-                out.assign(embd, embd + n);
-                ok = true;
-                break;
-            }
-        }
-        mtmd_input_chunks_free(chunks);
-        return ok;
-    }
-
     // decodes the buffered latents, the mimi decoder state carries over between calls
     bool flush_gen_wav() {
         if (feats_buf.empty()) {
@@ -1718,6 +1685,398 @@ private:
     bool   wav_header_sent = false;
 };
 
+// Chatterbox: the backbone emits s3 speech tokens, then one flow-matching mel decode and the NSF-iSTFT vocoder make the PCM
+// the multilingual variant runs cfg with an uncond sequence (text rows zeroed), paired as seq_id and n_seq_max / 2 + seq_id
+class chatterbox_gen_audio_pipeline : public mtmd_gen_audio_pipeline {
+public:
+    using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline;
+
+    void reset() override {
+        seq_id = 0;
+        seq_uncond = 0;
+        pos = 0;
+        ar_idx = 0;
+        codes_buf.clear();
+        audio_pcm.clear();
+        h_state_buf.clear();
+        out_buf.clear();
+        prompt_embd_buf.clear();
+        uncond_embd_buf.clear();
+        cond_logits.clear();
+        n_prompt = 0;
+        prompt_pos = 0;
+        decoded = false;
+        stream = false;
+        pcm_sent = 0;
+        wav_header_sent = false;
+        gen_state.reset();
+    }
+
+    int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override {
+        reset();
+        seq_id = inp->seq_id;
+
+        if (!ensure_cache()) {
+            return 1;
+        }
+
+        out_type = inp->out_type;
+        stream   = inp->stream;
+        if (!init_gen_state(inp->seed)) {
+            return 1;
+        }
+
+        if (mtl) {
+            // the uncond sequence lives half a seq space away from the cond one
+            const int32_t n_seq_max = (int32_t) llama_n_seq_max(lctx);
+            seq_uncond = (llama_seq_id) (n_seq_max / 2 + seq_id);
+            if (n_seq_max < 2 || seq_id >= n_seq_max / 2) {
+                LOG_ERR("mtmd_helper_gen_audio: multilingual chatterbox needs a paired cfg sequence "
+                        "(n_seq_max = %d, seq_id = %d)\n", n_seq_max, seq_id);
+                return 1;
+            }
+        }
+        llama_memory_t mem = llama_get_memory(lctx);
+        if (!llama_memory_seq_rm(mem, seq_id, 0, -1) || (mtl && !llama_memory_seq_rm(mem, seq_uncond, 0, -1))) {
+            LOG_ERR("mtmd_helper_gen_audio: cannot reset chatterbox sequences\n");
+            return 1;
+        }
+
+        // the encode gives the talker conditioning rows, set_ref keeps the flow decoder reference
+        if (!inp->speaker_ref) {
+            LOG_ERR("mtmd_helper_gen_audio: chatterbox requires a speaker reference or a voice\n");
+            return 1;
+        }
+        if (!encode_speaker(inp->speaker_ref, prompt_embd_buf) ||
+            mtmd_gen_state_set_ref(mctx, gen_state.get()) != 0) {
+            LOG_ERR("mtmd_helper_gen_audio: speaker reference encoding failed\n");
+            return 1;
+        }
+
+        const int n_e = n_embd;
+        auto push_row = [&](const float * src) {
+            prompt_embd_buf.insert(prompt_embd_buf.end(), src, src + n_e);
+        };
+        auto add_pos = [&](float * r, const std::vector<float> & tab, int idx) {
+            const float * p = tab.data() + (size_t) idx * n_e;
+            for (int j = 0; j < n_e; j++) {
+                r[j] += p[j];
+            }
+        };
+
+        // prompt rows: [conditioning, text, speech bos]
+        // no text preprocessing here, the [SPACE] substitution and the case folding are baked in at conversion
+        std::string txt(inp->prompt, inp->prompt_len);
+        if (mtl && inp->lang && inp->lang[0]) {
+            // the reference prefixes the language tag, e.g. "[fr]"
+            std::string lang = inp->lang;
+            std::transform(lang.begin(), lang.end(), lang.begin(), [](unsigned char c) { return (char) std::tolower(c); });
+            txt = "[" + lang + "]" + txt;
+        }
+        std::vector<llama_token> ids(txt.size() + 16);
+        int n_ids = llama_tokenize(vocab, txt.c_str(), (int32_t) txt.size(), ids.data(), (int32_t) ids.size(),
+                                   false, true);
+        if (n_ids < 1) {
+            LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n");
+            return 1;
+        }
+        ids.resize((size_t) n_ids);
+        if (mtl) {
+            // only the multilingual prompt wraps the text in start/stop tokens
+            ids.insert(ids.begin(), text_start);
+            ids.push_back(text_stop);
+            if (ids.size() > text_pos.size() / (size_t) n_e) {
+                LOG_ERR("mtmd_helper_gen_audio: prompt too long (%zu tokens, max %zu)\n", ids.size(), text_pos.size() / (size_t) n_e);
+                return 1;
+            }
+        }
+        const size_t text0 = prompt_embd_buf.size() / (size_t) n_e;
+        for (size_t i = 0; i < ids.size(); i++) {
+            push_row(tok_embd.data() + (size_t) ids[i] * n_e);
+            if (mtl) {
+                add_pos(prompt_embd_buf.data() + (text0 + i) * n_e, text_pos, (int) i);
+            }
+        }
+
+        // speech bos opens the AR stream
+        push_row(tok_embd.data() + (size_t) llama_vocab_bos(vocab) * n_e);
+        if (mtl) {
+            add_pos(prompt_embd_buf.data() + prompt_embd_buf.size() - n_e, speech_pos, 0);
+        }
+        ar_idx = 1;
+        n_prompt = (int) (prompt_embd_buf.size() / (size_t) n_e);
+
+        if (mtl) {
+            // cfg uncond prompt: text rows are zeroed but keep their learned positions, as in the reference
+            uncond_embd_buf = prompt_embd_buf;
+            for (size_t i = 0; i < ids.size(); i++) {
+                float * r = uncond_embd_buf.data() + (text0 + i) * n_e;
+                std::fill(r, r + n_e, 0.0f);
+                add_pos(r, text_pos, (int) i);
+            }
+        }
+
+        prompt_pos = 0;
+        return 0;
+    }
+
+    int32_t step_prompt(int32_t n_batch) override {
+        GGML_ASSERT(n_batch > 0);
+        // the uncond prompt follows the cond one on mtl
+        const int n_total = mtl ? 2 * n_prompt : n_prompt;
+        if (prompt_pos >= n_total) {
+            return 0;
+        }
+        const bool  uncond = prompt_pos >= n_prompt;
+        const int   p0     = uncond ? prompt_pos - n_prompt : prompt_pos;
+        const int32_t n_tokens_batch = std::min(n_batch, n_prompt - p0);
+        const bool  last   = p0 + n_tokens_batch == n_prompt;
+
+        decode_embd_batch batch((uncond ? uncond_embd_buf : prompt_embd_buf).data(), n_prompt, 1, n_embd);
+        batch.set_position_normal(0, uncond ? seq_uncond : seq_id);
+        if (last) {
+            batch.logits[n_prompt - 1] = 1;
+        }
+        if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.render(lctx, p0, n_tokens_batch)) != 0) {
+            LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n");
+            return -1;
+        }
+        prompt_pos += n_tokens_batch;
+
+        if (last && mtl) {
+            if (!uncond) {
+                cfg_read();
+            } else {
+                cfg_apply();
+            }
+        }
+
+        if (prompt_pos >= n_total) {
+            pos = n_prompt;
+            prompt_embd_buf.clear();
+            uncond_embd_buf.clear();
+            return 0;
+        }
+        return n_total - prompt_pos;
+    }
+
+    int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override {
+        GGML_UNUSED(h_state_in);
+        if (sampled < 0 || sampled >= llama_vocab_n_tokens(vocab)) {
+            LOG_ERR("mtmd_helper_gen_audio: chatterbox requires a token sampled from the backbone\n");
+            return 1;
+        }
+
+        // the speech stop token is the backbone eos, the multilingual speech positions end the stream too
+        if (llama_vocab_is_eog(vocab, sampled) || (mtl && (size_t) ar_idx >= speech_pos.size() / (size_t) n_embd)) {
+            *out_stop    = true;
+            *h_state_out = nullptr;
+            return 0;
+        }
+
+        // keep only the 6561 s3gen codes, dropping start/stop and oov ids
+        if (sampled >= speech_base && sampled - speech_base < 6561) {
+            codes_buf.push_back(sampled - speech_base);
+        }
+
+        std::vector<float> e(tok_embd.begin() + (size_t) sampled * n_embd,
+                             tok_embd.begin() + (size_t) (sampled + 1) * n_embd);
+        if (mtl) {
+            // the multilingual backbone takes the token row plus the learned speech position
+            const float * p = speech_pos.data() + (size_t) ar_idx * n_embd;
+            for (int j = 0; j < n_embd; j++) {
+                e[(size_t) j] += p[j];
+            }
+            ar_idx++;
+        }
+
+        if (!decode_one(e.data(), seq_id)) {
+            return 1;
+        }
+        if (mtl) {
+            // cfg: the sampled token feeds both branches
+            cfg_read();
+            if (!decode_one(e.data(), seq_uncond)) {
+                return 1;
+            }
+            cfg_apply();
+        }
+        pos++;
+
+        const float * h = llama_get_embeddings_ith(lctx, -1);
+        h_state_buf.assign(h, h + n_embd);
+        *h_state_out = h_state_buf.data();
+        return 0;
+    }
+
+    int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override {
+        *out_sample_rate = info.sample_rate;
+
+        if (!stream) {
+            if (!flush_gen_wav()) {
+                return 1;
+            }
+            if (out_n_samples) {
+                *out_n_samples = (int64_t) audio_pcm.size();
+            }
+            if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
+                *out_data     = (const char *) audio_pcm.data();
+                *out_data_len = audio_pcm.size() * sizeof(float);
+                return 0;
+            }
+            out_buf.clear();
+            if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) {
+                LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n");
+                return 1;
+            }
+            *out_data     = out_buf.data();
+            *out_data_len = out_buf.size();
+            return 0;
+        }
+
+        // streaming: the flow decode is one-shot, so all the audio arrives after flush()
+        const size_t n_new = audio_pcm.size() - pcm_sent;
+        if (out_n_samples) {
+            *out_n_samples = (int64_t) n_new;
+        }
+        if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) {
+            *out_data     = (const char *) (audio_pcm.data() + pcm_sent);
+            *out_data_len = n_new * sizeof(float);
+            pcm_sent = audio_pcm.size();
+            return 0;
+        }
+        out_buf.clear();
+        if (!wav_header_sent) {
+            write_wav16_header(out_buf, UINT32_MAX, info.sample_rate);
+            wav_header_sent = true;
+        }
+        append_wav16_pcm(out_buf, audio_pcm.data() + pcm_sent, n_new);
+        pcm_sent = audio_pcm.size();
+        *out_data     = out_buf.data();
+        *out_data_len = out_buf.size();
+        return 0;
+    }
+
+    int32_t flush() override {
+        return flush_gen_wav() ? 0 : 1;
+    }
+
+private:
+    bool decode_one(const float * e, llama_seq_id seq) {
+        decode_embd_batch batch(e, 1, 1, n_embd);
+        batch.set_position_normal(pos, seq);
+        batch.logits[0] = 1;
+        if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.render(lctx, 0, 1)) != 0) {
+            LOG_ERR("mtmd_helper_gen_audio: step decode failed\n");
+            return false;
+        }
+        return true;
+    }
+
+    // cfg of the reference: logits = cond + 0.5 * (cond - uncond)
+    // the mix goes into the last output row (the uncond one), the caller samples from it
+    void cfg_read() {
+        const float * c = llama_get_logits_ith(lctx, -1);
+        cond_logits.assign(c, c + llama_vocab_n_tokens(vocab));
+    }
+    void cfg_apply() {
+        float * u = llama_get_logits_ith(lctx, -1);
+        for (size_t i = 0; i < cond_logits.size(); i++) {
+            u[i] = cond_logits[i] + 0.5f * (cond_logits[i] - u[i]);
+        }
+    }
+
+    // one GEN_WAV call over the whole utterance, the flow decode is not streamed
+    bool flush_gen_wav() {
+        if (decoded || codes_buf.empty()) {
+            return true;
+        }
+        if (!mtl) {
+            // turbo appends a short silence tail before vocoding
+            codes_buf.insert(codes_buf.end(), 3, 4299);
+        }
+        mtmd_gen_inp inp = mtmd_gen_inp_default(mctx);
+        inp.type    = MTMD_GEN_PROCESS_TYPE_GEN_WAV;
+        inp.codes   = codes_buf.data();
+        inp.n_codes = codes_buf.size();
+        mtmd_gen_out out{};
+        if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) {
+            LOG_ERR("mtmd_helper_gen_audio: gen_wav decode failed\n");
+            return false;
+        }
+        audio_pcm.assign(out.audio, out.audio + out.n_samples);
+        decoded = true;
+        return true;
+    }
+
+    bool ensure_cache() {
+        if (!tok_embd.empty()) {
+            return true;
+        }
+        // fused vocab layout: [text 0..speech_base) then the speech tokens
+        speech_base = find_special_token(vocab, "<|speech_0|>");
+        if (speech_base == LLAMA_TOKEN_NULL) {
+            LOG_ERR("mtmd_helper_gen_audio: fused speech tokens not found in vocab\n");
+            return false;
+        }
+
+        const uint32_t n_tok_embd = llama_model_get_tok_embd(model, nullptr);
+        if (n_tok_embd != (uint32_t) llama_vocab_n_tokens(vocab) * (uint32_t) n_embd) {
+            LOG_ERR("mtmd_helper_gen_audio: unexpected token embedding size\n");
+            return false;
+        }
+        tok_embd.resize(n_tok_embd);
+        llama_model_get_tok_embd(model, tok_embd.data());
+
+        auto read_table = [&](const char * name, std::vector<float> & dst) {
+            const size_t n = mtmd_gen_audio_read_tensor(mctx, name, nullptr, 0);
+            dst.resize(n);
+            if (n == 0 || mtmd_gen_audio_read_tensor(mctx, name, dst.data(), n) != n) {
+                LOG_ERR("mtmd_helper_gen_audio: %s read failed\n", name);
+                return false;
+            }
+            return true;
+        };
+
+        // multilingual: the learned position tables
+        mtl = mtmd_gen_audio_read_tensor(mctx, "a.gen.t3.text_pos_emb", nullptr, 0) > 0;
+        if (mtl && (!read_table("a.gen.t3.text_pos_emb", text_pos) || !read_table("a.gen.t3.speech_pos_emb", speech_pos))) {
+            return false;
+        }
+        return true;
+    }
+
+    // vocab and mmproj data fixed across the whole session, read once
+    std::vector<float> tok_embd;
+    bool mtl = false;
+    std::vector<float> text_pos;
+    std::vector<float> speech_pos;
+    llama_token speech_base = LLAMA_TOKEN_NULL;
+    // start_text_token and stop_text_token of the reference, only the multilingual prompt uses them
+    static constexpr llama_token text_start = 255;
+    static constexpr llama_token text_stop  = 0;
+
+    // per-generation state, cleared by reset()
+    llama_seq_id seq_id     = 0;
+    llama_seq_id seq_uncond = 0;
+    llama_pos pos = 0;
+    int ar_idx = 0;
+    std::vector<float> prompt_embd_buf;
+    std::vector<float> uncond_embd_buf;
+    int n_prompt   = 0;
+    int prompt_pos = 0;
+    std::vector<float> cond_logits;
+    std::vector<int32_t> codes_buf;
+    bool decoded = false;
+    std::vector<float> audio_pcm;
+    std::vector<float> h_state_buf;
+    mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV;
+    std::vector<char> out_buf;
+    bool   stream          = false;
+    size_t pcm_sent        = 0; // samples already returned by get_output()
+    bool   wav_header_sent = false;
+};
+
 static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lctx, mtmd_context * mctx) {
     switch (mtmd_gen_audio_get_info(mctx).type) {
         case MTMD_GEN_AUDIO_TYPE_QWEN3TTS:
@@ -1728,6 +2087,8 @@ static std::unique_ptr<mtmd_gen_audio_pipeline> make_pipeline(llama_context * lc
             return std::unique_ptr<mtmd_gen_audio_pipeline>(new soprano_gen_audio_pipeline(lctx, mctx));
         case MTMD_GEN_AUDIO_TYPE_LFM3A:
             return std::unique_ptr<mtmd_gen_audio_pipeline>(new lfm3a_gen_audio_pipeline(lctx, mctx));
+        case MTMD_GEN_AUDIO_TYPE_CHATTERBOX:
+            return std::unique_ptr<mtmd_gen_audio_pipeline>(new chatterbox_gen_audio_pipeline(lctx, mctx));
         default:
             return nullptr;
     }

tools/mtmd/mtmd-helper.cpp

@@ -508,8 +508,6 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(const mtmd_context
 // Video input helpers
 //
 
-#ifdef MTMD_VIDEO
-
 // MSVC deprecates strdup(), use a std::string version directly
 static char * mtmd_strdup(const std::string & s) {
     size_t n = s.size() + 1;
@@ -520,15 +518,218 @@ static char * mtmd_strdup(const std::string & s) {
     return p;
 }
 
-struct mtmd_helper_video {
+mtmd_helper_video_seq::~mtmd_helper_video_seq() {
+    mtmd_bitmap_free(pending_frame);
+    mtmd_bitmap_free(pending_audio);
+}
+
+int32_t mtmd_helper_video_seq::ts_per_n_frames() const {
+    return tmpl == MTMD_VIDEO_TEMPLATE_LFM3 || tmpl == MTMD_VIDEO_TEMPLATE_LFM3_OMNI || tmpl == MTMD_VIDEO_TEMPLATE_QWEN3VL ? 2 : 0;
+}
+
+void mtmd_helper_video_seq::configure(const mtmd_video_config & cfg, const mtmd_helper_video_src & src, float fps_target_arg) {
+    tmpl = cfg.tmpl;
+    uint32_t width  = src.width;
+    uint32_t height = src.height;
+    float duration  = src.duration;
+    if (duration < 0.0f && src.n_frames > 0) {
+        duration = (float)src.n_frames / src.fps;
+    }
+
+    // the model's recommended sampling rate wins over the generic default
+    fps_target = cfg.fps > 0.0f ? cfg.fps
+               : fps_target_arg > 0.0f ? fps_target_arg : src.fps;
+
+    if (cfg.n_frames > 0 && duration > 0.0f) {
+        const int32_t n = cfg.n_frames;
+        target_frames = n;
+        fps_target    = (float) n / duration;
+
+        // same frames as the reference processors: round(linspace(0, n_frames - 1, n))
+        const double step = n > 1 ? (double) (src.n_frames - 1) / (n - 1) : 0.0;
+        frame_indices.resize(n);
+        for (int32_t i = 0; i < n; i++) {
+            frame_indices[i] = (int32_t) std::nearbyint(i * step);
+        }
+        src_fps      = src.fps;
+        n_src_frames = src.n_frames;
+
+        if (cfg.frame_width > 0 && cfg.frame_height > 0 &&
+                (cfg.frame_width != (int32_t) width || cfg.frame_height != (int32_t) height)) {
+            out_width  = cfg.frame_width;
+            out_height = cfg.frame_height;
+        }
+    }
+
+    if (out_width > 0 && out_height > 0) {
+        width  = (uint32_t) out_width;
+        height = (uint32_t) out_height;
+    }
+
+    info.width    = width;
+    info.height   = height;
+    info.fps      = fps_target;
+    LOG_DBG("%s: %ux%u fps=%.2f duration=%.2fs n_frames=%d target_frames=%d\n",
+            __func__, width, height, fps_target, duration, src.n_frames, target_frames);
+    info.n_frames = target_frames > 0 ? target_frames
+                  : duration > 0.0f ? (int32_t)(duration * fps_target + 0.5f) : -1;
+    frame_buf.resize((size_t)width * height * 3);
+}
+
+// the soundtrack of the group of frames that ends with frame idx, nullptr if the soundtrack has ended
+mtmd_bitmap * mtmd_helper_video_seq::soundtrack_slice(int32_t idx) const {
+    const int32_t ts_per_n_frames = this->ts_per_n_frames();
+    const int32_t group    = idx / ts_per_n_frames;
+    const int32_t n_groups = ((int32_t) frame_indices.size() + ts_per_n_frames - 1) / ts_per_n_frames;
+    auto group_start = [&](int32_t g) {
+        const double t = frame_indices[g * ts_per_n_frames] / src_fps;
+        return std::min(soundtrack.size(), (size_t) std::nearbyint(t * audio_rate));
+    };
+    const size_t begin = group == 0 ? 0 : group_start(group);
+    const size_t end   = group == n_groups - 1 ? soundtrack.size() : group_start(group + 1);
+    if (begin >= end) {
+        return nullptr;
+    }
+    mtmd_bitmap * slice = mtmd_bitmap_init_from_audio(end - begin, soundtrack.data() + begin);
+    if (!id.empty()) {
+        mtmd_bitmap_set_id(slice, (id + "+audio" + std::to_string(group)).c_str());
+    }
+    return slice;
+}
+
+mtmd_bitmap * mtmd_helper_video_seq::read_next_frame() {
+    if (!src_alive) return nullptr;
+
+    if (!frame_indices.empty()) {
+        if (next_index >= frame_indices.size()) {
+            return nullptr;
+        }
+        // skip to the wanted source frame; a repeated index reuses the frame in frame_buf
+        const int32_t want = frame_indices[next_index];
+        while (src_frame <= want) {
+            if (!read_raw_frame()) {
+                src_alive = false;
+                return nullptr;
+            }
+            src_frame++;
+        }
+        next_index++;
+    } else if (!read_raw_frame()) {
+        src_alive = false;
+        return nullptr;
+    }
+
+    LOG_DBG("%s: frame %d read OK\n", __func__, current_frame);
+    mtmd_bitmap * frame = mtmd_bitmap_init(info.width, info.height, frame_buf.data());
+    mtmd_bitmap_set_mergeable(frame, true);
+    mtmd_bitmap_set_is_video(frame, true);
+    if (!id.empty()) {
+        // each frame gets a unique id in the form of {hash}+{frame}, so that it can be identified in cache
+        std::string frame_id = id + "+" + std::to_string(current_frame);
+        mtmd_bitmap_set_id(frame, frame_id.c_str());
+    }
+    current_frame++;
+    return frame;
+}
+
+int32_t mtmd_helper_video_seq::read_next(mtmd_bitmap ** out_bitmap, char ** out_text) {
+    *out_bitmap = nullptr;
+    *out_text   = nullptr;
+
+    if (!pending_text.empty()) {
+        *out_text = mtmd_strdup(pending_text);
+        pending_text.clear();
+        return *out_text ? 0 : -2;
+    }
+
+    if (pending_frame) {
+        *out_bitmap  = pending_frame;
+        pending_frame = nullptr;
+        return 0;
+    }
+
+    if (pending_audio) {
+        *out_bitmap   = pending_audio;
+        pending_audio = nullptr;
+        return 0;
+    }
+
+    LOG_DBG("%s: proc_alive=%d start_emitted=%d current_frame=%d\n",
+            __func__, (int)src_alive, (int)start_emitted, current_frame);
+
+    if (!src_alive) {
+        return (current_frame == 0) ? -2 : -1;
+    }
+
+    if (!start_emitted) {
+        start_emitted = true;
+        // model templates have no prompt prefix, the video begin/end tokens are added by mtmd_tokenize
+        if (tmpl == MTMD_VIDEO_TEMPLATE_DEFAULT && !prompt_start.empty()) {
+            *out_text = mtmd_strdup(prompt_start);
+            return *out_text ? 0 : -2;
+        }
+    }
+
+    mtmd_bitmap * frame = read_next_frame();
+    if (!frame) {
+        return -1;
+    }
+
+    const int32_t ts_per_n_frames = this->ts_per_n_frames();
+    if (ts_per_n_frames > 0) {
+        // emit the group timestamp before the group's first frame; the value is the
+        // midpoint time of the group, matching HF processors that stamp each temporal group
+        const int32_t idx = current_frame - 1; // read_next_frame() already incremented
+        // the soundtrack of a group follows its last frame
+        const bool group_ends = idx % ts_per_n_frames == ts_per_n_frames - 1 || idx + 1 == (int32_t) frame_indices.size();
+        if (!soundtrack.empty() && !frame_indices.empty() && group_ends) {
+            pending_audio = soundtrack_slice(idx);
+        }
+        if (idx % ts_per_n_frames == 0) {
+            char ts_buf[64];
+            double ts_s = ((double)idx + 0.5*(double)(ts_per_n_frames - 1)) / info.fps;
+            if (!frame_indices.empty()) {
+                // stamp from the source frames; a short last group repeats its last frame
+                const size_t first = idx;
+                const size_t last  = std::min(first + ts_per_n_frames - 1, frame_indices.size() - 1);
+                ts_s = (frame_indices[first] / src_fps + frame_indices[last] / src_fps) / 2.0;
+            }
+            snprintf(ts_buf, sizeof(ts_buf), "<%.1f seconds>", ts_s);
+            pending_frame = frame;
+            *out_text = mtmd_strdup(ts_buf);
+            return *out_text ? 0 : -2;
+        }
+        *out_bitmap = frame;
+        return 0;
+    }
+
+    *out_bitmap = frame;
+
+    if (timestamp_interval_ms > 0) {
+        // current_frame was already incremented by read_next_frame(); undo for elapsed calc
+        float elapsed_ms = (float)(current_frame - 1) / info.fps * 1000.0f;
+        if (elapsed_ms >= next_timestamp_ms) {
+            char ts_buf[32];
+            float elapsed_s = elapsed_ms / 1000.0f;
+            int   minutes   = (int)(elapsed_s / 60);
+            float seconds   = elapsed_s - minutes * 60.0f;
+            snprintf(ts_buf, sizeof(ts_buf), "[%dm%.2fs]", minutes, seconds);
+            pending_text = ts_buf;
+            next_timestamp_ms += (float)timestamp_interval_ms;
+        }
+    }
+
+    return 0;
+}
+
+#ifdef MTMD_VIDEO
+
+struct mtmd_helper_video : mtmd_helper_video_seq {
     const mtmd_context * mctx;
     std::string path;
     std::vector<uint8_t> input_buf; // non-empty when initialized from buffer
     std::string ffmpeg_bin;
     std::string ffprobe_bin;
-    float fps_target = 0.0f;
-    mtmd_helper_video_info info = {};
-    std::string id; // hash of the input video
 
     // RAII wrapper for managing subprocess
     struct subprocess_handle {
@@ -596,46 +797,15 @@ struct mtmd_helper_video {
     };
 
     subprocess_handle sp;
-    int32_t current_frame = 0;
-
-    std::string prompt_start         = "Video:";
-    int32_t     timestamp_interval_ms = 5000; // emit a timestamp text every N ms (0 = disabled)
-    float       next_timestamp_ms     = 0.0f; // next elapsed-ms threshold at which to emit
 
     // apply mtmd_get_video_config(); false for still images decoded via ffmpeg
-    bool                use_video_config = false;
-    mtmd_video_template tmpl = MTMD_VIDEO_TEMPLATE_DEFAULT;
-
-    int32_t     target_frames = 0;   // exact frame count to emit (0 = no limit)
-    int32_t     out_width     = 0;   // frame size from mtmd_get_video_config (0 = source size)
-    int32_t     out_height    = 0;
-
-    // source frames to emit, as picked by the reference processors (empty = use the fps filter)
-    std::vector<int32_t> frame_indices;
-    size_t      next_index = 0;
-    int32_t     src_frame  = 0;      // source frames read from ffmpeg so far
-    double      src_fps    = 0.0;
-
-    std::vector<uint8_t> frame_buf;
-    std::string   pending_text;          // text queued to be returned before the next frame
-    mtmd_bitmap * pending_frame = nullptr; // frame queued to be returned after its group timestamp
-
-    // soundtrack: after each group of frames comes the audio from its first frame to the next group's
-    double             video_start_s = 0.0;
-    int32_t            n_src_frames  = 0;
-    std::vector<float> soundtrack;
-    mtmd_bitmap *      pending_audio = nullptr;
-    bool          start_emitted = false;
+    bool   use_video_config = false;
+    double video_start_s    = 0.0; // start time of the video stream, to align the soundtrack
 
     bool is_buf_input() const {
         return !input_buf.empty();
     }
 
-    // number of frames per timestamp group (0 = timestamp every timestamp_interval_ms)
-    int32_t ts_per_n_frames() const {
-        return tmpl == MTMD_VIDEO_TEMPLATE_LFM3 || tmpl == MTMD_VIDEO_TEMPLATE_LFM3_OMNI ? 2 : 0;
-    }
-
     bool probe(float fps_target_arg) {
         // cache: makes stdin seekable; without it ffprobe reports width/height 0
         // for mp4 files whose moov atom sits at the end
@@ -671,11 +841,7 @@ struct mtmd_helper_video {
             probe_sp.start_feeder(input_buf);
         }
 
-        uint32_t width  = 0;
-        uint32_t height = 0;
-        float orig_fps = 0.0f;
-        float duration = -1.0f;
-        int32_t n_frames_orig = -1;
+        mtmd_helper_video_src src;
         char line[256];
         FILE * fp = probe_sp.stdout_pipe();
 
@@ -689,15 +855,15 @@ struct mtmd_helper_video {
             if (nl) *nl = '\0';
 
             if (strcmp(key, "width") == 0) {
-                width = (uint32_t)atoi(val);
+                src.width = (uint32_t)atoi(val);
             } else if (strcmp(key, "height") == 0) {
-                height = (uint32_t)atoi(val);
+                src.height = (uint32_t)atoi(val);
             } else if (strcmp(key, "r_frame_rate") == 0) {
-                orig_fps = parse_rational(val);
+                src.fps = parse_rational(val);
             } else if (strcmp(key, "nb_frames") == 0 && strcmp(val, "N/A") != 0) {
-                n_frames_orig = atoi(val);
+                src.n_frames = atoi(val);
             } else if (strcmp(key, "duration") == 0 && strcmp(val, "N/A") != 0) {
-                duration = (float)atof(val);
+                src.duration = (float)atof(val);
             } else if (strcmp(key, "start_time") == 0 && strcmp(val, "N/A") != 0) {
                 video_start_s = atof(val);
             }
@@ -705,58 +871,15 @@ struct mtmd_helper_video {
 
         probe_sp.stop();
 
-        if (width == 0 || height == 0 || orig_fps <= 0.0f) {
+        if (src.width == 0 || src.height == 0 || src.fps <= 0.0f) {
             return false;
         }
 
-        if (duration < 0.0f && n_frames_orig > 0) {
-            duration = (float)n_frames_orig / orig_fps;
-        }
-
         mtmd_video_config cfg{MTMD_VIDEO_TEMPLATE_DEFAULT, 0.0f, 0, 0, 0};
         if (use_video_config) {
-            cfg  = mtmd_get_video_config(mctx, width, height, n_frames_orig, orig_fps);
-            tmpl = cfg.tmpl;
+            cfg = mtmd_get_video_config(mctx, src.width, src.height, src.n_frames, src.fps);
         }
-
-        // the model's recommended sampling rate wins over the generic default
-        fps_target = cfg.fps > 0.0f ? cfg.fps
-                   : fps_target_arg > 0.0f ? fps_target_arg : orig_fps;
-
-        if (cfg.n_frames > 0 && duration > 0.0f) {
-            const int32_t n = cfg.n_frames;
-            target_frames = n;
-            fps_target    = (float) n / duration;
-
-            // same frames as the reference processors: round(linspace(0, n_frames - 1, n))
-            const double step = n > 1 ? (double) (n_frames_orig - 1) / (n - 1) : 0.0;
-            frame_indices.resize(n);
-            for (int32_t i = 0; i < n; i++) {
-                frame_indices[i] = (int32_t) std::nearbyint(i * step);
-            }
-            src_fps      = orig_fps;
-            n_src_frames = n_frames_orig;
-
-            if (cfg.frame_width > 0 && cfg.frame_height > 0 &&
-                    (cfg.frame_width != (int32_t) width || cfg.frame_height != (int32_t) height)) {
-                out_width  = cfg.frame_width;
-                out_height = cfg.frame_height;
-            }
-        }
-
-        if (out_width > 0 && out_height > 0) {
-            width  = (uint32_t) out_width;
-            height = (uint32_t) out_height;
-        }
-
-        info.width    = width;
-        info.height   = height;
-        info.fps      = fps_target;
-        LOG_DBG("%s: %ux%u fps=%.2f duration=%.2fs n_frames=%d target_frames=%d\n",
-                __func__, width, height, fps_target, duration, n_frames_orig, target_frames);
-        info.n_frames = target_frames > 0 ? target_frames
-                      : duration > 0.0f ? (int32_t)(duration * fps_target + 0.5f) : -1;
-        frame_buf.resize((size_t)width * height * 3);
+        configure(cfg, src, fps_target_arg);
         return true;
     }
 
@@ -800,31 +923,10 @@ struct mtmd_helper_video {
             LOG_WRN("%s: the video has no soundtrack\n", __func__);
             return;
         }
+        audio_rate = rate;
         soundtrack.assign(pcm.begin() + start, pcm.begin() + std::min(pcm.size(), start + len));
     }
 
-    // the soundtrack of the group of frames that ends with frame idx, nullptr if the soundtrack has ended
-    mtmd_bitmap * soundtrack_slice(int32_t idx) const {
-        const int32_t ts_per_n_frames = this->ts_per_n_frames();
-        const int32_t group    = idx / ts_per_n_frames;
-        const int32_t n_groups = ((int32_t) frame_indices.size() + ts_per_n_frames - 1) / ts_per_n_frames;
-        const int     rate     = mtmd_get_audio_sample_rate(mctx);
-        auto group_start = [&](int32_t g) {
-            const double t = frame_indices[g * ts_per_n_frames] / src_fps;
-            return std::min(soundtrack.size(), (size_t) std::nearbyint(t * rate));
-        };
-        const size_t begin = group == 0 ? 0 : group_start(group);
-        const size_t end   = group == n_groups - 1 ? soundtrack.size() : group_start(group + 1);
-        if (begin >= end) {
-            return nullptr;
-        }
-        mtmd_bitmap * slice = mtmd_bitmap_init_from_audio(end - begin, soundtrack.data() + begin);
-        if (!id.empty()) {
-            mtmd_bitmap_set_id(slice, (id + "+audio" + std::to_string(group)).c_str());
-        }
-        return slice;
-    }
-
     bool start_ffmpeg(float seek_seconds) {
         char seek_buf[64];
         char fps_buf[64];
@@ -917,6 +1019,7 @@ struct mtmd_helper_video {
             sp.start_feeder(input_buf);
         }
 
+        src_alive = sp.alive;
         return sp.alive;
     }
 
@@ -925,7 +1028,7 @@ struct mtmd_helper_video {
     }
 
     // read one raw frame from ffmpeg into frame_buf
-    bool read_raw_frame() {
+    bool read_raw_frame() override {
         FILE * fp = sp.stdout_pipe();
         const size_t frame_size = (size_t)info.width * info.height * 3;
         LOG_DBG("%s: reading frame %d, expecting %zu bytes (%ux%u)\n",
@@ -938,7 +1041,6 @@ struct mtmd_helper_video {
                 // clean EOF only if no bytes read yet; partial frame is an error
                 LOG_DBG("%s: fread returned 0 after %zu/%zu bytes (ferror=%d)\n",
                         __func__, total_read, frame_size, ferror(fp));
-                sp.alive = false;
                 return false;
             }
             total_read += n;
@@ -946,129 +1048,6 @@ struct mtmd_helper_video {
         return true;
     }
 
-    mtmd_bitmap * read_next_frame() {
-        if (!sp.alive) return nullptr;
-
-        if (!frame_indices.empty()) {
-            if (next_index >= frame_indices.size()) {
-                return nullptr;
-            }
-            // skip to the wanted source frame; a repeated index reuses the frame in frame_buf
-            const int32_t want = frame_indices[next_index];
-            while (src_frame <= want) {
-                if (!read_raw_frame()) {
-                    return nullptr;
-                }
-                src_frame++;
-            }
-            next_index++;
-        } else if (!read_raw_frame()) {
-            return nullptr;
-        }
-
-        LOG_DBG("%s: frame %d read OK\n", __func__, current_frame);
-        mtmd_bitmap * frame = mtmd_bitmap_init(info.width, info.height, frame_buf.data());
-        mtmd_bitmap_set_mergeable(frame, true);
-        mtmd_bitmap_set_is_video(frame, true);
-        if (!id.empty()) {
-            // each frame gets a unique id in the form of {hash}+{frame}, so that it can be identified in cache
-            std::string frame_id = id + "+" + std::to_string(current_frame);
-            mtmd_bitmap_set_id(frame, frame_id.c_str());
-        }
-        current_frame++;
-        return frame;
-    }
-
-    int32_t read_next(mtmd_bitmap ** out_bitmap, char ** out_text) {
-        *out_bitmap = nullptr;
-        *out_text   = nullptr;
-
-        if (!pending_text.empty()) {
-            *out_text = mtmd_strdup(pending_text);
-            pending_text.clear();
-            return *out_text ? 0 : -2;
-        }
-
-        if (pending_frame) {
-            *out_bitmap  = pending_frame;
-            pending_frame = nullptr;
-            return 0;
-        }
-
-        if (pending_audio) {
-            *out_bitmap   = pending_audio;
-            pending_audio = nullptr;
-            return 0;
-        }
-
-        LOG_DBG("%s: proc_alive=%d start_emitted=%d current_frame=%d\n",
-                __func__, (int)sp.alive, (int)start_emitted, current_frame);
-
-        if (!sp.alive) {
-            return (current_frame == 0) ? -2 : -1;
-        }
-
-        if (!start_emitted) {
-            start_emitted = true;
-            // model templates have no prompt prefix, the video begin/end tokens are added by mtmd_tokenize
-            if (tmpl == MTMD_VIDEO_TEMPLATE_DEFAULT && !prompt_start.empty()) {
-                *out_text = mtmd_strdup(prompt_start);
-                return *out_text ? 0 : -2;
-            }
-        }
-
-        mtmd_bitmap * frame = read_next_frame();
-        if (!frame) {
-            return -1;
-        }
-
-        const int32_t ts_per_n_frames = this->ts_per_n_frames();
-        if (ts_per_n_frames > 0) {
-            // emit the group timestamp before the group's first frame; the value is the
-            // midpoint time of the group, matching HF processors that stamp each temporal group
-            const int32_t idx = current_frame - 1; // read_next_frame() already incremented
-            // the soundtrack of a group follows its last frame
-            const bool group_ends = idx % ts_per_n_frames == ts_per_n_frames - 1 || idx + 1 == (int32_t) frame_indices.size();
-            if (!soundtrack.empty() && !frame_indices.empty() && group_ends) {
-                pending_audio = soundtrack_slice(idx);
-            }
-            if (idx % ts_per_n_frames == 0) {
-                char ts_buf[64];
-                double ts_s = ((double)idx + 0.5*(double)(ts_per_n_frames - 1)) / info.fps;
-                if (!frame_indices.empty()) {
-                    // stamp from the source frames; a short last group repeats its last frame
-                    const size_t first = idx;
-                    const size_t last  = std::min(first + ts_per_n_frames - 1, frame_indices.size() - 1);
-                    ts_s = (frame_indices[first] / src_fps + frame_indices[last] / src_fps) / 2.0;
-                }
-                snprintf(ts_buf, sizeof(ts_buf), "<%.1f seconds>", ts_s);
-                pending_frame = frame;
-                *out_text = mtmd_strdup(ts_buf);
-                return *out_text ? 0 : -2;
-            }
-            *out_bitmap = frame;
-            return 0;
-        }
-
-        *out_bitmap = frame;
-
-        if (timestamp_interval_ms > 0) {
-            // current_frame was already incremented by read_next_frame(); undo for elapsed calc
-            float elapsed_ms = (float)(current_frame - 1) / info.fps * 1000.0f;
-            if (elapsed_ms >= next_timestamp_ms) {
-                char ts_buf[32];
-                float elapsed_s = elapsed_ms / 1000.0f;
-                int   minutes   = (int)(elapsed_s / 60);
-                float seconds   = elapsed_s - minutes * 60.0f;
-                snprintf(ts_buf, sizeof(ts_buf), "[%dm%.2fs]", minutes, seconds);
-                pending_text = ts_buf;
-                next_timestamp_ms += (float)timestamp_interval_ms;
-            }
-        }
-
-        return 0;
-    }
-
     static float parse_rational(const char * s) {
         int num = 0, den = 1;
         if (sscanf(s, "%d/%d", &num, &den) == 2 && den > 0) {
@@ -1233,8 +1212,6 @@ void mtmd_helper_video_free(mtmd_helper_video * ctx) {
 #ifdef MTMD_VIDEO
     if (!ctx) return;
     ctx->stop_ffmpeg();
-    mtmd_bitmap_free(ctx->pending_frame);
-    mtmd_bitmap_free(ctx->pending_audio);
     delete ctx;
 #else
     GGML_UNUSED(ctx);

tools/mtmd/mtmd-helper.h

@@ -318,6 +318,7 @@ MTMD_API int32_t mtmd_helper_gen_audio_flush(mtmd_helper_gen_audio * ctx);
 #ifdef __cplusplus
 #include <set>
 #include <memory>
+#include <string>
 #include <vector>
 
 namespace mtmd_helper {
@@ -408,6 +409,68 @@ inline std::vector<llama_logit_bias> audio_stream_logit_bias(const mtmd_context
 }
 
 } // namespace mtmd_helper
+
+//
+// test API (C++ only)
+// internals of mtmd_helper_video, do not use in applications
+//
+
+struct mtmd_helper_video_src {
+    uint32_t width    = 0;
+    uint32_t height   = 0;
+    float    fps      = 0.0f;
+    int32_t  n_frames = -1;
+    float    duration = -1.0f;
+};
+
+struct mtmd_helper_video_seq {
+    float fps_target = 0.0f;
+    mtmd_helper_video_info info = {};
+    std::string id;
+
+    bool    src_alive     = false;
+    int32_t current_frame = 0;
+
+    std::string prompt_start         = "Video:";
+    int32_t     timestamp_interval_ms = 5000;
+    float       next_timestamp_ms     = 0.0f;
+
+    mtmd_video_template tmpl = MTMD_VIDEO_TEMPLATE_DEFAULT;
+
+    int32_t     target_frames = 0;
+    int32_t     out_width     = 0;
+    int32_t     out_height    = 0;
+
+    std::vector<int32_t> frame_indices;
+    size_t      next_index = 0;
+    int32_t     src_frame  = 0;
+    double      src_fps    = 0.0;
+
+    std::vector<uint8_t> frame_buf;
+    std::string   pending_text;
+    mtmd_bitmap * pending_frame = nullptr;
+
+    int32_t            n_src_frames = 0;
+    int                audio_rate   = 0;
+    std::vector<float> soundtrack;
+    mtmd_bitmap *      pending_audio = nullptr;
+    bool          start_emitted = false;
+
+    mtmd_helper_video_seq() = default;
+    mtmd_helper_video_seq(const mtmd_helper_video_seq &) = delete;
+    mtmd_helper_video_seq & operator=(const mtmd_helper_video_seq &) = delete;
+    virtual ~mtmd_helper_video_seq();
+
+    virtual bool read_raw_frame() = 0;
+
+    void configure(const mtmd_video_config & cfg, const mtmd_helper_video_src & src, float fps_target_arg);
+
+    int32_t ts_per_n_frames() const;
+
+    mtmd_bitmap * soundtrack_slice(int32_t idx) const;
+    mtmd_bitmap * read_next_frame();
+    int32_t read_next(mtmd_bitmap ** out_bitmap, char ** out_text);
+};
 #endif
 
 #endif

tools/mtmd/mtmd-image.cpp

@@ -819,12 +819,12 @@ mtmd_image_preproc_out mtmd_image_preprocessor_dyn_size::preprocess(const clip_i
 }
 
 //
-// mtmd_image_preprocessor_lfm3vl
+// mtmd_image_preprocessor_qwen3vl
 //
 
 // rounds before the pixel checks (an edge may round to 0) and computes in double, as transformers does
-clip_image_size mtmd_image_preprocessor_lfm3vl::smart_resize(int width, int height, int factor, int64_t min_pixels, int64_t max_pixels,
-                                                             int n_frames, int temporal_factor) {
+clip_image_size mtmd_image_preprocessor_qwen3vl::smart_resize(int width, int height, int factor, int64_t min_pixels, int64_t max_pixels,
+                                                              int n_frames, int temporal_factor) {
     double h = height;
     double w = width;
     // the video variant first upscales edges below factor
@@ -850,18 +850,13 @@ clip_image_size mtmd_image_preprocessor_lfm3vl::smart_resize(int width, int heig
     return {std::max(factor, w_bar), std::max(factor, h_bar)};
 }
 
-clip_image_size mtmd_image_preprocessor_lfm3vl::target_size(const clip_image_size & original_size) const {
-    return smart_resize(original_size.width, original_size.height, hparams.patch_size * hparams.n_merge,
-                        hparams.image_min_pixels, hparams.image_max_pixels);
-}
-
-bool mtmd_image_preprocessor_lfm3vl::get_video_frame_size(int n_frames, int width, int height, clip_image_size & out) const {
-    // Lfm3VlVideoProcessor size: shortest_edge 4096, longest_edge 25165824, temporal_patch_size 2
+bool mtmd_image_preprocessor_qwen3vl::get_video_frame_size(int n_frames, int width, int height, clip_image_size & out) const {
+    // Qwen3VLVideoProcessor (and Lfm3VlVideoProcessor) size: shortest_edge 4096, longest_edge 25165824, temporal_patch_size 2
     out = smart_resize(width, height, hparams.patch_size * hparams.n_merge, 4096, 25165824, n_frames, 2);
     return true;
 }
 
-mtmd_image_preproc_out mtmd_image_preprocessor_lfm3vl::preprocess_video_frame(const clip_image_u8 & img) const {
+mtmd_image_preproc_out mtmd_image_preprocessor_qwen3vl::preprocess_video_frame(const clip_image_u8 & img) const {
     const int f = hparams.patch_size * hparams.n_merge;
     const auto size = img.get_size();
     return resize_to(img, {
@@ -870,7 +865,7 @@ mtmd_image_preproc_out mtmd_image_preprocessor_lfm3vl::preprocess_video_frame(co
     });
 }
 
-mtmd_image_preproc_out mtmd_image_preprocessor_lfm3vl::resize_to(const clip_image_u8 & img, const clip_image_size & size) const {
+mtmd_image_preproc_out mtmd_image_preprocessor_qwen3vl::resize_to(const clip_image_u8 & img, const clip_image_size & size) const {
     clip_image_u8 resized_image;
     img_tool::resize(img, resized_image, size,
                         hparams.image_resize_algo,
@@ -881,6 +876,15 @@ mtmd_image_preproc_out mtmd_image_preprocessor_lfm3vl::resize_to(const clip_imag
     return output;
 }
 
+//
+// mtmd_image_preprocessor_lfm3vl
+//
+
+clip_image_size mtmd_image_preprocessor_lfm3vl::target_size(const clip_image_size & original_size) const {
+    return smart_resize(original_size.width, original_size.height, hparams.patch_size * hparams.n_merge,
+                        hparams.image_min_pixels, hparams.image_max_pixels);
+}
+
 //
 // mtmd_image_preprocessor_glm5v
 //

tools/mtmd/mtmd-image.h

@@ -137,15 +137,15 @@ struct mtmd_image_preprocessor_fixed_size : mtmd_image_preprocessor {
 // this is used by models with native support for dynamic image size, for example: Qwen-VL, Pixtral, Kimi-VL, etc
 struct mtmd_image_preprocessor_dyn_size : mtmd_image_preprocessor {
     mtmd_image_preprocessor_dyn_size(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {}
+    mtmd_image_preprocessor_dyn_size(const clip_hparams & hparams) : mtmd_image_preprocessor(hparams) {}
     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;
     virtual clip_image_size target_size(const clip_image_size & original_size) const { return get_target_size(hparams, original_size); }
 };
 
-// LFM3-VL, same as dyn_size but sized exactly as the transformers Qwen2-VL smart_resize
-struct mtmd_image_preprocessor_lfm3vl : mtmd_image_preprocessor_dyn_size {
-    mtmd_image_preprocessor_lfm3vl(const clip_ctx * ctx) : mtmd_image_preprocessor_dyn_size(ctx) {}
-    clip_image_size target_size(const clip_image_size & original_size) const override;
+// Qwen3-VL, same as dyn_size for images, video frames are sized as Qwen3VLVideoProcessor
+struct mtmd_image_preprocessor_qwen3vl : mtmd_image_preprocessor_dyn_size {
+    using mtmd_image_preprocessor_dyn_size::mtmd_image_preprocessor_dyn_size;
     // the video helper sizes the frames (get_video_frame_size), they are only snapped to the grid
     mtmd_image_preproc_out preprocess_video_frame(const clip_image_u8 & img) const override;
     bool get_video_frame_size(int n_frames, int width, int height, clip_image_size & out) const override;
@@ -156,6 +156,12 @@ struct mtmd_image_preprocessor_lfm3vl : mtmd_image_preprocessor_dyn_size {
                                         int n_frames = 1, int temporal_factor = 1);
 };
 
+// LFM3-VL, same as qwen3vl but images are sized exactly as the transformers Qwen2-VL smart_resize
+struct mtmd_image_preprocessor_lfm3vl : mtmd_image_preprocessor_qwen3vl {
+    using mtmd_image_preprocessor_qwen3vl::mtmd_image_preprocessor_qwen3vl;
+    clip_image_size target_size(const clip_image_size & original_size) const override;
+};
+
 // GLM 5.3 flash, similar to dyn_size, but each edge is aligned up to patch_size*n_merge, and max budget is met by a search over the height with the width scaled proportionally
 struct mtmd_image_preprocessor_glm5v : mtmd_image_preprocessor {
     mtmd_image_preprocessor_glm5v(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {}

tools/mtmd/mtmd-internal.h

@@ -19,3 +19,10 @@ struct mtmd_internal_part {
 
 // [QWEN_VIDEO] merged parts are erased from `parts`, so one group always maps to one part
 std::vector<std::vector<const mtmd_bitmap *>> mtmd_group_mergeable_bitmaps(std::vector<mtmd_internal_part> & parts, int n_merge);
+
+struct mtmd_image_preprocessor;
+
+// mtmd_get_video_config() for a given template, without a context
+mtmd_video_config mtmd_video_config_init(mtmd_video_template tmpl, const mtmd_image_preprocessor * preproc,
+                                         int32_t src_width, int32_t src_height,
+                                         int32_t src_n_frames, float src_fps);

tools/mtmd/mtmd.cpp

@@ -176,6 +176,7 @@ struct mtmd_bitmap {
     std::string id; // optional user-defined id, for ex: can be set to image hash, useful for KV cache tracking
     bool is_audio = false; // true if the bitmap is audio
     bool is_voice = false;
+    std::shared_ptr<const clip_gen_ref> gen_ref; // gen-side reference of a predefined voice, if any
     bool mergeable = false; // [QWEN_VIDEO] set only on frames of the same video
     bool is_video  = false; // see mtmd_bitmap_set_is_video()
 
@@ -201,9 +202,14 @@ struct mtmd_bitmap {
         }
     }
 
-    static mtmd_bitmap * init_voice(const std::string & name, const float * rows, uint32_t n_embd, uint32_t n_rows) {
+    static mtmd_bitmap * init_voice(const std::string & name, const clip_gen_voice & voice, uint32_t n_embd) {
+        const float *  rows   = voice.rows.data();
+        const uint32_t n_rows = voice.rows.size() / n_embd;
         auto * bitmap = new mtmd_bitmap(nullptr, n_embd, n_rows);
         bitmap->is_voice = true;
+        if (!voice.ref.tokens.empty()) {
+            bitmap->gen_ref = std::make_shared<const clip_gen_ref>(voice.ref);
+        }
         bitmap->id = name;
         bitmap->data.assign((const unsigned char *) rows, (const unsigned char *) (rows + (size_t) n_embd * n_rows));
         return bitmap;
@@ -348,6 +354,7 @@ struct mtmd_audio_track {
     clip_image_f32_batch batch_f32; // the mel of the whole track
     uint32_t n_tokens = 0;
     std::vector<float> embd;        // filled by the first encode
+    std::shared_ptr<const clip_gen_ref> gen_ref; // gen-side reference of a predefined voice, if any
 };
 
 struct mtmd_audio_tokens {
@@ -564,6 +571,11 @@ mtmd_context_params mtmd_context_params_default() {
     return params;
 }
 
+// chatterbox multilingual has a perceiver in its conditioning encoder (speaker encoding context)
+static bool cbx_is_mtl(clip_ctx * ctx_a) {
+    return clip_cbx_read_tensor(ctx_a, "a.cenc.perceiver.pre_attention_query", nullptr, 0) > 0;
+}
+
 struct mtmd_context {
     struct clip_ctx * ctx_v; // vision
     struct clip_ctx * ctx_a; // audio
@@ -572,6 +584,9 @@ struct mtmd_context {
     // generation context
     struct clip_ctx * ctx_gen_a; // audio
 
+    // gen-side reference data of the last encoded chunk (chatterbox), see mtmd_gen_state_set_ref()
+    clip_gen_ref gen_ref;
+
     bool print_timings;
     int n_threads;
     std::string media_marker;
@@ -790,9 +805,15 @@ struct mtmd_context {
                         image_preproc = std::make_unique<mtmd_image_preprocessor_minicpmv>(ctx_v);
                     }
                 } break;
+            case PROJECTOR_TYPE_QWEN3VL:
+                {
+                    // <|vision_start|> ... (image embeddings) ... <|vision_end|>
+                    img_beg = "<|vision_start|>";
+                    img_end = "<|vision_end|>";
+                    image_preproc = std::make_unique<mtmd_image_preprocessor_qwen3vl>(ctx_v);
+                } break;
             case PROJECTOR_TYPE_QWEN2VL:
             case PROJECTOR_TYPE_QWEN25VL:
-            case PROJECTOR_TYPE_QWEN3VL:
             case PROJECTOR_TYPE_LING3VL:
             case PROJECTOR_TYPE_MIMOVL:
                 {
@@ -1153,6 +1174,10 @@ struct mtmd_context {
                 {
                     audio_preproc = std::make_unique<mtmd_audio_preprocessor_parakeet>(ctx_a);
                 } break;
+            case PROJECTOR_TYPE_NEMOTRON_TRANSCRIBE:
+                {
+                    audio_preproc = std::make_unique<mtmd_audio_preprocessor_nemotron_transcribe>(ctx_a);
+                } break;
             case PROJECTOR_TYPE_GEMMA4UA:
                 {
                     aud_beg = "<|audio>";
@@ -1180,6 +1205,10 @@ struct mtmd_context {
                 {
                     audio_preproc = std::make_unique<mtmd_audio_preprocessor_pockettts>(ctx_a);
                 } break;
+            case PROJECTOR_TYPE_CHATTERBOX_SPKENC:
+                {
+                    audio_preproc = std::make_unique<mtmd_audio_preprocessor_chatterbox_ref>(ctx_a, cbx_is_mtl(ctx_a));
+                } break;
             default:
                 throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj));
         }
@@ -1841,6 +1870,38 @@ struct mtmd_tokenizer {
                 }
             }
 
+            // one chunk for the five entries of a chatterbox clip, one token per talker conditioning row
+            if (clip_get_projector_type(ctx->ctx_a) == PROJECTOR_TYPE_CHATTERBOX_SPKENC) {
+                GGML_ASSERT(mel_spec_chunks.size() == 5);
+                // mtl: spkr, perceiver query rows, emotion; turbo: spkr, t3 speech tokens
+                const size_t n_query  = clip_cbx_read_tensor(ctx->ctx_a, "a.cenc.perceiver.pre_attention_query", nullptr, 0) / ctx->n_embd_out();
+                const size_t n_tokens = cbx_is_mtl(ctx->ctx_a) ? 2 + n_query : 1 + (size_t) mel_spec_chunks[2].n_len / 4;
+
+                clip_image_f32_batch batch_f32;
+                batch_f32.is_audio = true;
+                for (auto & mel_spec : mel_spec_chunks) {
+                    GGML_ASSERT(mel_spec.n_len <= INT32_MAX && mel_spec.n_len >= 0);
+                    GGML_ASSERT(mel_spec.n_mel <= INT32_MAX && mel_spec.n_mel >= 0);
+                    clip_image_f32 mel_f32;
+                    mel_f32.set_size({(int) mel_spec.n_len, (int) mel_spec.n_mel},
+                                     mel_spec.data.empty(), /* is_audio */ true);
+                    mel_f32.cpy_buf(mel_spec.data);
+                    batch_f32.entries.push_back(std::move(mel_f32));
+                }
+
+                mtmd_audio_tokens_ptr audio_tokens(new mtmd_audio_tokens);
+                audio_tokens->n_tokens  = (uint32_t) n_tokens;
+                audio_tokens->batch_f32 = std::move(batch_f32);
+                audio_tokens->id        = bitmap->id;
+
+                mtmd_input_chunk chunk{
+                    MTMD_INPUT_CHUNK_TYPE_AUDIO,
+                    {}, // text tokens
+                    nullptr, // image tokens
+                    std::move(audio_tokens),
+                };
+                cur.entries.emplace_back(std::move(chunk));
+            } else
             // consider each mel_spec as a separate audio chunk
             // TODO: maybe support batching, but this may come with memory cost
             for (auto & mel_spec : mel_spec_chunks) {
@@ -1894,6 +1955,7 @@ struct mtmd_tokenizer {
         track->n_tokens = bitmap.ny;
         const auto & buf = bitmap.get_ro_buf();
         track->embd.assign((const float *) buf.data(), (const float *) (buf.data() + buf.size()));
+        track->gen_ref = bitmap.gen_ref;
 
         mtmd_audio_tokens_ptr audio_tokens(new mtmd_audio_tokens);
         audio_tokens->n_tokens = bitmap.ny;
@@ -2185,7 +2247,10 @@ static int32_t mtmd_encode_impl(mtmd_context * ctx, const mtmd_image_tokens * im
     return ok ? 0 : 1;
 }
 
+static int32_t cbx_ref_encode(mtmd_context * ctx, const mtmd_audio_tokens * audio_tokens, std::vector<float> & out_embd);
+
 static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk * chunk, std::vector<float> & out_embd) {
+    ctx->gen_ref = {}; // see mtmd_gen_state_set_ref()
     if (chunk->type == MTMD_INPUT_CHUNK_TYPE_TEXT) {
         LOG_WRN("mtmd_encode_chunk has no effect for text chunks\n");
         return 0;
@@ -2230,8 +2295,14 @@ static int32_t mtmd_encode_chunk_impl(mtmd_context * ctx, const mtmd_input_chunk
             const size_t off = (size_t) chunk->tokens_audio->row_begin * n_mmproj_embd;
             out_embd.assign(track->embd.begin() + off,
                             track->embd.begin() + off + (size_t) chunk->tokens_audio->n_tokens * n_mmproj_embd);
+            if (track->gen_ref) {
+                ctx->gen_ref = *track->gen_ref;
+            }
             return 0;
         }
+        if (clip_get_projector_type(ctx->ctx_a) == PROJECTOR_TYPE_CHATTERBOX_SPKENC) {
+            return cbx_ref_encode(ctx, chunk->tokens_audio.get(), out_embd);
+        }
         out_embd.resize((size_t)chunk->tokens_audio->n_tokens * ctx->n_embd_out());
         bool ok = clip_image_batch_encode(
             ctx->ctx_a,
@@ -2298,6 +2369,10 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
             info.type = MTMD_GEN_AUDIO_TYPE_LFM3A;
             info.sample_rate = 24000;
             break;
+        case PROJECTOR_TYPE_CHATTERBOX:
+            info.type = MTMD_GEN_AUDIO_TYPE_CHATTERBOX;
+            info.sample_rate = 24000;
+            break;
         default:
             info.type = MTMD_GEN_AUDIO_TYPE_NONE;
             break;
@@ -2305,6 +2380,355 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) {
     return info;
 }
 
+size_t mtmd_gen_audio_read_tensor(mtmd_context * ctx, const char * name, float * out, size_t n_max) {
+    if (!ctx->ctx_gen_a) {
+        return 0;
+    }
+    return clip_cbx_read_tensor(ctx->ctx_gen_a, name, out, n_max);
+}
+
+// raw 192-dim CAMPPlus x-vector of the fbank entry
+static bool cbx_ref_xvec(mtmd_context * ctx, const clip_image_f32 & fbank, std::vector<float> & xvec) {
+    clip_image_f32 mel_img;
+    mel_img.set_size(fbank.get_size(), false, true);
+    mel_img.cpy_buf(fbank.get_ro_buf());
+    clip_image_f32_batch batch;
+    batch.is_audio = true;
+    batch.entries.push_back(std::move(mel_img));
+    xvec.resize(192);
+    if (!clip_image_batch_encode(ctx->ctx_a, ctx->n_threads, &batch, xvec)) {
+        LOG_ERR("%s: speaker encoder failed\n", __func__);
+        return false;
+    }
+    return true;
+}
+
+// s3 speech tokens of a log-mel entry, rows optionally receives their speech embedding rows
+static int32_t cbx_ref_tokenize(mtmd_context * ctx, const clip_image_f32 & mel,
+                            std::vector<int32_t> & codes, std::vector<float> * rows) {
+    clip_image_f32 mel_img;
+    mel_img.set_size(mel.get_size(), false, true);
+    mel_img.cpy_buf(mel.get_ro_buf());
+
+    clip_image_f32_batch batch;
+    batch.is_audio = true;
+    batch.entries.push_back(std::move(mel_img));
+
+    std::vector<int32_t> out_codes;
+    std::vector<float>   out_embd;
+
+    clip_encode_params params;
+    params.imgs        = &batch;
+    params.n_threads   = ctx->n_threads;
+    params.gen_process = CLIP_GEN_PROCESS_TOKENIZE;
+    params.out_codes     = &out_codes;
+    params.out_code_embd = rows ? &out_embd : nullptr;
+
+    if (!clip_encode(ctx->ctx_gen_a, &params)) {
+        LOG_ERR("%s: clip_encode failed (tokenize)\n", __func__);
+        return 1;
+    }
+
+    codes = std::move(out_codes);
+    if (rows) {
+        *rows = std::move(out_embd);
+    }
+    return 0;
+}
+
+// talker conditioning rows: the speaker row, then on mtl the perceiver block and the emotion row
+static int32_t cbx_ref_cond(mtmd_context * ctx, const clip_image_f32 & ve_mel,
+                        const std::vector<float> & pse, std::vector<float> & out_rows) {
+    clip_ctx * ctx_clip = ctx->ctx_a;
+    // n_expect 0 accepts any non-empty size
+    auto read_t = [&](const char * name, std::vector<float> & v, size_t n_expect) -> bool {
+        const size_t n = clip_cbx_read_tensor(ctx_clip, name, nullptr, 0);
+        if (n == 0 || (n_expect > 0 && n != n_expect)) {
+            return false;
+        }
+        v.resize(n);
+        return clip_cbx_read_tensor(ctx_clip, name, v.data(), n) == n;
+    };
+
+    // voice encoder (embeds_from_wavs): 3-layer lstm per 160-frame partial, the normalized partial embeddings are averaged
+    const std::vector<float> & mel = ve_mel.get_ro_buf();
+    const int n_mel     = 40;
+    const int n_partial = 160;
+    const int n_h       = 256;
+    const int step      = (int) lround((16000.0 / 1.3) / n_partial); // reference rate 1.3
+    if (ve_mel.nx() < n_partial || mel.size() < (size_t) ve_mel.nx() * n_mel) {
+        LOG_ERR("%s: voice encoder mel too short\n", __func__);
+        return 1;
+    }
+    const int n_wins = (ve_mel.nx() - n_partial) / step + 1;
+
+    std::vector<float> w_ih[3], w_hh[3], b_ih[3], b_hh[3];
+    std::vector<float> w_proj, b_proj;
+    for (int l = 0; l < 3; l++) {
+        const std::string s = std::to_string(l);
+        const size_t n_in = l == 0 ? n_mel : n_h;
+        if (!read_t(("a.ve.lstm.weight_ih_l" + s).c_str(), w_ih[l], 4 * n_h * n_in) ||
+            !read_t(("a.ve.lstm.weight_hh_l" + s).c_str(), w_hh[l], 4 * n_h * n_h) ||
+            !read_t(("a.ve.lstm.bias_ih_l" + s).c_str(), b_ih[l], 4 * n_h) ||
+            !read_t(("a.ve.lstm.bias_hh_l" + s).c_str(), b_hh[l], 4 * n_h)) {
+            LOG_ERR("%s: model has no voice encoder\n", __func__);
+            return 1;
+        }
+    }
+    if (!read_t("a.ve.proj.weight", w_proj, n_h * n_h) || !read_t("a.ve.proj.bias", b_proj, n_h)) {
+        LOG_ERR("%s: model has no voice encoder projection\n", __func__);
+        return 1;
+    }
+
+    std::vector<float> ve(n_h, 0.0f);
+    std::vector<float> h((size_t) 3 * n_h), c((size_t) 3 * n_h), x(n_h), g((size_t) 4 * n_h);
+    for (int p = 0; p < n_wins; p++) {
+        std::fill(h.begin(), h.end(), 0.0f);
+        std::fill(c.begin(), c.end(), 0.0f);
+        for (int t = 0; t < n_partial; t++) {
+            const float * in = mel.data() + (size_t) (p * step + t) * n_mel;
+            int n_in = n_mel;
+            for (int l = 0; l < 3; l++) {
+                float * hl = h.data() + (size_t) l * n_h;
+                float * cl = c.data() + (size_t) l * n_h;
+                for (int j = 0; j < 4 * n_h; j++) {
+                    double acc = b_ih[l][(size_t) j] + b_hh[l][(size_t) j];
+                    const float * wi = w_ih[l].data() + (size_t) j * n_in;
+                    for (int i = 0; i < n_in; i++) {
+                        acc += (double) wi[i] * in[i];
+                    }
+                    const float * wh = w_hh[l].data() + (size_t) j * n_h;
+                    for (int i = 0; i < n_h; i++) {
+                        acc += (double) wh[i] * hl[i];
+                    }
+                    g[(size_t) j] = (float) acc;
+                }
+                // torch gate order: input, forget, cell, output
+                for (int i = 0; i < n_h; i++) {
+                    const float gi = 1.0f / (1.0f + expf(-g[(size_t) i]));
+                    const float gf = 1.0f / (1.0f + expf(-g[(size_t) i + n_h]));
+                    const float gc = tanhf(g[(size_t) i + 2 * n_h]);
+                    const float go = 1.0f / (1.0f + expf(-g[(size_t) i + 3 * n_h]));
+                    cl[i] = gf * cl[i] + gi * gc;
+                    x[(size_t) i] = go * tanhf(cl[i]);
+                }
+                memcpy(hl, x.data(), (size_t) n_h * sizeof(float));
+                in   = hl;
+                n_in = n_h;
+            }
+        }
+        // projected, relu'd, normalized partial embedding
+        std::vector<float> e(n_h);
+        double norm = 0.0;
+        for (int o = 0; o < n_h; o++) {
+            double acc = b_proj[(size_t) o];
+            const float * w = w_proj.data() + (size_t) o * n_h;
+            const float * hl = h.data() + (size_t) 2 * n_h;
+            for (int i = 0; i < n_h; i++) {
+                acc += (double) w[i] * hl[i];
+            }
+            e[(size_t) o] = (float) std::max(acc, 0.0);
+            norm += (double) e[(size_t) o] * e[(size_t) o];
+        }
+        norm = sqrt(norm);
+        for (int o = 0; o < n_h; o++) {
+            ve[(size_t) o] += (float) (e[(size_t) o] / norm);
+        }
+    }
+    double norm = 0.0;
+    for (float v : ve) {
+        norm += (double) v * v;
+    }
+    norm = sqrt(norm);
+    for (float & v : ve) {
+        v = (float) (v / norm);
+    }
+
+    // speaker projection row
+    const int n_e = (int) ctx->n_embd_out();
+    std::vector<float> w_spkr, b_spkr;
+    if (!read_t("a.cenc.spkr_enc.weight", w_spkr, (size_t) n_e * n_h) || !read_t("a.cenc.spkr_enc.bias", b_spkr, n_e)) {
+        LOG_ERR("%s: model has no speaker conditioning projection\n", __func__);
+        return 1;
+    }
+    std::vector<float> rows((size_t) n_e);
+    for (int o = 0; o < n_e; o++) {
+        double acc = b_spkr[(size_t) o];
+        const float * w = w_spkr.data() + (size_t) o * n_h;
+        for (int i = 0; i < n_h; i++) {
+            acc += (double) w[i] * ve[(size_t) i];
+        }
+        rows[(size_t) o] = (float) acc;
+    }
+
+    // the perceiver runs its shared attention block as cross then self attention
+    std::vector<float> query;
+    if (read_t("a.cenc.perceiver.pre_attention_query", query, 0)) {
+        const int    n_head = 4;
+        const size_t n_ee   = (size_t) n_e * n_e;
+        if (pse.empty() || pse.size() % n_e != 0 || query.size() % n_e != 0 || n_e % n_head != 0) {
+            LOG_ERR("%s: reference speech embeddings required for the perceiver\n", __func__);
+            return 1;
+        }
+        std::vector<float> ln_w, ln_b, wq, bq, wk, bk, wv, bv, wo, bo, emo;
+        if (!read_t("a.cenc.perceiver.attn.norm.weight", ln_w, n_e) || !read_t("a.cenc.perceiver.attn.norm.bias", ln_b, n_e) ||
+            !read_t("a.cenc.perceiver.attn.to_q.weight", wq, n_ee) || !read_t("a.cenc.perceiver.attn.to_q.bias", bq, n_e) ||
+            !read_t("a.cenc.perceiver.attn.to_k.weight", wk, n_ee) || !read_t("a.cenc.perceiver.attn.to_k.bias", bk, n_e) ||
+            !read_t("a.cenc.perceiver.attn.to_v.weight", wv, n_ee) || !read_t("a.cenc.perceiver.attn.to_v.bias", bv, n_e) ||
+            !read_t("a.cenc.perceiver.attn.proj_out.weight", wo, n_ee) || !read_t("a.cenc.perceiver.attn.proj_out.bias", bo, n_e) ||
+            !read_t("a.cenc.emotion_adv_fc.weight", emo, n_e)) {
+            LOG_ERR("%s: model has an incomplete perceiver\n", __func__);
+            return 1;
+        }
+        const int d_head = n_e / n_head;
+
+        auto layer_norm = [&](const float * in, float * out) {
+            double mean = 0.0, var = 0.0;
+            for (int i = 0; i < n_e; i++) {
+                mean += in[i];
+            }
+            mean /= n_e;
+            for (int i = 0; i < n_e; i++) {
+                var += (in[i] - mean) * (in[i] - mean);
+            }
+            const double sd = sqrt(var / n_e + 1e-5);
+            for (int i = 0; i < n_e; i++) {
+                out[i] = (float) ((in[i] - mean) / sd * ln_w[(size_t) i] + ln_b[(size_t) i]);
+            }
+        };
+        auto linear = [&](const std::vector<float> & w, const std::vector<float> & b,
+                          const std::vector<float> & in, int n_rows, std::vector<float> & out) {
+            out.resize((size_t) n_rows * n_e);
+            for (int r = 0; r < n_rows; r++) {
+                for (int o = 0; o < n_e; o++) {
+                    double acc = b[(size_t) o];
+                    const float * wr = w.data() + (size_t) o * n_e;
+                    const float * ir = in.data() + (size_t) r * n_e;
+                    for (int i = 0; i < n_e; i++) {
+                        acc += (double) wr[i] * ir[i];
+                    }
+                    out[(size_t) r * n_e + o] = (float) acc;
+                }
+            }
+        };
+        auto attn_block = [&](const std::vector<float> & x1, int n1,
+                              const std::vector<float> & x2, int n2, std::vector<float> & out) {
+            std::vector<float> nx1((size_t) n1 * n_e), nx2((size_t) n2 * n_e);
+            for (int r = 0; r < n1; r++) {
+                layer_norm(x1.data() + (size_t) r * n_e, nx1.data() + (size_t) r * n_e);
+            }
+            for (int r = 0; r < n2; r++) {
+                layer_norm(x2.data() + (size_t) r * n_e, nx2.data() + (size_t) r * n_e);
+            }
+            std::vector<float> q, k, v;
+            linear(wq, bq, nx1, n1, q);
+            linear(wk, bk, nx2, n2, k);
+            linear(wv, bv, nx2, n2, v);
+
+            std::vector<float> ctxt((size_t) n1 * n_e);
+            std::vector<double> sc((size_t) n2);
+            for (int hd = 0; hd < n_head; hd++) {
+                const int off = hd * d_head;
+                for (int t = 0; t < n1; t++) {
+                    double mx = -1e30;
+                    for (int s = 0; s < n2; s++) {
+                        double acc = 0.0;
+                        for (int i = 0; i < d_head; i++) {
+                            acc += (double) q[(size_t) t * n_e + off + i] * k[(size_t) s * n_e + off + i];
+                        }
+                        sc[(size_t) s] = acc / sqrt((double) d_head);
+                        mx = std::max(mx, sc[(size_t) s]);
+                    }
+                    double sum = 0.0;
+                    for (int s = 0; s < n2; s++) {
+                        sc[(size_t) s] = exp(sc[(size_t) s] - mx);
+                        sum += sc[(size_t) s];
+                    }
+                    for (int i = 0; i < d_head; i++) {
+                        double acc = 0.0;
+                        for (int s = 0; s < n2; s++) {
+                            acc += sc[(size_t) s] * v[(size_t) s * n_e + off + i];
+                        }
+                        ctxt[(size_t) t * n_e + off + i] = (float) (acc / sum);
+                    }
+                }
+            }
+            linear(wo, bo, ctxt, n1, out);
+            for (size_t i = 0; i < out.size(); i++) {
+                out[i] += x1[i];
+            }
+        };
+
+        const int n_q = (int) (query.size() / n_e);
+        std::vector<float> pre, p32;
+        attn_block(query, n_q, pse, (int) (pse.size() / (size_t) n_e), pre);
+        attn_block(pre, n_q, pre, n_q, p32);
+
+        rows.insert(rows.end(), p32.begin(), p32.end());
+        const float exaggeration = 0.5f; // reference default
+        for (int i = 0; i < n_e; i++) {
+            rows.push_back(emo[(size_t) i] * exaggeration);
+        }
+    }
+
+    out_rows = std::move(rows);
+    return 0;
+}
+
+// talker conditioning rows go to out_embd, the flow decoder reference to ctx->gen_ref
+static int32_t cbx_ref_encode(mtmd_context * ctx, const mtmd_audio_tokens * audio_tokens, std::vector<float> & out_embd) {
+    if (!ctx->ctx_gen_a) {
+        LOG_ERR("%s: reference encoding needs the audio generation context\n", __func__);
+        return 1;
+    }
+    const auto & entries = audio_tokens->batch_f32.entries;
+    GGML_ASSERT(entries.size() == 5);
+
+    std::vector<float> xvec;
+    if (!cbx_ref_xvec(ctx, entries[0], xvec)) {
+        return 1;
+    }
+
+    std::vector<int32_t> flow_codes, t3_codes;
+    std::vector<float> t3_rows;
+    if (cbx_ref_tokenize(ctx, entries[1], flow_codes, nullptr) != 0 ||
+        cbx_ref_tokenize(ctx, entries[2], t3_codes, &t3_rows) != 0) {
+        return 1;
+    }
+
+    const size_t n_feat = (size_t) entries[3].nx() * 80;
+    if ((size_t) entries[3].nx() != 2 * flow_codes.size()) {
+        LOG_ERR("%s: reference mel length %d does not match %zu codes\n",
+                __func__, entries[3].nx(), flow_codes.size());
+        return 1;
+    }
+
+    // [spkr, perceiver block] on mtl, [spkr, reference speech rows] on turbo
+    const bool mtl = cbx_is_mtl(ctx->ctx_a);
+    std::vector<float> pse;
+    if (mtl) {
+        pse = std::move(t3_rows);
+    }
+    std::vector<float> rows;
+    if (cbx_ref_cond(ctx, entries[4], pse, rows) != 0) {
+        return 1;
+    }
+    if (!mtl) {
+        rows.insert(rows.end(), t3_rows.begin(), t3_rows.end());
+    }
+    if (rows.size() != (size_t) audio_tokens->n_tokens * (size_t) ctx->n_embd_out()) {
+        LOG_ERR("%s: unexpected number of conditioning rows\n", __func__);
+        return 1;
+    }
+
+    ctx->gen_ref.tokens = std::move(flow_codes);
+    ctx->gen_ref.feat.assign(entries[3].get_ro_buf().begin(), entries[3].get_ro_buf().begin() + n_feat);
+    ctx->gen_ref.spk = std::move(xvec);
+
+    out_embd = std::move(rows);
+    return 0;
+}
+
 const char * mtmd_gen_audio_voice_name(const mtmd_context * ctx, int32_t idx) {
     if (!ctx->ctx_gen_a || idx < 0) {
         return nullptr;
@@ -2326,8 +2750,7 @@ mtmd_bitmap * mtmd_gen_audio_get_voice(const mtmd_context * ctx, const char * na
     if (it == voices.end()) {
         return nullptr;
     }
-    const uint32_t n_embd = clip_n_mmproj_embd(ctx->ctx_gen_a);
-    return mtmd_bitmap::init_voice(it->first, it->second.data(), n_embd, it->second.size() / n_embd);
+    return mtmd_bitmap::init_voice(it->first, it->second, clip_n_mmproj_embd(ctx->ctx_gen_a));
 }
 
 mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) {
@@ -2378,6 +2801,9 @@ struct mtmd_gen_state {
     std::vector<float>   out_embd;  // next-step hidden state fed back to backbone (GEN_CODE)
     std::vector<float>   out_audio; // decoded PCM samples for the current frame (GEN_WAV)
 
+    // voice reference for GEN_WAV, from a speaker clip or a predefined voice (chatterbox)
+    clip_gen_ref ref;
+
     // lfm3a: the codec graph yields STFT frames, the ISTFT runs on the host
     std::unique_ptr<mtmd_audio_streaming_istft> istft;
 
@@ -2484,6 +2910,19 @@ static int32_t mtmd_gen_wav_lfm3a(mtmd_context * ctx, mtmd_gen_state * state, co
     return 0;
 }
 
+int32_t mtmd_gen_state_set_ref(mtmd_context * ctx, mtmd_gen_state * state) {
+    if (!state || state->ctx_clip != ctx->ctx_gen_a) {
+        LOG_ERR("%s: invalid gen state\n", __func__);
+        return 1;
+    }
+    if (ctx->gen_ref.tokens.empty()) {
+        LOG_ERR("%s: the last encode produced no reference data\n", __func__);
+        return 1;
+    }
+    state->ref = ctx->gen_ref;
+    return 0;
+}
+
 static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state * state, const mtmd_gen_inp * inp, mtmd_gen_out * out) {
     clip_ctx * ctx_clip = ctx->ctx_gen_a;
     if (!ctx_clip) {
@@ -2565,6 +3004,44 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state *
     }
 
     // MTMD_GEN_PROCESS_TYPE_GEN_WAV
+    if (clip_get_projector_type(ctx_clip) == PROJECTOR_TYPE_CHATTERBOX) {
+        if (!inp->codes || inp->n_codes == 0) {
+            LOG_ERR("%s: codes required for gen_wav\n", __func__);
+            return 1;
+        }
+        if (state->ref.tokens.empty()) {
+            LOG_ERR("%s: gen_wav needs a voice reference, see mtmd_gen_state_set_ref()\n", __func__);
+            return 1;
+        }
+        std::vector<int32_t> in_codes(inp->codes, inp->codes + inp->n_codes);
+
+        // the batch entry is unused, present to satisfy the encode interface
+        clip_image_f32 dummy;
+        dummy.set_size({1, 1}, false, true);
+        dummy.cpy_buf(std::vector<float>(1, 0.0f));
+        clip_image_f32_batch batch;
+        batch.is_audio = true;
+        batch.entries.push_back(std::move(dummy));
+
+        clip_encode_params params;
+        params.imgs        = &batch;
+        params.n_threads   = ctx->n_threads;
+        params.gen_process = CLIP_GEN_PROCESS_TTS;
+        params.codes       = &in_codes;
+        params.out_audio   = &state->out_audio;
+        params.ref         = &state->ref;
+        params.gen_state   = state->clip_state;
+
+        if (!clip_encode(ctx_clip, &params)) {
+            LOG_ERR("%s: clip_encode failed (gen_wav)\n", __func__);
+            return 1;
+        }
+
+        out->audio     = state->out_audio.data();
+        out->n_samples = state->out_audio.size();
+        return 0;
+    }
+
     const bool has_codes = inp->codes && inp->n_codes > 0;
     const bool has_feats = inp->feats && inp->n_feats > 0;
     if (has_codes == has_feats) {
@@ -2788,36 +3265,51 @@ int mtmd_get_audio_sample_rate(const mtmd_context * ctx) {
     return clip_get_hparams(ctx->ctx_a)->audio_sample_rate;
 }
 
-mtmd_video_config mtmd_get_video_config(const mtmd_context * ctx,
-                                        int32_t src_width, int32_t src_height,
-                                        int32_t src_n_frames, float src_fps) {
-    mtmd_video_config cfg{MTMD_VIDEO_TEMPLATE_DEFAULT, 0.0f, 0, 0, 0};
-    if (!ctx->ctx_v || clip_get_projector_type(ctx->ctx_v) != PROJECTOR_TYPE_LFM3VL) {
+mtmd_video_config mtmd_video_config_init(mtmd_video_template tmpl, const mtmd_image_preprocessor * preproc,
+                                         int32_t src_width, int32_t src_height,
+                                         int32_t src_n_frames, float src_fps) {
+    mtmd_video_config cfg{tmpl, 0.0f, 0, 0, 0};
+    if (tmpl == MTMD_VIDEO_TEMPLATE_DEFAULT) {
         return cfg;
     }
 
-    // matches Lfm3VlVideoProcessor: 2 fps sampling with a [4, 768] frame clamp
+    // matches Qwen3VLVideoProcessor and Lfm3VlVideoProcessor: 2 fps sampling with a [4, 768] frame clamp
     // LFM3-Omni samples at most 384 frames and hears the soundtrack
-    const bool omni = ctx->video_has_soundtrack();
-    cfg.tmpl = omni ? MTMD_VIDEO_TEMPLATE_LFM3_OMNI : MTMD_VIDEO_TEMPLATE_LFM3;
-    cfg.fps  = 2.0f;
+    cfg.fps = 2.0f;
     if (src_n_frames <= 0 || src_fps <= 0.0f) {
         return cfg;
     }
 
     int32_t n = (int32_t)((double) src_n_frames / src_fps * cfg.fps);
-    n = std::clamp(n, 4, omni ? 384 : 768);
+    n = std::clamp(n, 4, tmpl == MTMD_VIDEO_TEMPLATE_LFM3_OMNI ? 384 : 768);
     n = std::min(n, src_n_frames);
     cfg.n_frames = n;
 
     clip_image_size size;
-    if (ctx->image_preproc && ctx->image_preproc->get_video_frame_size(n, src_width, src_height, size)) {
+    if (preproc && preproc->get_video_frame_size(n, src_width, src_height, size)) {
         cfg.frame_width  = size.width;
         cfg.frame_height = size.height;
     }
     return cfg;
 }
 
+mtmd_video_config mtmd_get_video_config(const mtmd_context * ctx,
+                                        int32_t src_width, int32_t src_height,
+                                        int32_t src_n_frames, float src_fps) {
+    mtmd_video_template tmpl = MTMD_VIDEO_TEMPLATE_DEFAULT;
+    switch (ctx->ctx_v ? clip_get_projector_type(ctx->ctx_v) : PROJECTOR_TYPE_UNKNOWN) {
+        case PROJECTOR_TYPE_LFM3VL:
+            tmpl = ctx->video_has_soundtrack() ? MTMD_VIDEO_TEMPLATE_LFM3_OMNI : MTMD_VIDEO_TEMPLATE_LFM3;
+            break;
+        case PROJECTOR_TYPE_QWEN3VL:
+            tmpl = MTMD_VIDEO_TEMPLATE_QWEN3VL;
+            break;
+        default:
+            break;
+    }
+    return mtmd_video_config_init(tmpl, ctx->image_preproc.get(), src_width, src_height, src_n_frames, src_fps);
+}
+
 const char * mtmd_get_marker(const mtmd_context * ctx) {
     return ctx->media_marker.c_str();
 }

tools/mtmd/mtmd.h

@@ -63,6 +63,7 @@ enum mtmd_video_template {
     MTMD_VIDEO_TEMPLATE_DEFAULT,   // "Video:" prefix, then a "[XmY.YYs]" timestamp every N ms
     MTMD_VIDEO_TEMPLATE_LFM3,      // "<X.X seconds>" before every group of 2 frames
     MTMD_VIDEO_TEMPLATE_LFM3_OMNI, // same as LFM3, then the soundtrack of each group, as an audio bitmap after its frames
+    MTMD_VIDEO_TEMPLATE_QWEN3VL,   // "<X.X seconds>" before every group of 2 frames, no markers around the whole video
 };
 
 // opaque types
@@ -465,6 +466,7 @@ enum mtmd_gen_audio_type {
     MTMD_GEN_AUDIO_TYPE_POCKETTTS,
     MTMD_GEN_AUDIO_TYPE_SOPRANO,
     MTMD_GEN_AUDIO_TYPE_LFM3A,
+    MTMD_GEN_AUDIO_TYPE_CHATTERBOX,
 };
 
 struct mtmd_gen_audio_info {
@@ -476,6 +478,8 @@ struct mtmd_gen_audio_info {
 
 MTMD_API struct mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx);
 
+// read a named tensor of the gen audio model as F32, returns the element count (0 if not found), out may be null to query the size
+MTMD_API size_t mtmd_gen_audio_read_tensor(mtmd_context * ctx, const char * name, float * out, size_t n_max);
 // name of the idx-th predefined voice, nullptr when idx is out of range
 MTMD_API const char * mtmd_gen_audio_voice_name(const mtmd_context * ctx, int32_t idx);
 
@@ -543,6 +547,10 @@ MTMD_API void             mtmd_gen_state_free(mtmd_gen_state * state);
 // clear the decoder state (cold start), keep the RNG stream
 MTMD_API void             mtmd_gen_state_reset(mtmd_gen_state * state);
 
+// copy the gen-side reference data of the last mtmd_encode_chunk() call into the state (speaker clip or predefined voice)
+// returns non-zero if the last encode produced no such data
+MTMD_API int32_t          mtmd_gen_state_set_ref(mtmd_context * ctx, mtmd_gen_state * state);
+
 // defaults tuned for the loaded pipeline, callers override only what they care about
 MTMD_API struct mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx);
 

tools/tts/README.md

@@ -121,3 +121,36 @@ The mmproj holds the audio encoder (for ASR), the depthformer and the codec deco
 python convert_hf_to_gguf.py path/to/lfm3-audio --outfile lfm3-audio.gguf
 python convert_hf_to_gguf.py path/to/lfm3-audio --mmproj --outfile mmproj-lfm3-audio.gguf
 ```
+
+## Chatterbox
+
+Two variants are supported: [Chatterbox Turbo](https://huggingface.co/ResembleAI/chatterbox-turbo) (English) and [Chatterbox Multilingual V3](https://huggingface.co/ResembleAI/chatterbox) (23 languages).
+
+Available params:
+- `--tts-speaker-file` should point to a speaker reference audio file (wav, mp3). Optional, the predefined voice is used without it
+- `--tts-voice` can be `default`, the voice shipped with the model (`conds.pt`)
+- `--tts-lang` sets the language of the multilingual variant, for example `fr` (default: none). It is not used by Turbo
+
+Example usage:
+
+```sh
+llama-tts -m chatterbox-mtl.gguf \
+    -mm mmproj-chatterbox-mtl.gguf \
+    -p "Bonjour tout le monde" \
+    --tts-lang fr \
+    --tts-speaker-file speaker.mp3 \
+    --output out.wav
+```
+
+Note: the multilingual variant runs classifier-free guidance with a second sequence. `llama-tts` sets this up by itself, but `llama-server` does not support it yet.
+
+**Note for GGUF conversion:**
+
+The upstream repositories ship raw checkpoints without `config.json`. Download the files below into a directory, then convert it as usual (the variant is detected from the talker file):
+- Turbo: `t3_turbo_v1.safetensors`, `s3gen_meanflow.safetensors`, `ve.safetensors`, `conds.pt` and the GPT-2 tokenizer files (`vocab.json`, `merges.txt`, `added_tokens.json`, `special_tokens_map.json`, `tokenizer_config.json`)
+- Multilingual V3: `t3_mtl23ls_v3.safetensors`, `s3gen_v3.safetensors`, `ve.safetensors`, `conds.pt`, `mtl_tokenizer.json`
+
+```sh
+python convert_hf_to_gguf.py path/to/chatterbox --outfile chatterbox.gguf
+python convert_hf_to_gguf.py path/to/chatterbox --mmproj --outfile mmproj-chatterbox.gguf
+```

tools/tts/tts.cpp

@@ -69,6 +69,11 @@ int main(int argc, char ** argv) {
     // always enable embd, so that we can pass hidden states to the audio generation helper
     params.embedding = true;
 
+    // for chatterbox multilingual: its cfg pairs slot i with sequence n_parallel + i in a unified kv cache
+    // TODO: this applies to every model, find a better way for a pipeline to request extra sequences
+    params.n_parallel *= 2;
+    params.kv_unified  = true;
+
     llama_backend_init();
     llama_numa_init(params.numa);
 

tools/ui/.env.example

@@ -1,3 +1,21 @@
 VITE_PUBLIC_APP_NAME='llama-ui'
 VITE_DEBUG='true'
 VITE_PUBLIC_SERVER_ORIGIN='http://localhost:8033'
+
+# Demo videos (npm run demo:pr / demo:docs), see demo/README.md
+# Where demo scripts and videos are stored; unset keeps videos in demo/videos
+# DEMO_STORAGE='hf://buckets/<owner>/<bucket>'
+# DEMO_STORAGE='file:///path/to/dir'
+# Public base URL of DEMO_STORAGE, used to link videos in PR descriptions
+# DEMO_STORAGE_PUBLIC_URL='https://...'
+# HF dataset for demo conversations (llama.app traces): `demo:traces` pushes the
+# built-in ones there, and demos can import any <name>.jsonl from it
+# DEMO_TRACES_DATASET='<owner>/<dataset>'
+# DEMO_GITHUB_REPO='ggml-org/llama.cpp'
+# DEMO_SERVER_ORIGIN='http://localhost:9931'
+# DEMO_LLAMA_SERVER='../../build/bin/llama-server'
+# DEMO_COLOR_SCHEME='light'
+# Device pixel ratio demos render and record at (1440x900 viewport)
+# DEMO_SCALE='2'
+# Longer side of the GIF preview, in pixels; it shrinks further to stay under 5 MB
+# DEMO_PREVIEW_WIDTH='800'

tools/ui/README.md

@@ -710,6 +710,8 @@ npm run test:server   # Server-side unit tests
 npm run test:ui       # Storybook visual tests
 ```
 
+Locally, the browser tests (e2e, client and Storybook) run in your installed Google Chrome with its sandbox on, since Chrome updates itself and Playwright's bundled Chromium lags behind security fixes. CI (`CI` set) uses the bundled Chromium from `npx playwright install`.
+
 ### Storybook Development
 
 ```bash

tools/ui/demo/.gitignore

@@ -0,0 +1,5 @@
+.runs/
+test-results/
+videos/
+*.demo.ts
+/traces/

tools/ui/demo/README.md

@@ -0,0 +1,133 @@
+# Demo videos
+
+Scripted Playwright walkthroughs of the UI, exported as mp4 with an animated GIF preview, for pull requests and documentation. The harness lives here; the demo scripts don't:
+
+- **PR demos** live in the PR description, as a `ts demo` code block.
+- **Docs demos** live in the storage target, under `scripts/docs/`.
+
+```
+demo/
+├── cli.ts      demo:pr, demo:docs and demo:traces
+├── storage.ts  storage targets (hf://, file://)
+├── helpers.ts  `test` with a `demo` fixture: cursor, captions, human-like timing
+├── recorder.ts screencast recording and mp4 export
+├── fixtures/   fixture MCP servers, seeded settings and demo conversations
+└── videos/     output (git-ignored)
+```
+
+## Setup
+
+### Tools
+
+| Tool                                                 | Used for                                                                                                                                                                                                          |
+| ---------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| Node 22.18+                                          | runs `demo/cli.ts` directly (TypeScript type stripping)                                                                                                                                                           |
+| `npx playwright install chromium`                    | the browser demos record in, with its sandbox on                                                                                                                                                                  |
+| `ffmpeg`                                             | mp4 encoding and the GIF preview                                                                                                                                                                                  |
+| `gh`, logged in                                      | reads the PR description and writes the `## Demo` section, so you need edit rights on the PR                                                                                                                      |
+| `hf`, logged in with a write token (`hf auth login`) | uploads to the bucket and the traces dataset                                                                                                                                                                      |
+| `llama-server`                                       | `build/bin/llama-server`, or one already running at `DEMO_SERVER_ORIGIN` (default `:9931`). It runs in router mode, so demos can pick any model in the local cache; scripts that send messages need at least one. |
+
+`gh` is only needed for `demo:pr`, and `hf` only with the bucket or dataset below. Without them, the video stays in `demo/videos/`.
+
+Videos are recorded through Chromium's screencast at `DEMO_SCALE` (2x by default, so 2880x1800) and 60 fps, starting once the page has settled.
+
+### Hugging Face bucket (videos, previews, scripts)
+
+The bucket must be public, since GitHub loads the preview from it:
+
+```bash
+hf buckets create <owner>/llama-ui-demos
+```
+
+Then set it in `tools/ui/.env` (see `.env.example`):
+
+```bash
+DEMO_STORAGE='hf://buckets/<owner>/llama-ui-demos'
+DEMO_STORAGE_PUBLIC_URL='https://huggingface.co/buckets/<owner>/llama-ui-demos/resolve'
+```
+
+Layout: `videos/pr/<number>/<sha>.{mp4,gif,demo.ts}`, `videos/docs/<name>.{mp4,gif}` and `scripts/docs/*.demo.ts` (the docs scripts). Without `DEMO_STORAGE_PUBLIC_URL`, videos are uploaded but not linked in the PR. `DEMO_STORAGE='file:///some/dir'` stores them in a local folder instead.
+
+### Hugging Face dataset (conversations, optional)
+
+Only needed for demos that import conversations other than the built-in ones, or to publish the built-in ones with `npm run demo:traces`. It can be private: it is read with your `hf` login.
+
+```bash
+hf repos create <owner>/llama-ui-demo-traces --repo-type dataset --private
+```
+
+```bash
+DEMO_TRACES_DATASET='<owner>/llama-ui-demo-traces'
+```
+
+### Other storage
+
+Only Hugging Face buckets and `file://` folders are supported for now. S3, R2 or similar would be one more class in `storage.ts` (`upload` and `downloadDir`) plus the matching `DEMO_STORAGE_PUBLIC_URL`.
+
+## Recording a pull request
+
+Add the script to the PR description:
+
+````md
+<details>
+<summary>Demo script</summary>
+
+```ts demo
+test('demo', async ({ demo, page }) => {
+	await demo.click(page.getByRole('button', { exact: true, name: 'Settings' }));
+	await demo.caption('What this shows');
+});
+```
+
+</details>
+````
+
+Then run:
+
+```bash
+npm run demo:pr -- 31303
+```
+
+A PR by someone else runs its code and demo script on your machine, with your `gh` and `hf` logins, so `demo:pr` refuses it until you've read its diff and demo block and pass the head commit you read: `npm run demo:pr -- 31303 --trust <sha>`. If the PR gets new commits, it refuses again.
+
+It checks out the PR head in a temporary worktree, records the script against it, uploads `videos/pr/<number>/<sha>.{mp4,gif,demo.ts}` and adds or updates a `## Demo` section in the PR description with the preview linking to the video.
+
+A large PR can show its changes in several shorter videos: each `test()` in the block is recorded as its own video (`<sha>-1.mp4`, `<sha>-2.mp4`, ...), and the `## Demo` section shows their previews in a table, each headed by its test's title: up to four portrait (phone) videos to a row, or two landscape ones. Helpers and `test.use` at the top of the block are shared by all of them.
+
+````ts
+test('Header and safe areas', async ({ demo, page }) => {
+	// ...
+});
+
+test('Sidebar gestures', async ({ demo, page }) => {
+	// ...
+});
+``` GitHub has no API for attachments, so for an inline player drag `demo/videos/pr-<number>.mp4` into the description instead.
+
+## Recording docs
+
+```bash
+npm run demo:docs
+````
+
+Records every `scripts/docs/*.demo.ts` from the storage target and uploads `videos/docs/<name>.{mp4,gif}`.
+
+## Writing a demo
+
+Scripts get `test` and `expect` from `helpers.ts` without importing them. `demo.click`, `demo.hover` and `demo.type` glide a visible cursor to the target before acting, with timing varied like a person's (seeded per demo, so re-takes match); `demo.scroll` turns the mouse wheel over a target; `demo.zoom` pinches in on a control or a point to show a detail too small to read in the frame, holds, and pinches back out (with the fingertips drawn on a touch device); `demo.caption` shows a caption at the top of the frame. Use `test.use({ settings, startPath })` to seed other settings or start on another route.
+
+## Conversations
+
+Demos that need existing chats import them with `test.use({ traces: ['long-chat'] })`, through Settings > Import/Export before the video starts. Traces use the export format (llama.app JSONL), so the Hub shows them as Llama traces:
+
+- `fixtures/traces.ts` generates the built-in ones: `long-chat` (36 tall exchanges) and `agentic-debugging` (long reasoning, shell output and code blocks).
+- Any other name is downloaded as `<name>.jsonl` from `DEMO_TRACES_DATASET`, e.g. an exported conversation uploaded there.
+
+`npm run demo:traces` writes the built-in traces to `demo/traces/` and pushes them to `DEMO_TRACES_DATASET`.
+
+## Trying a script locally
+
+To try a script locally, put it in `demo/` as `<name>.demo.ts`, importing from `./helpers.ts`, and run `npm run demo:record -- demo/<name>`. `npm run demo:codegen` opens Playwright codegen against `:5180` for a first draft.
+
+The config starts, or reuses when already running, the fixture MCP servers on `:8765`, `llama-server` on `:9931` (router mode, `--tools all`) and the UI's vite dev server.

tools/ui/demo/cli.ts

@@ -0,0 +1,411 @@
+/**
+ * Records demo videos for pull requests and docs.
+ *
+ *   node demo/cli.ts pr <number>   record the ```ts demo block from a PR description, a video per test
+ *                                  (--trust <sha> to record someone else's PR at that head commit)
+ *   node demo/cli.ts docs          record every docs/*.demo.ts script in DEMO_STORAGE
+ *   node demo/cli.ts traces        write the built-in traces, and push them to DEMO_TRACES_DATASET
+ *
+ * Settings come from tools/ui/.env (see .env.example).
+ */
+
+import { DEMO_TRACES, generateTrace } from './fixtures/traces.ts';
+import { type DemoStorage, getDemoStorage, getPublicUrl } from './storage.ts';
+import { execFileSync } from 'node:child_process';
+import { createHash } from 'node:crypto';
+import {
+	existsSync,
+	mkdirSync,
+	readdirSync,
+	readFileSync,
+	rmSync,
+	statSync,
+	symlinkSync,
+	writeFileSync
+} from 'node:fs';
+import { createServer } from 'node:net';
+import { tmpdir } from 'node:os';
+import { basename, join, relative, resolve } from 'node:path';
+
+const DEMO_DIR = import.meta.dirname;
+const UI_DIR = resolve(DEMO_DIR, '..');
+const RUNS_DIR = join(DEMO_DIR, '.runs');
+const VIDEOS_DIR = join(DEMO_DIR, 'videos');
+const TRACES_DIR = join(DEMO_DIR, 'traces');
+const DEMO_BLOCK = /```(?:ts|typescript)[ \t]+demo[ \t]*\r?\n([\s\S]*?)```/;
+const HELPERS_IMPORT = /^import\s[^;]*from\s+['"][^'"]*helpers(?:\.ts)?['"];?\s*$/gm;
+const SECTION_START = '<!-- demo:start -->';
+const SECTION_END = '<!-- demo:end -->';
+/** GitHub's image proxy refuses larger images, so a bigger preview shows as a broken image. */
+const PREVIEW_MAX_BYTES = 5 * 1024 * 1024;
+const envFile = join(UI_DIR, '.env');
+
+if (existsSync(envFile)) process.loadEnvFile(envFile);
+
+const REPO = process.env.DEMO_GITHUB_REPO ?? 'ggml-org/llama.cpp';
+
+function run(
+	cmd: string,
+	args: string[],
+	opts: { cwd?: string; env?: NodeJS.ProcessEnv } = {}
+): string {
+	return execFileSync(cmd, args, {
+		cwd: opts.cwd ?? UI_DIR,
+		encoding: 'utf8',
+		env: opts.env
+	}).trim();
+}
+
+/** Writes a demo script into .runs/, importing the harness helpers. */
+function writeDemoScript(name: string, source: string): string {
+	const file = join(RUNS_DIR, `${name}.demo.ts`);
+	const helpers = relative(RUNS_DIR, join(DEMO_DIR, 'helpers.ts'));
+	const body = source.replace(HELPERS_IMPORT, '').trim();
+
+	mkdirSync(RUNS_DIR, { recursive: true });
+	writeFileSync(file, `import { expect, test } from './${helpers}';\n\n${body}\n`);
+
+	return file;
+}
+
+async function freePort(): Promise<number> {
+	return new Promise((done, fail) => {
+		const server = createServer();
+
+		server.once('error', fail);
+		server.listen(0, () => {
+			const { port } = server.address() as { port: number };
+
+			server.close(() => done(port));
+		});
+	});
+}
+
+/** A video recorded from one `test()` of a script. */
+interface Video {
+	mp4: string;
+	title: string;
+}
+
+/** Records the given scripts; returns a video for each of their tests, in order. */
+async function record(scripts: string[], uiDir = UI_DIR): Promise<Video[]> {
+	const port = await freePort();
+	const manifests = scripts.map((s) => join(VIDEOS_DIR, `${basename(s, '.demo.ts')}.videos.jsonl`));
+
+	for (const manifest of manifests) rmSync(manifest, { force: true });
+
+	execFileSync(
+		'npx',
+		[
+			'playwright',
+			'test',
+			'--config=demo/playwright.config.ts',
+			...scripts.map((s) => relative(UI_DIR, s))
+		],
+		{
+			cwd: UI_DIR,
+			env: { ...process.env, DEMO_UI_DIR: uiDir, DEMO_UI_PORT: String(port) },
+			stdio: 'inherit'
+		}
+	);
+
+	// helpers.ts lists each video it encodes in the script's manifest
+	return manifests.flatMap((manifest) =>
+		readFileSync(manifest, 'utf8')
+			.trim()
+			.split('\n')
+			.map((line) => JSON.parse(line) as Video)
+	);
+}
+
+/**
+ * Animated GIF preview that GitHub renders inline, its longer side DEMO_PREVIEW_WIDTH (so a
+ * phone's portrait video is that tall), shrunk until it fits under PREVIEW_MAX_BYTES.
+ */
+function makePreview(mp4: string): string {
+	const gif = mp4.replace(/\.mp4$/, '.gif');
+
+	for (let side = Number(process.env.DEMO_PREVIEW_WIDTH ?? 800); ; side = Math.round(side * 0.8)) {
+		execFileSync('ffmpeg', [
+			'-y',
+			'-loglevel',
+			'error',
+			'-i',
+			mp4,
+			'-vf',
+			`fps=10,scale='if(gt(iw,ih),${side},-2)':'if(gt(iw,ih),-2,${side})':flags=lanczos,split[a][b];[a]palettegen=stats_mode=diff[p];[b][p]paletteuse=dither=bayer:bayer_scale=5:diff_mode=rectangle`,
+			'-loop',
+			'0',
+			gif
+		]);
+
+		if (statSync(gif).size <= PREVIEW_MAX_BYTES || side < 320) return gif;
+	}
+}
+
+function upload(storage: DemoStorage, files: Record<string, string>): void {
+	for (const [key, file] of Object.entries(files)) {
+		storage.upload(file, key);
+		console.log(`uploaded ${key}`);
+	}
+}
+
+/**
+ * Public URL with a content hash, so a re-recorded file at the same key gets a new
+ * URL: GitHub's image proxy caches previews by URL.
+ */
+function versionedUrl(key: string, file: string): null | string {
+	const url = getPublicUrl(key);
+	const version = createHash('sha1').update(readFileSync(file)).digest('hex').slice(0, 10);
+
+	return url && `${url}?v=${version}`;
+}
+
+/** Whether a video is taller than wide, like a phone's. */
+function isPortrait(mp4: string): boolean {
+	const [width, height] = run('ffprobe', [
+		'-v',
+		'error',
+		'-select_streams',
+		'v:0',
+		'-show_entries',
+		'stream=width,height',
+		'-of',
+		'csv=p=0',
+		mp4
+	])
+		.split(',')
+		.map(Number);
+
+	return height > width;
+}
+
+/**
+ * One preview, or a table of them when the script has several tests, each headed by its
+ * test's title: up to four portrait videos to a row, or two landscape ones, so each stays
+ * large enough to read.
+ */
+function demoSection(
+	links: { gif: string; mp4: string; title: string }[],
+	sha: string,
+	columns: number
+): string {
+	const preview = ({ gif, mp4, title }: (typeof links)[number]) =>
+		`[![${title.replace(/[[\]|]/g, '')}](${gif})](${mp4})`;
+	const cell = (text: string) => text.replace(/\|/g, '\\|');
+	const rows = Array.from({ length: Math.ceil(links.length / columns) }, (_, i) =>
+		links.slice(i * columns, (i + 1) * columns)
+	);
+	const pad = (cells: string[]) => [...cells, ...Array(columns - cells.length).fill(' ')];
+	const body =
+		links.length === 1
+			? [preview({ ...links[0], title: 'Demo' })]
+			: rows.flatMap((row, i) => [
+					// a table has one header row, so later rows title their videos in bold
+					`| ${pad(row.map((l) => (i === 0 ? cell(l.title) : `**${cell(l.title)}**`))).join(' | ')} |`,
+					...(i === 0 ? [`|${Array(columns).fill('---').join('|')}|`] : []),
+					`| ${pad(row.map(preview)).join(' | ')} |`
+				]);
+	const note =
+		links.length === 1 ? 'the preview for the full video' : 'a preview for its full video';
+
+	return [
+		SECTION_START,
+		...body,
+		'',
+		`Recorded at ${sha.slice(0, 10)}; click ${note}.`,
+		SECTION_END
+	].join('\n');
+}
+
+/** Replaces the managed demo section, or adds a `## Demo` section before the details. */
+function withDemoSection(body: string, section: string): string {
+	const start = body.indexOf(SECTION_START);
+	const end = body.indexOf(SECTION_END);
+
+	if (start !== -1 && end > start) {
+		return body.slice(0, start) + section + body.slice(end + SECTION_END.length);
+	}
+
+	const block = `## Demo\n\n${section}\n\n`;
+	const anchor = body.search(/^## (Additional information|Requirements)\b/m);
+
+	return anchor === -1
+		? `${body.trimEnd()}\n\n${block}`
+		: body.slice(0, anchor) + block + body.slice(anchor);
+}
+
+async function recordPr(number: string, trustedSha?: string): Promise<void> {
+	const pr = JSON.parse(
+		run('gh', ['pr', 'view', number, '--repo', REPO, '--json', 'author,body,headRefOid'])
+	) as { author: { login: string }; body: string; headRefOid: string };
+	const source = DEMO_BLOCK.exec(pr.body)?.[1];
+
+	if (!source) throw new Error(`PR #${number} has no \`\`\`ts demo block in its description`);
+
+	const sha = pr.headRefOid;
+
+	// the PR's code and demo script run on this machine, with the gh and hf logins,
+	// so someone else's PR only runs at a head commit you've read and passed to --trust
+	if (pr.author.login !== run('gh', ['api', 'user', '--jq', '.login'])) {
+		if (!trustedSha) {
+			throw new Error(
+				`PR #${number} is by ${pr.author.login}, and recording it runs its code here: ` +
+					`read its diff and demo block, then run again with --trust ${sha.slice(0, 10)}`
+			);
+		}
+
+		if (trustedSha.length < 7 || !sha.startsWith(trustedSha)) {
+			throw new Error(`PR #${number} head is ${sha.slice(0, 10)}, not the trusted ${trustedSha}`);
+		}
+	}
+
+	const worktree = join(tmpdir(), `llama-ui-demo-pr-${number}`);
+
+	run('git', ['fetch', '--no-tags', `https://github.com/${REPO}.git`, `pull/${number}/head`]);
+
+	if (run('git', ['rev-parse', 'FETCH_HEAD']) !== sha) {
+		throw new Error(`PR #${number} head moved while fetching, run again`);
+	}
+
+	rmSync(worktree, { force: true, recursive: true });
+	run('git', ['worktree', 'prune']);
+	run('git', ['worktree', 'add', '--detach', worktree, sha]);
+	symlinkSync(join(UI_DIR, 'node_modules'), join(worktree, 'tools/ui/node_modules'));
+
+	const name = `pr-${number}`;
+	const script = writeDemoScript(name, source);
+
+	try {
+		const videos = (await record([script], join(worktree, 'tools/ui'))).map((video) => ({
+			...video,
+			gif: makePreview(video.mp4)
+		}));
+		const storage = getDemoStorage();
+
+		for (const { gif, mp4 } of videos) console.log(`\nvideo:   ${mp4}\npreview: ${gif}`);
+
+		if (!storage) {
+			console.log('DEMO_STORAGE is not set: drag the mp4 into the PR description.');
+
+			return;
+		}
+
+		const prefix = `videos/pr/${number}/${sha.slice(0, 10)}`;
+		// a script with several tests gets <sha>-1.mp4, <sha>-2.mp4 and so on
+		const keys = videos.map((_, i) => (videos.length === 1 ? prefix : `${prefix}-${i + 1}`));
+
+		upload(storage, { [`${prefix}.demo.ts`]: script });
+		videos.forEach(({ gif, mp4 }, i) => {
+			upload(storage, { [`${keys[i]}.gif`]: gif, [`${keys[i]}.mp4`]: mp4 });
+		});
+
+		if (!getPublicUrl(prefix)) {
+			console.log('DEMO_STORAGE_PUBLIC_URL is not set: not linking the video in the PR.');
+
+			return;
+		}
+
+		const links = videos.map(({ gif, mp4, title }, i) => ({
+			gif: versionedUrl(`${keys[i]}.gif`, gif) ?? '',
+			mp4: versionedUrl(`${keys[i]}.mp4`, mp4) ?? '',
+			title
+		}));
+
+		execFileSync('gh', ['pr', 'edit', number, '--repo', REPO, '--body-file', '-'], {
+			cwd: UI_DIR,
+			input: withDemoSection(
+				pr.body,
+				demoSection(links, sha, Math.min(links.length, isPortrait(videos[0].mp4) ? 4 : 2))
+			),
+			stdio: ['pipe', 'inherit', 'inherit']
+		});
+		console.log(`linked the demo in https://github.com/${REPO}/pull/${number}`);
+	} finally {
+		rmSync(script, { force: true });
+		run('git', ['worktree', 'remove', '--force', worktree]);
+	}
+}
+
+async function recordDocs(): Promise<void> {
+	const storage = getDemoStorage();
+
+	if (!storage) throw new Error('DEMO_STORAGE is not set: docs demo scripts are read from it');
+
+	const downloaded = join(RUNS_DIR, 'docs-source');
+
+	rmSync(downloaded, { force: true, recursive: true });
+
+	if (!storage.downloadDir('scripts/docs', downloaded)) {
+		throw new Error('no scripts/docs/*.demo.ts in DEMO_STORAGE');
+	}
+
+	const scripts = readdirSync(downloaded)
+		.filter((f) => f.endsWith('.demo.ts'))
+		.map((f) =>
+			writeDemoScript(basename(f, '.demo.ts'), readFileSync(join(downloaded, f), 'utf8'))
+		);
+
+	try {
+		const videos = await record(scripts);
+
+		for (const { mp4 } of videos) {
+			const name = basename(mp4, '.mp4');
+
+			upload(storage, {
+				[`videos/docs/${name}.gif`]: makePreview(mp4),
+				[`videos/docs/${name}.mp4`]: mp4
+			});
+		}
+	} finally {
+		rmSync(RUNS_DIR, { force: true, recursive: true });
+	}
+}
+
+/** Writes the built-in traces to demo/traces/ and uploads them to the traces dataset. */
+function writeTraces(): void {
+	rmSync(TRACES_DIR, { force: true, recursive: true });
+	mkdirSync(TRACES_DIR, { recursive: true });
+
+	for (const name of Object.keys(DEMO_TRACES)) {
+		writeFileSync(join(TRACES_DIR, `${name}.jsonl`), `${generateTrace(name)}\n`);
+		console.log(`wrote demo/traces/${name}.jsonl`);
+	}
+
+	const dataset = process.env.DEMO_TRACES_DATASET;
+
+	if (!dataset) {
+		console.log('DEMO_TRACES_DATASET is not set: not uploading.');
+
+		return;
+	}
+
+	execFileSync(
+		'hf',
+		[
+			'upload',
+			dataset,
+			TRACES_DIR,
+			'.',
+			'--repo-type',
+			'dataset',
+			'--commit-message',
+			'Update demo traces'
+		],
+		{ stdio: 'inherit' }
+	);
+}
+
+const [command, arg, ...flags] = process.argv.slice(2);
+const trustIndex = flags.indexOf('--trust');
+
+if (command === 'pr' && arg) {
+	await recordPr(arg, trustIndex === -1 ? undefined : flags[trustIndex + 1]);
+} else if (command === 'docs') {
+	await recordDocs();
+} else if (command === 'traces') {
+	writeTraces();
+} else {
+	console.error('usage: node demo/cli.ts pr <number> [--trust <sha>] | docs | traces');
+	process.exit(1);
+}

tools/ui/demo/fixtures/demo-config.ts

@@ -0,0 +1,49 @@
+/** Shared fixture data for demo recordings (used by both Node and the browser seed). */
+
+export const DEMO_MCP_PORT = 8765;
+
+export interface DemoMcpServer {
+	id: string;
+	name: string;
+	title: string;
+	tools: { name: string; description: string; result: string }[];
+}
+
+export const DEMO_MCP_SERVERS: DemoMcpServer[] = [
+	{
+		id: 'weather',
+		name: 'weather',
+		title: 'Weather',
+		tools: [
+			{
+				description: 'Get the weather forecast for a city',
+				name: 'get_forecast',
+				result: 'Sunny, 22°C'
+			},
+			{
+				description: 'List active weather alerts for a region',
+				name: 'get_alerts',
+				result: 'No active alerts'
+			}
+		]
+	},
+	{
+		id: 'notes',
+		name: 'notes',
+		title: 'Notes',
+		tools: [
+			{ description: 'Search notes by keyword', name: 'search_notes', result: 'No notes found' },
+			{ description: 'Create a new note', name: 'create_note', result: 'Note created' },
+			{ description: 'List recent notes', name: 'list_notes', result: 'No notes yet' }
+		]
+	}
+];
+
+/** `mcpServers` setting value pointing the UI at the fixture servers. */
+export function demoMcpServerSettings(): object[] {
+	return DEMO_MCP_SERVERS.map((server) => ({
+		enabled: true,
+		id: `demo-${server.id}`,
+		url: `http://localhost:${DEMO_MCP_PORT}/${server.id}/mcp`
+	}));
+}

tools/ui/demo/fixtures/mcp-servers.ts

@@ -0,0 +1,91 @@
+/**
+ * Fixture MCP servers for demo recordings.
+ *
+ * Serves a few stateless Streamable HTTP MCP servers with canned tools, so
+ * demos show MCP groups without depending on real third-party servers.
+ *
+ * Usage: node demo/fixtures/mcp-servers.ts
+ */
+
+import { DEMO_MCP_PORT, DEMO_MCP_SERVERS, type DemoMcpServer } from './demo-config.ts';
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
+import { createServer, type IncomingMessage, type ServerResponse } from 'node:http';
+import { z } from 'zod';
+
+function buildServer(fixture: DemoMcpServer): McpServer {
+	const server = new McpServer({ name: fixture.name, title: fixture.title, version: '1.0.0' });
+
+	for (const tool of fixture.tools) {
+		server.registerTool(
+			tool.name,
+			{ description: tool.description, inputSchema: { query: z.string().optional() } },
+			async () => ({ content: [{ text: tool.result, type: 'text' }] })
+		);
+	}
+
+	return server;
+}
+
+function setCorsHeaders(res: ServerResponse): void {
+	res.setHeader('Access-Control-Allow-Origin', '*');
+	res.setHeader('Access-Control-Allow-Headers', '*');
+	res.setHeader('Access-Control-Allow-Methods', 'GET, POST, DELETE, OPTIONS');
+	res.setHeader('Access-Control-Expose-Headers', 'mcp-session-id, mcp-protocol-version');
+	// the UI runs with COEP require-corp
+	res.setHeader('Cross-Origin-Resource-Policy', 'cross-origin');
+}
+
+async function readBody(req: IncomingMessage): Promise<unknown> {
+	const chunks: Buffer[] = [];
+
+	for await (const chunk of req) chunks.push(chunk as Buffer);
+
+	const raw = Buffer.concat(chunks).toString();
+
+	return raw ? JSON.parse(raw) : undefined;
+}
+
+const httpServer = createServer(async (req, res) => {
+	setCorsHeaders(res);
+
+	if (req.method === 'OPTIONS') {
+		res.writeHead(204).end();
+
+		return;
+	}
+
+	const path = new URL(req.url ?? '/', 'http://localhost').pathname;
+	const fixture = DEMO_MCP_SERVERS.find((s) => path === `/${s.id}/mcp`);
+
+	if (!fixture) {
+		res.writeHead(path === '/health' ? 200 : 404).end();
+
+		return;
+	}
+
+	if (req.method !== 'POST') {
+		// stateless: no standalone SSE stream or session teardown
+		res.writeHead(405).end();
+
+		return;
+	}
+
+	const server = buildServer(fixture);
+	const transport = new StreamableHTTPServerTransport({
+		enableJsonResponse: true,
+		sessionIdGenerator: undefined
+	});
+
+	res.on('close', () => {
+		transport.close();
+		server.close();
+	});
+
+	await server.connect(transport);
+	await transport.handleRequest(req, res, await readBody(req));
+});
+
+httpServer.listen(DEMO_MCP_PORT, () => {
+	console.log(`demo MCP servers listening on http://localhost:${DEMO_MCP_PORT}`);
+});

tools/ui/demo/fixtures/traces.ts

@@ -0,0 +1,312 @@
+/**
+ * Synthetic conversations for demos, in the UI's export format (llama.app JSONL
+ * traces): a `session` record followed by one `message` record per message.
+ * Deterministic, so a trace and the demos recorded from it don't drift. Messages
+ * carry no model, so the UI doesn't flag one the server doesn't have.
+ */
+
+import { createHash } from 'node:crypto';
+
+const HARNESS = 'llama.app';
+const START = Date.UTC(2026, 9, 1, 9, 0, 0);
+
+interface ToolCall {
+	function: { arguments: string; name: string };
+	id: string;
+	type: 'function';
+}
+
+interface TraceMessage {
+	children: string[];
+	content: string;
+	convId: string;
+	id: string;
+	parent: null | string;
+	reasoningContent?: string;
+	role: 'assistant' | 'system' | 'tool' | 'user';
+	timestamp: number;
+	toolCallId?: string;
+	toolCalls?: ToolCall[];
+	type: 'root' | 'text';
+}
+
+/** Builds one linear conversation, message by message. */
+class Trace {
+	private calls = 0;
+	private readonly id: string;
+	private readonly messages: TraceMessage[] = [];
+	private readonly name: string;
+
+	assistant(content: string, reasoning?: string): this {
+		return this.push({ content, reasoningContent: reasoning, role: 'assistant' });
+	}
+
+	constructor(name: string) {
+		this.name = name;
+		this.id = uuid(name, 'conv');
+		this.push({ content: '', role: 'system', type: 'root' });
+	}
+
+	toJsonl(): string {
+		const last = this.messages.at(-1)!;
+		const session = {
+			currNode: last.id,
+			harness: HARNESS,
+			id: this.id,
+			lastModified: last.timestamp,
+			name: this.name,
+			type: 'session'
+		};
+
+		return [session, ...this.messages.map((message) => ({ message, type: 'message' }))]
+			.map((record) => JSON.stringify(record))
+			.join('\n');
+	}
+
+	/** An assistant turn calling one tool, followed by the tool's result. */
+	tool(name: string, args: Record<string, unknown>, result: string, reasoning?: string): this {
+		const id = uuid(this.name, `call-${this.calls++}`).replaceAll('-', '');
+
+		this.push({
+			content: '',
+			reasoningContent: reasoning,
+			role: 'assistant',
+			toolCalls: [{ function: { arguments: JSON.stringify(args), name }, id, type: 'function' }]
+		});
+
+		return this.push({ content: result, role: 'tool', toolCallId: id });
+	}
+
+	user(content: string): this {
+		return this.push({ content, role: 'user' });
+	}
+
+	private push(message: Partial<TraceMessage> & Pick<TraceMessage, 'content' | 'role'>): this {
+		const parent = this.messages.at(-1);
+		const index = this.messages.length;
+		const entry: TraceMessage = {
+			children: [],
+			convId: this.id,
+			id: uuid(this.name, `msg-${index}`),
+			parent: parent?.id ?? null,
+			timestamp: START + index * 20_000,
+			type: 'text',
+			...message
+		};
+
+		parent?.children.push(entry.id);
+		this.messages.push(entry);
+
+		return this;
+	}
+}
+
+function uuid(...parts: string[]): string {
+	const hex = createHash('sha1').update(parts.join(':')).digest('hex');
+
+	return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-4${hex.slice(13, 16)}-a${hex.slice(17, 20)}-${hex.slice(20, 32)}`;
+}
+
+const STEPS = [
+	['set up the project', 'cargo new forecast && cd forecast', 'Cargo.toml', 'toml'],
+	['parse command line arguments', 'cargo add clap --features derive', 'src/cli.rs', 'rust'],
+	['call the weather API', 'cargo add reqwest --features json,blocking', 'src/api.rs', 'rust'],
+	['deserialize the response', 'cargo add serde --features derive', 'src/model.rs', 'rust'],
+	['cache responses on disk', 'cargo add directories', 'src/cache.rs', 'rust'],
+	['print a table', 'cargo add comfy-table', 'src/table.rs', 'rust'],
+	['handle errors', 'cargo add anyhow', 'src/main.rs', 'rust'],
+	['add colors', 'cargo add owo-colors', 'src/style.rs', 'rust'],
+	['read a config file', 'cargo add toml', 'src/config.rs', 'rust'],
+	['write tests', 'cargo add --dev assert_cmd predicates', 'tests/cli.rs', 'rust'],
+	['set up CI', 'mkdir -p .github/workflows', '.github/workflows/ci.yml', 'yaml'],
+	['publish a release', 'cargo install cargo-dist', 'dist-workspace.toml', 'toml']
+] as const;
+
+function codeFor(step: number, file: string, lang: string): string {
+	const lines =
+		lang === 'rust'
+			? [
+					`// ${file}`,
+					'use anyhow::{Context, Result};',
+					'',
+					`pub fn step_${step}(input: &str) -> Result<Vec<String>> {`,
+					'    let mut out = Vec::new();',
+					'    for (i, line) in input.lines().enumerate() {',
+					'        let trimmed = line.trim();',
+					'        if trimmed.is_empty() {',
+					'            continue;',
+					'        }',
+					'        out.push(format!("{i:>3}: {trimmed}"));',
+					'    }',
+					`    Ok(out).context("step ${step} failed")`,
+					'}'
+				]
+			: lang === 'yaml'
+				? [
+						`# ${file}`,
+						'on: [push, pull_request]',
+						'jobs:',
+						'  test:',
+						'    runs-on: ubuntu-latest'
+					]
+				: [`# ${file}`, `name = "forecast"`, `version = "0.${step}.0"`, 'edition = "2024"'];
+
+	return `\`\`\`${lang}\n${lines.join('\n')}\n\`\`\``;
+}
+
+/** A long chat with tall messages, for scrolling and landing demos. */
+function longChat(): Trace {
+	const trace = new Trace('Building a weather CLI in Rust');
+
+	for (let i = 0; i < 36; i++) {
+		const [goal, command, file, lang] = STEPS[i % STEPS.length];
+		const round = Math.floor(i / STEPS.length);
+		const ask =
+			round === 0 ? `How do I ${goal}?` : `Can we revisit how we ${goal}? (pass ${round + 1})`;
+
+		trace.user(ask);
+		trace.assistant(
+			[
+				`To ${goal}, start from a clean working tree so the change stays easy to review.`,
+				'',
+				'```sh',
+				command,
+				'```',
+				'',
+				`Then add \`${file}\`:`,
+				'',
+				codeFor(i + 1, file, lang),
+				'',
+				'A few things to keep in mind:',
+				'',
+				`- Keep \`${file}\` focused on one job, so the next step only touches what it needs.`,
+				'- Run `cargo clippy` after each step; it catches most mistakes early.',
+				'- Commit once it builds, so you can always go back to a working version.',
+				'',
+				i % 3 === 0
+					? `Next we can look at how to ${STEPS[(i + 1) % STEPS.length][0]}.`
+					: 'Let me know if anything here does not build.'
+			].join('\n'),
+			`The user wants to ${goal}. Give the command, the file and a short checklist.`
+		);
+	}
+
+	return trace;
+}
+
+function testLog(failing: boolean): string {
+	const lines = ['running 48 tests'];
+
+	for (let i = 1; i <= 48; i++) {
+		const ok = !failing || (i !== 17 && i !== 31);
+
+		lines.push(
+			`test cache::tests::entry_${String(i).padStart(2, '0')} ... ${ok ? 'ok' : 'FAILED'}`
+		);
+	}
+
+	if (failing) {
+		lines.push(
+			'',
+			'failures:',
+			'',
+			'---- cache::tests::entry_17 stdout ----',
+			"thread 'cache::tests::entry_17' panicked at src/cache.rs:88:9:",
+			'assertion `left == right` failed: entry expired too early',
+			'  left: Some(1759312800)',
+			' right: Some(1759316400)',
+			'',
+			'---- cache::tests::entry_31 stdout ----',
+			"thread 'cache::tests::entry_31' panicked at src/cache.rs:88:9:",
+			'assertion `left == right` failed: entry expired too early',
+			'',
+			'test result: FAILED. 46 passed; 2 failed; 0 ignored; 0 measured; 0 filtered out'
+		);
+	} else {
+		lines.push('', 'test result: ok. 48 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out');
+	}
+
+	return `${lines.join('\n')}\n\n[exit code: ${failing ? 101 : 0}]`;
+}
+
+function reasoning(topic: string, paragraphs: number): string {
+	const out: string[] = [];
+
+	for (let i = 0; i < paragraphs; i++) {
+		out.push(
+			`${i === 0 ? `Looking at ${topic}.` : 'Another angle:'} The cache stores an expiry timestamp per entry, ` +
+				'and the failing tests are the ones created right at an hour boundary. If the expiry is computed ' +
+				'from local time while the check uses UTC, entries created near the boundary would expire an hour ' +
+				`early. Step ${i + 1}: confirm by reading how \`expires_at\` is built, then compare with \`is_fresh\`.`
+		);
+	}
+
+	return out.join('\n\n');
+}
+
+/** An agentic session with long reasoning, shell output and code, for nested scroller demos. */
+function agenticDebugging(): Trace {
+	const trace = new Trace('Fix the flaky cache tests');
+	const fixed = Array.from({ length: 60 }, (_, i) =>
+		i === 0
+			? 'use chrono::{DateTime, Duration, Utc};'
+			: i === 30
+				? '        let expires_at = Utc::now() + self.ttl; // was Local::now()'
+				: `        // cache.rs line ${i + 1}: entry handling stays the same`
+	).join('\n');
+
+	trace.user('The cache tests fail on CI every few runs. Can you find out why and fix it?');
+
+	for (let round = 0; round < 3; round++) {
+		trace.tool(
+			'exec_shell_command',
+			{ command: 'cargo test cache -- --test-threads=1' },
+			testLog(true),
+			reasoning('the failing cache tests', 4 + round)
+		);
+		trace.tool(
+			'exec_shell_command',
+			{ command: "grep -n 'expires_at\\|is_fresh' src/cache.rs" },
+			Array.from(
+				{ length: 30 },
+				(_, i) => `${80 + i}:        let expires_at = Local::now() + self.ttl; // entry ${i}`
+			).join('\n') + '\n\n[exit code: 0]',
+			reasoning('where the expiry is computed', 3)
+		);
+		trace.assistant(
+			[
+				`Round ${round + 1}: the expiry uses local time while the freshness check uses UTC.`,
+				'',
+				'```rust',
+				fixed,
+				'```',
+				'',
+				'I switched the expiry to `Utc::now()` so both sides use the same clock.'
+			].join('\n'),
+			reasoning('the fix', 5)
+		);
+		trace.user(round < 2 ? 'Still flaky on CI, can you check again?' : 'Run the tests once more.');
+	}
+
+	trace.tool(
+		'exec_shell_command',
+		{ command: 'cargo test cache' },
+		testLog(false),
+		reasoning('the final test run', 2)
+	);
+	trace.assistant(
+		'All 48 cache tests pass now, including the two that failed at the hour boundary.'
+	);
+
+	return trace;
+}
+
+/** Built-in traces, by name. */
+export const DEMO_TRACES: Record<string, () => Trace> = {
+	'agentic-debugging': agenticDebugging,
+	'long-chat': longChat
+};
+
+export function generateTrace(name: string): null | string {
+	return DEMO_TRACES[name]?.().toJsonl() ?? null;
+}

tools/ui/demo/helpers.ts

@@ -0,0 +1,500 @@
+import { demoMcpServerSettings } from './fixtures/demo-config.ts';
+import { generateTrace } from './fixtures/traces.ts';
+import { ScreencastRecorder } from './recorder.ts';
+import { type CDPSession, type Locator, type Page, test as base } from '@playwright/test';
+import { execFileSync } from 'node:child_process';
+import { appendFileSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs';
+import { basename, join, resolve } from 'node:path';
+
+const VIDEOS_DIR = resolve(import.meta.dirname, 'videos');
+/** How long a demo waits for the app to be ready before it starts recording. */
+const APP_READY_TIMEOUT_MS = 90000;
+/** How far below the viewport's top a pinch has to start for Chromium to zoom at all. */
+const ZOOM_TOP_INSET_PX = 64;
+
+/**
+ * Draws a fake cursor and click ripples, since recorded videos don't include
+ * the OS pointer, and a fingertip per touch, so a pinch shows as two fingers.
+ * Runs in the page via addInitScript.
+ */
+function installCursorOverlay(): void {
+	const mount = () => {
+		const style = document.createElement('style');
+
+		style.textContent = `
+			#demo-cursor { position: fixed; top: 0; left: 0; z-index: 2147483647; pointer-events: none;
+				width: 22px; height: 22px; margin: -3px 0 0 -3px; transition: transform 80ms ease-out; }
+			#demo-cursor.down { transform: scale(0.85); }
+			.demo-ripple { position: fixed; z-index: 2147483646; pointer-events: none; width: 36px; height: 36px;
+				margin: -18px 0 0 -18px; border-radius: 50%; background: rgba(59, 130, 246, 0.35);
+				animation: demo-ripple 450ms ease-out forwards; }
+			@keyframes demo-ripple { from { transform: scale(0.3); opacity: 1; } to { transform: scale(1.4); opacity: 0; } }
+			#demo-caption { position: fixed; left: 50%; top: 72px; z-index: 2147483647; pointer-events: none;
+				transform: translate(-50%, -8px); opacity: 0; transition: opacity 200ms ease, transform 200ms ease;
+				max-width: 80%; padding: 10px 18px; border-radius: 10px; background: rgba(17, 17, 17, 0.88);
+				color: #fff; font: 500 18px/1.4 system-ui, sans-serif; text-align: center; }
+			#demo-caption.visible { opacity: 1; transform: translate(-50%, 0); }
+			.demo-finger { position: fixed; z-index: 2147483646; pointer-events: none; width: 40px; height: 40px;
+				margin: -20px 0 0 -20px; border-radius: 50%; background: rgba(59, 130, 246, 0.3);
+				border: 2px solid rgba(59, 130, 246, 0.55); transition: opacity 200ms ease; }
+		`;
+		document.head.appendChild(style);
+
+		const cursor = document.createElement('div');
+
+		cursor.id = 'demo-cursor';
+		cursor.innerHTML = `<svg viewBox="0 0 24 24" width="22" height="22"><path d="M4 2l16 10.5-7 1.5-3.5 6.5z"
+			fill="#111" stroke="#fff" stroke-width="1.5" stroke-linejoin="round"/></svg>`;
+		document.documentElement.appendChild(cursor);
+
+		const caption = document.createElement('div');
+
+		caption.id = 'demo-caption';
+		document.documentElement.appendChild(caption);
+
+		document.addEventListener(
+			'mousemove',
+			(e) => (cursor.style.translate = `${e.clientX}px ${e.clientY}px`),
+			true
+		);
+		document.addEventListener(
+			'mousedown',
+			(e) => {
+				cursor.classList.add('down');
+				const ripple = document.createElement('div');
+
+				ripple.className = 'demo-ripple';
+				ripple.style.left = `${e.clientX}px`;
+				ripple.style.top = `${e.clientY}px`;
+				document.documentElement.appendChild(ripple);
+				setTimeout(() => ripple.remove(), 500);
+			},
+			true
+		);
+		document.addEventListener('mouseup', () => cursor.classList.remove('down'), true);
+
+		const fingers: HTMLDivElement[] = [];
+		const placeFingers = (e: TouchEvent) => {
+			Array.from(e.touches).forEach((touch, i) => {
+				const finger = (fingers[i] ??= Object.assign(document.createElement('div'), {
+					className: 'demo-finger'
+				}));
+
+				finger.style.left = `${touch.clientX}px`;
+				finger.style.top = `${touch.clientY}px`;
+				finger.style.opacity = '1';
+				document.documentElement.appendChild(finger);
+			});
+			fingers.slice(e.touches.length).forEach((finger) => (finger.style.opacity = '0'));
+		};
+
+		for (const type of ['touchstart', 'touchmove', 'touchend', 'touchcancel']) {
+			document.addEventListener(type, placeFingers as EventListener, {
+				capture: true,
+				passive: true
+			});
+		}
+	};
+
+	if (document.readyState === 'loading') {
+		document.addEventListener('DOMContentLoaded', mount);
+	} else {
+		mount();
+	}
+}
+
+/** Seeds UI settings once per demo, so later reloads keep what the demo changed. */
+function seedSettings(config: Record<string, unknown>): void {
+	const key = 'LlamaUi.config';
+
+	if (localStorage.getItem(key) === null) {
+		localStorage.setItem(key, JSON.stringify(config));
+	}
+}
+
+/** Seeded randomness, so a demo moves the same way on every recording. */
+function createRandom(seed: string): () => number {
+	let state = [...seed].reduce((h, c) => Math.imul(h ^ c.charCodeAt(0), 2654435761), 1);
+
+	return () => {
+		state = (state + 0x6d2b79f5) | 0;
+
+		let t = Math.imul(state ^ (state >>> 15), 1 | state);
+
+		t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
+
+		return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
+	};
+}
+
+const easeInOut = (t: number) => (t < 0.5 ? 4 * t * t * t : 1 - (-2 * t + 2) ** 3 / 2);
+
+/**
+ * Drives the page like a person: the cursor glides along a slight arc with
+ * Fitts-like durations, settles briefly before pressing, and typing and pauses
+ * vary a little.
+ */
+export class Demo {
+	private cdp: CDPSession | undefined;
+	private readonly page: Page;
+	private readonly random: () => number;
+	private x = 0;
+	private y = 0;
+
+	/** Shows a caption at the top of the frame for `ms`, clear of the chat form. */
+	async caption(text: string, ms = 2200): Promise<void> {
+		await this.page.evaluate((t) => {
+			const el = document.getElementById('demo-caption');
+
+			if (!el) return;
+
+			el.textContent = t;
+			el.classList.add('visible');
+		}, text);
+		await this.page.waitForTimeout(ms);
+		await this.page.evaluate(() =>
+			document.getElementById('demo-caption')?.classList.remove('visible')
+		);
+		await this.page.waitForTimeout(250);
+	}
+
+	async click(target: Locator, ms = 400): Promise<void> {
+		await this.moveTo(target);
+		await this.page.waitForTimeout(this.between(80, 160));
+		await this.page.mouse.down();
+		await this.page.waitForTimeout(this.between(50, 90));
+		await this.page.mouse.up();
+		await this.pause(ms);
+	}
+
+	constructor(page: Page, seed: string) {
+		this.page = page;
+		this.random = createRandom(seed);
+
+		const viewport = page.viewportSize();
+
+		this.x = (viewport?.width ?? 0) / 2;
+		this.y = (viewport?.height ?? 0) / 2;
+	}
+
+	async hover(target: Locator, ms = 350): Promise<void> {
+		await this.moveTo(target);
+		await this.pause(ms);
+	}
+
+	/** Glides the cursor to a point near the center of `target`. */
+	async moveTo(target: Locator): Promise<void> {
+		await target.scrollIntoViewIfNeeded();
+		const box = await target.boundingBox();
+
+		if (!box) throw new Error(`demo: ${target} has no bounding box`);
+
+		const x = box.x + box.width * this.between(0.4, 0.6);
+		const y = box.y + box.height * this.between(0.4, 0.6);
+		const dx = x - this.x;
+		const dy = y - this.y;
+		const distance = Math.hypot(dx, dy);
+		const duration = 110 + 55 * Math.log2(distance / 40 + 1);
+		// bow the path a little to one side, like a wrist turning
+		const bow = distance * this.between(-0.08, 0.08);
+		const fromX = this.x;
+		const fromY = this.y;
+		const started = Date.now();
+
+		for (;;) {
+			const t = Math.min((Date.now() - started) / duration, 1);
+			const p = easeInOut(t);
+			const arc = Math.sin(Math.PI * p) * bow;
+
+			this.x = fromX + dx * p - (dy / (distance || 1)) * arc;
+			this.y = fromY + dy * p + (dx / (distance || 1)) * arc;
+			await this.page.mouse.move(this.x, this.y);
+
+			if (t === 1) break;
+
+			await this.page.waitForTimeout(8);
+		}
+	}
+
+	/** Waits about `ms`, varied by up to 15% so pauses don't feel metronomic. */
+	async pause(ms = 800): Promise<void> {
+		await this.page.waitForTimeout(ms * this.between(0.85, 1.15));
+	}
+
+	async press(key: string, ms = 600): Promise<void> {
+		await this.page.keyboard.press(key);
+		await this.pause(ms);
+	}
+
+	/** Scrolls the wheel over `target` by `dy` pixels (negative scrolls up), in small steps like a trackpad. */
+	async scroll(target: Locator, dy: number, ms = 500): Promise<void> {
+		await this.moveTo(target);
+
+		const step = Math.sign(dy) * 40;
+
+		for (let done = 0; Math.abs(done) < Math.abs(dy); done += step) {
+			await this.page.mouse.wheel(0, step);
+			await this.page.waitForTimeout(16);
+		}
+
+		await this.pause(ms);
+	}
+
+	/** Types quickly with uneven key timing, a little slower after spaces and punctuation. */
+	async type(target: Locator, text: string, ms = 500): Promise<void> {
+		await this.click(target, 200);
+
+		for (const char of text) {
+			await this.page.keyboard.type(char);
+			await this.page.waitForTimeout(
+				this.between(15, 45) + (/[\s.,!?]/.test(char) ? this.between(20, 60) : 0)
+			);
+		}
+
+		await this.pause(ms);
+	}
+
+	/**
+	 * Pinches in on `target`, a control or a point in the viewport, holds about `ms`, and pinches
+	 * back out to 1x: a two-finger zoom onto a detail too small to read in the frame. On a touch
+	 * device the fingertips show.
+	 */
+	async zoom(target: Locator | { x: number; y: number }, scale = 2.5, ms = 1400): Promise<void> {
+		const point = 'x' in target ? target : await this.center(target);
+		// a pinch anchored in the viewport's top band does not zoom at all
+		const at = { x: point.x, y: Math.max(point.y, ZOOM_TOP_INSET_PX) };
+		const viewport = this.page.viewportSize() ?? { height: 0, width: 0 };
+		const middle = { x: viewport.width / 2, y: viewport.height / 2 };
+
+		await this.pinch(at, scale, 1400);
+
+		// One near the bottom edge overshoots its ratio, and a pinch there again misses too. A small
+		// pinch from the middle evens it out, drawing the detail toward the middle, so it stays in view
+		const reached = await this.pageScale();
+
+		if (Math.abs(reached / scale - 1) > 0.03) await this.pinch(middle, scale / reached, 1400);
+
+		await this.liftFingers();
+		await this.pause(ms);
+
+		// back to exactly 1x, which the next gesture needs, in one quick pinch from the middle of
+		// the viewport, and another for what is left when one lands short
+		for (let tries = 0; ; tries++) {
+			const current = await this.pageScale();
+
+			if (current <= 1.001) break;
+
+			if (tries === 4) throw new Error(`demo: zoom stuck at ${current.toFixed(3)}x`);
+
+			await this.pinch(middle, 1 / current, 1800);
+		}
+
+		await this.liftFingers();
+		await this.pause(200);
+	}
+
+	private between(min: number, max: number): number {
+		return min + (max - min) * this.random();
+	}
+
+	private async center(target: Locator): Promise<{ x: number; y: number }> {
+		const box = await target.boundingBox();
+
+		if (!box) throw new Error(`demo: ${target} has no bounding box`);
+
+		return { x: box.x + box.width / 2, y: box.y + box.height / 2 };
+	}
+
+	/** A pinch that reaches the screen's edge can end without its last touchend reaching the page. */
+	private async liftFingers(): Promise<void> {
+		await this.page.evaluate(() =>
+			document
+				.querySelectorAll<HTMLElement>('.demo-finger')
+				.forEach((finger) => (finger.style.opacity = '0'))
+		);
+	}
+
+	private pageScale(): Promise<number> {
+		return this.page.evaluate(() => window.visualViewport?.scale ?? 1);
+	}
+
+	/**
+	 * A two-finger pinch by `factor` around `at`, in viewport pixels. While zoomed in, the gesture
+	 * reads its position in page-scale units, so it is divided by the scale the page is at.
+	 */
+	private async pinch(at: { x: number; y: number }, factor: number, speed: number): Promise<void> {
+		const cdp = await this.session();
+		const current = await this.pageScale();
+
+		await cdp.send('Input.synthesizePinchGesture', {
+			gestureSourceType: 'touch',
+			relativeSpeed: speed,
+			scaleFactor: factor,
+			x: at.x / current,
+			y: at.y / current
+		});
+	}
+
+	private async session(): Promise<CDPSession> {
+		this.cdp ??= await this.page.context().newCDPSession(this.page);
+
+		return this.cdp;
+	}
+}
+
+interface DemoOptions {
+	/** Path the demo starts from. */
+	startPath: string;
+	/** Extra UI settings merged into the seeded config. */
+	settings: Record<string, unknown>;
+	/**
+	 * Conversations imported before the demo starts: built-in traces from
+	 * fixtures/traces.ts, or `<name>.jsonl` files from DEMO_TRACES_DATASET.
+	 */
+	traces: string[];
+}
+
+/** Reads a trace by name, generating built-in ones and downloading the rest. */
+function loadTrace(name: string, dir: string): string {
+	const generated = generateTrace(name);
+
+	if (generated) return generated;
+
+	const dataset = process.env.DEMO_TRACES_DATASET;
+
+	if (!dataset)
+		throw new Error(`demo: no built-in trace "${name}" and DEMO_TRACES_DATASET is not set`);
+
+	execFileSync(
+		'hf',
+		['download', dataset, `${name}.jsonl`, '--repo-type', 'dataset', '--local-dir', dir],
+		{ stdio: 'inherit' }
+	);
+
+	return readFileSync(join(dir, `${name}.jsonl`), 'utf8');
+}
+
+/** Imports conversations through Settings > Import/Export, like a user would. */
+async function importTraces(page: Page, file: string): Promise<void> {
+	// the expanded sidebar adds the shortcut to the name
+	const settings = page.getByRole('button', { name: /^Settings(\s|$)/ }).first();
+	const reachable = (target: Locator) =>
+		target.click({ timeout: 2000, trial: true }).then(
+			() => true,
+			() => false
+		);
+	// a collapsed sidebar can sit under the page: on a phone the page's menu key slides it aside,
+	// on a desktop its logo expands it
+	const keys = [
+		page.locator('xpath=//button[@aria-label="Expand navigation"][not(ancestor::aside)]'),
+		page.locator('aside button[aria-label="Expand navigation"]')
+	];
+
+	for (const key of keys) {
+		if (await reachable(settings)) break;
+
+		if ((await key.count()) > 0 && (await reachable(key.first()))) await key.first().click();
+	}
+
+	await settings.click();
+
+	const dialog = page.getByRole('dialog');
+
+	await dialog.getByText('Import/Export', { exact: true }).filter({ visible: true }).click();
+
+	const chooser = page.waitForEvent('filechooser');
+
+	await dialog.getByRole('button', { name: 'Import conversations' }).click();
+	await (await chooser).setFiles(file);
+	await page.getByRole('button', { name: /^Import \(\d+\)$/ }).click();
+	await page.getByRole('heading', { name: /^Imported \d+ conversation/ }).waitFor();
+
+	// the app remembers the sidebar the import expanded; a demo starts with it collapsed, as on a
+	// first visit
+	await page.evaluate(() => {
+		for (const key of Object.keys(localStorage)) {
+			if (key.endsWith('.sidebarExpanded')) localStorage.removeItem(key);
+		}
+	});
+}
+
+export const test = base.extend<DemoOptions & { demo: Demo }>({
+	demo: async ({ page }, use, testInfo) => {
+		await use(new Demo(page, `${testInfo.file}:${testInfo.title}`));
+	},
+	page: async ({ page, settings, startPath, traces }, use, testInfo) => {
+		await page.addInitScript(installCursorOverlay);
+		await page.addInitScript(seedSettings, {
+			mcpServers: JSON.stringify(demoMcpServerSettings()),
+			...settings
+		});
+
+		if (traces.length > 0) {
+			const dir = testInfo.outputPath('traces');
+			const file = join(dir, 'import.jsonl');
+
+			mkdirSync(dir, { recursive: true });
+			writeFileSync(file, traces.map((name) => loadTrace(name, dir)).join('\n'));
+			await page.goto('/');
+			await page
+				.getByRole('textbox', { name: 'Type a message...' })
+				.waitFor({ timeout: APP_READY_TIMEOUT_MS });
+			await importTraces(page, file);
+		}
+
+		await page.goto(startPath);
+
+		// a hash route after the import doesn't reload, which the conversations list needs
+		if (traces.length > 0) await page.reload();
+
+		// a fresh checkout's dev server optimizes its dependencies on the first load and reloads the
+		// page, which can take well past an action's timeout; the recording has not started yet
+		await page
+			.getByRole('textbox', { name: 'Type a message...' })
+			.waitFor({ timeout: APP_READY_TIMEOUT_MS });
+
+		const viewport = page.viewportSize();
+
+		if (viewport) await page.mouse.move(viewport.width / 2, viewport.height / 2);
+
+		// let MCP connections and fonts settle before the recording starts
+		await page.waitForTimeout(1500);
+
+		const recorder = new ScreencastRecorder(testInfo.outputPath('frames'));
+
+		await recorder.start(page);
+		await use(page);
+		await recorder.stop();
+		await page.close();
+
+		if (testInfo.status !== testInfo.expectedStatus) return;
+
+		const file = basename(testInfo.file).replace(/\.demo\.ts$/, '');
+		const name = testInfo.title === 'demo' ? file : `${file}--${slugify(testInfo.title)}`;
+		const mp4 = resolve(VIDEOS_DIR, `${name}.mp4`);
+
+		mkdirSync(VIDEOS_DIR, { recursive: true });
+		recorder.encode(mp4);
+		// the CLI reads which videos a script made, in order, with their titles
+		appendFileSync(
+			resolve(VIDEOS_DIR, `${file}.videos.jsonl`),
+			`${JSON.stringify({ mp4, title: testInfo.title })}\n`
+		);
+		console.log(`demo video: ${mp4}`);
+	},
+
+	settings: [{}, { option: true }],
+
+	startPath: ['/', { option: true }],
+
+	traces: [[], { option: true }]
+});
+
+export { expect } from '@playwright/test';
+
+function slugify(text: string): string {
+	return text
+		.toLowerCase()
+		.replace(/[^a-z0-9]+/g, '-')
+		.replace(/^-|-$/g, '');
+}

tools/ui/demo/playwright.config.ts

@@ -0,0 +1,87 @@
+import { DEMO_MCP_PORT } from './fixtures/demo-config.ts';
+import { defineConfig, devices } from '@playwright/test';
+import { existsSync } from 'node:fs';
+import { resolve } from 'node:path';
+
+const ENV_FILE = resolve(import.meta.dirname, '../.env');
+
+if (existsSync(ENV_FILE)) process.loadEnvFile(ENV_FILE);
+
+/**
+ * Demo video recordings.
+ *
+ * Env (also read from tools/ui/.env):
+ *   DEMO_UI_DIR         tools/ui checkout to record (default: this one), e.g. a PR worktree
+ *   DEMO_UI_PORT        port for that checkout's vite dev server (default: 5180)
+ *   DEMO_SERVER_ORIGIN  llama-server the UI proxies to (default: http://localhost:9931)
+ *   DEMO_LLAMA_SERVER   llama-server binary started when none runs there (default: build/bin)
+ *   DEMO_COLOR_SCHEME   light | dark (default: light)
+ */
+
+// Playwright clears outputDir when a run starts, so concurrent recordings each get
+// their own; workers inherit the env, so they resolve the same directory
+process.env.DEMO_RUN_ID ??= String(process.pid);
+
+const UI_DIR = resolve(process.env.DEMO_UI_DIR ?? resolve(import.meta.dirname, '..'));
+const UI_PORT = Number(process.env.DEMO_UI_PORT ?? 5180);
+const SERVER_ORIGIN = process.env.DEMO_SERVER_ORIGIN ?? 'http://localhost:9931';
+const SERVER_PORT = Number(new URL(SERVER_ORIGIN).port || 80);
+const LLAMA_SERVER =
+	process.env.DEMO_LLAMA_SERVER ?? resolve(import.meta.dirname, '../../../build/bin/llama-server');
+
+export const VIEWPORT = { height: 900, width: 1440 };
+/** Device pixel ratio the demos render and record at. */
+const SCALE = Number(process.env.DEMO_SCALE ?? 2);
+
+export default defineConfig({
+	expect: {
+		timeout: 10000
+	},
+	fullyParallel: false,
+	outputDir: `test-results/${process.env.DEMO_RUN_ID}`,
+	projects: [
+		{
+			name: 'chromium',
+			use: { ...devices['Desktop Chrome'], deviceScaleFactor: SCALE, viewport: VIEWPORT }
+		}
+	],
+	reporter: 'line',
+	retries: 0,
+	testDir: '.',
+	testMatch: ['**/*.demo.ts'],
+	timeout: 120000,
+	use: {
+		actionTimeout: 10000,
+		baseURL: `http://localhost:${UI_PORT}`,
+		colorScheme: process.env.DEMO_COLOR_SCHEME === 'dark' ? 'dark' : 'light',
+		// headless Chromium screencasts at 1x unless the scale is forced at launch; the sandbox
+		// is on since a PR demo runs that PR's code (Playwright turns it off by default). Not the
+		// installed Chrome: its headless mode screencasts a window that ignores the viewport
+		launchOptions: { args: [`--force-device-scale-factor=${SCALE}`], chromiumSandbox: true },
+		// recorded through the screencast in recorder.ts, at the device pixel ratio
+		video: 'off',
+		viewport: VIEWPORT
+	},
+	webServer: [
+		{
+			command: `node ${resolve(import.meta.dirname, 'fixtures/mcp-servers.ts')}`,
+			reuseExistingServer: true,
+			url: `http://localhost:${DEMO_MCP_PORT}/health`
+		},
+		{
+			command: `${LLAMA_SERVER} --host 127.0.0.1 --port ${SERVER_PORT} --tools all`,
+			reuseExistingServer: true,
+			timeout: 60000,
+			url: `${SERVER_ORIGIN}/health`
+		},
+		{
+			command: `npx vite dev --port ${UI_PORT} --strictPort`,
+			cwd: UI_DIR,
+			env: { VITE_PUBLIC_SERVER_ORIGIN: SERVER_ORIGIN },
+			reuseExistingServer: true,
+			timeout: 120000,
+			url: `http://localhost:${UI_PORT}`
+		}
+	],
+	workers: 1
+});

tools/ui/demo/recorder.ts

@@ -0,0 +1,109 @@
+/**
+ * Records a page through Chromium's screencast instead of Playwright's video,
+ * which encodes VP8 at about 1 Mbit/s at 1x scale. Frames arrive at the page's
+ * device pixel ratio, each with its timestamp, and are encoded once into an
+ * H.264 mp4 (what GitHub accepts in PR descriptions).
+ */
+
+import type { CDPSession, Page } from '@playwright/test';
+import { execFileSync } from 'node:child_process';
+import { mkdirSync, rmSync, writeFileSync } from 'node:fs';
+import { writeFile } from 'node:fs/promises';
+import { join } from 'node:path';
+
+const FPS = 60;
+const JPEG_QUALITY = 92;
+
+interface Frame {
+	file: string;
+	/** Seconds, on the same clock as Date.now(). */
+	time: number;
+}
+
+export class ScreencastRecorder {
+	private readonly dir: string;
+	private readonly frames: Frame[] = [];
+	private readonly pending: Promise<void>[] = [];
+	private session: CDPSession | undefined;
+	private stoppedAt = 0;
+
+	constructor(dir: string) {
+		this.dir = dir;
+	}
+
+	/** Encodes the frames, each held until the next one arrived (the screencast only sends changes). */
+	encode(mp4: string): void {
+		if (this.frames.length === 0) throw new Error('demo: the screencast captured no frames');
+
+		const list = this.frames.flatMap((frame, i) => {
+			const end = this.frames[i + 1]?.time ?? this.stoppedAt;
+
+			return [`file '${frame.file}'`, `duration ${Math.max(end - frame.time, 1 / FPS).toFixed(4)}`];
+		});
+		const listFile = join(this.dir, 'frames.txt');
+
+		// the concat demuxer ignores the last entry's duration unless the file is repeated
+		writeFileSync(listFile, [...list, `file '${this.frames.at(-1)!.file}'`, ''].join('\n'));
+		execFileSync(
+			'ffmpeg',
+			[
+				'-y',
+				'-loglevel',
+				'error',
+				'-f',
+				'concat',
+				'-safe',
+				'0',
+				'-i',
+				listFile,
+				'-vf',
+				`fps=${FPS},scale=trunc(iw/2)*2:trunc(ih/2)*2`,
+				'-c:v',
+				'libx264',
+				'-preset',
+				'slow',
+				'-crf',
+				'18',
+				'-pix_fmt',
+				'yuv420p',
+				'-movflags',
+				'+faststart',
+				mp4
+			],
+			{ stdio: 'inherit' }
+		);
+		rmSync(this.dir, { force: true, recursive: true });
+	}
+
+	async start(page: Page): Promise<void> {
+		const size = page.viewportSize();
+		const scale = await page.evaluate(() => window.devicePixelRatio);
+
+		rmSync(this.dir, { force: true, recursive: true });
+		mkdirSync(this.dir, { recursive: true });
+
+		this.session = await page.context().newCDPSession(page);
+		this.session.on('Page.screencastFrame', ({ data, metadata, sessionId }) => {
+			const file = join(this.dir, `${String(this.frames.length).padStart(6, '0')}.jpg`);
+
+			// acknowledged first: Chromium sends the next frame only after the ack
+			this.session?.send('Page.screencastFrameAck', { sessionId }).catch(() => {});
+			this.frames.push({ file, time: metadata.timestamp ?? Date.now() / 1000 });
+			this.pending.push(writeFile(file, Buffer.from(data, 'base64')));
+		});
+		await this.session.send('Page.startScreencast', {
+			everyNthFrame: 1,
+			format: 'jpeg',
+			maxHeight: size ? size.height * scale : undefined,
+			maxWidth: size ? size.width * scale : undefined,
+			quality: JPEG_QUALITY
+		});
+	}
+
+	async stop(): Promise<void> {
+		this.stoppedAt = Date.now() / 1000;
+		await this.session?.send('Page.stopScreencast').catch(() => {});
+		await this.session?.detach().catch(() => {});
+		await Promise.all(this.pending);
+	}
+}

tools/ui/demo/storage.ts

@@ -0,0 +1,92 @@
+/**
+ * Storage targets for demo scripts and videos, picked by the scheme of the
+ * DEMO_STORAGE URI:
+ *
+ *   hf://buckets/<owner>/<bucket>[/prefix]   Hugging Face bucket, via the `hf` CLI login
+ *   file:///some/dir                         local or mounted directory
+ */
+
+import { execFileSync } from 'node:child_process';
+import { copyFileSync, cpSync, existsSync, mkdirSync } from 'node:fs';
+import { dirname, join } from 'node:path';
+import { fileURLToPath } from 'node:url';
+
+export interface DemoStorage {
+	/** Copies every file under remote `dir` into local `target`; returns false when it has none. */
+	downloadDir(dir: string, target: string): boolean;
+	/** Copies local `file` to remote `key`. */
+	upload(file: string, key: string): void;
+}
+
+function joinKey(base: string, key: string): string {
+	return `${base.replace(/\/+$/, '')}/${key}`;
+}
+
+// plain fields: the CLI runs under Node's type stripping, which has no parameter properties
+class HfStorage implements DemoStorage {
+	private readonly base: string;
+
+	constructor(base: string) {
+		this.base = base;
+	}
+
+	downloadDir(dir: string, target: string): boolean {
+		mkdirSync(target, { recursive: true });
+
+		try {
+			execFileSync('hf', ['buckets', 'sync', joinKey(this.base, dir), target], {
+				stdio: 'inherit'
+			});
+		} catch {
+			return false;
+		}
+
+		return true;
+	}
+
+	upload(file: string, key: string): void {
+		execFileSync('hf', ['buckets', 'cp', file, joinKey(this.base, key)], { stdio: 'inherit' });
+	}
+}
+
+class FileStorage implements DemoStorage {
+	private readonly root: string;
+
+	constructor(root: string) {
+		this.root = root;
+	}
+
+	downloadDir(dir: string, target: string): boolean {
+		const source = join(this.root, dir);
+
+		if (!existsSync(source)) return false;
+
+		cpSync(source, target, { recursive: true });
+
+		return true;
+	}
+
+	upload(file: string, key: string): void {
+		const target = join(this.root, key);
+
+		mkdirSync(dirname(target), { recursive: true });
+		copyFileSync(file, target);
+	}
+}
+
+export function getDemoStorage(uri = process.env.DEMO_STORAGE): DemoStorage | null {
+	if (!uri) return null;
+
+	if (uri.startsWith('hf://')) return new HfStorage(uri);
+
+	if (uri.startsWith('file://')) return new FileStorage(fileURLToPath(uri));
+
+	throw new Error(`DEMO_STORAGE: unsupported scheme in "${uri}" (use hf:// or file://)`);
+}
+
+/** Public URL of a stored key, when DEMO_STORAGE_PUBLIC_URL is set. */
+export function getPublicUrl(key: string): string | null {
+	const base = process.env.DEMO_STORAGE_PUBLIC_URL;
+
+	return base ? joinKey(base, key) : null;
+}