graph: store the conv state of every rollback arch in one set_rows

The rollback slots of the causal conv history were written with one
copy per slot in qwen4exp, glm5-next, kimi-k3, bailingmoe3, lfm and
mamba2, while delta_net already wrote them with a single set_rows.
build_conv_state_store now does this for all of them: one copy without
rollback, otherwise one set_rows of the overlapping windows into the
min(n, n_rs_seq + 1) slots a rollback can read.

The slot order {n_rs_seq, ..., 0} is a constant of the recurrent
memory, kept in its own buffer next to the state, and each graph takes
a view of its tail. Its size no longer depends on the ubatch, so the
reserved graph covers every ubatch shape, and the graph input that
held it is gone.

A model that is not causal never rolls back, so the context clamps
n_rs_seq to 0 for it, like for an arch without rollback support.

.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

conversion/__init__.py

@@ -223,6 +223,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "ModernBertForSequenceClassification": "bert",
     "ModernBertModel": "bert",
     "NanbeigeForCausalLM": "nanbeige",
+    "Nemotron3_5TranscribeForConditionalGeneration": "nemotron",
     "NemotronForCausalLM": "nemotron",
     "NemotronHForCausalLM": "nemotron",
     "NemotronHPuzzleForCausalLM": "nemotron",
@@ -380,6 +381,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

@@ -1277,6 +1277,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 +1326,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 +1335,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 +1352,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 +1364,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])
@@ -3384,7 +3386,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/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)

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(
@@ -700,9 +726,39 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             }
             GGML_ABORT("fatal error");
         }
+        // permuted view: split the view dim that spans the src split axis, other dims must be inside one element of it or outside the span
+        if (ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0]) && axis >= 0 && axis < GGML_MAX_DIMS-1) {
+            const size_t nb_axis = tensor->src[0]->nb[axis];
+            const size_t nb_next = tensor->src[0]->nb[axis+1];
+            for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
+                if (tensor->ne[dim] == 1 || tensor->nb[dim] < nb_axis || tensor->nb[dim]*tensor->ne[dim] != nb_next) {
+                    continue;
+                }
+                bool others_ok = true;
+                for (int i = 0; i < GGML_MAX_DIMS; i++) {
+                    if (i == dim || tensor->ne[i] == 1) {
+                        continue;
+                    }
+                    others_ok = others_ok && (tensor->nb[i]*tensor->ne[i] <= tensor->nb[dim] || tensor->nb[i] % nb_next == 0);
+                }
+                if (others_ok) {
+                    return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1};
+                }
+            }
+        }
         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};
     };
@@ -1000,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-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

@@ -1457,8 +1457,8 @@ static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vec
         lut_size = 4*16;
         break;
     case GGML_TYPE_NVFP4:
-        // Same kvalues budget as MXFP4 plus ue4m3_fp32_lut[128] (types.glsl, DATA_A_NVFP4).
-        lut_size = 4*16 + 128u * (uint32_t)sizeof(float);
+        // Same kvalues budget as MXFP4 plus ue4m3_fp32_lut[256] (types.glsl, DATA_A_NVFP4).
+        lut_size = 4*16 + 256u * (uint32_t)sizeof(float);
         break;
     default:
         break;
@@ -1592,7 +1592,7 @@ static bool ggml_vk_matmul_cm1_int_shmem_support(const vk_device& device, const
         total += 16u * (uint32_t)sizeof(int8_t);         // cm1_kvalues[16]
     }
     if (src0_type == GGML_TYPE_NVFP4 && !device->ocp_fp4) {
-        total += 128u * (uint32_t)sizeof(float);         // ue4m3_fp32_lut[128]
+        total += 256u * (uint32_t)sizeof(float);         // ue4m3_fp32_lut[256]
     }
     if (mul_mat_id) {
         total += BN * 2u * (uint32_t)sizeof(uint16_t);   // row_ids[BN] (u16vec2)
@@ -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-vulkan/vulkan-shaders/mul_mmq_cm1.comp

@@ -173,7 +173,7 @@ void main() {
         cm1_kvalues[gl_LocalInvocationIndex] = kvalues_mxfp4_const[gl_LocalInvocationIndex];
     }
 #if !defined(USE_OCP_FP4)
-    for (uint i = gl_LocalInvocationIndex; i < 128u; i += BLOCK_SIZE) {
+    for (uint i = gl_LocalInvocationIndex; i < 256u; i += BLOCK_SIZE) {
         ue4m3_fp32_lut[i] = ue4m3_to_fp32_build(i);
     }
 #endif

ggml/src/ggml-vulkan/vulkan-shaders/types.glsl

@@ -1921,8 +1921,9 @@ shared int8_t kvalues_mxfp4[16];
 #endif
 
 #if defined(DATA_A_NVFP4) && !defined(USE_OCP_FP4)
-// UE4M3 scale in NVFP4 blocks use only 7 bits; sign (bit 7) is always zero.
-shared float ue4m3_fp32_lut[128];
+// Cover all scale bytes accepted by ggml_validate_row_data.
+// ue4m3_to_fp32_build() masks the exponent and mantissa, but checks the full byte for 0 and 0x7F.
+shared float ue4m3_fp32_lut[256];
 
 float ue4m3_to_fp32_build(uint u) {
     if (u == 0u || u == 127u) {
@@ -1947,7 +1948,7 @@ void init_iq_shmem(uvec3 wgsize)
         kvalues_mxfp4[i] = kvalues_mxfp4_const[i];
     }
 #if defined(DATA_A_NVFP4)
-    for (uint i = gl_LocalInvocationIndex.x; i < 128u; i += wgsize.x) {
+    for (uint i = gl_LocalInvocationIndex.x; i < 256u; i += wgsize.x) {
         ue4m3_fp32_lut[i] = ue4m3_to_fp32_build(i);
     }
 #endif

ggml/src/ggml.c

@@ -1788,6 +1788,7 @@ static struct ggml_tensor * ggml_new_tensor_impl(
         enum   ggml_type      type,
         int                   n_dims,
         const int64_t       * ne,
+        const size_t        * nb,
         struct ggml_tensor  * view_src,
         size_t                view_offs) {
 
@@ -1814,8 +1815,6 @@ static struct ggml_tensor * ggml_new_tensor_impl(
         data_size *= ne[i];
     }
 
-    GGML_ASSERT(view_src == NULL || data_size == 0 || data_size + view_offs <= ggml_nbytes(view_src));
-
     void * data = view_src != NULL ? view_src->data : NULL;
     if (data != NULL) {
         data = (char *) data + view_offs;
@@ -1859,12 +1858,23 @@ static struct ggml_tensor * ggml_new_tensor_impl(
         result->ne[i] = ne[i];
     }
 
-    result->nb[0] = ggml_type_size(type);
-    result->nb[1] = result->nb[0]*(result->ne[0]/ggml_blck_size(type));
-    for (int i = 2; i < GGML_MAX_DIMS; i++) {
-        result->nb[i] = result->nb[i - 1]*result->ne[i - 1];
+    // a view built with strides takes them from the caller, any other tensor is contiguous
+    if (nb != NULL) {
+        for (int i = 0; i < GGML_MAX_DIMS; i++) {
+            result->nb[i] = nb[i];
+        }
+    } else {
+        result->nb[0] = ggml_type_size(type);
+        result->nb[1] = result->nb[0]*(result->ne[0]/ggml_blck_size(type));
+        for (int i = 2; i < GGML_MAX_DIMS; i++) {
+            result->nb[i] = result->nb[i - 1]*result->ne[i - 1];
+        }
     }
 
+    // a view stays inside its source, over the bytes its strides actually reach
+    const size_t view_size = ggml_nbytes(result);
+    GGML_ASSERT(view_src == NULL || view_size == 0 || view_size + view_offs <= ggml_nbytes(view_src));
+
     ctx->n_objects++;
 
     return result;
@@ -1875,7 +1885,7 @@ struct ggml_tensor * ggml_new_tensor(
         enum   ggml_type      type,
         int                   n_dims,
         const int64_t       * ne) {
-    return ggml_new_tensor_impl(ctx, type, n_dims, ne, NULL, 0);
+    return ggml_new_tensor_impl(ctx, type, n_dims, ne, NULL, NULL, 0);
 }
 
 struct ggml_tensor * ggml_new_tensor_1d(
@@ -1992,13 +2002,9 @@ struct ggml_tensor * ggml_format_name(struct ggml_tensor * tensor, const char *
 struct ggml_tensor * ggml_view_tensor(
         struct ggml_context * ctx,
         struct ggml_tensor  * src) {
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, src->type, GGML_MAX_DIMS, src->ne, src, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, src->type, GGML_MAX_DIMS, src->ne, src->nb, src, 0);
     ggml_format_name(result, "%s (view)", src->name);
 
-    for (int i = 0; i < GGML_MAX_DIMS; i++) {
-        result->nb[i] = src->nb[i];
-    }
-
     return result;
 }
 
@@ -2948,7 +2954,7 @@ static struct ggml_tensor * ggml_glu_impl(
     }
 
     int64_t ne[GGML_MAX_DIMS] = { a->ne[0] / 2 }; for (int i = 1; i < GGML_MAX_DIMS; i++) ne[i] = a->ne[i];
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b ? a->ne : ne, NULL, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b ? a->ne : ne, NULL, NULL, 0);
 
     ggml_set_op_params_i32(result, 0, (int32_t) op);
     ggml_set_op_params_i32(result, 1, (int32_t) swapped);
@@ -3757,7 +3763,7 @@ struct ggml_tensor * ggml_reshape(
     // as only the shape of b is relevant, and not its memory layout, b is allowed to be non contiguous.
     GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b));
 
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b->ne, a, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b->ne, NULL, a, 0);
     ggml_format_name(result, "%s (reshaped)", a->name);
 
     result->op     = GGML_OP_RESHAPE;
@@ -3774,7 +3780,7 @@ struct ggml_tensor * ggml_reshape_1d(
     GGML_ASSERT(ggml_nelements(a) == ne0);
 
     const int64_t ne[1] = { ne0 };
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 1, ne, a, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 1, ne, NULL, a, 0);
     ggml_format_name(result, "%s (reshaped)", a->name);
 
     result->op     = GGML_OP_RESHAPE;
@@ -3792,7 +3798,7 @@ struct ggml_tensor * ggml_reshape_2d(
     GGML_ASSERT(ggml_nelements(a) == ne0*ne1);
 
     const int64_t ne[2] = { ne0, ne1 };
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 2, ne, a, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 2, ne, NULL, a, 0);
     ggml_format_name(result, "%s (reshaped)", a->name);
 
     result->op     = GGML_OP_RESHAPE;
@@ -3811,7 +3817,7 @@ struct ggml_tensor * ggml_reshape_3d(
     GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2);
 
     const int64_t ne[3] = { ne0, ne1, ne2 };
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 3, ne, a, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 3, ne, NULL, a, 0);
     ggml_format_name(result, "%s (reshaped)", a->name);
 
     result->op     = GGML_OP_RESHAPE;
@@ -3831,7 +3837,7 @@ struct ggml_tensor * ggml_reshape_4d(
     GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2*ne3);
 
     const int64_t ne[4] = { ne0, ne1, ne2, ne3 };
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 4, ne, a, 0);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 4, ne, NULL, a, 0);
     ggml_format_name(result, "%s (reshaped)", a->name);
 
     result->op     = GGML_OP_RESHAPE;
@@ -3845,8 +3851,9 @@ static struct ggml_tensor * ggml_view_impl(
         struct ggml_tensor  * a,
         int                   n_dims,
         const int64_t       * ne,
+        const size_t        * nb,
         size_t                offset) {
-    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, n_dims, ne, a, offset);
+    struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, n_dims, ne, nb, a, offset);
     ggml_format_name(result, "%s (view)", a->name);
 
     ggml_set_op_params(result, &offset, sizeof(offset));
@@ -3864,7 +3871,7 @@ struct ggml_tensor * ggml_view_1d(
         struct ggml_tensor  * a,
         int64_t               ne0,
         size_t                offset) {
-    struct ggml_tensor * result = ggml_view_impl(ctx, a, 1, &ne0, offset);
+    struct ggml_tensor * result = ggml_view_impl(ctx, a, 1, &ne0, NULL, offset);
 
     return result;
 }
@@ -3879,12 +3886,9 @@ struct ggml_tensor * ggml_view_2d(
         size_t                nb1,
         size_t                offset) {
     const int64_t ne[2] = { ne0, ne1 };
+    const size_t  nb[4] = { ggml_type_size(a->type), nb1, nb1*ne1, nb1*ne1 };
 
-    struct ggml_tensor * result = ggml_view_impl(ctx, a, 2, ne, offset);
-
-    result->nb[1] = nb1;
-    result->nb[2] = result->nb[1]*ne1;
-    result->nb[3] = result->nb[2];
+    struct ggml_tensor * result = ggml_view_impl(ctx, a, 2, ne, nb, offset);
 
     return result;
 }
@@ -3901,12 +3905,9 @@ struct ggml_tensor * ggml_view_3d(
         size_t                nb2,
         size_t                offset) {
     const int64_t ne[3] = { ne0, ne1, ne2 };
+    const size_t  nb[4] = { ggml_type_size(a->type), nb1, nb2, nb2*ne2 };
 
-    struct ggml_tensor * result = ggml_view_impl(ctx, a, 3, ne, offset);
-
-    result->nb[1] = nb1;
-    result->nb[2] = nb2;
-    result->nb[3] = result->nb[2]*ne2;
+    struct ggml_tensor * result = ggml_view_impl(ctx, a, 3, ne, nb, offset);
 
     return result;
 }
@@ -3925,12 +3926,9 @@ struct ggml_tensor * ggml_view_4d(
         size_t                nb3,
         size_t                offset) {
     const int64_t ne[4] = { ne0, ne1, ne2, ne3 };
+    const size_t  nb[4] = { ggml_type_size(a->type), nb1, nb2, nb3 };
 
-    struct ggml_tensor * result = ggml_view_impl(ctx, a, 4, ne, offset);
-
-    result->nb[1] = nb1;
-    result->nb[2] = nb2;
-    result->nb[3] = nb3;
+    struct ggml_tensor * result = ggml_view_impl(ctx, a, 4, ne, nb, offset);
 
     return result;
 }
@@ -5389,10 +5387,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);
 

gguf-py/gguf/constants.py

@@ -506,6 +506,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"
@@ -1411,6 +1416,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()
@@ -2219,6 +2242,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 +2628,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,
@@ -6740,6 +6799,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

@@ -1712,6 +1712,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

@@ -2532,6 +2532,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 +2547,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 +2599,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 +2635,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 +2649,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 +2661,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 +2692,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 +2712,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 +2803,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 +2926,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
         ),

src/llama-arch.cpp

@@ -1390,7 +1390,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-context.cpp

@@ -201,7 +201,7 @@ llama_context::llama_context(
     }
 
     cparams.n_rs_seq = params.n_rs_seq;
-    if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) {
+    if (cparams.n_rs_seq > 0 && (!llm_arch_supports_rs_rollback(model.arch) || !hparams.causal_attn)) {
         LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n",
                         __func__, cparams.n_rs_seq);
         cparams.n_rs_seq = 0;

src/llama-graph.cpp

@@ -3994,6 +3994,63 @@ ggml_tensor * llm_graph_context::build_rs(
                     get_state_rows);
 }
 
+void llm_graph_context::build_conv_state_store(
+        llm_graph_input_rs * inp,
+        ggml_tensor * s,
+        ggml_tensor * conv_x,
+            int64_t   n_cols,
+             size_t   offs,
+                int   il) const {
+    const auto * mctx_cur = inp->mctx;
+
+    const int64_t n        = conv_x->ne[0] - n_cols; // n_seq_tokens
+    const int64_t channels = conv_x->ne[1];
+    const int64_t n_seqs   = conv_x->ne[2];
+
+    const size_t row_size  = s->nb[1];
+    const size_t cell_offs = mctx_cur->get_head()*row_size + offs;
+
+    if (cparams.n_rs_seq == 0) {
+        ggml_tensor * last = ggml_view_3d(ctx0, conv_x,
+                n_cols, channels, n_seqs,
+                conv_x->nb[1], conv_x->nb[2],
+                ggml_row_size(conv_x->type, n));
+
+        ggml_tensor * cell = ggml_view_2d(ctx0, s,
+                n_cols*channels, n_seqs,
+                row_size, cell_offs);
+
+        ggml_build_forward_expand(gf, ggml_cpy(ctx0, last, cell));
+        return;
+    }
+
+    // [TAG_RECURRENT_ROLLBACK_SPLITS]
+    // this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside
+    //   the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument
+    // a rollback of k tokens reads slot k, and the memory allows k < n only (see rs_max), so slots from n on are left unwritten
+    const int64_t K  = (int64_t) cparams.n_rs_seq + 1;
+    const int64_t Kw = std::min(n, K);
+
+    // {Kw-1, ..., 0}, the tail of the slot order the memory keeps for every ubatch
+    ggml_tensor * slot_order = mctx_cur->get_slot_l(il);
+    slot_order = ggml_view_1d(ctx0, slot_order, Kw, ggml_row_size(slot_order->type, K - Kw));
+
+    // window w starts at column n - (Kw - 1) + w (ascending time) and is rollback slot Kw - 1 - w,
+    // consecutive windows start one element apart, so they overlap
+    ggml_tensor * windows = ggml_view_4d(ctx0, conv_x,
+            n_cols, Kw, channels, n_seqs,
+            ggml_element_size(conv_x), conv_x->nb[1], conv_x->nb[2],
+            ggml_row_size(conv_x->type, n - (Kw - 1)));
+
+    // [n_cols, slot, channel, seq] view of the cells of this ubatch
+    ggml_tensor * slots = ggml_view_4d(ctx0, s,
+            n_cols, K, channels, n_seqs,
+            (size_t) mctx_cur->get_size()*row_size, ggml_row_size(s->type, n_cols), row_size,
+            cell_offs);
+
+    ggml_build_forward_expand(gf, ggml_set_rows(ctx0, slots, windows, slot_order));
+}
+
 ggml_tensor * llm_graph_context::build_rwkv_token_shift_load(
     llm_graph_input_rs * inp,
     const llama_ubatch & ubatch,

src/llama-graph.h

@@ -1433,6 +1433,16 @@ struct llm_graph_context {
                 int32_t   n_seqs,
             const llm_graph_get_rows_fn & get_state_rows = nullptr) const;
 
+    // store the conv history of the ubatch cells, the last n_cols columns of conv_x {n_cols + n_seq_tokens, channels, n_seqs},
+    // at byte offset offs of each row of s, and with n_rs_seq > 0 once per rollback slot of this ubatch
+    void build_conv_state_store(
+            llm_graph_input_rs * inp,
+            ggml_tensor * s,
+            ggml_tensor * conv_x,
+                int64_t   n_cols,
+                 size_t   offs,
+                    int   il) const;
+
     ggml_tensor * build_rwkv_token_shift_load(
         llm_graph_input_rs * inp,
         const llama_ubatch & ubatch,

src/llama-memory-recurrent.cpp

@@ -26,7 +26,7 @@ llama_memory_recurrent::llama_memory_recurrent(
                  uint32_t   mem_size,
                  uint32_t   n_seq_max,
                  uint32_t   n_rs_seq,
-    const layer_filter_cb & filter) : hparams(model.hparams), n_seq_max(n_seq_max), ctxs_bufs(no_alloc) {
+    const layer_filter_cb & filter) : hparams(model.hparams), n_seq_max(n_seq_max), ctxs_bufs(no_alloc), slot_bufs(no_alloc) {
     const int32_t n_layer = hparams.n_layer();
 
     head = 0;
@@ -46,6 +46,7 @@ llama_memory_recurrent::llama_memory_recurrent(
         }
     };
     std::map<ggml_backend_buffer_type_t, ggml_context_ptr, ggml_backend_buft_comparator> ctx_map;
+    std::map<ggml_backend_buffer_type_t, ggml_context_ptr, ggml_backend_buft_comparator> slot_ctx_map;
 
     // create a context for each buffer type
     auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * {
@@ -74,6 +75,7 @@ llama_memory_recurrent::llama_memory_recurrent(
     r_l.resize(n_layer);
     s_l.resize(n_layer);
     p_l.resize(n_layer);
+    slot_l.resize(n_layer);
 
     for (int i = 0; i < n_layer; i++) {
         if (filter && !filter(i)) {
@@ -113,6 +115,28 @@ llama_memory_recurrent::llama_memory_recurrent(
             ggml_format_name(p, "cache_ple_r_l%d", i);
             p_l[i] = p;
         }
+
+        if (n_rs_seq > 0) {
+            auto it = slot_ctx_map.find(buft);
+            if (it == slot_ctx_map.end()) {
+                ggml_init_params params = {
+                    /*.mem_size   =*/ size_t(n_layer*ggml_tensor_overhead()),
+                    /*.mem_buffer =*/ NULL,
+                    /*.no_alloc   =*/ true,
+                };
+
+                ggml_context * slot_ctx = ggml_init(params);
+                if (!slot_ctx) {
+                    throw std::runtime_error("failed to create ggml context for rs slots");
+                }
+
+                it = slot_ctx_map.emplace(buft, slot_ctx).first;
+            }
+
+            ggml_tensor * slot = ggml_new_tensor_1d(it->second.get(), GGML_TYPE_I32, n_rs_seq + 1);
+            ggml_format_name(slot, "cache_rs_slot_l%d", i);
+            slot_l[i] = slot;
+        }
     }
 
     // allocate the tensors and initialize the buffers to avoid NaNs in the padding
@@ -122,6 +146,23 @@ llama_memory_recurrent::llama_memory_recurrent(
         LLAMA_LOG_INFO("%s: %10s RS buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0);
     }
 
+    for (auto & [buft, ctx] : slot_ctx_map) {
+        slot_bufs.emplace(std::move(ctx), buft, "rs slots");
+    }
+
+    if (!slot_bufs.no_alloc) {
+        std::vector<int32_t> order(n_rs_seq + 1);
+        for (uint32_t i = 0; i <= n_rs_seq; ++i) {
+            order[i] = (int32_t) (n_rs_seq - i);
+        }
+
+        for (ggml_tensor * slot : slot_l) {
+            if (slot != nullptr) {
+                ggml_backend_tensor_set(slot, order.data(), 0, ggml_nbytes(slot));
+            }
+        }
+    }
+
     if (is_empty()) {
         if (n_rs_seq > 0) {
             n_rs_seq = 0;
@@ -432,7 +473,11 @@ void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) {
 }
 
 std::map<ggml_backend_buffer_type_t, size_t> llama_memory_recurrent::memory_breakdown() const {
-    return ctxs_bufs.memory_breakdown();
+    auto ret = ctxs_bufs.memory_breakdown();
+    for (const auto & [buft, size] : slot_bufs.memory_breakdown()) {
+        ret[buft] += size;
+    }
+    return ret;
 }
 
 llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {
@@ -1360,6 +1405,10 @@ ggml_tensor * llama_memory_recurrent_context::get_p_l(int32_t il) const {
     return mem->p_l[il];
 }
 
+ggml_tensor * llama_memory_recurrent_context::get_slot_l(int32_t il) const {
+    return mem->slot_l[il];
+}
+
 int32_t llama_memory_recurrent_context::s_copy(int i) const {
     const uint32_t cell_idx = i + mem->head;
     const int32_t  src0     = mem->cells[cell_idx].src0;

src/llama-memory-recurrent.h

@@ -118,6 +118,8 @@ public:
     std::vector<ggml_tensor *> s_l;
     // a second conv history that must stay replicated across devices, so it cannot share the r row
     std::vector<ggml_tensor *> p_l;
+    // I32 [n_rs_seq + 1] = {n_rs_seq, ..., 0}, the rollback slot of each conv window, the same for every ubatch
+    std::vector<ggml_tensor *> slot_l;
 
 private:
     //const llama_model & model;
@@ -128,6 +130,9 @@ private:
     // ggml contexts for the RS cache along with the allocated backend buffers:
     llama_memory_bufs ctxs_bufs;
 
+    // the constants of the cache, kept out of the state buffers so that clearing or restoring the state leaves them intact
+    llama_memory_bufs slot_bufs;
+
     // true if no layers - can happen if the layer filter removes all layers
     bool is_empty() const;
 
@@ -184,6 +189,7 @@ public:
     ggml_tensor * get_r_l(int32_t il) const;
     ggml_tensor * get_s_l(int32_t il) const;
     ggml_tensor * get_p_l(int32_t il) const;
+    ggml_tensor * get_slot_l(int32_t il) const;
 
     int32_t s_copy(int i) const;
 

src/llama-model.cpp

@@ -410,6 +410,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 +430,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 +443,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");
@@ -532,6 +536,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 +668,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 +765,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);

src/models/bailingmoe3.cpp

@@ -160,25 +160,11 @@ std::unique_ptr<llm_graph_context> llama_model_bailingmoe3::build_arch_graph(con
     return std::make_unique<graph>(*this, params);
 }
 
-static ggml_tensor * bailingmoe3_causal_conv1d(
-        ggml_cgraph * gf,
-        ggml_context * ctx0,
-        ggml_tensor * conv_states_all,
-        ggml_tensor * conv_state_all,
-        int64_t qkv,
-        ggml_tensor * x,
-        ggml_tensor * proj_w,
-        ggml_tensor * conv_w,
-        int64_t d_conv,
-        int64_t head_dim,
-        int64_t n_head,
-        int64_t n_seq_tokens,
-        int64_t n_seqs,
-        int64_t n_tokens,
-        int64_t cache_head,
-        uint32_t mem_size,
-        uint32_t n_rs_seq) {
-    const int64_t d_inner = head_dim * n_head;
+ggml_tensor * llama_model_bailingmoe3::graph::build_causal_conv1d(
+        llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+        int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+        int64_t d_conv, int64_t head_dim, int64_t n_head_kda, int64_t n_seq_tokens, int64_t n_seqs, int il) {
+    const int64_t d_inner = head_dim * n_head_kda;
     const int64_t conv_state_size = (d_conv - 1) * d_inner;
     const int64_t total_state_size = 3 * conv_state_size;
 
@@ -191,23 +177,13 @@ static ggml_tensor * bailingmoe3_causal_conv1d(
     x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
     ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0);
 
-    const int64_t K = (int64_t) n_rs_seq + 1;
-    const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
-
-    for (int64_t slot = 0; slot < n_written; ++slot) {
-        ggml_tensor * conv_snap = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
-                conv_x->nb[1], conv_x->nb[2], (conv_x->ne[0] - (d_conv - 1) - slot) * conv_x->nb[0]);
-        ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_snap,
-                ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
-                    (d_conv - 1) * ggml_element_size(conv_states_all),
-                    total_state_size * ggml_element_size(conv_states_all),
-                    ((slot * mem_size + cache_head) * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
-    }
+    build_conv_state_store(inp_rs, conv_states_all, conv_x, d_conv - 1,
+            ggml_row_size(conv_states_all->type, qkv * conv_state_size), il);
 
     ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
     ggml_tensor * out = ggml_ssm_conv(ctx0, conv_x, conv_weight);
     out = ggml_silu(ctx0, ggml_reshape_2d(ctx0, out, d_inner, n_tokens));
-    return ggml_reshape_4d(ctx0, out, head_dim, n_head, n_seq_tokens, n_seqs);
+    return ggml_reshape_4d(ctx0, out, head_dim, n_head_kda, n_seq_tokens, n_seqs);
 }
 
 llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph_params & params) :
@@ -253,20 +229,18 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
 
         if (hparams.is_recr(il)) {
             const auto * mctx_cur = inp_rs->mctx;
-            const auto cache_head = mctx_cur->get_head();
-            const auto mem_size = mctx_cur->get_size();
             ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
             ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
 
-            ggml_tensor * q = bailingmoe3_causal_conv1d(
-                    gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
-            ggml_tensor * k = bailingmoe3_causal_conv1d(
-                    gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
-            ggml_tensor * v = bailingmoe3_causal_conv1d(
-                    gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv,
-                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
+            ggml_tensor * q = build_causal_conv1d(
+                    inp_rs, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv,
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, il);
+            ggml_tensor * k = build_causal_conv1d(
+                    inp_rs, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv,
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, il);
+            ggml_tensor * v = build_causal_conv1d(
+                    inp_rs, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv,
+                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, il);
 
             ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
             gate = ggml_add(ctx0, gate, layer.ssm_dt_b);

src/models/delta-net-base.cpp

@@ -450,11 +450,6 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state(
         int64_t              conv_kernel_size,
         int64_t              conv_channels,
         int                  il) {
-    const auto * mctx_cur = inp->mctx;
-
-    const auto kv_head  = mctx_cur->get_head();
-    const auto mem_size = mctx_cur->get_size();
-
     const int64_t n_seqs = ubatch.n_seqs;
 
     ggml_tensor * conv_states = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs);
@@ -469,54 +464,7 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state(
     ggml_tensor * conv_input = ggml_concat(ctx0, conv_states, qkv_mixed, 0);
     cb(conv_input, "conv_input", il);
 
-    const int64_t row_count = (conv_kernel_size - 1) * conv_channels;
-
-    const size_t row_size  = ggml_row_size(conv_states_all->type, row_count);
-
-    if (cparams.n_rs_seq == 0) {
-        const int64_t s_idx  = conv_input->ne[0] - conv_states->ne[0];
-        const int64_t s_slot = 0;
-
-        ggml_tensor * conv_state_last =
-            ggml_view_3d(ctx0, conv_input,
-                    conv_kernel_size - 1, conv_channels, n_seqs,
-                    conv_input->nb[1], conv_input->nb[2],
-                    ggml_row_size(conv_input->type, s_idx));
-        cb(conv_state_last, "conv_state_last", il);
-
-        ggml_tensor * conv_state_update =
-            ggml_view_2d(ctx0, conv_states_all,
-                    row_count, n_seqs, conv_states_all->nb[1],
-                    (s_slot * mem_size + kv_head) * row_size);
-        cb(conv_state_update, "conv_state_update", il);
-
-        ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_state_last, conv_state_update));
-    } else {
-        // [TAG_RECURRENT_ROLLBACK_SPLITS]
-        // this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside
-        //   the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument
-
-        const int64_t K = (int64_t) cparams.n_rs_seq + 1;
-
-        for (int64_t t = 1; t <= K; ++t) {
-            const int64_t s_idx  = std::max<int64_t>(0, conv_input->ne[0] - conv_states->ne[0] - K + t);
-            const int64_t s_slot = K - t;
-
-            ggml_tensor * conv_state_last =
-                ggml_view_3d(ctx0, conv_input,
-                        conv_kernel_size - 1, conv_channels, n_seqs,
-                        conv_input->nb[1], conv_input->nb[2],
-                        ggml_row_size(conv_input->type, s_idx));
-
-            ggml_tensor * conv_state_update =
-                ggml_view_2d(ctx0,
-                        conv_states_all, row_count, n_seqs,
-                        conv_states_all->nb[1],
-                        (s_slot * mem_size + kv_head) * row_size);
-
-            ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_state_last, conv_state_update));
-        }
-    }
+    build_conv_state_store(inp, conv_states_all, conv_input, conv_kernel_size - 1, 0, il);
 
     return conv_input;
 }

src/models/glm5-next.cpp

@@ -191,13 +191,11 @@ std::unique_ptr<llm_graph_context> llama_model_glm5_next::build_arch_graph(const
 }
 
 // Causal conv1d over one of Q/K/V
-static ggml_tensor * glm5_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
-                                 ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
-                                 int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
-                                 int64_t d_conv, int64_t head_dim, int64_t n_head,
-                                 int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head,
-                                 int64_t mem_size, int64_t K_rs) {
-    const int64_t d_inner         = head_dim * n_head;
+ggml_tensor * llama_model_glm5_next::graph::build_kda_conv1d(
+        llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+        int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+        int64_t d_conv, int64_t head_dim, int64_t n_head_kda, int64_t n_seq_tokens, int64_t n_seqs, int il) {
+    const int64_t d_inner         = head_dim * n_head_kda;
     const int64_t conv_state_size = (d_conv - 1) * d_inner;
     const int64_t n_embd_r_total  = 3 * conv_state_size;
 
@@ -210,26 +208,15 @@ static ggml_tensor * glm5_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
     ggml_tensor * x_3d   = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
     ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0);
 
-    // group s holds the conv window s tokens back.
-    // [TAG_RECURRENT_ROLLBACK_SPLITS]: the last K_rs tokens must share one ubatch.
-    for (int64_t s = 0; s < K_rs; ++s) {
-        const int64_t s_idx = std::max<int64_t>(0, n_seq_tokens - s);
-        ggml_tensor * conv_x_s = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
-            conv_x->nb[1], conv_x->nb[2], s_idx * conv_x->nb[0]);
-        ggml_build_forward_expand(gf,
-            ggml_cpy(ctx0, conv_x_s,
-                ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
-                    (d_conv - 1)   * ggml_element_size(conv_states_all),
-                    n_embd_r_total * ggml_element_size(conv_states_all),
-                    ((s * mem_size + kv_head) * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
-    }
+    build_conv_state_store(inp_rs, conv_states_all, conv_x, d_conv - 1,
+        ggml_row_size(conv_states_all->type, qkv * conv_state_size), il);
 
     ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
     ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight);
     Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens);
     Xcur = ggml_silu(ctx0, Xcur);
 
-    return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs);
+    return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head_kda, n_seq_tokens, n_seqs);
 }
 
 
@@ -812,17 +799,13 @@ ggml_tensor * llama_model_glm5_next::graph::build_kda_layer(
         int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) {
 
     const auto * mctx_cur = inp_rs->mctx;
-    const auto   kv_head  = mctx_cur->get_head();
 
     ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
     ggml_tensor * conv_state_all  = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
 
-    const int64_t mem_size = mctx_cur->get_size();
-    const int64_t K_rs     = (int64_t) cparams.n_rs_seq + 1;
-
-    ggml_tensor * Qcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
-    ggml_tensor * Kcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
-    ggml_tensor * Vcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
+    ggml_tensor * Qcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, il);
+    ggml_tensor * Kcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, il);
+    ggml_tensor * Vcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, il);
     cb(Qcur, "kda_q_conv", il);
     cb(Kcur, "kda_k_conv", il);
     cb(Vcur, "kda_v_conv", il);

src/models/kimi-k3.cpp

@@ -1,6 +1,5 @@
 #include "models.h"
 
-#include <algorithm>
 #include "llama-memory-recurrent.h"
 
 //
@@ -353,13 +352,11 @@ llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_par
 //
 
 // causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use
-static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
-                                    ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
-                                    int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
-                                    int64_t d_conv, int64_t head_dim, int64_t n_head,
-                                    int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head,
-                                    int64_t mem_size, int64_t K_rs) {
-    const int64_t d_inner         = head_dim * n_head;
+ggml_tensor * llama_model_kimi_k3::graph::build_kda_conv1d(
+        llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+        int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+        int64_t d_conv, int64_t head_dim, int64_t n_head_kda, int64_t n_seq_tokens, int64_t n_seqs, int il) {
+    const int64_t d_inner         = head_dim * n_head_kda;
     const int64_t conv_state_size = (d_conv - 1) * d_inner;
     const int64_t n_embd_r_total  = 3 * conv_state_size;
 
@@ -372,26 +369,15 @@ static ggml_tensor * kimi_k3_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
     ggml_tensor * x_3d   = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
     ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0);
 
-    // group s holds the conv window s tokens back.
-    // [TAG_RECURRENT_ROLLBACK_SPLITS]: the last K_rs tokens must share one ubatch.
-    for (int64_t s = 0; s < K_rs; ++s) {
-        const int64_t s_idx = std::max<int64_t>(0, n_seq_tokens - s);
-        ggml_tensor * conv_x_s = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
-            conv_x->nb[1], conv_x->nb[2], s_idx * conv_x->nb[0]);
-        ggml_build_forward_expand(gf,
-            ggml_cpy(ctx0, conv_x_s,
-                ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
-                    (d_conv - 1)   * ggml_element_size(conv_states_all),
-                    n_embd_r_total * ggml_element_size(conv_states_all),
-                    ((s * mem_size + kv_head) * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
-    }
+    build_conv_state_store(inp_rs, conv_states_all, conv_x, d_conv - 1,
+        ggml_row_size(conv_states_all->type, qkv * conv_state_size), il);
 
     ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
     ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight);
     Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens);
     Xcur = ggml_silu(ctx0, Xcur);
 
-    return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs);
+    return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head_kda, n_seq_tokens, n_seqs);
 }
 
 ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer(
@@ -400,17 +386,13 @@ ggml_tensor * llama_model_kimi_k3::graph::build_kda_layer(
         int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) {
 
     const auto * mctx_cur = inp_rs->mctx;
-    const auto   kv_head  = mctx_cur->get_head();
 
     ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
     ggml_tensor * conv_state_all  = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
 
-    const int64_t mem_size = mctx_cur->get_size();
-    const int64_t K_rs     = (int64_t) cparams.n_rs_seq + 1;
-
-    ggml_tensor * Qcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
-    ggml_tensor * Kcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
-    ggml_tensor * Vcur = kimi_k3_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
+    ggml_tensor * Qcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, il);
+    ggml_tensor * Kcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, il);
+    ggml_tensor * Vcur = build_kda_conv1d(inp_rs, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, il);
     cb(Qcur, "kda_q_conv", il);
     cb(Kcur, "kda_k_conv", il);
     cb(Vcur, "kda_v_conv", il);

src/models/lfm-base.cpp

@@ -2,8 +2,6 @@
 
 #include "llama-memory-recurrent.h"
 
-#include <algorithm>
-
 llm_build_lfm_base::llm_build_lfm_base(const llm_graph_params & params) : llm_graph_context(params) {}
 
 ggml_tensor * llm_build_lfm_base::build_lfm_shortconv(const llama_model & model,
@@ -11,7 +9,6 @@ ggml_tensor * llm_build_lfm_base::build_lfm_shortconv(const llama_model & model,
                                                       llm_graph_input_rs * inp_recr,
                                                       ggml_tensor * cur,
                                                       int il) {
-    const uint32_t kv_head      = mctx_recr->get_head();
     const int64_t  n_seq_tokens = ubatch.n_seq_tokens;
     const int64_t  n_seqs       = ubatch.n_seqs;
     GGML_ASSERT(n_seqs != 0);
@@ -55,20 +52,7 @@ ggml_tensor * llm_build_lfm_base::build_lfm_shortconv(const llama_model & model,
     }
     GGML_ASSERT(bx->ne[0] > conv->ne[0]);
 
-    // write conv states: slot 0 = the final state, slot s = the state s tokens back (partial rollback)
-    const int64_t K         = hparams.causal_attn && cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1;
-    const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
-    const auto    mem_size  = mctx_recr->get_size();
-    const size_t  row_size  = ggml_row_size(conv_state->type, (int64_t) d_conv * n_embd);
-
-    for (int64_t slot = 0; slot < n_written; ++slot) {
-        auto * conv_snap = ggml_view_3d(ctx0, bx, d_conv, bx->ne[1], bx->ne[2], bx->nb[1], bx->nb[2],
-                                        (bx->ne[0] - d_conv - slot) * ggml_element_size(bx));
-        ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_snap,
-                                               ggml_view_2d(ctx0, conv_state, (int64_t) d_conv * n_embd, n_seqs,
-                                                            conv_state->nb[1],
-                                                            ((size_t) slot * mem_size + kv_head) * row_size)));
-    }
+    build_conv_state_store(inp_recr, conv_state, bx, d_conv, 0, il);
 
     auto * conv_kernel = model.layers[il].shortconv.conv;
     auto * conv_out    = ggml_ssm_conv(ctx0, bx, conv_kernel);

src/models/lfm3.cpp

@@ -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

@@ -206,19 +206,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp,
         // => {d_conv - 1 + n_seq_tokens, d_inner + 2*n_group*d_state, n_seqs}
         ggml_tensor * conv_x = ggml_concat(ctx0, conv, ggml_transpose(ctx0, xBC), 0);
 
-        const int64_t row_count = (d_conv - 1) * (d_inner + 2 * n_group * d_state);
-        const size_t  row_size  = ggml_row_size(conv_states_all->type, row_count);
-        const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
-
-        for (int64_t slot = 0; slot < n_written; ++slot) {
-            ggml_tensor * last_conv = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner + 2 * n_group * d_state, n_seqs,
-                                                   conv_x->nb[1], conv_x->nb[2], (n_seq_tokens - slot) * conv_x->nb[0]);
-
-            ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv,
-                                                   ggml_view_2d(ctx0, conv_states_all, row_count, n_seqs,
-                                                                conv_states_all->nb[1],
-                                                                ((size_t) slot * mem_size + kv_head) * row_size)));
-        }
+        build_conv_state_store(inp, conv_states_all, conv_x, d_conv - 1, 0, il);
 
         // 1D convolution
         // The equivalent is to make a self-overlapping view of conv_x

src/models/models.h

@@ -2143,6 +2143,11 @@ struct llama_model_bailingmoe3 : public llama_model_base {
         graph(const llama_model & model, const llm_graph_params & params);
 
         const llama_model & model;
+
+        // causal conv1d over one of Q/K/V, qkv selects which third of the conv state row is used
+        ggml_tensor * build_causal_conv1d(llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+                                          int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+                                          int64_t d_conv, int64_t head_dim, int64_t n_head_kda, int64_t n_seq_tokens, int64_t n_seqs, int il);
     };
 
     struct graph_mtp : public llm_graph_context {
@@ -3024,6 +3029,11 @@ struct llama_model_kimi_k3 : public llama_model_base {
                                       int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
                                       int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il);
 
+        // causal conv1d over one of Q/K/V, qkv selects which third of the conv state row is used
+        ggml_tensor * build_kda_conv1d(llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+                                       int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+                                       int64_t d_conv, int64_t head_dim, int64_t n_head_kda, int64_t n_seq_tokens, int64_t n_seqs, int il);
+
         ggml_tensor * build_mla_layer(ggml_tensor * cur, const llama_layer & layer,
                                       llm_graph_input_attn_k  * inp_attn_k,
                                       llm_graph_input_attn_kv * inp_attn_kv,
@@ -3096,6 +3106,11 @@ struct llama_model_glm5_next : public llama_model_base {
                                       int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
                                       int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il);
 
+        // causal conv1d over one of Q/K/V, qkv selects which third of the conv state row is used
+        ggml_tensor * build_kda_conv1d(llm_graph_input_rs * inp_rs, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+                                       int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+                                       int64_t d_conv, int64_t head_dim, int64_t n_head_kda, int64_t n_seq_tokens, int64_t n_seqs, int il);
+
         ggml_tensor * build_kpool_select(ggml_tensor * cur, ggml_tensor * qr, const llama_layer & layer,
                                          const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_kpool * inp_kpool, int il);
 

src/models/qwen4exp.cpp

@@ -1402,10 +1402,6 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(
         int64_t              state_cols,
         int64_t              channels,
         int                  il) {
-    const auto * mctx_cur = inp->mctx;
-
-    const auto kv_head = mctx_cur->get_head();
-
     const int64_t n_seqs    = ubatch.n_seqs;
     const int64_t row_total = conv_states_all->ne[0];
 
@@ -1423,28 +1419,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(
 
     ggml_tensor * conv_input = ggml_concat(ctx0, state, ggml_transpose(ctx0, x), 0);
 
-    // [TAG_RECURRENT_ROLLBACK_SPLITS] keep the last state_cols columns once per rollback slot,
-    // slot s ending s tokens earlier so a rollback of s tokens reads a history that never saw them
-    const size_t row_size = ggml_row_size(conv_states_all->type, row_total);
-    const uint32_t mem_size = mctx_cur->get_size();
-
-    const int64_t n_slots = (int64_t) cparams.n_rs_seq + 1;
-
-    for (int64_t slot = 0; slot < n_slots; ++slot) {
-        const int64_t s_idx = std::max<int64_t>(0, conv_input->ne[0] - state_cols - slot);
-
-        ggml_tensor * tail = ggml_view_3d(ctx0, conv_input,
-                state_cols, channels, n_seqs,
-                conv_input->nb[1], conv_input->nb[2],
-                ggml_row_size(conv_input->type, s_idx));
-
-        ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all,
-                state_cols * channels, n_seqs,
-                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));
-    }
+    build_conv_state_store(inp, conv_states_all, conv_input, state_cols, 0, il);
 
     return conv_input;
 }

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,6 +288,11 @@ 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_CHAMELEON) {
@@ -357,6 +361,10 @@ 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);
         }

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

@@ -47,6 +47,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,
@@ -647,6 +665,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"},

tools/mtmd/clip-model.h

@@ -165,6 +165,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 +298,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;
@@ -926,6 +934,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.cpp

@@ -405,13 +405,13 @@ ggml_tensor * clip_graph::build_vit(
 
                 if (layer.q_norm) {
                     GGML_ASSERT(layer.q_norm->ne[0] == Qcur->ne[0]);
-                    Qcur = build_norm(Qcur, layer.q_norm, NULL, norm_t, eps, il);
+                    Qcur = build_norm(Qcur, layer.q_norm, layer.q_norm_b, norm_t, eps, il);
                     cb(Qcur, "Qcur_norm", il);
                 }
 
                 if (layer.k_norm) {
                     GGML_ASSERT(layer.k_norm->ne[0] == Kcur->ne[0]);
-                    Kcur = build_norm(Kcur, layer.k_norm, NULL, norm_t, eps, il);
+                    Kcur = build_norm(Kcur, layer.k_norm, layer.k_norm_b, norm_t, eps, il);
                     cb(Kcur, "Kcur_norm", il);
                 }
 
@@ -437,11 +437,11 @@ ggml_tensor * clip_graph::build_vit(
 
                 if (!norm_per_head) {
                     if (layer.q_norm) {
-                        Qcur = build_norm(Qcur, layer.q_norm, NULL, norm_t, eps, il);
+                        Qcur = build_norm(Qcur, layer.q_norm, layer.q_norm_b, norm_t, eps, il);
                         cb(Qcur, "Qcur_norm", il);
                     }
                     if (layer.k_norm) {
-                        Kcur = build_norm(Kcur, layer.k_norm, NULL, norm_t, eps, il);
+                        Kcur = build_norm(Kcur, layer.k_norm, layer.k_norm_b, norm_t, eps, il);
                         cb(Kcur, "Kcur_norm", il);
                     }
                 }
@@ -452,11 +452,11 @@ ggml_tensor * clip_graph::build_vit(
 
                 if (norm_per_head) {
                     if (layer.q_norm) {
-                        Qcur = build_norm(Qcur, layer.q_norm, NULL, norm_t, eps, il);
+                        Qcur = build_norm(Qcur, layer.q_norm, layer.q_norm_b, norm_t, eps, il);
                         cb(Qcur, "Qcur_norm_per_head", il);
                     }
                     if (layer.k_norm) {
-                        Kcur = build_norm(Kcur, layer.k_norm, NULL, norm_t, eps, il);
+                        Kcur = build_norm(Kcur, layer.k_norm, layer.k_norm_b, norm_t, eps, il);
                         cb(Kcur, "Kcur_norm_per_head", il);
                     }
                 }
@@ -1242,6 +1242,10 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
             {
                 builder = std::make_unique<clip_graph_parakeet>(ctx, img);
             } break;
+        case PROJECTOR_TYPE_NEMOTRON_TRANSCRIBE:
+            {
+                builder = std::make_unique<clip_graph_nemotron_transcribe>(ctx, img, params ? params->enc_state : nullptr);
+            } break;
         case PROJECTOR_TYPE_GRANITE4_VISION:
             {
                 builder = std::make_unique<clip_graph_granite4_vision>(ctx, img);
@@ -1623,6 +1627,27 @@ struct clip_model_loader {
                         hparams.audio_window_len   = 400;
                         hparams.audio_hop_len      = 160;
                     } break;
+                case PROJECTOR_TYPE_NEMOTRON_TRANSCRIBE:
+                    {
+                        get_u32(KEY_AUDIO_SUBSMPL_FACTOR, hparams.subsampling_factor);
+                        get_u32(KEY_A_ROPE_DIMS,          hparams.audio_rope_dims);
+                        get_u32(KEY_A_SPK_N_BLOCK,        hparams.spk_n_layer);
+                        get_u32(KEY_A_SPK_N_HEAD,         hparams.spk_n_head);
+                        if (hparams.subsampling_factor != 8) {
+                            throw std::runtime_error(string_format("%s: unsupported subsampling factor %d\n", __func__, hparams.subsampling_factor));
+                        }
+                        if (hparams.spk_n_layer <= 0 || hparams.spk_n_layer > (int) gguf_get_n_tensors(ctx_gguf.get())) {
+                            throw std::runtime_error(string_format("%s: invalid speaker block count %d\n", __func__, hparams.spk_n_layer));
+                        }
+                        hparams.audio_chunk_len    = 0;
+                        hparams.audio_sample_rate  = 16000;
+                        hparams.audio_n_fft        = 512;
+                        hparams.audio_window_len   = 400;
+                        hparams.audio_hop_len      = 160;
+                        // one full window
+                        using nt = clip_graph_nemotron_transcribe;
+                        hparams.warmup_audio_size = (nt::window_left + nt::window_core + nt::window_right) * hparams.subsampling_factor;
+                    } break;
                 case PROJECTOR_TYPE_IDEFICS3:
                     {
                         // use default llava-uhd preprocessing params
@@ -4073,6 +4098,87 @@ struct clip_model_loader {
                     model.mm_model_mlp_2_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 1, "weight"));
                     model.mm_model_mlp_3_w = get_tensor(string_format(TN_MVLM_PROJ_MLP, 3, "weight"));
                 } break;
+            case PROJECTOR_TYPE_NEMOTRON_TRANSCRIBE:
+                {
+                    hparams.mel_filters = get_vector(TN_MEL_FILTERS);
+                    hparams.window      = get_vector(TN_WINDOW);
+
+                    // ASR expert: blocks, pre_ln and post_ln come from the generic loader
+                    model.inp_proj_w = get_tensor(string_format(TN_INP_PROJ, "weight"));
+                    for (int il = 0; il < hparams.n_layer; ++il) {
+                        auto & layer = model.layers[il];
+                        layer.q_norm_b = get_tensor(string_format(TN_ATTN_Q_NORM, prefix, il, "bias"));
+                        layer.k_norm_b = get_tensor(string_format(TN_ATTN_K_NORM, prefix, il, "bias"));
+                    }
+
+                    // speaker expert
+                    const char * sp = "a.spk";
+                    model.spk_inp_proj_w   = get_tensor(string_format(TN_A_SPK_INP_PROJ, "weight"));
+                    model.spk_pre_ln_w     = get_tensor(string_format(TN_LN_PRE,  sp, "weight"));
+                    model.spk_pre_ln_b     = get_tensor(string_format(TN_LN_PRE,  sp, "bias"));
+                    model.spk_post_ln_w    = get_tensor(string_format(TN_LN_POST, sp, "weight"));
+                    model.spk_post_ln_b    = get_tensor(string_format(TN_LN_POST, sp, "bias"));
+                    model.spk_layers.resize(hparams.spk_n_layer);
+                    for (int il = 0; il < hparams.spk_n_layer; ++il) {
+                        auto & layer = model.spk_layers[il];
+                        layer.ln_1_w    = get_tensor(string_format(TN_LN_1,        sp, il, "weight"));
+                        layer.ln_1_b    = get_tensor(string_format(TN_LN_1,        sp, il, "bias"));
+                        layer.qkv_w     = get_tensor(string_format(TN_ATTN_QKV,    sp, il, "weight"));
+                        layer.o_w       = get_tensor(string_format(TN_ATTN_OUTPUT, sp, il, "weight"));
+                        layer.o_b       = get_tensor(string_format(TN_ATTN_OUTPUT, sp, il, "bias"));
+                        layer.ln_2_w    = get_tensor(string_format(TN_LN_2,        sp, il, "weight"));
+                        layer.ln_2_b    = get_tensor(string_format(TN_LN_2,        sp, il, "bias"));
+                        layer.ff_up_w   = get_tensor(string_format(TN_FFN_UP,      sp, il, "weight"));
+                        layer.ff_up_b   = get_tensor(string_format(TN_FFN_UP,      sp, il, "bias"));
+                        layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN,    sp, il, "weight"));
+                        layer.ff_down_b = get_tensor(string_format(TN_FFN_DOWN,    sp, il, "bias"));
+                    }
+
+                    // Sortformer head, fusion
+                    model.spk_head_proj_w   = get_tensor(string_format(TN_A_SPK_HEAD, "proj",     "weight"));
+                    model.spk_head_proj_b   = get_tensor(string_format(TN_A_SPK_HEAD, "proj",     "bias"));
+                    model.spk_head_up_w     = get_tensor(string_format(TN_A_SPK_HEAD, "upsample", "weight"));
+                    model.spk_head_up_b     = get_tensor(string_format(TN_A_SPK_HEAD, "upsample", "bias"));
+                    model.spk_head_fc1_w    = get_tensor(string_format(TN_A_SPK_HEAD, "fc1",      "weight"));
+                    model.spk_head_fc1_b    = get_tensor(string_format(TN_A_SPK_HEAD, "fc1",      "bias"));
+                    model.spk_head_fc2_w    = get_tensor(string_format(TN_A_SPK_HEAD, "fc2",      "weight"));
+                    model.spk_head_fc2_b    = get_tensor(string_format(TN_A_SPK_HEAD, "fc2",      "bias"));
+                    model.fusion_asr_norm_w = get_tensor(string_format(TN_A_FUSION_NORM, "asr_norm", "weight"));
+                    model.fusion_asr_norm_b = get_tensor(string_format(TN_A_FUSION_NORM, "asr_norm", "bias"));
+                    model.fusion_spk_norm_w = get_tensor(string_format(TN_A_FUSION_NORM, "spk_norm", "weight"));
+                    model.fusion_spk_norm_b = get_tensor(string_format(TN_A_FUSION_NORM, "spk_norm", "bias"));
+                    model.fusion_spk_kernel = get_tensor(TN_A_FUSION_KERNEL);
+
+                    model.spk_sil_emb = get_vector(TN_A_SPK_SIL_EMB);
+
+                    // the head layouts and the widths the graph and the speaker state use must match the tensors
+                    hparams.spk_n_embd = model.spk_inp_proj_w->ne[1];
+                    const int64_t d_head     = hparams.n_head     > 0 ? hparams.n_embd     / hparams.n_head     : 0;
+                    const int64_t d_head_spk = hparams.spk_n_head > 0 ? hparams.spk_n_embd / hparams.spk_n_head : 0;
+                    const int64_t n_spk      = model.spk_head_fc2_w->ne[1];
+                    // the graph splits each fused QKV into three views of the layer width
+                    bool qkv_ok = true;
+                    for (const auto & layer : model.layers) {
+                        qkv_ok = qkv_ok && layer.qkv_w && layer.qkv_w->ne[0] == hparams.n_embd && layer.qkv_w->ne[1] == 3 * hparams.n_embd;
+                    }
+                    for (const auto & layer : model.spk_layers) {
+                        qkv_ok = qkv_ok && layer.qkv_w->ne[0] == hparams.spk_n_embd && layer.qkv_w->ne[1] == 3 * hparams.spk_n_embd;
+                    }
+                    if (!qkv_ok || model.inp_proj_w->ne[1] != hparams.n_embd || d_head <= 0 || d_head * hparams.n_head != hparams.n_embd ||
+                            (hparams.n_embd_head > 0 && hparams.n_embd_head != d_head) ||
+                            hparams.audio_rope_dims <= 0 || hparams.audio_rope_dims % 2 != 0 || hparams.audio_rope_dims > d_head ||
+                            d_head_spk <= 0 || d_head_spk % 2 != 0 || d_head_spk * hparams.spk_n_head != hparams.spk_n_embd ||
+                            (int64_t) model.spk_sil_emb.size() != hparams.spk_n_embd || n_spk <= 0 ||
+                            model.fusion_spk_kernel->ne[0] != n_spk || model.fusion_spk_kernel->ne[1] != hparams.n_embd ||
+                            hparams.mel_filters.size() != (size_t) hparams.n_mel_bins * (hparams.audio_n_fft / 2 + 1) ||
+                            hparams.window.size() != (size_t) hparams.audio_window_len) {
+                        throw std::runtime_error(string_format("%s: hparams do not match the tensor shapes\n", __func__));
+                    }
+
+                    // projection to the text model
+                    model.mm_fc_w = get_tensor(string_format(TN_MM_AUDIO_FC, "weight"));
+                    model.mm_fc_b = get_tensor(string_format(TN_MM_AUDIO_FC, "bias"));
+                } break;
             case PROJECTOR_TYPE_GRANITE_SPEECH:
                 {
                     model.inp_proj_w     = get_tensor(string_format(TN_INP_PROJ,    "weight"));
@@ -5013,6 +5119,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
                 }
                 n_patches = n;
             } break;
+        case PROJECTOR_TYPE_NEMOTRON_TRANSCRIBE:
         case PROJECTOR_TYPE_PARAKEET:
             {
                 n_patches = (img->nx() + (params.subsampling_factor - 1)) / params.subsampling_factor;
@@ -5100,6 +5207,9 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
     params.n_threads = n_threads;
     params.out_embd = &out_batch_embd;
 
+    if (ctx->proj_type() == PROJECTOR_TYPE_NEMOTRON_TRANSCRIBE) {
+        return clip_graph_nemotron_transcribe::encode_windows(ctx, ctx->model, &params);
+    }
     return clip_encode(ctx, &params);
 }
 
@@ -6593,6 +6703,20 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
 
                 set_lfm3a_window_inputs(base, n_pos, hparams.attn_window_size);
             } break;
+        case PROJECTOR_TYPE_NEMOTRON_TRANSCRIBE:
+            {
+                GGML_ASSERT(imgs.entries.size() == 1);
+                std::vector<int32_t> positions(ggml_nelements(get_inp_tensor("positions")));
+                for (size_t i = 0; i < positions.size(); i++) {
+                    positions[i] = (int32_t) i;
+                }
+                set_input_i32("positions", positions);
+                GGML_ASSERT(params->enc_state != nullptr);
+                set_input_f32("inp_norm", params->enc_state->mel_norm);
+                if (ggml_graph_get_tensor(gf, "spk_prev") != nullptr) {
+                    set_input_f32("spk_prev", params->enc_state->spk_prev);
+                }
+            } break;
         case PROJECTOR_TYPE_PARAKEET:
             {
                 GGML_ASSERT(imgs.entries.size() == 1);
@@ -6899,6 +7023,15 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
             GGML_ASSERT(t != nullptr);
             copy_gen_state_out(t, slots[i++]);
         }
+    } else if (params->enc_state != nullptr && model.proj_type == PROJECTOR_TYPE_NEMOTRON_TRANSCRIBE) {
+        clip_enc_state * st = params->enc_state;
+        ggml_tensor * chunk = ggml_graph_get_tensor(gf, "spk_chunk");
+        ggml_tensor * preds = ggml_graph_get_tensor(gf, "spk_preds");
+        GGML_ASSERT(chunk != nullptr && preds != nullptr);
+        st->spk_chunk.resize(ggml_nelements(chunk));
+        st->spk_preds.resize(ggml_nelements(preds));
+        ggml_backend_tensor_get(chunk, st->spk_chunk.data(), 0, ggml_nbytes(chunk));
+        ggml_backend_tensor_get(preds, st->spk_preds.data(), 0, ggml_nbytes(preds));
     } else if (params->enc_state != nullptr) {
         clip_enc_state * st = params->enc_state;
         ggml_tensor * tail = ggml_graph_get_tensor(gf, "out_stem_tail");
@@ -7057,6 +7190,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
         case PROJECTOR_TYPE_LLAMA4:
             return ctx->model.mm_model_proj->ne[1];
         case PROJECTOR_TYPE_QWEN2A:
+        case PROJECTOR_TYPE_NEMOTRON_TRANSCRIBE:
             return ctx->model.mm_fc_w->ne[1];
         case PROJECTOR_TYPE_QWEN3A:
             return ctx->model.mm_2_w->ne[1];

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;

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

@@ -1385,13 +1385,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 +1417,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 +1455,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 +1502,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

@@ -176,6 +176,9 @@ struct mtmd_audio_preprocessor_parakeet : mtmd_audio_preprocessor {
     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 +189,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.cpp

@@ -1153,6 +1153,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>";

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/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;
+}

tools/ui/package.json

@@ -18,6 +18,11 @@
 		"test": "npm run test:ui -- --run && npm run test:client -- --run && npm run test:unit -- --run && npm run test:e2e",
 		"test:e2e": "playwright test",
 		"test:e2e:pwa": "playwright test tests/e2e/pwa.e2e.ts",
+		"demo:codegen": "playwright codegen --target=playwright-test --viewport-size=1440,900 --output=demo/recording.demo.ts http://localhost:5180/",
+		"demo:docs": "node demo/cli.ts docs",
+		"demo:pr": "node demo/cli.ts pr",
+		"demo:record": "playwright test --config=demo/playwright.config.ts",
+		"demo:traces": "node demo/cli.ts traces",
 		"test:client": "vitest --project=client",
 		"test:unit": "vitest --project=unit",
 		"test:unit:pwa": "vitest --project=unit --run tests/unit/pwa.spec.ts",