model: add MTP support for Nemotron model (#26725)

* model: add MTP support for Nemotron Nano model

* model: add mtp_flags for nemotron model

* address review comments

.devops/rocm.Dockerfile

@@ -57,7 +57,6 @@ COPY --from=web /app/tools/ui/dist tools/ui/dist
 RUN HIPCXX="$(hipconfig -l)/clang" HIP_PATH="$(hipconfig -R)" \
     cmake -S . -B build \
         -DGGML_HIP=ON \
-        -DGGML_HIP_ROCWMMA_FATTN=ON \
         -DAMDGPU_TARGETS="$ROCM_DOCKER_ARCH" \
         -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON \
         -DCMAKE_BUILD_TYPE=Release -DLLAMA_BUILD_TESTS=OFF \

.github/workflows/build-cuda-ubuntu.yml

@@ -99,7 +99,6 @@ jobs:
         run: |
           cmake -B build -S . \
             -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" \
-            -DGGML_HIP_ROCWMMA_FATTN=ON \
             -DGPU_TARGETS="gfx1030" \
             -DGGML_HIP=ON
           cmake --build build --config Release -j $(nproc)

.github/workflows/build-cuda-windows.yml

@@ -150,7 +150,6 @@ jobs:
             -DLLAMA_BUILD_BORINGSSL=ON `
             -DROCM_DIR="${env:HIP_PATH}" `
             -DGGML_HIP=ON `
-            -DGGML_HIP_ROCWMMA_FATTN=ON `
             -DGPU_TARGETS="gfx1100"  `
             -DGGML_RPC=ON
           cmake --build build -j ${env:NUMBER_OF_PROCESSORS}

.github/workflows/build-sanitize.yml

@@ -15,6 +15,12 @@ on:
       '**/*.cpp'
     ]
 
+  pull_request:
+    types: [opened, synchronize, reopened]
+    paths: [
+      '.github/workflows/build-sanitize.yml'
+    ]
+
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}
   cancel-in-progress: true
@@ -28,19 +34,35 @@ env:
 
 jobs:
   ctest:
-    runs-on: [self-hosted, X64, CPU, Linux]
-
     continue-on-error: true
 
     strategy:
       matrix:
-        sanitizer: [ADDRESS, THREAD, UNDEFINED]
+        include:
+        - sanitizer: ADDRESS
+          machine: [self-hosted, X64, Linux]
+        # thread doesn't run properly on some self hosted machines, so run it on Github instead
+        - sanitizer: THREAD
+          machine: ubuntu-24.04
+        - sanitizer: UNDEFINED
+          machine: [self-hosted, X64, Linux]
+
+    runs-on: ${{ matrix.machine }}
 
     steps:
       - name: Clone
         id: checkout
         uses: actions/checkout@v6
 
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.21
+        if: ${{ matrix.sanitizer == 'THREAD' }}
+        with:
+          key: ctest-thread-ubuntu-24.04
+          variant: ccache
+          evict-old-files: 1d
+          save: ${{ github.event_name == 'push' && github.ref == 'refs/heads/master' }}
+
       # with UNDEFINED sanitizer, we have to build in Debug to avoid GCC 13 false-positive warnings
       - name: Build (undefined)
         id: cmake_build_undefined

.github/workflows/pr-draft-label.yml

@@ -0,0 +1,23 @@
+name: Convert PR to draft
+
+on:
+  pull_request_target:
+    types: [labeled]
+
+permissions:
+  pull-requests: write
+  issues: write
+  contents: write # required for "gh pr ready" command, see https://github.com/cli/cli/issues/8910
+
+jobs:
+  convert-to-draft:
+    if: github.event.label.name == 'draft' && github.event.pull_request.draft == false
+    runs-on: ubuntu-slim
+    steps:
+      - name: Convert PR to draft
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          PR_URL: ${{ github.event.pull_request.html_url }}
+        run: |
+          gh pr ready --undo "$PR_URL"
+          gh pr edit "$PR_URL" --remove-label draft

.github/workflows/release.yml

@@ -1229,7 +1229,6 @@ jobs:
             -DGPU_TARGETS="${{ matrix.gpu_targets }}" \
             -DGGML_HIP=ON \
             -DHIP_PLATFORM=amd \
-            -DGGML_HIP_ROCWMMA_FATTN=ON \
             -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} \
             ${{ env.CMAKE_ARGS }}
           cmake --build build --config Release -j $(nproc)
@@ -1353,7 +1352,6 @@ jobs:
             -DGGML_NATIVE=OFF `
             -DGGML_CPU=OFF `
             -DGPU_TARGETS="${{ matrix.gpu_targets }}" `
-            -DGGML_HIP_ROCWMMA_FATTN=ON `
             -DGGML_HIP=ON `
             -DHF_UI_VERSION=${{ needs.get-version.outputs.ui_version }} `
             -DLLAMA_BUILD_BORINGSSL=ON

README.md

@@ -12,7 +12,7 @@
 [![Docker](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml/badge.svg)](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml)
 [![Winget](https://github.com/ggml-org/llama.cpp/actions/workflows/winget.yml/badge.svg)](https://github.com/ggml-org/llama.cpp/actions/workflows/winget.yml)
 
-[manifesto](https://github.com/ggml-org/llama.cpp/discussions/205) / [ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/issues?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Abartowski1182%20OR%20author%3Ahipudding%20OR%20author%3AServeurpersoCom%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3A0cc4m%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3Aallozaur%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Alhez%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming)%20sort%3Aupdated-desc) / [dev branches](https://github.com/ggml-org/llama.cpp-dev/blob/master/README-features.md) / [compile times](https://github.com/ggml-org/llama.cpp-dev/blob/master/README-compile-times.md) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)
+[manifesto](https://github.com/ggml-org/llama.cpp/discussions/205) / [ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/issues?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Abartowski1182%20OR%20author%3Ahipudding%20OR%20author%3AServeurpersoCom%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3A0cc4m%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3Aallozaur%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Alhez%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming)%20sort%3Aupdated-desc) / [compile times](https://github.com/ggml-org/llama.cpp-dev/blob/master/README-compile-times.md) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291)
 
 </div>
 

ci/run.sh

@@ -92,7 +92,7 @@ if [ ! -z ${GG_BUILD_CUDA} ]; then
 fi
 
 if [ ! -z ${GG_BUILD_ROCM} ]; then
-    CMAKE_EXTRA="${CMAKE_EXTRA} -DCMAKE_HIP_COMPILER=$(hipconfig -l)/clang -DGGML_HIP=ON -DGGML_HIP_ROCWMMA_FATTN=ON"
+    CMAKE_EXTRA="${CMAKE_EXTRA} -DCMAKE_HIP_COMPILER=$(hipconfig -l)/clang -DGGML_HIP=ON"
     if [ -z ${GG_BUILD_AMDGPU_TARGETS} ]; then
         echo "Missing GG_BUILD_AMDGPU_TARGETS, please set it to your GPU architecture (e.g. gfx90a, gfx1100, etc.)"
         exit 1

conversion/__init__.py

@@ -103,6 +103,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
     "GraniteMoeForCausalLM": "granite",
     "GraniteMoeHybridForCausalLM": "granite",
     "GraniteMoeSharedForCausalLM": "granite",
+    "GraniteSwitchForCausalLM": "granite",
     "GraniteSpeechForConditionalGeneration": "granite",
     "GraniteSpeechPlusForConditionalGeneration": "granite",
     "Grok1ForCausalLM": "grok",

conversion/granite.py

@@ -123,6 +123,166 @@ class GraniteMoeModel(GraniteModel):
         yield from super().modify_tensors(data_torch, name, bid)
 
 
+@ModelBase.register("GraniteSwitchForCausalLM")
+class GraniteSwitchModel(GraniteMoeModel):
+    """Dense, all-attention Granite with N per-token embedded LoRA adapters, stacked
+    over the adapter dim with a zero adapter at slot 0 (N = num_adapters + 1)."""
+    model_arch = gguf.MODEL_ARCH.GRANITE_SWITCH
+
+    # permute q/k per-slice below (NORM-rope layout), not via the parent's auto-permute
+    undo_permute = False
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # the weightless switch reserves one cache slot: one fewer block than num_hidden_layers
+        self.block_count = self.block_count - 1
+        self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+        self._n_adapters = int(self.hparams["num_adapters"])
+        self._max_lora_rank = int(self.hparams["max_lora_rank"])
+        self._n_slots = self._n_adapters + 1  # +1 for the zero slot at index 0
+
+        n_head = int(self.hparams["num_attention_heads"])
+        n_kv_head = int(self.hparams["num_key_value_heads"])
+        head_dim = (
+            self.hparams.get("projection_head_dim")
+            or self.hparams.get("head_dim")
+            or (self.hparams["hidden_size"] // n_head)
+        )
+        self._n_head = n_head
+        self._n_kv_head = n_kv_head
+        self._head_dim = int(head_dim)
+        self._q_size = n_head * self._head_dim
+        self._kv_size = n_kv_head * self._head_dim
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+
+        # dense: pin expert_used_count to 0 (config carries a leftover num_experts_per_tok)
+        if not self.hparams.get("num_local_experts"):
+            self.gguf_writer.add_expert_used_count(0)
+
+        self.gguf_writer.add_adapter_count(self._n_adapters)
+        self.gguf_writer.add_adapter_lora_rank(self._max_lora_rank)
+        self.gguf_writer.add_adapter_token_ids_activate(self.hparams["adapter_token_ids"])
+        self.gguf_writer.add_adapter_token_ids_substitute(self.hparams["adapter_substitute_token_ids"])
+        router_gain = float(self.hparams.get("control_token_gain", 15.0))
+        self.gguf_writer.add_adapter_router_gain(router_gain)
+        logger.info("gguf: (graniteswitch) num_adapters=%s max_lora_rank=%s n_slots=%s router_gain=%s", self._n_adapters, self._max_lora_rank, self._n_slots, router_gain)
+
+    def _lora_a(self, data: Tensor) -> Tensor:
+        # on-disk A: [n_adapters, 1, max_rank, in] -> [n_adapters+1, max_rank, in]
+        a = data.squeeze(1)
+        zero = torch.zeros_like(a[:1])
+        return torch.cat([zero, a], dim=0).contiguous()
+
+    def _lora_b(self, data: Tensor, permute_n_head: int | None = None) -> Tensor:
+        # on-disk B: [n_adapters, 1, out, max_rank] -> [n_adapters+1, out, max_rank]
+        b = data.squeeze(1)
+        if permute_n_head is not None:
+            # permute each adapter's B output rows to match the permuted q/k base
+            b = torch.stack([self.permute(b[i], permute_n_head, permute_n_head) for i in range(b.shape[0])], dim=0)
+        zero = torch.zeros_like(b[:1])
+        return torch.cat([zero, b], dim=0).contiguous()
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        T = gguf.MODEL_TENSOR
+
+        # skip the weightless switch + control-token buffers (rebuilt at load time)
+        bare = name.split(".")[-1]
+        if (
+            name.startswith("model.switch.") or name.startswith("switch.")
+            or bare in ("adapter_token_ids", "control_to_substitute_lut")
+        ):
+            return
+
+        if "self_attn.qkv_proj" in name:
+            if name.endswith("base_layer.weight"):
+                # fused [q|k|v] rows: permute q/k row-blocks for ggml's NORM-rope layout
+                q, k, v = data_torch.split([self._q_size, self._kv_size, self._kv_size], dim=0)
+                q = self.permute(q, self._n_head, self._n_head)
+                k = self.permute(k, self._n_kv_head, self._n_kv_head)
+                fused = torch.cat([q, k, v], dim=0)
+                yield (self.format_tensor_name(T.ATTN_QKV, bid), fused)
+                return
+            if "lora_A_slices." in name:
+                slot = int(name.rsplit(".", 1)[1])
+                key = {0: T.ATTN_Q, 1: T.ATTN_K, 2: T.ATTN_V}[slot]
+                yield (self.format_tensor_name(key, bid, suffix=".lora_a"), self._lora_a(data_torch))
+                return
+            if "lora_B_slices." in name:
+                slot = int(name.rsplit(".", 1)[1])
+                key, ph = {
+                    0: (T.ATTN_Q, self._n_head),
+                    1: (T.ATTN_K, self._n_kv_head),
+                    2: (T.ATTN_V, None),
+                }[slot]
+                yield (self.format_tensor_name(key, bid, suffix=".lora_b"), self._lora_b(data_torch, ph))
+                return
+            raise ValueError(f"Unexpected qkv_proj tensor: {name}")
+
+        if "self_attn.o_proj" in name:
+            if name.endswith("base_layer.weight"):
+                yield (self.format_tensor_name(T.ATTN_OUT, bid), data_torch)
+                return
+            if name.endswith("lora_A"):
+                yield (self.format_tensor_name(T.ATTN_OUT, bid, suffix=".lora_a"), self._lora_a(data_torch))
+                return
+            if name.endswith("lora_B"):
+                yield (self.format_tensor_name(T.ATTN_OUT, bid, suffix=".lora_b"), self._lora_b(data_torch))
+                return
+            raise ValueError(f"Unexpected o_proj tensor: {name}")
+
+        if "shared_mlp.input_linear" in name:
+            ffn = self.hparams["shared_intermediate_size"]
+            if name.endswith("base_layer.weight"):
+                gate, up = data_torch.split([ffn, ffn], dim=0)
+                yield (self.format_tensor_name(T.FFN_GATE, bid), gate)
+                yield (self.format_tensor_name(T.FFN_UP, bid), up)
+                return
+            if "lora_A_slices." in name:
+                slot = int(name.rsplit(".", 1)[1])
+                key = {0: T.FFN_GATE, 1: T.FFN_UP}[slot]
+                yield (self.format_tensor_name(key, bid, suffix=".lora_a"), self._lora_a(data_torch))
+                return
+            if "lora_B_slices." in name:
+                slot = int(name.rsplit(".", 1)[1])
+                key = {0: T.FFN_GATE, 1: T.FFN_UP}[slot]
+                yield (self.format_tensor_name(key, bid, suffix=".lora_b"), self._lora_b(data_torch))
+                return
+            raise ValueError(f"Unexpected shared_mlp.input_linear tensor: {name}")
+
+        if "shared_mlp.output_linear" in name:
+            if name.endswith("base_layer.weight"):
+                yield (self.format_tensor_name(T.FFN_DOWN, bid), data_torch)
+                return
+            if name.endswith("lora_A"):
+                yield (self.format_tensor_name(T.FFN_DOWN, bid, suffix=".lora_a"), self._lora_a(data_torch))
+                return
+            if name.endswith("lora_B"):
+                yield (self.format_tensor_name(T.FFN_DOWN, bid, suffix=".lora_b"), self._lora_b(data_torch))
+                return
+            raise ValueError(f"Unexpected shared_mlp.output_linear tensor: {name}")
+
+        if bid is not None and ".layers." in name and (
+            "input_layernorm" in name or "post_attention_layernorm" in name
+        ):
+            key = T.ATTN_NORM if "input_layernorm" in name else T.FFN_NORM
+            yield (self.format_tensor_name(key, bid), data_torch)
+            return
+
+        if name in ("model.embed_tokens.weight", "embed_tokens.weight"):
+            yield (self.format_tensor_name(T.TOKEN_EMBD), data_torch)
+            return
+        if name in ("model.norm.weight", "norm.weight"):
+            yield (self.format_tensor_name(T.OUTPUT_NORM), data_torch)
+            return
+        if name == "lm_head.weight":
+            return  # tied to token_embd
+
+        raise ValueError(f"graniteswitch: unhandled tensor {name!r} (bid={bid})")
+
+
 @ModelBase.register("GraniteMoeHybridForCausalLM", "BambaForCausalLM")
 class GraniteHybridModel(Mamba2Model, GraniteMoeModel):
     """GraniteHybrid is a hybrid SSM + Attention model that uses Mamba2 SSM

conversion/nemotron.py

@@ -197,6 +197,7 @@ class NemotronHModel(GraniteHybridModel):
     """Hybrid mamba2/attention model from NVIDIA"""
     model_arch = gguf.MODEL_ARCH.NEMOTRON_H
     is_moe: bool = False
+    supports_mtp_export = True
 
     def __init__(self, *args, **kwargs):
         # We have to determine the correct model architecture (MoE vs non-MoE) before
@@ -236,6 +237,25 @@ class NemotronHModel(GraniteHybridModel):
             self._ssm_layers = [i for i, val in enumerate(pattern) if val == "mamba"]
             self._mlp_layers = [i for i, val in enumerate(pattern) if val == "moe"]
 
+        # `--no-mtp` drops it entirely; `--mtp` exports only the MTP head
+        self._mtp_bid: int | None = None
+        if self.is_moe and not self.no_mtp:
+            n_nextn = self.hparams.get("num_nextn_predict_layers", 0) or 0
+            if n_nextn > 0:
+                assert n_nextn == 1, (
+                    "NemotronH MTP conversion currently supports num_nextn_predict_layers == 1"
+                )
+                self._mtp_bid = self.block_count
+                self.block_count += 1
+                # The folded MTP block carries both an attention sub-layer and a
+                # MoE sub-layer, so register it as both so the per-layer metadata arrays cover it
+                self._attn_layers.append(self._mtp_bid)
+                self._mlp_layers.append(self._mtp_bid)
+                self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+        if self.mtp_only and self._mtp_bid is None:
+            raise ValueError("--mtp was requested, but this model does not contain a supported MTP head")
+
     def get_attn_layers(self):
         pattern = self.hparams.get("hybrid_override_pattern") or self.hparams.get("layers_block_type")
         if pattern is None:
@@ -246,6 +266,36 @@ class NemotronHModel(GraniteHybridModel):
 
         return [i for i, val in enumerate(pattern) if val == "attention"]
 
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        name, gen = item
+        if name.startswith("mtp."):
+            # --no-mtp: drop the MTP head entirely
+            if cls.no_mtp:
+                return None
+        elif cls.mtp_only:
+            # --mtp: export the MTP head plus the tensors it shares with the target model
+            keep = name in (
+                "backbone.embeddings.weight",
+                "backbone.norm_f.weight",
+                "lm_head.weight",
+            )
+            if not keep:
+                return None
+        return super().filter_tensors((name, gen))
+
+    def prepare_metadata(self, vocab_only: bool):
+        from_dir = self.fname_out.is_dir()
+        super().prepare_metadata(vocab_only=vocab_only)
+
+        if not self.mtp_only or not from_dir:
+            return
+        output_type: str = self.ftype.name.partition("_")[2]
+        fname_default: str = gguf.naming_convention(
+            self.metadata.name, self.metadata.basename, self.metadata.finetune,
+            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
+        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
+
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
 
@@ -284,6 +334,10 @@ class NemotronHModel(GraniteHybridModel):
             if (latent_size := self.hparams.get("moe_latent_size")) is not None:
                 self.gguf_writer.add_moe_latent_size(latent_size)
 
+        # MTP head: number of trailing NextN blocks
+        if self._mtp_bid is not None:
+            self.gguf_writer.add_nextn_predict_layers(self.hparams["num_nextn_predict_layers"])
+
     def set_vocab(self):
         # The NemotronH config uses pattern characters (e.g. '-') that may not
         # be supported by the installed transformers version. AutoTokenizer
@@ -350,15 +404,24 @@ class NemotronHModel(GraniteHybridModel):
         if not self.is_moe:
             self.gguf_writer.add_add_bos_token(True)
 
-    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
-        if self.is_moe and bid is not None:
-            # Skip Multi-Token Prediction (MTP) tensors. These are used for
-            # for speculative decoding but we don't include them in this model
-            # conversion. See https://github.com/ggml-org/llama.cpp/pull/18886
-            if name.startswith("mtp."):
-                logger.info(f"gguf: Skipping MTP (Speculative) layer: {name}")
-                return
+    _MTP_SPECIAL_RENAMES = {
+        "mtp.layers.0.enorm.weight":           "model.layers.{bid}.enorm.weight",
+        "mtp.layers.0.hnorm.weight":           "model.layers.{bid}.hnorm.weight",
+        "mtp.layers.0.eh_proj.weight":         "model.layers.{bid}.eh_proj.weight",
+        "mtp.layers.1.norm.weight":            "model.layers.{bid}.post_attention_layernorm.weight",
+        "mtp.layers.1.final_layernorm.weight": "model.layers.{bid}.shared_head.norm.weight",
+    }
 
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        #   mtp.layers.0: NextN input fusion + attention
+        #   mtp.layers.1: MoE + final head norm
+        if self._mtp_bid is not None and name.startswith(("mtp.layers.0.", "mtp.layers.1.")):
+            suffix = name.split(".", 3)[3]
+            bid = self._mtp_bid
+            renamed = self._MTP_SPECIAL_RENAMES.get(name)
+            name = renamed.format(bid=bid) if renamed else f"backbone.layers.{bid}.{suffix}"
+
+        if self.is_moe and bid is not None:
             if name.endswith("mixer.gate.e_score_correction.bias"):
                 yield from ModelBase.modify_tensors(self, data_torch, name, bid)
                 return

ggml/src/ggml-cpu/spacemit/ime.cpp

@@ -195,6 +195,7 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
                     case GGML_TYPE_Q4_K:
                     case GGML_TYPE_Q6_K:
                     case GGML_TYPE_Q8_0:
+                    case GGML_TYPE_Q5_0:
                     case GGML_TYPE_Q5_1:
                     case GGML_TYPE_Q5_K:
                         //case GGML_TYPE_MXFP4:
@@ -214,6 +215,7 @@ template <typename BLOC_TYPE, int64_t INTER_SIZE, int64_t NB_COLS> class tensor_
                     case GGML_TYPE_Q4_K:
                     case GGML_TYPE_Q6_K:
                     case GGML_TYPE_Q8_0:
+                    case GGML_TYPE_Q5_0:
                     case GGML_TYPE_Q5_1:
                     case GGML_TYPE_Q5_K:
                         //case GGML_TYPE_MXFP4:

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

@@ -2651,6 +2651,52 @@ static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope,
     return true;
 }
 
+static bool ggml_cuda_should_fuse_rms_norm_mul_rope(const ggml_tensor * rms_norm,
+                                                    const ggml_tensor * mul,
+                                                    const ggml_tensor * rope) {
+    if (rms_norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || rope->op != GGML_OP_ROPE) {
+        return false;
+    }
+
+    if (rms_norm->src[0]->type != GGML_TYPE_F32 || rms_norm->type != GGML_TYPE_F32 ||
+        mul->src[0]->type != GGML_TYPE_F32 || mul->src[1]->type != GGML_TYPE_F32 ||
+        mul->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32) {
+        return false;
+    }
+
+    if (rope->src[0] != mul) {
+        return false;
+    }
+
+    //if rms norm is the B operand, then we don't handle broadcast
+    if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) {
+        return false;
+    }
+
+    if (!ggml_are_same_shape(rms_norm, mul)) {
+        return false;
+    }
+
+    //rms_norm kernel assumes contiguous rows
+    if (!ggml_is_contiguous_rows(rms_norm->src[0]) ||
+        !ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) {
+        return false;
+    }
+
+    // the fused kernel handles the norm/neox rope modes only
+    const int mode = ((const int32_t *) rope->op_params)[2];
+    if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) {
+        return false;
+    }
+
+    const int n_dims = ((const int32_t *) rope->op_params)[1];
+    if (n_dims % 2 != 0 || rope->src[0]->ne[0] % 2 != 0) {
+        return false;
+    }
+
+    return true;
+}
+
 // match gated_delta_net + the strided cpy that scatters its state snapshots into the cache
 // (slot i -> rollback group i, slot 0 newest), so the kernel can write them and skip the cpy.
 static int ggml_cuda_try_gdn_cache_fusion(
@@ -2980,6 +3026,36 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
         }
     }
 
+    std::initializer_list<enum ggml_op> rms_norm_mul_rope_ops          = { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE };
+    std::initializer_list<enum ggml_op> rms_norm_mul_rope_set_rows_ops = { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS };
+
+    if (is_equal(rms_norm_mul_rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) {
+        const ggml_tensor * rms_norm = cgraph->nodes[node_idx];
+        const ggml_tensor * mul      = cgraph->nodes[node_idx + 1];
+        const ggml_tensor * rope     = cgraph->nodes[node_idx + 2];
+        const ggml_tensor * view     = cgraph->nodes[node_idx + 3];
+        const ggml_tensor * set_rows = cgraph->nodes[node_idx + 4];
+
+        if (ggml_check_edges(cgraph, node_idx, {{1, 0, 0}, {2, 0, 1}, {3, 0, 2}, {4, 0, 3}}) &&
+            ggml_cuda_should_fuse_rms_norm_mul_rope(rms_norm, mul, rope) &&
+            ggml_cuda_should_fuse_rope_set_rows(rope, view, set_rows)) {
+            int out_nodes[] = { node_idx + 4 };
+            return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1);
+        }
+    }
+
+    if (is_equal(rms_norm_mul_rope_ops, ops) && ggml_can_fuse(cgraph, node_idx, ops)) {
+        const ggml_tensor * rms_norm = cgraph->nodes[node_idx];
+        const ggml_tensor * mul      = cgraph->nodes[node_idx + 1];
+        const ggml_tensor * rope     = cgraph->nodes[node_idx + 2];
+
+        if (ggml_cuda_should_fuse_rms_norm_mul_rope(rms_norm, mul, rope)) {
+            int out_nodes[] = { node_idx + 2 };
+            return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1);
+        }
+        return false;
+    }
+
     std::initializer_list<enum ggml_op> rope_set_rows_ops = { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS };
 
     if (is_equal(rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) {
@@ -2988,7 +3064,8 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph *                cgraph,
         const ggml_tensor * set_rows = cgraph->nodes[node_idx + 2];
 
         if (ggml_cuda_should_fuse_rope_set_rows(rope, view, set_rows)) {
-            return true;
+            int out_nodes[] = { node_idx + 2 };
+            return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1);
         }
     }
 
@@ -3840,6 +3917,16 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
         return fused_node_count - 1;
     }
 
+    if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) {
+        ggml_cuda_op_rms_norm_mul_rope_fused(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2], cgraph->nodes[i + 4]);
+        return 4;
+    }
+
+    if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE }, {})) {
+        ggml_cuda_op_rms_norm_mul_rope_fused(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2], nullptr);
+        return 2;
+    }
+
     if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }, {})) {
         ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i + 1], cgraph->nodes[i + 2]);
         return 2;

ggml/src/ggml-cuda/rope.cu

@@ -670,3 +670,238 @@ void ggml_cuda_op_rope_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
 void ggml_cuda_op_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * rope, ggml_tensor * set_rows) {
     ggml_cuda_op_rope_impl<true>(ctx, rope, set_rows);
 }
+
+// fused RMS_NORM + MUL + ROPE (+ VIEW + SET_ROWS)
+// one block per row: block_reduce gives the norm scale, then each thread applies mul and rope to the elements it owns
+template <int block_size, bool has_ff, typename D>
+static __global__ void rms_norm_mul_rope_f32(
+        const float * x, D * dst, const int ncols,
+        const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t s1, const int64_t s2, const int64_t s3,
+        const float eps,
+        const float * mul,
+        const int64_t mul_s01, const int64_t mul_s02, const int64_t mul_s03,
+        const uint3 mul_ncols_packed, const uint3 mul_nrows_packed,
+        const uint3 mul_nchannels_packed, const uint3 mul_nsamples_packed,
+        const int n_dims, const int32_t * pos,
+        const float freq_scale, const float ext_factor, const float attn_factor,
+        const rope_corr_dims corr_dims, const float theta_scale,
+        const float * freq_factors,
+        const int64_t * row_indices, const int set_rows_stride,
+        const bool is_neox) {
+    ggml_cuda_pdl_lc();
+    const int row     = blockIdx.x;
+    const int channel = blockIdx.y;
+    const int sample  = blockIdx.z;
+    const int tid     = threadIdx.x;
+
+    x += sample*s03 + channel*s02 + row*s01;
+
+    const uint32_t mul_row     = fastmodulo(row,     mul_nrows_packed);
+    const uint32_t mul_channel = fastmodulo(channel, mul_nchannels_packed);
+    const uint32_t mul_sample  = fastmodulo(sample,  mul_nsamples_packed);
+    mul += mul_sample*mul_s03 + mul_channel*mul_s02 + mul_row*mul_s01;
+
+    float tmp = 0.0f;
+
+    ggml_cuda_pdl_sync();
+    for (int col = tid; col < ncols; col += block_size) {
+        const float xi = x[col];
+        tmp += xi * xi;
+    }
+
+    extern __shared__ float s_sum[];
+    tmp = block_reduce<block_reduce_method::SUM, block_size>(tmp, s_sum);
+
+    const float scale = rsqrtf(tmp/ncols + eps);
+
+    int64_t idst = sample*s3 + channel*s2 + row*s1;
+    if (set_rows_stride != 0) {
+        idst = row*s1 + row_indices[channel]*set_rows_stride;
+    }
+    dst += idst;
+
+    for (int i0 = 2*tid; i0 < ncols; i0 += 2*block_size) {
+        int ix0;
+        int ix1;
+        if (is_neox && i0 < n_dims) {
+            ix0 = i0/2;
+            ix1 = i0/2 + n_dims/2;
+        } else {
+            ix0 = i0 + 0;
+            ix1 = i0 + 1;
+        }
+
+        const float x0 = scale * x[ix0] * mul[fastmodulo(ix0, mul_ncols_packed)];
+        const float x1 = scale * x[ix1] * mul[fastmodulo(ix1, mul_ncols_packed)];
+
+        if (i0 >= n_dims) {
+            dst[ix0] = ggml_cuda_cast<D>(x0);
+            dst[ix1] = ggml_cuda_cast<D>(x1);
+            continue;
+        }
+
+        const float theta_base  = pos[channel]*powf(theta_scale, i0/2.0f);
+        const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f;
+
+        float cos_theta;
+        float sin_theta;
+        rope_yarn<true>(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta);
+
+        dst[ix0] = ggml_cuda_cast<D>(x0*cos_theta - x1*sin_theta);
+        dst[ix1] = ggml_cuda_cast<D>(x0*sin_theta + x1*cos_theta);
+    }
+}
+
+template <typename D>
+static void rms_norm_mul_rope_cuda(
+        const float * x, D * dst,
+        const int ncols, const int nrows, const int nchannels, const int nsamples,
+        const int64_t s01, const int64_t s02, const int64_t s03,
+        const int64_t s1, const int64_t s2, const int64_t s3,
+        const float eps,
+        const float * mul,
+        const int64_t mul_s01, const int64_t mul_s02, const int64_t mul_s03,
+        const uint32_t mul_ncols, const uint32_t mul_nrows,
+        const uint32_t mul_nchannels, const uint32_t mul_nsamples,
+        const int n_dims, const int32_t * pos,
+        const float freq_scale, const float freq_base, const float ext_factor, const float attn_factor,
+        const rope_corr_dims corr_dims,
+        const float * freq_factors,
+        const int64_t * row_indices, const int set_rows_stride,
+        const bool is_neox, cudaStream_t stream) {
+    GGML_ASSERT(ncols % 2 == 0);
+
+    const dim3 blocks_num(nrows, nchannels, nsamples);
+
+    const float theta_scale = powf(freq_base, -2.0f/n_dims);
+
+    const uint3 mul_ncols_packed     = init_fastdiv_values(mul_ncols);
+    const uint3 mul_nrows_packed     = init_fastdiv_values(mul_nrows);
+    const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels);
+    const uint3 mul_nsamples_packed  = init_fastdiv_values(mul_nsamples);
+
+    if (ncols < 1024) {
+        const dim3 block_dims(256, 1, 1);
+        const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32*sizeof(float), stream};
+        if (freq_factors == nullptr) {
+            ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<256, false, D>, launch_params,
+                x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
+                mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
+                n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
+                freq_factors, row_indices, set_rows_stride, is_neox);
+        } else {
+            ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<256, true, D>, launch_params,
+                x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
+                mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
+                n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
+                freq_factors, row_indices, set_rows_stride, is_neox);
+        }
+    } else {
+        const dim3 block_dims(1024, 1, 1);
+        const ggml_cuda_kernel_launch_params launch_params = {blocks_num, block_dims, 32*sizeof(float), stream};
+        if (freq_factors == nullptr) {
+            ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<1024, false, D>, launch_params,
+                x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
+                mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
+                n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
+                freq_factors, row_indices, set_rows_stride, is_neox);
+        } else {
+            ggml_cuda_kernel_launch(rms_norm_mul_rope_f32<1024, true, D>, launch_params,
+                x, dst, ncols, s01, s02, s03, s1, s2, s3, eps, mul, mul_s01, mul_s02, mul_s03,
+                mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed,
+                n_dims, pos, freq_scale, ext_factor, attn_factor, corr_dims, theta_scale,
+                freq_factors, row_indices, set_rows_stride, is_neox);
+        }
+    }
+}
+
+void ggml_cuda_op_rms_norm_mul_rope_fused(ggml_backend_cuda_context & ctx,
+        ggml_tensor * rms_norm, ggml_tensor * mul, ggml_tensor * rope, ggml_tensor * set_rows) {
+    const ggml_tensor * x = rms_norm->src[0];
+    const ggml_tensor * mul_src = mul->src[0] == rms_norm ? mul->src[1] : mul->src[0];
+
+    float eps = 0.0f;
+    memcpy(&eps, rms_norm->op_params, sizeof(float));
+    GGML_ASSERT(eps >= 0.0f);
+
+    GGML_ASSERT(x->type == GGML_TYPE_F32);
+    GGML_ASSERT(mul_src->type == GGML_TYPE_F32);
+    GGML_ASSERT(rope->type == GGML_TYPE_F32);
+
+    void *          dst_d           = rope->data;
+    ggml_type       dst_type        = rope->type;
+    const int64_t * row_indices     = nullptr;
+    int             set_rows_stride = 0;
+
+    if (set_rows != nullptr) {
+        dst_d           = set_rows->data;
+        dst_type        = set_rows->type;
+        row_indices     = (const int64_t *) set_rows->src[1]->data;
+        set_rows_stride = set_rows->nb[1] / ggml_type_size(set_rows->type);
+    }
+
+    const int n_dims     = ((const int32_t *) rope->op_params)[1];
+    const int mode       = ((const int32_t *) rope->op_params)[2];
+    const int n_ctx_orig = ((const int32_t *) rope->op_params)[4];
+
+    float freq_base;
+    float freq_scale;
+    float ext_factor;
+    float attn_factor;
+    float beta_fast;
+    float beta_slow;
+
+    memcpy(&freq_base,   (const int32_t *) rope->op_params +  5, sizeof(float));
+    memcpy(&freq_scale,  (const int32_t *) rope->op_params +  6, sizeof(float));
+    memcpy(&ext_factor,  (const int32_t *) rope->op_params +  7, sizeof(float));
+    memcpy(&attn_factor, (const int32_t *) rope->op_params +  8, sizeof(float));
+    memcpy(&beta_fast,   (const int32_t *) rope->op_params +  9, sizeof(float));
+    memcpy(&beta_slow,   (const int32_t *) rope->op_params + 10, sizeof(float));
+
+    const bool is_neox = mode & GGML_ROPE_TYPE_NEOX;
+
+    const int32_t * pos = (const int32_t *) rope->src[1]->data;
+
+    const float * freq_factors = rope->src[2] != nullptr ? (const float *) rope->src[2]->data : nullptr;
+
+    rope_corr_dims corr_dims;
+    ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims.v);
+
+    const size_t ts0 = ggml_type_size(x->type);
+    GGML_ASSERT(x->nb[0] == ts0);
+    const int64_t s01 = x->nb[1] / ts0;
+    const int64_t s02 = x->nb[2] / ts0;
+    const int64_t s03 = x->nb[3] / ts0;
+
+    const size_t ts_mul = ggml_type_size(mul_src->type);
+    GGML_ASSERT(mul_src->nb[0] == ts_mul);
+    const int64_t mul_s01 = mul_src->nb[1] / ts_mul;
+    const int64_t mul_s02 = mul_src->nb[2] / ts_mul;
+    const int64_t mul_s03 = mul_src->nb[3] / ts_mul;
+
+    const size_t ts_dst = ggml_type_size(rope->type);
+    const int64_t s1 = rope->nb[1] / ts_dst;
+    const int64_t s2 = rope->nb[2] / ts_dst;
+    const int64_t s3 = rope->nb[3] / ts_dst;
+
+    cudaStream_t stream = ctx.stream();
+
+    if (dst_type == GGML_TYPE_F32) {
+        rms_norm_mul_rope_cuda((const float *) x->data, (float *) dst_d,
+            x->ne[0], x->ne[1], x->ne[2], x->ne[3], s01, s02, s03, s1, s2, s3, eps,
+            (const float *) mul_src->data, mul_s01, mul_s02, mul_s03,
+            mul_src->ne[0], mul_src->ne[1], mul_src->ne[2], mul_src->ne[3],
+            n_dims, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims,
+            freq_factors, row_indices, set_rows_stride, is_neox, stream);
+    } else if (dst_type == GGML_TYPE_F16) {
+        rms_norm_mul_rope_cuda((const float *) x->data, (half *) dst_d,
+            x->ne[0], x->ne[1], x->ne[2], x->ne[3], s01, s02, s03, s1, s2, s3, eps,
+            (const float *) mul_src->data, mul_s01, mul_s02, mul_s03,
+            mul_src->ne[0], mul_src->ne[1], mul_src->ne[2], mul_src->ne[3],
+            n_dims, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims,
+            freq_factors, row_indices, set_rows_stride, is_neox, stream);
+    } else {
+        GGML_ABORT("fatal error");
+    }
+}

ggml/src/ggml-cuda/rope.cuh

@@ -7,3 +7,5 @@ void ggml_cuda_op_rope(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
 void ggml_cuda_op_rope_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
 
 void ggml_cuda_op_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * set_rows);
+
+void ggml_cuda_op_rms_norm_mul_rope_fused(ggml_backend_cuda_context & ctx, ggml_tensor * rms_norm, ggml_tensor * mul, ggml_tensor * rope, ggml_tensor * set_rows);

ggml/src/ggml-webgpu/ggml-webgpu-shader-lib.hpp

@@ -3221,17 +3221,17 @@ class ggml_webgpu_shader_lib {
         auto push_type_defines = [&](const char * prefix, ggml_type type) {
             std::string s_prefix = prefix;
             if (type == GGML_TYPE_F32) {
-                defines.push_back(s_prefix + "_F32");
+                defines.push_back(s_prefix + "=f32");
             } else if (type == GGML_TYPE_F16) {
-                defines.push_back(s_prefix + "_F16");
+                defines.push_back(s_prefix + "=f16");
             } else {
                 GGML_ABORT("Unsupported type for CONV_2D shader");
             }
         };
 
-        push_type_defines("WEIGHT", key.weight_type);
-        push_type_defines("INPUT", key.input_type);
-        push_type_defines("OUTPUT", key.output_type);
+        push_type_defines("WEIGHT_TYPE", key.weight_type);
+        push_type_defines("INPUT_TYPE", key.input_type);
+        push_type_defines("OUTPUT_TYPE", key.output_type);
 
         defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
 
@@ -3263,17 +3263,18 @@ class ggml_webgpu_shader_lib {
         auto push_type_defines = [&](const char * prefix, ggml_type type) {
             std::string s_prefix = prefix;
             if (type == GGML_TYPE_F32) {
-                defines.push_back(s_prefix + "_F32");
+                defines.push_back(s_prefix + "=f32");
             } else if (type == GGML_TYPE_F16) {
-                defines.push_back(s_prefix + "_F16");
+                defines.push_back(s_prefix + "=f16");
             } else {
-                GGML_ABORT("Unsupported type for CONV_2D_DW shader");
+                GGML_ABORT("Unsupported type for CONV_2D shader");
             }
         };
 
-        push_type_defines("WEIGHT", key.weight_type);
-        push_type_defines("INPUT", key.input_type);
-        push_type_defines("OUTPUT", key.output_type);
+        push_type_defines("WEIGHT_TYPE", key.weight_type);
+        push_type_defines("INPUT_TYPE", key.input_type);
+        push_type_defines("OUTPUT_TYPE", key.output_type);
+
         if (whcn) {
             defines.push_back("WHCN");
         }
@@ -3304,16 +3305,16 @@ class ggml_webgpu_shader_lib {
         auto push_type_defines = [&](const char * prefix, ggml_type type) {
             std::string s_prefix = prefix;
             if (type == GGML_TYPE_F32) {
-                defines.push_back(s_prefix + "_F32");
+                defines.push_back(s_prefix + "=f32");
             } else if (type == GGML_TYPE_F16) {
-                defines.push_back(s_prefix + "_F16");
+                defines.push_back(s_prefix + "=f16");
             } else {
                 GGML_ABORT("Unsupported type for IM2COL shader");
             }
         };
 
-        push_type_defines("INPUT", key.input_type);
-        push_type_defines("OUTPUT", key.output_type);
+        push_type_defines("INPUT_TYPE", key.input_type);
+        push_type_defines("OUTPUT_TYPE", key.output_type);
 
         defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
 

ggml/src/ggml-webgpu/ggml-webgpu.cpp

@@ -930,7 +930,6 @@ static webgpu_encoded_op ggml_webgpu_solve_tri(webgpu_context & ctx,
 
         (uint32_t) src1->ne[0],
         (uint32_t) dst->ne[2],
-        (uint32_t) dst->ne[3],
     };
 
     std::vector<wgpu::BindGroupEntry> entries = {
@@ -1039,7 +1038,6 @@ static webgpu_encoded_op ggml_webgpu_conv_2d_dw(webgpu_context & ctx,
 
         (uint32_t) ggml_nelements(dst),
         (uint32_t) dst->ne[2],
-        (uint32_t) dst->ne[3],
         (uint32_t) dst->ne[0],
         (uint32_t) dst->ne[1],
         (uint32_t) src1->ne[0],
@@ -1328,7 +1326,6 @@ static webgpu_encoded_op ggml_webgpu_ssm_scan(webgpu_context & ctx,
         (uint32_t) src0->ne[2],
         (uint32_t) src4->ne[1],
         (uint32_t) src1->ne[2],
-        (uint32_t) src1->ne[3],
         (uint32_t) ggml_nelements(src1),
     };
 
@@ -1921,25 +1918,20 @@ static bool ggml_webgpu_flash_attn_use_vec_path(const webgpu_global_context & gl
                                                 const ggml_tensor *           K,
                                                 const ggml_tensor *           V) {
     const size_t storage_offset_alignment = global_ctx->capabilities.limits.minStorageBufferOffsetAlignment;
-    const bool   k_float_vec4_aligned     = (K->type != GGML_TYPE_F16 && K->type != GGML_TYPE_F32) ||
-                                            ggml_webgpu_flash_attn_float_vec4_aligned(K, storage_offset_alignment);
-    const bool   v_float_vec4_aligned     = (V->type != GGML_TYPE_F16 && V->type != GGML_TYPE_F32) ||
-                                            ggml_webgpu_flash_attn_float_vec4_aligned(V, storage_offset_alignment);
-    const bool   k_vec_type_supported =
-        K->type == GGML_TYPE_F32 || K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q4_0 || K->type == GGML_TYPE_Q8_0;
-    const bool v_vec_type_supported =
-        V->type == GGML_TYPE_F32 || V->type == GGML_TYPE_F16 || V->type == GGML_TYPE_Q4_0 || V->type == GGML_TYPE_Q8_0;
-    const uint32_t k_vec_head_align         = (K->type == GGML_TYPE_F32 || K->type == GGML_TYPE_F16) ?
-                                                  GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH :
-                                                  (uint32_t) ggml_blck_size(K->type);
-    const uint32_t v_vec_head_align         = (V->type == GGML_TYPE_F32 || V->type == GGML_TYPE_F16) ?
-                                                  GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH :
-                                                  (uint32_t) ggml_blck_size(V->type);
-    const bool     kv_vec_head_dims_aligned = Q->ne[0] % k_vec_head_align == 0 && V->ne[0] % v_vec_head_align == 0;
+
+    const bool k_float_vec4_aligned = (K->type != GGML_TYPE_F16 && K->type != GGML_TYPE_F32) ||
+                                      ggml_webgpu_flash_attn_float_vec4_aligned(K, storage_offset_alignment);
+    const bool v_float_vec4_aligned = (V->type != GGML_TYPE_F16 && V->type != GGML_TYPE_F32) ||
+                                      ggml_webgpu_flash_attn_float_vec4_aligned(V, storage_offset_alignment);
+
+    const uint32_t k_vec_head_align =
+        ggml_is_quantized(K->type) ? ggml_blck_size(K->type) : GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH;
+    const uint32_t v_vec_head_align =
+        ggml_is_quantized(V->type) ? ggml_blck_size(V->type) : GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH;
+    const bool kv_vec_head_dims_aligned = Q->ne[0] % k_vec_head_align == 0 && V->ne[0] % v_vec_head_align == 0;
 
     return global_ctx->capabilities.supports_subgroups && (Q->ne[1] < GGML_WEBGPU_FLASH_ATTN_VEC_MAX_SEQ_LEN) &&
-           kv_vec_head_dims_aligned && k_vec_type_supported && v_vec_type_supported && k_float_vec4_aligned &&
-           v_float_vec4_aligned;
+           kv_vec_head_dims_aligned && k_float_vec4_aligned && v_float_vec4_aligned;
 }
 
 static ggml_webgpu_flash_attn_op ggml_webgpu_flash_attn_prepare(webgpu_context & ctx,
@@ -2514,7 +2506,6 @@ static webgpu_encoded_op ggml_webgpu_concat(webgpu_context & ctx,
                                      (uint32_t) dst->ne[0],
                                      (uint32_t) dst->ne[1],
                                      (uint32_t) dst->ne[2],
-                                     (uint32_t) dst->ne[3],
                                      dim,
                                      (uint32_t) src0->ne[dim] };
 
@@ -2610,7 +2601,6 @@ static std::optional<webgpu_encoded_op> ggml_webgpu_rms_norm_mul(webgpu_context
         (uint32_t) dst->ne[0],
         (uint32_t) dst->ne[1],
         (uint32_t) dst->ne[2],
-        (uint32_t) dst->ne[3],
         ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(rn_dst, 0))  // epsilon, treated as f32 in the shader
     };
 
@@ -2666,7 +2656,6 @@ static webgpu_encoded_op ggml_webgpu_row_norm(webgpu_context & ctx, ggml_tensor
         (uint32_t) src->ne[0],
         (uint32_t) src->ne[1],
         (uint32_t) src->ne[2],
-        (uint32_t) src->ne[3],
         ggml_webgpu_u32_from_f32(ggml_get_op_params_f32(dst, 0))  // epsilon, treated as f32 in the shader
     };
 
@@ -2925,7 +2914,6 @@ static webgpu_encoded_op ggml_webgpu_soft_max(webgpu_context & ctx,
         (uint32_t) (dst->nb[1] / ggml_type_size(dst->type)),
         (uint32_t) (dst->nb[2] / ggml_type_size(dst->type)),
         (uint32_t) (dst->nb[3] / ggml_type_size(dst->type)),
-        (uint32_t) ggml_nelements(dst),
         (uint32_t) src0->ne[0],
         (uint32_t) src0->ne[1],
         (uint32_t) src0->ne[2],

ggml/src/ggml-webgpu/wgsl-shaders/concat.wgsl

@@ -18,7 +18,6 @@ struct Params {
     ne0: u32,
     ne1: u32,
     ne2: u32,
-    ne3: u32,
 
     dim: u32,
     src0_nedim: u32

ggml/src/ggml-webgpu/wgsl-shaders/conv2d.wgsl

@@ -2,25 +2,11 @@
 enable f16;
 
 @group(0) @binding(0)
-#if defined(WEIGHT_F32)
-var<storage, read_write> weights: array<f32>;
-#elif defined(WEIGHT_F16)
-var<storage, read_write> weights: array<f16>;
-#endif
-
+var<storage, read_write> weights: array<WEIGHT_TYPE>;
 @group(0) @binding(1)
-#if defined(INPUT_F32)
-var<storage, read_write> input: array<f32>;
-#elif defined(INPUT_F16)
-var<storage, read_write> input: array<f16>;
-#endif
-
+var<storage, read_write> input: array<INPUT_TYPE>;
 @group(0) @binding(2)
-#if defined(OUTPUT_F32)
-var<storage, read_write> output: array<f32>;
-#elif defined(OUTPUT_F16)
-var<storage, read_write> output: array<f16>;
-#endif
+var<storage, read_write> output: array<OUTPUT_TYPE>;
 
 struct Params {
     offset_w: u32,
@@ -50,30 +36,6 @@ struct Params {
 @group(0) @binding(3)
 var<uniform> params: Params;
 
-fn load_weight(idx: u32) -> f32 {
-    #if defined(WEIGHT_F32)
-        return weights[idx];
-    #elif defined(WEIGHT_F16)
-        return f32(weights[idx]);
-    #endif
-}
-
-fn load_input(idx: u32) -> f32 {
-    #if defined(INPUT_F32)
-        return input[idx];
-    #elif defined(INPUT_F16)
-        return f32(input[idx]);
-    #endif
-}
-
-fn store_output(idx: u32, val: f32) {
-    #if defined(OUTPUT_F32)
-        output[idx] = val;
-    #elif defined(OUTPUT_F16)
-        output[idx] = f16(val);
-    #endif
-}
-
 fn ceil_div_u32(x: u32, y: u32) -> u32 {
     return (x + y - 1) / y;
 }
@@ -136,7 +98,7 @@ fn main(
     // entire receptive field is out of bounds
     if (kw_begin >= kw_end || kh_begin >= kh_end) {
         let out_idx = params.offset_o + ow * params.so0 + oh * params.so1 + oc * params.so2 + n * params.so3;
-        store_output(out_idx, 0.0);
+        output[out_idx] = OUTPUT_TYPE(0.0);
         return;
     }
 
@@ -155,11 +117,11 @@ fn main(
                 let iw = u32(ow_base + i32(kw * params.d0));
                 let w_idx = w_row_base + kw * params.sw0;
                 let in_idx = in_row_base + iw * params.si0;
-                sum += load_weight(w_idx) * load_input(in_idx);
+                sum += f32(weights[w_idx]) * f32(input[in_idx]);
             }
         }
     }
 
     let out_idx = params.offset_o + ow * params.so0 + oh * params.so1 + oc * params.so2 + n * params.so3;
-    store_output(out_idx, sum);
+    output[out_idx] = OUTPUT_TYPE(sum);
 }

ggml/src/ggml-webgpu/wgsl-shaders/conv2d_dw.wgsl

@@ -6,25 +6,11 @@ enable f16;
 // weight (src0) is [KW,KH,1,C]; output matches the input layout.
 
 @group(0) @binding(0)
-#if defined(WEIGHT_F32)
-var<storage, read_write> weights: array<f32>;
-#elif defined(WEIGHT_F16)
-var<storage, read_write> weights: array<f16>;
-#endif
-
+var<storage, read_write> weights: array<WEIGHT_TYPE>;
 @group(0) @binding(1)
-#if defined(INPUT_F32)
-var<storage, read_write> input: array<f32>;
-#elif defined(INPUT_F16)
-var<storage, read_write> input: array<f16>;
-#endif
-
+var<storage, read_write> input: array<INPUT_TYPE>;
 @group(0) @binding(2)
-#if defined(OUTPUT_F32)
-var<storage, read_write> output: array<f32>;
-#elif defined(OUTPUT_F16)
-var<storage, read_write> output: array<f16>;
-#endif
+var<storage, read_write> output: array<OUTPUT_TYPE>;
 
 struct Params {
     offset_w: u32,
@@ -33,7 +19,6 @@ struct Params {
 
     ne: u32,
     channels: u32,
-    batches: u32,
     dst_w: u32, dst_h: u32,
     src_w: u32, src_h: u32,
     knl_w: u32, knl_h: u32,
@@ -46,28 +31,6 @@ struct Params {
 @group(0) @binding(3)
 var<uniform> params: Params;
 
-fn load_weight(idx: u32) -> f32 {
-    #if defined(WEIGHT_F32)
-        return weights[idx];
-    #elif defined(WEIGHT_F16)
-        return f32(weights[idx]);
-    #endif
-}
-fn load_input(idx: u32) -> f32 {
-    #if defined(INPUT_F32)
-        return input[idx];
-    #elif defined(INPUT_F16)
-        return f32(input[idx]);
-    #endif
-}
-fn store_output(idx: u32, val: f32) {
-    #if defined(OUTPUT_F32)
-        output[idx] = val;
-    #elif defined(OUTPUT_F16)
-        output[idx] = f16(val);
-    #endif
-}
-
 #if defined(WHCN)
 // Input/output/kernel contiguous in [W, H, C, N] order (kernel [KW,KH,C]).
 fn conv_2d_dw(idx: u32) -> f32 {
@@ -89,8 +52,8 @@ fn conv_2d_dw(idx: u32) -> f32 {
         for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) {
             let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x;
             if (src_x < 0 || src_x >= i32(params.src_w)) { continue; }
-            let v = load_input(src_i + u32(src_y) * params.src_w + u32(src_x));
-            let k = load_weight(knl_i + ky * params.knl_w + kx);
+            let v = f32(input[src_i + u32(src_y) * params.src_w + u32(src_x)]);
+            let k = f32(weights[knl_i + ky * params.knl_w + kx]);
             sum += v * k;
         }
     }
@@ -117,8 +80,8 @@ fn conv_2d_dw(idx: u32) -> f32 {
         for (var kx: u32 = 0u; kx < params.knl_w; kx += 1u) {
             let src_x = i32(dst_x) * params.stride_x + i32(kx) * params.dilation_x - params.pad_x;
             if (src_x < 0 || src_x >= i32(params.src_w)) { continue; }
-            let v = load_input(src_i + u32(src_y) * src_row + u32(src_x) * params.channels + c);
-            let k = load_weight(params.offset_w + ky * knl_row + kx * params.channels + c);
+            let v = f32(input[src_i + u32(src_y) * src_row + u32(src_x) * params.channels + c]);
+            let k = f32(weights[params.offset_w + ky * knl_row + kx * params.channels + c]);
             sum += v * k;
         }
     }
@@ -133,5 +96,5 @@ fn main(
 ) {
     let idx = gid.x + (num_wg.x * u32(WG_SIZE)) * gid.y;
     if (idx >= params.ne) { return; }
-    store_output(params.offset_o + idx, conv_2d_dw(idx));
+    output[params.offset_o + idx] = OUTPUT_TYPE(conv_2d_dw(idx));
 }

ggml/src/ggml-webgpu/wgsl-shaders/flash_attn.wgsl

@@ -7,32 +7,18 @@ enable chromium_experimental_subgroup_matrix;
 #define BYTE_HELPERS
 #include "common_decls.tmpl"
 
-#ifdef K_F32
-#define K_TYPE f32
-#elif defined(K_Q4_0) || defined(K_Q8_0)
-#define K_TYPE u32
-#else
-#define K_TYPE f16
-#endif
-
-#ifdef V_F32
-#define V_TYPE f32
-#elif defined(V_Q4_0) || defined(V_Q8_0)
-#define V_TYPE u32
-#else
-#define V_TYPE f16
-#endif
+#define FLASH_ATTN_SCALAR_KV
+#include "flash_attn_decls.tmpl"
 
 // Default values
+// The actual values are defined in shader-lib.
 #define HEAD_DIM_QK 64
 #define HEAD_DIM_V 64
-
 // The number of rows/columns/k in a subgroup matrix. MxK * KxN = MxN
 // Note that the "K" here does not correspond to the K in attention's Q/K/V, it's just the common dimension.
 #define SG_MAT_M 8
 #define SG_MAT_N 8
 #define SG_MAT_K 8
-
 // Each workgroup processes one subgroup matrix of Q rows
 #define Q_TILE SG_MAT_M
 #define KV_TILE 16
@@ -41,104 +27,13 @@ enable chromium_experimental_subgroup_matrix;
 // Number of subgroup-matrix-width blocks that span the KV tile. SG_MAT_N must divide KV_TILE.
 #define KV_BLOCKS (KV_TILE / SG_MAT_N)
 
-struct Params {
-    offset_q: u32,
-    offset_k: u32,
-    offset_v: u32,
-    offset_mask: u32,
-    offset_sinks: u32,
-    offset_dst: u32,
-
-    // shapes of Q/K/V
-    n_heads: u32,
-    seq_len_q: u32,
-    seq_len_kv: u32,
-
-    // strides (in elements)
-    stride_q1: u32,
-    stride_q2: u32,
-    stride_q3: u32,
-    stride_k1: u32,
-    stride_k2: u32,
-    stride_k3: u32,
-    stride_v1: u32,
-    stride_v2: u32,
-    stride_v3: u32,
-    stride_mask3: u32,
-
-    // repeat factors for K/V, e.g., MHA vs. MQA vs. GQA
-    q_per_kv: u32,
-
-    // softmax params
-    scale: f32,
-    max_bias: f32,
-    logit_softcap: f32,
-    n_head_log2: f32,
-    m0: f32,
-    m1: f32,
-};
-
-@group(0) @binding(0) var<storage, read_write> Q: array<f32>;
-#ifdef KV_OVERLAP
-@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
-#define V K
-#else
-@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
-@group(0) @binding(2) var<storage, read_write> V: array<V_TYPE>;
-#endif
-
-#if defined(MASK) && defined(SINKS)
-#ifdef KV_OVERLAP
-@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
-@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
-#define DST_BINDING 4
-#define PARAMS_BINDING 5
-#else
-@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
-@group(0) @binding(4) var<storage, read_write> sinks: array<f32>;
-#define DST_BINDING 5
-#define PARAMS_BINDING 6
-#endif
-#elif defined(MASK)
-#ifdef KV_OVERLAP
-@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
-#define DST_BINDING 3
-#define PARAMS_BINDING 4
-#else
-@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
-#define DST_BINDING 4
-#define PARAMS_BINDING 5
-#endif
-#elif defined(SINKS)
-#ifdef KV_OVERLAP
-@group(0) @binding(2) var<storage, read_write> sinks: array<f32>;
-#define DST_BINDING 3
-#define PARAMS_BINDING 4
-#else
-@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
-#define DST_BINDING 4
-#define PARAMS_BINDING 5
-#endif
-#else
-#ifdef KV_OVERLAP
-#define DST_BINDING 2
-#define PARAMS_BINDING 3
-#else
-#define DST_BINDING 3
-#define PARAMS_BINDING 4
-#endif
-#endif
-
-@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<vec4<f32>>;
-@group(0) @binding(PARAMS_BINDING) var<uniform> params: Params;
-
-// Just a very small float value.
-const FLOAT_MIN: f32 = -1.0e9;
-
 // The number of Q rows processed per workgroup
 var<workgroup> q_shmem: array<f16, Q_TILE * HEAD_DIM_QK>;
 
 #if !defined(K_DIRECT) || !defined(V_DIRECT)
+#define STAGING_SHMEM kv_shmem
+#define STAGING_OUT_TYPE f16
+#include "flash_attn_staging.tmpl"
 const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
 // we can reuse the same shmem for K and V since we only need one at a time
 var<workgroup> kv_shmem: array<f16, kv_shmem_size>;
@@ -175,50 +70,6 @@ fn calc_softmax_term(kv_idx: u32, q_tile_row: u32, slope: f32) -> f32 {
     return v;
 }
 
-fn load_f32x4(buf: ptr<storage, array<vec4<f32>>, read_write>, scalar_index: u32) -> vec4<f32> {
-    return (*buf)[scalar_index >> 2u];
-}
-
-fn load_kx4(buf: ptr<storage, array<vec4<K_TYPE>>, read_write>, scalar_index: u32) -> vec4<K_TYPE> {
-    return (*buf)[scalar_index >> 2u];
-}
-
-#if !defined(K_DIRECT) || !defined(V_DIRECT)
-#define QUANT_SHMEM kv_shmem
-#define QUANT_OUT_TYPE f16
-#include "flash_attn_quant_staging.tmpl"
-
-#if !defined(K_DIRECT) && !defined(K_Q4_0) && !defined(K_Q8_0)
-fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
-    for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE) {
-        let k_row = elem_idx / HEAD_DIM_QK;
-        let k_col = elem_idx % HEAD_DIM_QK;
-        let global_k_row = kv_tile + k_row;
-        let global_k_row_offset = k_head_offset + global_k_row * params.stride_k1;
-        kv_shmem[elem_idx] = f16(select(
-            0.0,
-            K[global_k_row_offset + k_col],
-            global_k_row < params.seq_len_kv && k_col < HEAD_DIM_QK));
-    }
-}
-#endif
-
-#if !defined(V_DIRECT) && !defined(V_Q4_0) && !defined(V_Q8_0)
-fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
-    for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE) {
-        let v_row = elem_idx / HEAD_DIM_V;
-        let v_col = elem_idx % HEAD_DIM_V;
-        let global_v_row = kv_tile + v_row;
-        let global_v_row_offset = v_head_offset + global_v_row * params.stride_v1;
-        kv_shmem[elem_idx] = f16(select(
-            0.0,
-            V[global_v_row_offset + v_col],
-            global_v_row < params.seq_len_kv && v_col < HEAD_DIM_V));
-    }
-}
-#endif
-#endif
-
 @compute @workgroup_size(WG_SIZE)
 fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
     @builtin(local_invocation_id) local_id: vec3<u32>,

ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_decls.tmpl

@@ -0,0 +1,134 @@
+#ifdef Q_F32
+#define Q_TYPE f32
+#else
+#define Q_TYPE f16
+#endif
+
+#ifdef K_F32
+#define K_TYPE f32
+#elif defined(K_Q4_0) || defined(K_Q8_0)
+#define K_TYPE u32
+#else
+#define K_TYPE f16
+#endif
+
+#ifdef V_F32
+#define V_TYPE f32
+#elif defined(V_Q4_0) || defined(V_Q8_0)
+#define V_TYPE u32
+#else
+#define V_TYPE f16
+#endif
+
+#ifdef DST_F32
+#define DST_TYPE f32
+#else
+#define DST_TYPE f16
+#endif
+
+#if defined(FLASH_ATTN_SCALAR_KV) || defined(K_Q4_0) || defined(K_Q8_0)
+#define K_STORAGE_TYPE K_TYPE
+#else
+#define K_STORAGE_TYPE vec4<K_TYPE>
+#endif
+
+#if defined(FLASH_ATTN_SCALAR_KV) || defined(V_Q4_0) || defined(V_Q8_0)
+#define V_STORAGE_TYPE V_TYPE
+#else
+#define V_STORAGE_TYPE vec4<V_TYPE>
+#endif
+
+// Just a very small float value.
+const FLOAT_MIN: f32 = -1.0e9;
+
+struct Params {
+    offset_q: u32,
+    offset_k: u32,
+    offset_v: u32,
+    offset_mask: u32,
+    offset_sinks: u32,
+    offset_dst: u32,
+
+    // shapes of Q/K/V
+    n_heads: u32,
+    seq_len_q: u32,
+    seq_len_kv: u32,
+
+    // strides (in elements)
+    stride_q1: u32,
+    stride_q2: u32,
+    stride_q3: u32,
+    stride_k1: u32,
+    stride_k2: u32,
+    stride_k3: u32,
+    stride_v1: u32,
+    stride_v2: u32,
+    stride_v3: u32,
+    stride_mask3: u32,
+
+    // repeat factors for K/V, e.g., MHA vs. MQA vs. GQA
+    q_per_kv: u32,
+
+    // softmax params
+    scale: f32,
+    max_bias: f32,
+    logit_softcap: f32,
+    n_head_log2: f32,
+    m0: f32,
+    m1: f32,
+
+#ifdef FLASH_ATTN_VEC_SPLIT
+#ifdef BLK
+    blk_base: u32,
+    blk_nblk0: u32,
+    blk_nblk1: u32,
+#endif
+
+    tmp_data_base: u32,
+    tmp_stats_base: u32,
+    nwg: u32,
+#endif
+};
+
+@group(0) @binding(0) var<storage, read_write> Q: array<Q_TYPE>;
+@group(0) @binding(1) var<storage, read_write> K: array<K_STORAGE_TYPE>;
+#ifdef KV_OVERLAP
+#define V K
+#define MASK_BINDING 2
+#else
+@group(0) @binding(2) var<storage, read_write> V: array<V_STORAGE_TYPE>;
+#define MASK_BINDING 3
+#endif // KV_OVERLAP
+
+#ifdef MASK
+@group(0) @binding(MASK_BINDING) var<storage, read_write> mask: array<f16>;
+#define SINKS_BINDING (MASK_BINDING + 1)
+#else
+#define SINKS_BINDING MASK_BINDING
+#endif
+
+#ifdef SINKS
+@group(0) @binding(SINKS_BINDING) var<storage, read_write> sinks: array<f32>;
+#define BLK_BINDING (SINKS_BINDING + 1)
+#else
+#define BLK_BINDING SINKS_BINDING
+#endif
+
+#ifdef FLASH_ATTN_VEC_SPLIT
+#ifdef BLK
+@group(0) @binding(BLK_BINDING) var<storage, read_write> blk: array<u32>;
+#define TMP_BINDING (BLK_BINDING + 1)
+#else
+#define TMP_BINDING BLK_BINDING
+#endif
+
+@group(0) @binding(TMP_BINDING) var<storage, read_write> tmp: array<f32>;
+#define DST_BINDING (TMP_BINDING + 1)
+#else
+#define DST_BINDING BLK_BINDING
+#endif // FLASH_ATTN_VEC_SPLIT
+
+@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<vec4<DST_TYPE>>;
+
+#define PARAMS_BINDING (DST_BINDING + 1)
+@group(0) @binding(PARAMS_BINDING) var<uniform> params: Params;

ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_quant_staging.tmpl

@@ -1,83 +0,0 @@
-#include "quant_inner_loops.tmpl"
-
-#define BLOCK_SIZE 32
-#define BLOCKS_K ((HEAD_DIM_QK + BLOCK_SIZE - 1) / BLOCK_SIZE)
-#define BLOCKS_V ((HEAD_DIM_V + BLOCK_SIZE - 1) / BLOCK_SIZE)
-
-#if defined(K_Q4_0)
-#define K_NQ 16
-#define K_BLOCK_SIZE_BYTES 18u
-#define K_BYTES_PER_THREAD 8u
-#define K_BYTES_PER_INNER_LOOP 4u
-#elif defined(K_Q8_0)
-#define K_NQ 16
-#define K_BLOCK_SIZE_BYTES 34u
-#define K_BYTES_PER_THREAD 16u
-#define K_BYTES_PER_INNER_LOOP 4u
-#endif
-
-#if defined(V_Q4_0)
-#define V_NQ 16
-#define V_BLOCK_SIZE_BYTES 18u
-#define V_BYTES_PER_THREAD 8u
-#define V_BYTES_PER_INNER_LOOP 4u
-#elif defined(V_Q8_0)
-#define V_NQ 16
-#define V_BLOCK_SIZE_BYTES 34u
-#define V_BYTES_PER_THREAD 16u
-#define V_BYTES_PER_INNER_LOOP 4u
-#endif
-
-#if defined(K_Q4_0) || defined(K_Q8_0)
-fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
-    for (var elem_idx = local_x * K_NQ; elem_idx < kv_count * HEAD_DIM_QK; elem_idx += WG_SIZE * K_NQ) {
-        let blck_idx = elem_idx / BLOCK_SIZE;
-        let block_offset = (elem_idx % BLOCK_SIZE) / K_NQ;
-        let k_row = blck_idx / BLOCKS_K;
-        let global_k_row = kv_tile + k_row;
-        let block_k = blck_idx % BLOCKS_K;
-        let row_offset = k_row * HEAD_DIM_QK;
-        let global_block_idx = k_head_offset + global_k_row * params.stride_k1 + block_k;
-        let block_byte_base = global_block_idx * K_BLOCK_SIZE_BYTES;
-        let d = f16_from_u16(load_k_u16_at(block_byte_base));
-        let thread_byte_offset = block_offset * K_BYTES_PER_THREAD;
-        let shmem_idx = row_offset + block_k * BLOCK_SIZE + thread_byte_offset;
-        for (var j = 0u; j < K_BYTES_PER_THREAD / K_BYTES_PER_INNER_LOOP; j += 1u) {
-            let q_byte_offset = block_byte_base + 2u + thread_byte_offset + j * K_BYTES_PER_INNER_LOOP;
-            let q_packed = load_k_u32_at(q_byte_offset);
-#if defined(K_Q4_0)
-            dequant_q4_0_packed_to_shmem(q_packed, d, shmem_idx + j * K_BYTES_PER_INNER_LOOP);
-#elif defined(K_Q8_0)
-            dequant_q8_0_packed_to_shmem(q_packed, d, shmem_idx + j * K_BYTES_PER_INNER_LOOP);
-#endif
-        }
-    }
-}
-#endif
-
-#if defined(V_Q4_0) || defined(V_Q8_0)
-fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
-    for (var elem_idx = local_x * V_NQ; elem_idx < kv_count * HEAD_DIM_V; elem_idx += WG_SIZE * V_NQ) {
-        let blck_idx = elem_idx / BLOCK_SIZE;
-        let block_offset = (elem_idx % BLOCK_SIZE) / V_NQ;
-        let v_row = blck_idx / BLOCKS_V;
-        let global_v_row = kv_tile + v_row;
-        let block_k = blck_idx % BLOCKS_V;
-        let row_offset = v_row * HEAD_DIM_V;
-        let global_block_idx = v_head_offset + global_v_row * params.stride_v1 + block_k;
-        let block_byte_base = global_block_idx * V_BLOCK_SIZE_BYTES;
-        let d = f16_from_u16(load_v_u16_at(block_byte_base));
-        let thread_byte_offset = block_offset * V_BYTES_PER_THREAD;
-        let shmem_idx = row_offset + block_k * BLOCK_SIZE + thread_byte_offset;
-        for (var j = 0u; j < V_BYTES_PER_THREAD / V_BYTES_PER_INNER_LOOP; j += 1u) {
-            let q_byte_offset = block_byte_base + 2u + thread_byte_offset + j * V_BYTES_PER_INNER_LOOP;
-            let q_packed = load_v_u32_at(q_byte_offset);
-#if defined(V_Q4_0)
-            dequant_q4_0_packed_to_shmem(q_packed, d, shmem_idx + j * V_BYTES_PER_INNER_LOOP);
-#elif defined(V_Q8_0)
-            dequant_q8_0_packed_to_shmem(q_packed, d, shmem_idx + j * V_BYTES_PER_INNER_LOOP);
-#endif
-        }
-    }
-}
-#endif

ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_staging.tmpl

@@ -0,0 +1,136 @@
+#if defined(K_Q4_0) || defined(K_Q8_0) || defined(V_Q4_0) || defined(V_Q8_0)
+#define QUANT_SHMEM STAGING_SHMEM
+#define QUANT_OUT_TYPE STAGING_OUT_TYPE
+#include "quant_inner_loops.tmpl"
+#undef QUANT_SHMEM
+#undef QUANT_OUT_TYPE
+#define BLOCK_SIZE 32
+#define BLOCKS_K ((HEAD_DIM_QK + BLOCK_SIZE - 1) / BLOCK_SIZE)
+#define BLOCKS_V ((HEAD_DIM_V + BLOCK_SIZE - 1) / BLOCK_SIZE)
+#endif
+
+#if defined(K_Q4_0)
+#define K_NQ 16
+#define K_BLOCK_SIZE_BYTES 18u
+#define K_BYTES_PER_THREAD 8u
+#define K_BYTES_PER_INNER_LOOP 4u
+#define DEQUANT_K_PACKED_TO_SHMEM dequant_q4_0_packed_to_shmem
+#elif defined(K_Q8_0)
+#define K_NQ 16
+#define K_BLOCK_SIZE_BYTES 34u
+#define K_BYTES_PER_THREAD 16u
+#define K_BYTES_PER_INNER_LOOP 4u
+#define DEQUANT_K_PACKED_TO_SHMEM dequant_q8_0_packed_to_shmem
+#endif
+
+#if defined(V_Q4_0)
+#define V_NQ 16
+#define V_BLOCK_SIZE_BYTES 18u
+#define V_BYTES_PER_THREAD 8u
+#define V_BYTES_PER_INNER_LOOP 4u
+#define DEQUANT_V_PACKED_TO_SHMEM dequant_q4_0_packed_to_shmem
+#elif defined(V_Q8_0)
+#define V_NQ 16
+#define V_BLOCK_SIZE_BYTES 34u
+#define V_BYTES_PER_THREAD 16u
+#define V_BYTES_PER_INNER_LOOP 4u
+#define DEQUANT_V_PACKED_TO_SHMEM dequant_q8_0_packed_to_shmem
+#endif
+
+#ifndef K_DIRECT
+fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
+#if defined(K_Q4_0) || defined(K_Q8_0)
+    for (var elem_idx = local_x * K_NQ; elem_idx < kv_count * HEAD_DIM_QK; elem_idx += WG_SIZE * K_NQ) {
+        let blck_idx = elem_idx / BLOCK_SIZE;
+        let block_offset = (elem_idx % BLOCK_SIZE) / K_NQ;
+        let k_row = blck_idx / BLOCKS_K;
+        let global_k_row = kv_tile + k_row;
+        let block_k = blck_idx % BLOCKS_K;
+        let row_offset = k_row * HEAD_DIM_QK;
+        let global_block_idx = k_head_offset + global_k_row * params.stride_k1 + block_k;
+        let block_byte_base = global_block_idx * K_BLOCK_SIZE_BYTES;
+        let d = f16_from_u16(load_k_u16_at(block_byte_base));
+        let thread_byte_offset = block_offset * K_BYTES_PER_THREAD;
+        let shmem_idx = row_offset + block_k * BLOCK_SIZE + thread_byte_offset;
+        for (var j = 0u; j < K_BYTES_PER_THREAD / K_BYTES_PER_INNER_LOOP; j += 1u) {
+            let q_byte_offset = block_byte_base + 2u + thread_byte_offset + j * K_BYTES_PER_INNER_LOOP;
+            let q_packed = load_k_u32_at(q_byte_offset);
+            DEQUANT_K_PACKED_TO_SHMEM(q_packed, d, shmem_idx + j * K_BYTES_PER_INNER_LOOP);
+        }
+    }
+#elif defined(FLASH_ATTN_SCALAR_KV)
+    for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE) {
+        let k_row = elem_idx / HEAD_DIM_QK;
+        let k_col = elem_idx % HEAD_DIM_QK;
+        let global_k_row = kv_tile + k_row;
+        let global_k_row_offset = k_head_offset + global_k_row * params.stride_k1;
+        STAGING_SHMEM[elem_idx] = STAGING_OUT_TYPE(select(
+            0.0,
+            K[global_k_row_offset + k_col],
+            global_k_row < params.seq_len_kv && k_col < HEAD_DIM_QK));
+    }
+#else
+    for (var vec_idx_local = local_x; vec_idx_local < kv_count * Q_CHUNKS; vec_idx_local += WG_SIZE) {
+        let kv_local = vec_idx_local / Q_CHUNKS;
+        let chunk = vec_idx_local % Q_CHUNKS;
+        let global_k_row = kv_tile + kv_local;
+        let k_vec_index = (k_head_offset + global_k_row * params.stride_k1 + chunk * 4u) >> 2u;
+        let k4 = K[k_vec_index];
+        let kv_off = kv_local * HEAD_DIM_QK + chunk * 4u;
+        STAGING_SHMEM[kv_off + 0u] = STAGING_OUT_TYPE(k4.x);
+        STAGING_SHMEM[kv_off + 1u] = STAGING_OUT_TYPE(k4.y);
+        STAGING_SHMEM[kv_off + 2u] = STAGING_OUT_TYPE(k4.z);
+        STAGING_SHMEM[kv_off + 3u] = STAGING_OUT_TYPE(k4.w);
+    }
+#endif
+}
+#endif // !defined(K_DIRECT)
+
+#ifndef V_DIRECT
+fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
+#if defined(V_Q4_0) || defined(V_Q8_0)
+    for (var elem_idx = local_x * V_NQ; elem_idx < kv_count * HEAD_DIM_V; elem_idx += WG_SIZE * V_NQ) {
+        let blck_idx = elem_idx / BLOCK_SIZE;
+        let block_offset = (elem_idx % BLOCK_SIZE) / V_NQ;
+        let v_row = blck_idx / BLOCKS_V;
+        let global_v_row = kv_tile + v_row;
+        let block_k = blck_idx % BLOCKS_V;
+        let row_offset = v_row * HEAD_DIM_V;
+        let global_block_idx = v_head_offset + global_v_row * params.stride_v1 + block_k;
+        let block_byte_base = global_block_idx * V_BLOCK_SIZE_BYTES;
+        let d = f16_from_u16(load_v_u16_at(block_byte_base));
+        let thread_byte_offset = block_offset * V_BYTES_PER_THREAD;
+        let shmem_idx = row_offset + block_k * BLOCK_SIZE + thread_byte_offset;
+        for (var j = 0u; j < V_BYTES_PER_THREAD / V_BYTES_PER_INNER_LOOP; j += 1u) {
+            let q_byte_offset = block_byte_base + 2u + thread_byte_offset + j * V_BYTES_PER_INNER_LOOP;
+            let q_packed = load_v_u32_at(q_byte_offset);
+            DEQUANT_V_PACKED_TO_SHMEM(q_packed, d, shmem_idx + j * V_BYTES_PER_INNER_LOOP);
+        }
+    }
+#elif defined(FLASH_ATTN_SCALAR_KV)
+    for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE) {
+        let v_row = elem_idx / HEAD_DIM_V;
+        let v_col = elem_idx % HEAD_DIM_V;
+        let global_v_row = kv_tile + v_row;
+        let global_v_row_offset = v_head_offset + global_v_row * params.stride_v1;
+        STAGING_SHMEM[elem_idx] = STAGING_OUT_TYPE(select(
+            0.0,
+            V[global_v_row_offset + v_col],
+            global_v_row < params.seq_len_kv && v_col < HEAD_DIM_V));
+    }
+#else
+    for (var vec_idx_local = local_x; vec_idx_local < kv_count * V_CHUNKS; vec_idx_local += WG_SIZE) {
+        let kv_local = vec_idx_local / V_CHUNKS;
+        let chunk = vec_idx_local % V_CHUNKS;
+        let global_v_row = kv_tile + kv_local;
+        let v_vec_index = (v_head_offset + global_v_row * params.stride_v1 + chunk * 4u) >> 2u;
+        let v4 = V[v_vec_index];
+        let kv_off = kv_local * HEAD_DIM_V + chunk * 4u;
+        STAGING_SHMEM[kv_off + 0u] = STAGING_OUT_TYPE(v4.x);
+        STAGING_SHMEM[kv_off + 1u] = STAGING_OUT_TYPE(v4.y);
+        STAGING_SHMEM[kv_off + 2u] = STAGING_OUT_TYPE(v4.z);
+        STAGING_SHMEM[kv_off + 3u] = STAGING_OUT_TYPE(v4.w);
+    }
+#endif
+}
+#endif // !defined(V_DIRECT)

ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_tile.wgsl

@@ -3,192 +3,32 @@ enable subgroups;
 
 #define BYTE_HELPERS
 #include "common_decls.tmpl"
+#include "flash_attn_decls.tmpl"
 
-#ifdef Q_F16
-#define Q_TYPE f16
-#else
-#define Q_TYPE f32
-#endif
-
-#ifdef K_F32
-#define K_TYPE f32
-#elif defined(K_Q4_0) || defined(K_Q8_0)
-#define K_TYPE u32
-#else
-#define K_TYPE f16
-#endif
-
-#ifdef V_F32
-#define V_TYPE f32
-#elif defined(V_Q4_0) || defined(V_Q8_0)
-#define V_TYPE u32
-#else
-#define V_TYPE f16
-#endif
-
-#ifdef DST_F16
-#define DST_TYPE f16
-#else
-#define DST_TYPE f32
-#endif
-
+// Default values
+// The actual values are defined in shader-lib.
 #define HEAD_DIM_QK 64
 #define HEAD_DIM_V 64
 #define Q_TILE 4
 #define KV_TILE 64
 #define WG_SIZE 128
-#ifndef MIN_SUBGROUP_SIZE
-#define MIN_SUBGROUP_SIZE MAX_SUBGROUP_SIZE
-#endif
 
-struct Params {
-    offset_q: u32,
-    offset_k: u32,
-    offset_v: u32,
-    offset_mask: u32,
-    offset_sinks: u32,
-    offset_dst: u32,
-
-    n_heads: u32,
-    seq_len_q: u32,
-    seq_len_kv: u32,
-
-    stride_q1: u32,
-    stride_q2: u32,
-    stride_q3: u32,
-    stride_k1: u32,
-    stride_k2: u32,
-    stride_k3: u32,
-    stride_v1: u32,
-    stride_v2: u32,
-    stride_v3: u32,
-    stride_mask3: u32,
-
-    q_per_kv: u32,
-
-    scale: f32,
-    max_bias: f32,
-    logit_softcap: f32,
-    n_head_log2: f32,
-    m0: f32,
-    m1: f32,
-};
-
-@group(0) @binding(0) var<storage, read_write> Q: array<Q_TYPE>;
-#ifdef KV_OVERLAP
-#if defined(K_Q4_0) || defined(K_Q8_0)
-@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
-#else
-@group(0) @binding(1) var<storage, read_write> K: array<vec4<K_TYPE>>;
-#endif
-#define V K
-#else
-#if defined(K_Q4_0) || defined(K_Q8_0)
-@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
-#else
-@group(0) @binding(1) var<storage, read_write> K: array<vec4<K_TYPE>>;
-#endif
-#if defined(V_Q4_0) || defined(V_Q8_0)
-@group(0) @binding(2) var<storage, read_write> V: array<V_TYPE>;
-#else
-@group(0) @binding(2) var<storage, read_write> V: array<vec4<V_TYPE>>;
-#endif
-#endif
-
-#if defined(MASK) && defined(SINKS)
-#ifdef KV_OVERLAP
-@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
-@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
-#define DST_BINDING 4
-#define PARAMS_BINDING 5
-#else
-@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
-@group(0) @binding(4) var<storage, read_write> sinks: array<f32>;
-#define DST_BINDING 5
-#define PARAMS_BINDING 6
-#endif
-#elif defined(MASK)
-#ifdef KV_OVERLAP
-@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
-#define DST_BINDING 3
-#define PARAMS_BINDING 4
-#else
-@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
-#define DST_BINDING 4
-#define PARAMS_BINDING 5
-#endif
-#elif defined(SINKS)
-#ifdef KV_OVERLAP
-@group(0) @binding(2) var<storage, read_write> sinks: array<f32>;
-#define DST_BINDING 3
-#define PARAMS_BINDING 4
-#else
-@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
-#define DST_BINDING 4
-#define PARAMS_BINDING 5
-#endif
-#else
-#ifdef KV_OVERLAP
-#define DST_BINDING 2
-#define PARAMS_BINDING 3
-#else
-#define DST_BINDING 3
-#define PARAMS_BINDING 4
-#endif
-#endif
-
-@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<vec4<DST_TYPE>>;
-@group(0) @binding(PARAMS_BINDING) var<uniform> params: Params;
-
-const FLOAT_MIN: f32 = -1.0e9;
 const Q_CHUNKS: u32 = HEAD_DIM_QK / 4u;
 const V_CHUNKS: u32 = HEAD_DIM_V / 4u;
 const SCORE_REGS_PER_LANE: u32 = (KV_TILE + MIN_SUBGROUP_SIZE - 1u) / MIN_SUBGROUP_SIZE;
 const OUT_REGS_PER_LANE: u32 = (V_CHUNKS + MIN_SUBGROUP_SIZE - 1u) / MIN_SUBGROUP_SIZE;
+
+#if !defined(K_DIRECT) || !defined(V_DIRECT)
+#define STAGING_SHMEM kv_shmem
+#define STAGING_OUT_TYPE f16
+#include "flash_attn_staging.tmpl"
 const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
+var<workgroup> kv_shmem: array<f16, kv_shmem_size>;
+#endif
 
 var<workgroup> q_shmem: array<Q_TYPE, Q_TILE * HEAD_DIM_QK>;
-var<workgroup> kv_shmem: array<f16, kv_shmem_size>;
 var<workgroup> p_shmem: array<f16, Q_TILE * KV_TILE>;
 
-#define QUANT_SHMEM kv_shmem
-#define QUANT_OUT_TYPE f16
-#include "flash_attn_quant_staging.tmpl"
-
-#if !defined(K_Q4_0) && !defined(K_Q8_0)
-fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
-    for (var vec_idx_local = local_x; vec_idx_local < kv_count * Q_CHUNKS; vec_idx_local += WG_SIZE) {
-        let kv_local = vec_idx_local / Q_CHUNKS;
-        let chunk = vec_idx_local % Q_CHUNKS;
-        let global_k_row = kv_tile + kv_local;
-        let k_vec_index = (k_head_offset + global_k_row * params.stride_k1 + chunk * 4u) >> 2u;
-        let k4 = K[k_vec_index];
-        let kv_off = kv_local * HEAD_DIM_QK + chunk * 4u;
-        kv_shmem[kv_off + 0u] = f16(k4.x);
-        kv_shmem[kv_off + 1u] = f16(k4.y);
-        kv_shmem[kv_off + 2u] = f16(k4.z);
-        kv_shmem[kv_off + 3u] = f16(k4.w);
-    }
-}
-#endif
-
-#if !defined(V_Q4_0) && !defined(V_Q8_0)
-fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
-    for (var vec_idx_local = local_x; vec_idx_local < kv_count * V_CHUNKS; vec_idx_local += WG_SIZE) {
-        let kv_local = vec_idx_local / V_CHUNKS;
-        let chunk = vec_idx_local % V_CHUNKS;
-        let global_v_row = kv_tile + kv_local;
-        let v_vec_index = (v_head_offset + global_v_row * params.stride_v1 + chunk * 4u) >> 2u;
-        let v4 = V[v_vec_index];
-        let kv_off = kv_local * HEAD_DIM_V + chunk * 4u;
-        kv_shmem[kv_off + 0u] = f16(v4.x);
-        kv_shmem[kv_off + 1u] = f16(v4.y);
-        kv_shmem[kv_off + 2u] = f16(v4.z);
-        kv_shmem[kv_off + 3u] = f16(v4.w);
-    }
-}
-#endif
-
 @compute @workgroup_size(WG_SIZE)
 fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
         @builtin(local_invocation_id) local_id: vec3<u32>,

ggml/src/ggml-webgpu/wgsl-shaders/flash_attn_vec_split.wgsl

@@ -4,200 +4,35 @@ enable subgroups;
 
 #define BYTE_HELPERS
 #include "common_decls.tmpl"
+#define FLASH_ATTN_VEC_SPLIT
+#include "flash_attn_decls.tmpl"
 
-#ifdef K_F32
-#define K_TYPE f32
-#elif defined(K_Q4_0) || defined(K_Q8_0)
-#define K_TYPE u32
-#else
-#define K_TYPE f16
-#endif
-
-#ifdef V_F32
-#define V_TYPE f32
-#elif defined(V_Q4_0) || defined(V_Q8_0)
-#define V_TYPE u32
-#else
-#define V_TYPE f16
-#endif
-
-#ifdef Q_F16
-#define Q_TYPE f16
-#else
-#define Q_TYPE f32
-#endif
-
-#ifdef DST_F16
-#define DST_TYPE f16
-#else
-#define DST_TYPE f32
-#endif
-
+// Default values
+// The actual values are defined in shader-lib.
 #define HEAD_DIM_QK 64
 #define HEAD_DIM_V 64
-
-#define KV_GRANULARITY 8
 #define KV_TILE 16
 #define WG_SIZE 64
 
-#define KV_BLOCKS (KV_TILE / KV_GRANULARITY)
-
-struct Params {
-    offset_q: u32,
-    offset_k: u32,
-    offset_v: u32,
-    offset_mask: u32,
-    offset_sinks: u32,
-    offset_dst: u32,
-
-    // shapes of Q/K/V
-    n_heads: u32,
-    seq_len_q: u32,
-    seq_len_kv: u32,
-
-    // strides (in elements)
-    stride_q1: u32,
-    stride_q2: u32,
-    stride_q3: u32,
-    stride_k1: u32,
-    stride_k2: u32,
-    stride_k3: u32,
-    stride_v1: u32,
-    stride_v2: u32,
-    stride_v3: u32,
-    stride_mask3: u32,
-
-    // repeat factors for K/V, e.g., MHA vs. MQA vs. GQA
-    q_per_kv: u32,
-
-    // softmax params
-    scale: f32,
-    max_bias: f32,
-    logit_softcap: f32,
-    n_head_log2: f32,
-    m0: f32,
-    m1: f32,
-
-#ifdef BLK
-    blk_base: u32,
-    blk_nblk0: u32,
-    blk_nblk1: u32,
-#endif
-
-    tmp_data_base: u32,
-    tmp_stats_base: u32,
-    nwg: u32,
-};
-
-@group(0) @binding(0) var<storage, read_write> Q: array<Q_TYPE>;
-#ifdef KV_OVERLAP
-#if defined(K_Q4_0) || defined(K_Q8_0)
-@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
-#else
-@group(0) @binding(1) var<storage, read_write> K: array<vec4<K_TYPE>>;
-#endif
-#define V K
-#else
-#if defined(K_Q4_0) || defined(K_Q8_0)
-@group(0) @binding(1) var<storage, read_write> K: array<K_TYPE>;
-#else
-@group(0) @binding(1) var<storage, read_write> K: array<vec4<K_TYPE>>;
-#endif
-#if defined(V_Q4_0) || defined(V_Q8_0)
-@group(0) @binding(2) var<storage, read_write> V: array<V_TYPE>;
-#else
-@group(0) @binding(2) var<storage, read_write> V: array<vec4<V_TYPE>>;
-#endif
-#endif
-#if defined(MASK) && defined(SINKS)
-#ifdef KV_OVERLAP
-@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
-@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
-#ifdef BLK
-#define BLK_BINDING 4
-#define TMP_BINDING 5
-#define DST_BINDING 6
-#define PARAMS_BINDING 7
-#else
-#define TMP_BINDING 4
-#define DST_BINDING 5
-#define PARAMS_BINDING 6
-#endif
-#else
-@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
-@group(0) @binding(4) var<storage, read_write> sinks: array<f32>;
-#ifdef BLK
-#define BLK_BINDING 5
-#define TMP_BINDING 6
-#define DST_BINDING 7
-#define PARAMS_BINDING 8
-#else
-#define TMP_BINDING 5
-#define DST_BINDING 6
-#define PARAMS_BINDING 7
-#endif
-#endif
-#elif defined(MASK)
-#ifdef KV_OVERLAP
-@group(0) @binding(2) var<storage, read_write> mask: array<f16>;
-#ifdef BLK
-#define BLK_BINDING 3
-#define TMP_BINDING 4
-#define DST_BINDING 5
-#define PARAMS_BINDING 6
-#else
-#define TMP_BINDING 3
-#define DST_BINDING 4
-#define PARAMS_BINDING 5
-#endif
-#else
-@group(0) @binding(3) var<storage, read_write> mask: array<f16>;
-#ifdef BLK
-#define BLK_BINDING 4
-#define TMP_BINDING 5
-#define DST_BINDING 6
-#define PARAMS_BINDING 7
-#else
-#define TMP_BINDING 4
-#define DST_BINDING 5
-#define PARAMS_BINDING 6
-#endif
-#endif
-#elif defined(SINKS)
-#ifdef KV_OVERLAP
-@group(0) @binding(2) var<storage, read_write> sinks: array<f32>;
-#define TMP_BINDING 3
-#define DST_BINDING 4
-#define PARAMS_BINDING 5
-#else
-@group(0) @binding(3) var<storage, read_write> sinks: array<f32>;
-#define TMP_BINDING 4
-#define DST_BINDING 5
-#define PARAMS_BINDING 6
-#endif
-#else
-#ifdef KV_OVERLAP
-#define TMP_BINDING 2
-#define DST_BINDING 3
-#define PARAMS_BINDING 4
-#else
-#define TMP_BINDING 3
-#define DST_BINDING 4
-#define PARAMS_BINDING 5
-#endif
-#endif
-
-#ifdef BLK
-@group(0) @binding(BLK_BINDING) var<storage, read_write> blk: array<u32>;
-#endif
-@group(0) @binding(TMP_BINDING) var<storage, read_write> tmp: array<f32>;
-@group(0) @binding(DST_BINDING) var<storage, read_write> dst: array<vec4<DST_TYPE>>;
-@group(0) @binding(PARAMS_BINDING) var<uniform> params: Params;
-
-// Just a very small float value.
-const FLOAT_MIN: f32 = -1.0e9;
+const Q_CHUNKS: u32 = HEAD_DIM_QK / 4u;
+const V_CHUNKS: u32 = HEAD_DIM_V / 4u;
 const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
 
+#if defined(K_DIRECT) || defined(V_DIRECT)
+// Shared memory for scale factor (d) in quantized K/V. Multiple threads use the same value,
+// so caching it is more efficient, even on the direct path.
+var<workgroup> d_shmem: array<f32, kv_shmem_size / 32>;
+#endif
+
+// K/V shared memory handling
+#if !defined(K_DIRECT) || !defined(V_DIRECT)
+#define STAGING_SHMEM kv_shmem
+#define STAGING_OUT_TYPE f32
+#include "flash_attn_staging.tmpl"
+// we can reuse the same shmem for K and V since we only need one at a time
+var<workgroup> kv_shmem: array<f32, kv_shmem_size>;
+#endif
+
 var<workgroup> q_shmem: array<f32, HEAD_DIM_QK>;
 var<workgroup> o_shmem: array<f32, HEAD_DIM_V>;
 // note that we reuse the same storage for both since we only need one at a time
@@ -208,59 +43,6 @@ var<workgroup> inter_shmem: array<f32, KV_TILE>;
 var<workgroup> mask_shmem: array<f32, KV_TILE>;
 #endif
 
-#if defined(K_DIRECT) || defined(V_DIRECT)
-// Shared memory for scale factor (d) in quantized K/V. Multiple threads use the same value,
-// so caching it is more efficient, even on the direct path.
-var<workgroup> d_shmem: array<f32, kv_shmem_size / 32>;
-#endif
-
-// K/V shared memory handling
-#if !defined(K_DIRECT) || !defined(V_DIRECT)
-
-// we can reuse the same shmem for K and V since we only need one at a time
-var<workgroup> kv_shmem: array<f32, kv_shmem_size>;
-
-#define QUANT_SHMEM kv_shmem
-#define QUANT_OUT_TYPE f32
-#include "flash_attn_quant_staging.tmpl"
-
-#if !defined(K_DIRECT) && !defined(K_Q4_0) && !defined(K_Q8_0)
-fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
-    for (var elem_idx = local_x * 4u; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE * 4u) {
-        let k_row = elem_idx / HEAD_DIM_QK;
-        let k_col = elem_idx % HEAD_DIM_QK;
-        let global_k_row = kv_tile + k_row;
-        let global_k_row_offset = k_head_offset + global_k_row * params.stride_k1;
-        let in_bounds = global_k_row < params.seq_len_kv && (k_col + 3u) < HEAD_DIM_QK;
-        let vec_idx = (global_k_row_offset + k_col) >> 2u;
-        let k4 = select(vec4<K_TYPE>(0.0), K[vec_idx], in_bounds);
-        kv_shmem[elem_idx + 0u] = f32(k4.x);
-        kv_shmem[elem_idx + 1u] = f32(k4.y);
-        kv_shmem[elem_idx + 2u] = f32(k4.z);
-        kv_shmem[elem_idx + 3u] = f32(k4.w);
-    }
-}
-#endif
-
-#if !defined(V_DIRECT) && !defined(V_Q4_0) && !defined(V_Q8_0)
-fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
-    for (var elem_idx = local_x * 4u; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE * 4u) {
-        let v_row = elem_idx / HEAD_DIM_V;
-        let v_col = elem_idx % HEAD_DIM_V;
-        let global_v_row = kv_tile + v_row;
-        let global_v_row_offset = v_head_offset + global_v_row * params.stride_v1;
-        let in_bounds = global_v_row < params.seq_len_kv && (v_col + 3u) < HEAD_DIM_V;
-        let vec_idx = (global_v_row_offset + v_col) >> 2u;
-        let v4 = select(vec4<V_TYPE>(0.0), V[vec_idx], in_bounds);
-        kv_shmem[elem_idx + 0u] = f32(v4.x);
-        kv_shmem[elem_idx + 1u] = f32(v4.y);
-        kv_shmem[elem_idx + 2u] = f32(v4.z);
-        kv_shmem[elem_idx + 3u] = f32(v4.w);
-    }
-}
-#endif
-#endif // !defined(K_DIRECT) || !defined(V_DIRECT)
-
 // Storage for row max and exp sum during online softmax
 fn calc_softmax_term(kv_idx: u32, slope: f32, has_bias: bool, apply_mask: bool) -> f32 {
     var v = select(FLOAT_MIN,

ggml/src/ggml-webgpu/wgsl-shaders/im2col.wgsl

@@ -1,19 +1,9 @@
-#include "common_decls.tmpl"
 enable f16;
 
 @group(0) @binding(0)
-#if defined(INPUT_F32)
-var<storage, read_write> input: array<f32>;
-#elif defined(INPUT_F16)
-var<storage, read_write> input: array<f16>;
-#endif
-
+var<storage, read_write> input: array<INPUT_TYPE>;
 @group(0) @binding(1)
-#if defined(OUTPUT_F32)
-var<storage, read_write> output: array<f32>;
-#elif defined(OUTPUT_F16)
-var<storage, read_write> output: array<f16>;
-#endif
+var<storage, read_write> output: array<OUTPUT_TYPE>;
 
 struct Params {
     offset_i: u32,
@@ -38,22 +28,6 @@ struct Params {
 @group(0) @binding(2)
 var<uniform> params: Params;
 
-fn load_input(idx: u32) -> f32 {
-    #if defined(INPUT_F32)
-        return input[idx];
-    #elif defined(INPUT_F16)
-        return f32(input[idx]);
-    #endif
-}
-
-fn store_output(idx: u32, val: f32) {
-    #if defined(OUTPUT_F32)
-        output[idx] = val;
-    #elif defined(OUTPUT_F16)
-        output[idx] = f16(val);
-    #endif
-}
-
 @compute @workgroup_size(WG_SIZE)
 fn main(
     @builtin(global_invocation_id) gid: vec3<u32>,
@@ -90,12 +64,14 @@ fn main(
     let iw_i32 = i32(ow * params.s0 + kw * params.d0) - i32(params.p0);
     let ih_i32 = i32(oh * params.s1 + kh * params.d1) - i32(params.p1);
 
+    let output_idx = params.offset_o + k * params.so0 + ow * params.so1 + oh * params.so2 + n * params.so3;
+
     if (iw_i32 >= 0 && iw_i32 < i32(params.IW) && ih_i32 >= 0 && ih_i32 < i32(params.IH)) {
         let iw = u32(iw_i32);
         let ih = u32(ih_i32);
         let in_idx = params.offset_i + iw * params.si0 + ih * params.si1 + ic * params.si2 + n * params.si3;
-        store_output(params.offset_o + k * params.so0 + ow * params.so1 + oh * params.so2 + n * params.so3, load_input(in_idx));
+        output[output_idx] = OUTPUT_TYPE(input[in_idx]);
     } else {
-        store_output(params.offset_o + k * params.so0 + ow * params.so1 + oh * params.so2 + n * params.so3, 0.0);
+        output[output_idx] = OUTPUT_TYPE(0.0);
     }
 }

ggml/src/ggml-webgpu/wgsl-shaders/rms_norm_mul.wgsl

@@ -88,7 +88,6 @@ struct Params {
     ne0: u32,
     ne1: u32,
     ne2: u32,
-    ne3: u32,
 
     eps: f32
 };

ggml/src/ggml-webgpu/wgsl-shaders/row_norm.wgsl

@@ -31,7 +31,6 @@ struct Params {
     ne0: u32,
     ne1: u32,
     ne2: u32,
-    ne3: u32,
 
     eps: f32
 };

ggml/src/ggml-webgpu/wgsl-shaders/soft_max.wgsl

@@ -27,7 +27,6 @@ struct Params {
     stride_dst3: u32,
 
     // shape of src0/dst
-    ne: u32,
     ne0: u32,
     ne1: u32,
     ne2: u32,
@@ -43,71 +42,38 @@ struct Params {
     m1: f32,
 };
 
-@group(0) @binding(0)
+#define SRC_BINDING 0
+@group(0) @binding(SRC_BINDING)
 var<storage, read_write> src: array<f32>;
 
 #ifdef HAS_MASK
-#ifdef HAS_SINK
-@group(0) @binding(1)
+#define MASK_BINDING SRC_BINDING + 1
+@group(0) @binding(MASK_BINDING)
 var<storage, read_write> mask: array<MaskType>;
-@group(0) @binding(2)
-var<storage, read_write> sinks: array<f32>;
-
-#ifdef INPLACE
-@group(0) @binding(3)
-var<uniform> params: Params;
-
 #else
-@group(0) @binding(3)
-var<storage, read_write> dst: array<f32>;
-@group(0) @binding(4)
-var<uniform> params: Params;
+#define MASK_BINDING SRC_BINDING
 #endif
 
-#else
-@group(0) @binding(1)
-var<storage, read_write> mask: array<MaskType>;
-
-#ifdef INPLACE
-@group(0) @binding(2)
-var<uniform> params: Params;
-
-#else
-@group(0) @binding(2)
-var<storage, read_write> dst: array<f32>;
-@group(0) @binding(3)
-var<uniform> params: Params;
-#endif
-#endif
-
-#else
 #ifdef HAS_SINK
-@group(0) @binding(1)
+#define SINKS_BINDING MASK_BINDING + 1
+@group(0) @binding(SINKS_BINDING)
 var<storage, read_write> sinks: array<f32>;
+#else
+#define SINKS_BINDING MASK_BINDING
+#endif
+
+#define DST_BINDING SINKS_BINDING + 1
+@group(0) @binding(DST_BINDING)
+var<storage, read_write> dst: array<f32>;
 
 #ifdef INPLACE
-@group(0) @binding(2)
-var<uniform> params: Params;
-
+#define PARAMS_BINDING DST_BINDING
 #else
-@group(0) @binding(2)
-var<storage, read_write> dst: array<f32>;
-@group(0) @binding(3)
-var<uniform> params: Params;
+#define PARAMS_BINDING (DST_BINDING + 1)
 #endif
 
-#else
-#ifdef INPLACE
-@group(0) @binding(1)
+@group(0) @binding(PARAMS_BINDING)
 var<uniform> params: Params;
-#else
-@group(0) @binding(1)
-var<storage, read_write> dst: array<f32>;
-@group(0) @binding(2)
-var<uniform> params: Params;
-#endif
-#endif
-#endif
 
 #ifdef INPLACE
 fn inter_value(i: u32) -> f32 {
@@ -242,4 +208,3 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
         col += WG_SIZE;
     }
 }
-

ggml/src/ggml-webgpu/wgsl-shaders/solve_tri.wgsl

@@ -29,7 +29,6 @@ struct Params {
 
     k: u32,
     ne2: u32,
-    ne3: u32,
 };
 
 @group(0) @binding(3)

ggml/src/ggml-webgpu/wgsl-shaders/ssm_scan.wgsl

@@ -39,7 +39,6 @@ struct Params {
     n_head: u32,
     n_group: u32,
     n_seq_tokens: u32,
-    n_seqs: u32,
 
     y_elems: u32,
 };

gguf-py/gguf/constants.py

@@ -164,6 +164,13 @@ class Keys:
         NORM_BEFORE_RESIDUAL              = "{arch}.norm_before_residual"
         NORM_BEFORE_FC                    = "{arch}.norm_before_fc"
 
+    class Adapters:
+        COUNT                = "{arch}.adapters.count"
+        TOKEN_IDS_ACTIVATE   = "{arch}.adapters.token_ids_activate"
+        TOKEN_IDS_SUBSTITUTE = "{arch}.adapters.token_ids_substitute"
+        LORA_RANK            = "{arch}.adapters.lora_rank"
+        ROUTER_GAIN          = "{arch}.adapters.router_gain"
+
     class Attention:
         HEAD_COUNT                   = "{arch}.attention.head_count"
         HEAD_COUNT_KV                = "{arch}.attention.head_count_kv"
@@ -527,6 +534,7 @@ class MODEL_ARCH(IntEnum):
     GRANITE          = auto()
     GRANITE_MOE      = auto()
     GRANITE_HYBRID   = auto()
+    GRANITE_SWITCH   = auto()
     CHAMELEON        = auto()
     WAVTOKENIZER_DEC = auto()
     PLM              = auto()
@@ -1198,6 +1206,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.GRANITE:          "granite",
     MODEL_ARCH.GRANITE_MOE:      "granitemoe",
     MODEL_ARCH.GRANITE_HYBRID:   "granitehybrid",
+    MODEL_ARCH.GRANITE_SWITCH:   "graniteswitch",
     MODEL_ARCH.CHAMELEON:        "chameleon",
     MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec",
     MODEL_ARCH.PLM:              "plm",
@@ -3837,6 +3846,12 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_DOWN_SHEXP,
         MODEL_TENSOR.FFN_UP_SHEXP,
         MODEL_TENSOR.FFN_EXP_PROBS_B,
+        # NextN/MTP (draft head)
+        MODEL_TENSOR.ATTN_POST_NORM,
+        MODEL_TENSOR.NEXTN_EH_PROJ,
+        MODEL_TENSOR.NEXTN_ENORM,
+        MODEL_TENSOR.NEXTN_HNORM,
+        MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
     ],
     MODEL_ARCH.EXAONE: [
         MODEL_TENSOR.TOKEN_EMBD,
@@ -3972,6 +3987,21 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_DOWN,
         MODEL_TENSOR.FFN_UP,
     ],
+    MODEL_ARCH.GRANITE_SWITCH: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_QKV,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+    ],
     MODEL_ARCH.CHAMELEON: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,

gguf-py/gguf/gguf_writer.py

@@ -906,6 +906,21 @@ class GGUFWriter:
     def add_embedding_scale(self, value: float) -> None:
         self.add_float32(Keys.LLM.EMBEDDING_SCALE.format(arch=self.arch), value)
 
+    def add_adapter_count(self, count: int) -> None:
+        self.add_uint32(Keys.Adapters.COUNT.format(arch=self.arch), count)
+
+    def add_adapter_token_ids_activate(self, ids: Sequence[int]) -> None:
+        self.add_array(Keys.Adapters.TOKEN_IDS_ACTIVATE.format(arch=self.arch), ids)
+
+    def add_adapter_token_ids_substitute(self, ids: Sequence[int]) -> None:
+        self.add_array(Keys.Adapters.TOKEN_IDS_SUBSTITUTE.format(arch=self.arch), ids)
+
+    def add_adapter_lora_rank(self, rank: int) -> None:
+        self.add_uint32(Keys.Adapters.LORA_RANK.format(arch=self.arch), rank)
+
+    def add_adapter_router_gain(self, gain: float) -> None:
+        self.add_float32(Keys.Adapters.ROUTER_GAIN.format(arch=self.arch), gain)
+
     def add_wkv_head_size(self, size: int) -> None:
         self.add_uint32(Keys.WKV.HEAD_SIZE.format(arch=self.arch), size)
 

scripts/sync_vendor.py

@@ -5,7 +5,7 @@ import os
 import sys
 import subprocess
 
-HTTPLIB_VERSION = "refs/tags/v0.52.0"
+HTTPLIB_VERSION = "refs/tags/v0.53.0"
 
 vendor = {
     "https://github.com/nlohmann/json/releases/latest/download/json.hpp":     "vendor/nlohmann/json.hpp",

src/llama-arch.cpp

@@ -100,6 +100,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
     { LLM_ARCH_GRANITE,          "granite"          },
     { LLM_ARCH_GRANITE_MOE,      "granitemoe"       },
     { LLM_ARCH_GRANITE_HYBRID,   "granitehybrid"    },
+    { LLM_ARCH_GRANITE_SWITCH,   "graniteswitch"    },
     { LLM_ARCH_CHAMELEON,        "chameleon"        },
     { LLM_ARCH_WAVTOKENIZER_DEC, "wavtokenizer-dec" },
     { LLM_ARCH_PLM,              "plm"              },
@@ -220,6 +221,11 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
     { LLM_KV_TIME_DECAY_EXTRA_DIM,              "%s.time_decay_extra_dim"              },
     { LLM_KV_RESIDUAL_SCALE,                    "%s.residual_scale"                    },
     { LLM_KV_EMBEDDING_SCALE,                   "%s.embedding_scale"                   },
+    { LLM_KV_ADAPTER_COUNT,                     "%s.adapters.count"                    },
+    { LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE,        "%s.adapters.token_ids_activate"       },
+    { LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE,      "%s.adapters.token_ids_substitute"     },
+    { LLM_KV_ADAPTER_LORA_RANK,                 "%s.adapters.lora_rank"                },
+    { LLM_KV_ADAPTER_ROUTER_GAIN,               "%s.adapters.router_gain"              },
     { LLM_KV_TOKEN_SHIFT_COUNT,                 "%s.token_shift_count"                 },
     { LLM_KV_INTERLEAVE_MOE_LAYER_STEP,         "%s.interleave_moe_layer_step"         },
     { LLM_KV_FULL_ATTENTION_INTERVAL,           "%s.full_attention_interval"           },

src/llama-arch.h

@@ -105,6 +105,7 @@ enum llm_arch {
     LLM_ARCH_GRANITE,
     LLM_ARCH_GRANITE_MOE,
     LLM_ARCH_GRANITE_HYBRID,
+    LLM_ARCH_GRANITE_SWITCH,
     LLM_ARCH_CHAMELEON,
     LLM_ARCH_WAVTOKENIZER_DEC,
     LLM_ARCH_PLM,
@@ -225,6 +226,11 @@ enum llm_kv {
     LLM_KV_TIME_DECAY_EXTRA_DIM,
     LLM_KV_RESIDUAL_SCALE,
     LLM_KV_EMBEDDING_SCALE,
+    LLM_KV_ADAPTER_COUNT,
+    LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE,
+    LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE,
+    LLM_KV_ADAPTER_LORA_RANK,
+    LLM_KV_ADAPTER_ROUTER_GAIN,
     LLM_KV_TOKEN_SHIFT_COUNT,
     LLM_KV_INTERLEAVE_MOE_LAYER_STEP,
     LLM_KV_FULL_ATTENTION_INTERVAL,

src/llama-context.cpp

@@ -3602,8 +3602,9 @@ llama_context * llama_init_from_model(
                        model->hparams.pooling_type, params.pooling_type);
     }
 
+    // router_layer >= 0 means n_layer_nextn is repurposed for a router layer, not real MTP
     if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
-        model->hparams.n_layer_nextn == 0) {
+        (model->hparams.n_layer_nextn == 0 || model->hparams.router_layer >= 0)) {
         LLAMA_LOG_WARN("%s: context type MTP requested but model doesn't contain MTP layers\n", __func__);
         return nullptr;
     }

src/llama-hparams.cpp

@@ -277,6 +277,16 @@ bool llama_hparams::has_kv(uint32_t il) const {
     return true;
 }
 
+bool llama_hparams::has_rope(uint32_t il) const {
+    // the router layer stores adapter routing signal, not positional info,
+    // so it must not be RoPE-shifted
+    if (router_layer >= 0 && (int32_t) il == router_layer) {
+        return false;
+    }
+
+    return true;
+}
+
 uint32_t llama_hparams::n_layer() const {
     return n_layer_all - n_layer_nextn;
 }

src/llama-hparams.h

@@ -53,6 +53,10 @@ struct llama_hparams {
     uint32_t n_embd;
     uint32_t n_layer_all;
     uint32_t n_layer_nextn = 0;
+
+    // granite-switch: index of the single-head "router" KV layer that encodes
+    // per-token adapter selection. -1 when the model has no such layer.
+    int32_t  router_layer = -1;
     uint32_t n_expert = 0;
     uint32_t n_expert_used = 0;
     uint32_t n_rel_attn_bkts = 0;
@@ -371,6 +375,8 @@ struct llama_hparams {
 
     bool has_kv(uint32_t il) const;
 
+    bool has_rope(uint32_t il) const;
+
     // number of effective layers (excludes nextn layers)
     uint32_t n_layer() const;
 

src/llama-kv-cache.cpp

@@ -1931,6 +1931,10 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co
     for (const auto & layer : layers) {
         const uint32_t il = layer.il;
 
+        if (!hparams.has_rope(il)) {
+            continue;
+        }
+
         const int64_t n_head_kv    = hparams.n_head_kv(il);
         const int64_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);
 

src/llama-model-loader.cpp

@@ -937,10 +937,11 @@ static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w
             } break;
         case GGML_OP_MUL_MAT_ID:
             {
-                const int n_expert_used = hparams.n_expert_used;
-                GGML_ASSERT(n_expert_used > 0);
-                ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_expert_used, 512);
-                ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512);
+                // Used for either MoE expert routing or embedded adapter routing
+                const int n_ids_used = hparams.router_layer >= 0 ? 1 : hparams.n_expert_used;
+                GGML_ASSERT(n_ids_used > 0);
+                ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_ids_used, 512);
+                ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_ids_used, 512);
                 op_tensor = ggml_mul_mat_id(ctx, w, b, ids);
             } break;
         case GGML_OP_ADD:
@@ -1123,15 +1124,14 @@ struct ggml_tensor * llama_model_loader::create_tensor(
             return nullptr;
         }
 
-        // tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID
+        // tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID;
+        // embedded-adapter ".lora_a"/".lora_b" tensors are always used with GGML_OP_MUL_MAT_ID
         ggml_op op;
-        bool bias = tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0;
-        if (bias) {
-            if (info.op == GGML_OP_MUL_MAT_ID) {
-                op = GGML_OP_ADD_ID;
-            } else {
-                op = GGML_OP_ADD;
-            }
+        if (tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0) {
+            op = info.op == GGML_OP_MUL_MAT_ID ? GGML_OP_ADD_ID : GGML_OP_ADD;
+        } else if (hparams.router_layer >= 0 && tn.suffix != nullptr &&
+                (strcmp(tn.suffix, "lora_a") == 0 || strcmp(tn.suffix, "lora_b") == 0)) {
+            op = GGML_OP_MUL_MAT_ID;
         } else {
             op = info.op;
         }

src/llama-model-saver.cpp

@@ -213,7 +213,7 @@ void llama_model_saver::add_kv_from_model() {
     add_kv(LLM_KV_FEED_FORWARD_LENGTH,               hparams.n_ff_arr, true);
     add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,        hparams.n_ff_exp);
     add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp);
-    add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_chexp);
+    add_kv(LLM_KV_EXPERT_CHUNK_FEED_FORWARD_LENGTH,  hparams.n_ff_chexp);
     add_kv(LLM_KV_SWIGLU_CLAMP_EXP,                  hparams.swiglu_clamp_exp);
     add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP,                hparams.swiglu_clamp_shexp);
     add_kv(LLM_KV_USE_PARALLEL_RESIDUAL,             hparams.use_par_res);

src/llama-model.cpp

@@ -234,6 +234,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
             return new llama_model_granite(params);
         case LLM_ARCH_GRANITE_MOE:
             return new llama_model_granite_moe(params);
+        case LLM_ARCH_GRANITE_SWITCH:
+            return new llama_model_granite_switch(params);
         case LLM_ARCH_MINICPM:
             return new llama_model_minicpm(params);
         case LLM_ARCH_GRANITE_HYBRID:
@@ -1912,6 +1914,7 @@ void llama_model::print_info() const {
                 arch == LLM_ARCH_GRANITE ||
                 arch == LLM_ARCH_GRANITE_MOE ||
                 arch == LLM_ARCH_GRANITE_HYBRID ||
+                arch == LLM_ARCH_GRANITE_SWITCH ||
                 arch == LLM_ARCH_NEMOTRON_H_MOE) {
             LLAMA_LOG_INFO("%s: f_embedding_scale     = %f\n", __func__, hparams.f_embedding_scale);
             LLAMA_LOG_INFO("%s: f_residual_scale      = %f\n", __func__, hparams.f_residual_scale);
@@ -2228,6 +2231,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                     params.ctx_type == LLAMA_CONTEXT_TYPE_MTP &&
                     (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE);
 
+                const bool mtp_on_hybrid_nemotron =
+                    params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE;
+
                 if (llm_arch_is_recurrent(arch)) {
                     res = new llama_memory_recurrent(
                             *this,
@@ -2238,7 +2244,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                             cparams.n_seq_max,
                             cparams.n_rs_seq,
                             nullptr);
-                } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen) {
+                } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen && !mtp_on_hybrid_nemotron) {
                     // The main difference between hybrid architectures is the
                     // layer filters, so pick the right one here
                     llama_memory_hybrid::layer_filter_cb filter_attn = nullptr;
@@ -2319,7 +2325,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                         };
                     }
 
-                    if (mtp_on_hybrid_qwen) {
+                    if (mtp_on_hybrid_qwen || mtp_on_hybrid_nemotron) {
                         filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
                     }
 
@@ -2596,6 +2602,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
         case LLM_ARCH_GRANITE:
         case LLM_ARCH_GRANITE_MOE:
         case LLM_ARCH_GRANITE_HYBRID:
+        case LLM_ARCH_GRANITE_SWITCH:
         case LLM_ARCH_CHAMELEON:
         case LLM_ARCH_BAILINGMOE:
         case LLM_ARCH_NEO_BERT:

src/llama-model.h

@@ -223,6 +223,24 @@ struct llama_layer_nextn {
     struct ggml_tensor * shared_head_norm      = nullptr;
 };
 
+struct llama_layer_switch_lora {
+    struct ggml_tensor * a_q    = nullptr;
+    struct ggml_tensor * b_q    = nullptr;
+    struct ggml_tensor * a_k    = nullptr;
+    struct ggml_tensor * b_k    = nullptr;
+    struct ggml_tensor * a_v    = nullptr;
+    struct ggml_tensor * b_v    = nullptr;
+    struct ggml_tensor * a_o    = nullptr;
+    struct ggml_tensor * b_o    = nullptr;
+
+    struct ggml_tensor * a_gate = nullptr;
+    struct ggml_tensor * b_gate = nullptr;
+    struct ggml_tensor * a_up   = nullptr;
+    struct ggml_tensor * b_up   = nullptr;
+    struct ggml_tensor * a_down = nullptr;
+    struct ggml_tensor * b_down = nullptr;
+};
+
 struct llama_layer {
     // normalization
     struct ggml_tensor * attn_norm       = nullptr;
@@ -533,6 +551,8 @@ struct llama_layer {
     struct llama_layer_shortconv shortconv;
 
     struct llama_layer_nextn nextn;
+
+    struct llama_layer_switch_lora switch_lora;
 };
 
 struct llama_device {

src/models/granite-switch.cpp

@@ -0,0 +1,426 @@
+#include "models.h"
+
+#include <cmath>
+
+void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) {
+    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+    ml.get_key(LLM_KV_LOGIT_SCALE,                 hparams.f_logit_scale);
+    ml.get_key(LLM_KV_RESIDUAL_SCALE,              hparams.f_residual_scale, false);
+    ml.get_key(LLM_KV_EMBEDDING_SCALE,             hparams.f_embedding_scale, false);
+    ml.get_key(LLM_KV_ATTENTION_SCALE,             hparams.f_attention_scale, false);
+
+    bool rope_finetuned = true;
+    ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);
+    hparams.rope_finetuned = rope_finetuned;
+
+    switch (hparams.n_layer()) {
+        case 40: type = hparams.n_embd == 4096 ? LLM_TYPE_8B : LLM_TYPE_3B; break;
+        case 64: type = LLM_TYPE_30B; break;
+        default: type = LLM_TYPE_UNKNOWN;
+    }
+
+    ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false);
+
+    ml.get_key(LLM_KV_ADAPTER_COUNT,     n_adapters);
+    ml.get_key(LLM_KV_ADAPTER_LORA_RANK, max_lora_rank);
+    ml.get_key(LLM_KV_ADAPTER_ROUTER_GAIN, router_gain, /* required */ false);
+
+    // bound counts that size tensors
+    if (n_adapters > 4096) {
+        throw std::runtime_error(format("graniteswitch: invalid adapter count %u", n_adapters));
+    }
+    if (max_lora_rank > 4096) {
+        throw std::runtime_error(format("graniteswitch: invalid lora rank %u", max_lora_rank));
+    }
+
+    std::vector<llama_token> token_ids;
+    std::vector<llama_token> substitute_ids;
+    ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_ACTIVATE,   token_ids);
+    ml.get_arr(LLM_KV_ADAPTER_TOKEN_IDS_SUBSTITUTE, substitute_ids);
+
+    if (token_ids.size() != n_adapters || substitute_ids.size() != n_adapters) {
+        throw std::runtime_error(format(
+            "graniteswitch: adapter token id arrays (%zu activate, %zu substitute) do not match adapter count %u",
+            token_ids.size(), substitute_ids.size(), n_adapters));
+    }
+
+    adapter_token_to_slot.clear();
+    adapter_token_to_substitute.clear();
+    for (uint32_t i = 0; i < n_adapters; ++i) {
+        // adapter i -> stacked slot i+1 (slot 0 is the base/zero delta)
+        adapter_token_to_slot[token_ids[i]]       = (int32_t) (i + 1);
+        adapter_token_to_substitute[token_ids[i]] = substitute_ids[i];
+    }
+
+    // extra single-head attention layer at the END (index n_real) holds the router
+    // K/V. reusing n_layer_nextn keeps n_layer() == n_real, so the regular layers
+    // keep their indices and the KV cache shift/defrag skips the router layer.
+    // n_layer_nextn is repurposed here (no MTP): it leaks as 1 into the
+    // llama_model_n_layer_nextn() getter and a re-saved nextn_predict_layers
+    const uint32_t n_real = hparams.n_layer();
+    if (n_real >= LLAMA_MAX_LAYERS) {
+        throw std::runtime_error(format("graniteswitch: block count %u exceeds LLAMA_MAX_LAYERS", n_real));
+    }
+    hparams.router_layer  = (int32_t) n_real;
+    hparams.n_layer_all   = n_real + 1;
+    hparams.n_layer_nextn = 1;
+
+    hparams.n_head_arr[n_real]    = 1;
+    hparams.n_head_kv_arr[n_real] = 1;
+    hparams.n_ff_arr[n_real]      = 0;
+}
+
+void llama_model_granite_switch::load_arch_tensors(llama_model_loader &) {
+    LLAMA_LOAD_LOCALS;
+
+    const int64_t n_slots     = (int64_t) n_adapters + 1; // slot 0 = base/zero delta
+    const int64_t n_rank      = (int64_t) max_lora_rank;
+    const int64_t n_embd_q    = n_embd_head_k * n_head;
+    const int64_t n_embd_kv   = n_embd_k_gqa;
+
+    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+    // substitute ids index tok_embd rows directly; range-check against n_vocab
+    for (const auto & kv : adapter_token_to_substitute) {
+        const llama_token sub = kv.second;
+        if (sub < 0 || (int64_t) sub >= n_vocab) {
+            throw std::runtime_error(format(
+                "graniteswitch: substitute token id %d out of range [0, %d)", sub, (int) n_vocab));
+        }
+    }
+
+    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);
+    if (output == NULL) {
+        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+    }
+
+    for (int i = 0; i < n_layer; ++i) {
+        auto & layer = layers[i];
+
+        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+
+        layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, n_embd_q + 2*n_embd_kv}, 0);
+        layer.wo   = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_q, n_embd}, 0);
+
+        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+        layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd,   n_ff}, 0);
+        layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  n_ff, n_embd}, 0);
+        layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, 0);
+
+        auto & sl = layer.switch_lora;
+
+        sl.a_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_a", i), {n_embd,  n_rank, n_slots}, 0);
+        sl.b_q = create_tensor(tn(LLM_TENSOR_ATTN_Q, "lora_b", i), {n_rank, n_embd_q, n_slots}, 0);
+        sl.a_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_a", i), {n_embd,  n_rank, n_slots}, 0);
+        sl.b_k = create_tensor(tn(LLM_TENSOR_ATTN_K, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0);
+        sl.a_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_a", i), {n_embd,  n_rank, n_slots}, 0);
+        sl.b_v = create_tensor(tn(LLM_TENSOR_ATTN_V, "lora_b", i), {n_rank, n_embd_kv, n_slots}, 0);
+
+        sl.a_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_a", i), {n_embd_q, n_rank, n_slots}, 0);
+        sl.b_o = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "lora_b", i), {n_rank,   n_embd, n_slots}, 0);
+
+        sl.a_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_a", i), {n_embd, n_rank, n_slots}, 0);
+        sl.b_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "lora_b", i), {n_rank,  n_ff,  n_slots}, 0);
+        sl.a_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "lora_a", i), {n_embd, n_rank, n_slots}, 0);
+        sl.b_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "lora_b", i), {n_rank,  n_ff,  n_slots}, 0);
+        sl.a_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_a", i), {  n_ff, n_rank, n_slots}, 0);
+        sl.b_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "lora_b", i), {n_rank, n_embd, n_slots}, 0);
+    }
+}
+
+class llm_graph_input_switch : public llm_graph_input_i {
+public:
+    llm_graph_input_switch(const llama_model_granite_switch & smodel) : smodel(smodel) {}
+    virtual ~llm_graph_input_switch() = default;
+
+    void set_input(const llama_ubatch * ubatch) override;
+
+    ggml_tensor * sub_tokens  = nullptr; // I32 [n_tokens] adapter-substituted token ids
+    ggml_tensor * router_ksig = nullptr; // F32 [n_tokens] router K signal (+/-gain)
+    ggml_tensor * router_vval = nullptr; // F32 [n_tokens] router V value (adapter slot / 0)
+    ggml_tensor * router_q    = nullptr; // F32 [n_tokens] router Q value (constant 1.0)
+
+    const llama_model_granite_switch & smodel;
+};
+
+// K dim-0 is +gain for an adapter token, -gain otherwise; the causal softmax then
+// lets a single visible adapter token dominate so the readback recovers its slot.
+void llm_graph_input_switch::set_input(const llama_ubatch * ubatch) {
+    if (!ubatch->token) {
+        return;
+    }
+
+    const int64_t n_tokens = ubatch->n_tokens;
+
+    std::vector<int32_t> sub (n_tokens);
+    std::vector<float>   ksig(n_tokens);
+    std::vector<float>   vval(n_tokens);
+    std::vector<float>   q   (n_tokens, 1.0f);
+
+    for (int64_t i = 0; i < n_tokens; ++i) {
+        const llama_token tok = ubatch->token[i];
+
+        const auto it = smodel.adapter_token_to_slot.find(tok);
+        if (it != smodel.adapter_token_to_slot.end()) {
+            ksig[i] = +smodel.router_gain;
+            vval[i] = (float) it->second;
+        } else {
+            ksig[i] = -smodel.router_gain;
+            vval[i] = 0.0f;
+        }
+
+        const auto sit = smodel.adapter_token_to_substitute.find(tok);
+        sub[i] = (sit != smodel.adapter_token_to_substitute.end())
+            ? (int32_t) sit->second
+            : (int32_t) tok;
+    }
+
+    ggml_backend_tensor_set(sub_tokens,  sub.data(),  0, n_tokens*ggml_element_size(sub_tokens));
+    ggml_backend_tensor_set(router_ksig, ksig.data(), 0, n_tokens*ggml_element_size(router_ksig));
+    ggml_backend_tensor_set(router_vval, vval.data(), 0, n_tokens*ggml_element_size(router_vval));
+    ggml_backend_tensor_set(router_q,    q.data(),    0, n_tokens*ggml_element_size(router_q));
+}
+
+std::unique_ptr<llm_graph_context> llama_model_granite_switch::build_arch_graph(const llm_graph_params & params) const {
+    return std::make_unique<graph>(*this, params);
+}
+
+// per-token switched LoRA delta: B_a*(A_a*x), adapter selected per token via ids.
+// cur: {n_in, n_tokens}, ids: {n_tokens} -> {n_out, n_tokens}
+ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_delta(
+          ggml_tensor * lora_a,
+          ggml_tensor * lora_b,
+          ggml_tensor * cur,
+          ggml_tensor * ids) {
+    const int64_t n_in     = cur->ne[0];
+    const int64_t n_tokens = cur->ne[1];
+
+    ggml_tensor * x    = ggml_reshape_3d(ctx0, cur, n_in, 1, n_tokens);
+    ggml_tensor * ids2 = ggml_reshape_2d(ctx0, ids, 1, n_tokens);
+
+    ggml_tensor * a = ggml_mul_mat_id(ctx0, lora_a, x, ids2); // {max_rank, 1, n_tokens}
+    ggml_tensor * d = ggml_mul_mat_id(ctx0, lora_b, a, ids2); // {n_out,    1, n_tokens}
+
+    return ggml_reshape_2d(ctx0, d, d->ne[0], n_tokens);
+}
+
+ggml_tensor * llama_model_granite_switch::graph::build_switched_lora_mm(
+          ggml_tensor * w,
+          ggml_tensor * lora_a,
+          ggml_tensor * lora_b,
+          ggml_tensor * cur,
+          ggml_tensor * ids) {
+    ggml_tensor * base  = ggml_mul_mat(ctx0, w, cur);
+    ggml_tensor * delta = build_switched_lora_delta(lora_a, lora_b, cur, ids);
+    return ggml_add(ctx0, base, delta);
+}
+
+llama_model_granite_switch::graph::graph(
+    const llama_model & model,
+    const llm_graph_params & params)
+    : llm_graph_context(params) {
+
+    const auto & smodel = static_cast<const llama_model_granite_switch &>(model);
+
+    // TODO: support raw embedding input (multimodal / pre-embedded tokens) when needed
+    GGML_ASSERT(ubatch.token && "granite-switch requires token input");
+
+    const int64_t n_embd_head = hparams.n_embd_head_v();
+    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+    GGML_ASSERT(n_embd_head == n_rot);
+
+    auto inp_switch = std::make_unique<llm_graph_input_switch>(smodel);
+    inp_switch->sub_tokens  = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+    inp_switch->router_ksig = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
+    inp_switch->router_vval = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
+    inp_switch->router_q    = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_tokens);
+    ggml_set_input(inp_switch->sub_tokens);
+    ggml_set_input(inp_switch->router_ksig);
+    ggml_set_input(inp_switch->router_vval);
+    ggml_set_input(inp_switch->router_q);
+    ggml_tensor * sub_tokens  = inp_switch->sub_tokens;
+    ggml_tensor * router_ksig = inp_switch->router_ksig;
+    ggml_tensor * router_vval = inp_switch->router_vval;
+    ggml_tensor * router_q    = inp_switch->router_q;
+    res->add_input(std::move(inp_switch));
+
+    // embed the substituted ids directly; build_inp_embd would embed the raw tokens
+    ggml_tensor * inpL = ggml_get_rows(ctx0, model.tok_embd, sub_tokens);
+    if (hparams.f_embedding_scale != 0.0f) {
+        inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale);
+    }
+    cb(inpL, "inp_embd", -1);
+
+    ggml_tensor * inp_pos = nullptr;
+    if (hparams.rope_finetuned) {
+        inp_pos = build_inp_pos();
+    }
+    auto * inp_attn = build_attn_inp_kv();
+
+    // single causal head at layer R recovers the adapter index in-graph: only dim 0
+    // carries signal (Q[0]=1, K[0]=+/-gain, V[0]=slot/0), the rest is zero-padded.
+    const int R = hparams.router_layer;
+    GGML_ASSERT(R >= 0);
+    auto router_lane = [&](ggml_tensor * sig1d) {
+        ggml_tensor * t = ggml_reshape_3d(ctx0, sig1d, 1, 1, n_tokens);
+        return ggml_pad(ctx0, t, (int) n_embd_head - 1, 0, 0, 0);
+    };
+    ggml_tensor * Qr = router_lane(router_q);
+    ggml_tensor * Kr = router_lane(router_ksig);
+    ggml_tensor * Vr = router_lane(router_vval);
+
+    ggml_tensor * router_out = build_attn(inp_attn,
+            nullptr, nullptr, nullptr,
+            Qr, Kr, Vr, nullptr, nullptr, nullptr, /*kq_scale=*/1.0f, /*il=*/R);
+    cb(router_out, "router_out", R);
+
+    // row 0 of router_out is the attended slot; clamp+round to an I32 index
+    ggml_tensor * slot_f = ggml_cont(ctx0,
+        ggml_view_2d(ctx0, router_out, 1, n_tokens, router_out->nb[1], 0));
+    slot_f = ggml_reshape_1d(ctx0, slot_f, n_tokens);
+    slot_f = ggml_clamp(ctx0, slot_f, 0.0f, (float) smodel.n_adapters);
+    slot_f = ggml_round(ctx0, slot_f);
+    ggml_tensor * adapter_ids = ggml_cast(ctx0, slot_f, GGML_TYPE_I32);
+    cb(adapter_ids, "adapter_ids", -1);
+
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    ggml_tensor * cur;
+
+    for (int il = 0; il < n_layer; ++il) {
+        ggml_tensor * inpSA = inpL;
+
+        cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
+        cb(cur, "attn_norm", il);
+
+        cur = build_attention_layer(cur, inp_pos, adapter_ids, inp_attn, model, n_embd_head, il);
+
+        if (il == n_layer - 1 && inp_out_ids) {
+            cur   = ggml_get_rows(ctx0, cur,   inp_out_ids);
+            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+            // keep adapter_ids aligned to the kept rows (2D round-trip for get_rows)
+            const int64_t n_out = inp_out_ids->ne[0];
+            adapter_ids = ggml_get_rows(ctx0,
+                ggml_reshape_2d(ctx0, adapter_ids, 1, adapter_ids->ne[0]), inp_out_ids);
+            adapter_ids = ggml_reshape_1d(ctx0, adapter_ids, n_out);
+        }
+
+        cur = build_layer_ffn(cur, inpSA, adapter_ids, model, il);
+
+        inpL = cur;
+    }
+
+    cur = inpL;
+
+    cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+    cb(cur, "result_norm", -1);
+    res->t_embd = cur;
+
+    cur = build_lora_mm(model.output, cur, model.output_s);
+
+    cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale);
+    cb(cur, "result_output", -1);
+    res->t_logits = cur;
+
+    ggml_build_forward_expand(gf, cur);
+}
+
+ggml_tensor * llama_model_granite_switch::graph::build_attention_layer(
+          ggml_tensor             * cur,
+          ggml_tensor             * inp_pos,
+          ggml_tensor             * adapter_ids,
+          llm_graph_input_attn_kv * inp_attn,
+    const llama_model             & model,
+    const int64_t                 n_embd_head,
+    const int                     il) {
+
+    const auto & layer = model.layers[il];
+    const auto & sl    = layer.switch_lora;
+
+    const int64_t n_head    = hparams.n_head(il);
+    const int64_t n_head_kv = hparams.n_head_kv(il);
+
+    ggml_tensor * qkv = ggml_mul_mat(ctx0, layer.wqkv, cur);
+    cb(qkv, "wqkv", il);
+
+    const int64_t n_embd_q  = n_embd_head * n_head;
+    const int64_t n_embd_kv = n_embd_head * n_head_kv;
+
+    // slice fused qkv into Q/K/V, made contiguous so LoRA deltas can be added
+    ggml_tensor * Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_q,  qkv->ne[1], qkv->nb[1], 0));
+    ggml_tensor * Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], n_embd_q*ggml_element_size(qkv)));
+    ggml_tensor * Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd_kv, qkv->ne[1], qkv->nb[1], (n_embd_q + n_embd_kv)*ggml_element_size(qkv)));
+
+    Qcur = ggml_add(ctx0, Qcur, build_switched_lora_delta(sl.a_q, sl.b_q, cur, adapter_ids));
+    Kcur = ggml_add(ctx0, Kcur, build_switched_lora_delta(sl.a_k, sl.b_k, cur, adapter_ids));
+    Vcur = ggml_add(ctx0, Vcur, build_switched_lora_delta(sl.a_v, sl.b_v, cur, adapter_ids));
+
+    Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head,    n_tokens);
+    Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
+    Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);
+
+    if (hparams.rope_finetuned) {
+        ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
+        Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors,
+                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+                ext_factor, attn_factor, beta_fast, beta_slow);
+        Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, rope_factors,
+                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+                ext_factor, attn_factor, beta_fast, beta_slow);
+    }
+    cb(Qcur, "Qcur", il);
+    cb(Kcur, "Kcur", il);
+    cb(Vcur, "Vcur", il);
+
+    const float kq_scale = hparams.f_attention_scale == 0.0f
+        ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
+
+    // wo = nullptr so build_attn returns concatenated heads; o-proj is switched below
+    ggml_tensor * attn = build_attn(inp_attn,
+            nullptr, nullptr, nullptr,
+            Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+    cb(attn, "attn_pre_o", il);
+
+    cur = build_switched_lora_mm(layer.wo, sl.a_o, sl.b_o, attn, adapter_ids);
+    cb(cur, "attn_out", il);
+    return cur;
+}
+
+ggml_tensor * llama_model_granite_switch::graph::build_layer_ffn(
+          ggml_tensor       * cur,
+          ggml_tensor       * inpSA,
+          ggml_tensor       * adapter_ids,
+    const llama_model       & model,
+    const int                 il) {
+
+    const auto & layer = model.layers[il];
+    const auto & sl    = layer.switch_lora;
+
+    if (hparams.f_residual_scale) {
+        cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
+    }
+    ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+    cb(ffn_inp, "ffn_inp", il);
+
+    cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
+    cb(cur, "ffn_norm", il);
+
+    ggml_tensor * g = build_switched_lora_mm(layer.ffn_gate, sl.a_gate, sl.b_gate, cur, adapter_ids);
+    ggml_tensor * u = build_switched_lora_mm(layer.ffn_up,   sl.a_up,   sl.b_up,   cur, adapter_ids);
+    g = ggml_silu(ctx0, g);
+    ggml_tensor * gu = ggml_mul(ctx0, g, u);
+    cur = build_switched_lora_mm(layer.ffn_down, sl.a_down, sl.b_down, gu, adapter_ids);
+    cb(cur, "ffn_out", il);
+
+    if (hparams.f_residual_scale) {
+        cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);
+    }
+    cur = ggml_add(ctx0, cur, ffn_inp);
+
+    cur = build_cvec(cur, il);
+    cb(cur, "l_out", il);
+
+    return cur;
+}

src/models/models.h

@@ -1461,6 +1461,10 @@ struct llama_model_nemotron_h_moe : public llama_model_nemotron_h {
 
     using graph = llama_model_nemotron_h::graph;
 
+    struct graph_mtp : public llm_graph_context {
+        graph_mtp(const llama_model & model, const llm_graph_params & params);
+    };
+
     std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
 };
 
@@ -1596,6 +1600,56 @@ struct llama_model_granite_moe : public llama_model_base {
 };
 
 
+struct llama_model_granite_switch : public llama_model_base {
+    llama_model_granite_switch(const struct llama_model_params & params) : llama_model_base(params) {}
+    void load_arch_hparams(llama_model_loader & ml) override;
+    void load_arch_tensors(llama_model_loader & ml) override;
+
+    uint32_t n_adapters    = 0;
+    uint32_t max_lora_rank = 0;
+    float    router_gain   = 15.0f;
+
+    std::unordered_map<llama_token, int32_t>     adapter_token_to_slot;
+    std::unordered_map<llama_token, llama_token> adapter_token_to_substitute;
+
+    struct graph : public llm_graph_context {
+        graph(const llama_model & model, const llm_graph_params & params);
+
+    private:
+        ggml_tensor * build_switched_lora_delta(
+                  ggml_tensor * lora_a,
+                  ggml_tensor * lora_b,
+                  ggml_tensor * cur,
+                  ggml_tensor * ids);
+
+        ggml_tensor * build_switched_lora_mm(
+                  ggml_tensor * w,
+                  ggml_tensor * lora_a,
+                  ggml_tensor * lora_b,
+                  ggml_tensor * cur,
+                  ggml_tensor * ids);
+
+        ggml_tensor * build_attention_layer(
+                  ggml_tensor             * cur,
+                  ggml_tensor             * inp_pos,
+                  ggml_tensor             * adapter_ids,
+                  llm_graph_input_attn_kv * inp_attn,
+            const llama_model             & model,
+            const int64_t                 n_embd_head,
+            const int                     il);
+
+        ggml_tensor * build_layer_ffn(
+                  ggml_tensor       * cur,
+                  ggml_tensor       * inpSA,
+                  ggml_tensor       * adapter_ids,
+            const llama_model       & model,
+            const int                 il);
+    };
+
+    std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
+
 struct llama_model_minicpm : public llama_model_base {
     llama_model_minicpm(const struct llama_model_params & params) : llama_model_base(params) {}
     void load_arch_hparams(llama_model_loader & ml) override;

src/models/nemotron-h-moe.cpp

@@ -1,6 +1,156 @@
 #include "models.h"
 
 std::unique_ptr<llm_graph_context> llama_model_nemotron_h_moe::build_arch_graph(const llm_graph_params & params) const {
+    if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
+        return std::make_unique<graph_mtp>(*this, params);
+    }
     return std::make_unique<graph>(*this, params);
 }
 
+// MTP draft head for Nemotron-H MoE
+llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)
+    : llm_graph_context(params) {
+    GGML_ASSERT(hparams.n_layer_nextn == 1 && "NEMOTRON_H_MOE MTP currently supports a single MTP block");
+
+    const int64_t n_embd_head = hparams.n_embd_head_v();
+    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+
+    const int il = hparams.n_layer();
+    const auto & layer = model.layers[il];
+
+    GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm);
+    GGML_ASSERT(layer.ffn_gate_inp);
+
+    // token embedding weights
+    ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;
+    GGML_ASSERT(tok_embd_w != nullptr && "NEMOTRON_H_MOE MTP requires token embeddings");
+
+    auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);
+
+    inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+    ggml_set_input(inp->tokens);
+
+    inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);
+    ggml_set_input(inp->embd);
+
+    ggml_tensor * tok_embd;
+    if (ubatch.token) {
+        tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);
+    } else {
+        tok_embd = inp->embd;
+    }
+    cb(tok_embd, "mtp_tok_embd", il);
+
+    inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
+    ggml_set_input(inp->h);
+    ggml_set_name(inp->h, "mtp_h_input");
+
+    ggml_tensor * h_embd = inp->h;
+
+    res->add_input(std::move(inp));
+
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    // attention fills KV over all tokens, but the MoE is position-wise: gather output rows before
+    // it to save FFN compute (unless unmasked embeddings_nextn needs the full-length hidden state)
+    const bool emit_h_nextn    = cparams.embeddings_nextn;
+    const bool crop_before_ffn = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked);
+
+    auto * inp_attn = build_attn_inp_kv();
+
+    ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
+    cb(h_norm, "mtp_hnorm", il);
+    ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
+    cb(e_norm, "mtp_enorm", il);
+
+    ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);
+    cb(concat, "mtp_concat", il);
+
+    ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);
+    cb(cur, "mtp_eh_proj", il);
+
+    // dense NoPE attention sub-layer (mtp.layers.0)
+    ggml_tensor * inpSA = cur;
+    cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+    cb(cur, "mtp_attn_norm", il);
+
+    {
+        auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);
+        const float kq_scale = hparams.f_attention_scale == 0.0f
+            ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
+        cur = build_attn(inp_attn, layer.wo, layer.wo_b, layer.wo_s,
+                Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+        cb(cur, "mtp_attn_out", il);
+    }
+
+    cur = ggml_add(ctx0, cur, inpSA);
+    cb(cur, "mtp_attn_residual", il);
+
+    // gather the output rows here so the MoE FFN below only runs on the positions we keep
+    if (crop_before_ffn) {
+        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+    }
+
+    // MoE FFN sub-layer (mtp.layers.1)
+    ggml_tensor * ffn_residual = cur;
+    cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);
+    cb(cur, "mtp_attn_post_norm", il);
+
+    {
+        ggml_tensor * router_logits = build_lora_mm(layer.ffn_gate_inp, cur);
+        cb(router_logits, "mtp_ffn_moe_logits", il);
+
+        ggml_tensor * moe_out =
+            build_moe_ffn(cur,
+                layer.ffn_gate_inp,
+                layer.ffn_up_exps,
+                nullptr, // no gate
+                layer.ffn_down_exps,
+                layer.ffn_exp_probs_b,
+                n_expert, n_expert_used,
+                LLM_FFN_RELU_SQR, hparams.expert_weights_norm,
+                hparams.expert_weights_scale,
+                LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID,
+                il,
+                router_logits, nullptr,
+                layer.ffn_up_exps_s,
+                nullptr, // no gate
+                layer.ffn_down_exps_s);
+        cb(moe_out, "mtp_ffn_moe_out", il);
+
+        ggml_tensor * ffn_shexp = build_ffn(cur,
+                layer.ffn_up_shexp,   NULL, layer.ffn_up_shexp_s,
+                NULL,                 NULL, NULL,
+                layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,
+                NULL,
+                LLM_FFN_RELU_SQR, LLM_FFN_PAR, il);
+        cb(ffn_shexp, "mtp_ffn_shexp", il);
+
+        cur = ggml_add(ctx0, moe_out, ffn_shexp);
+        cb(cur, "mtp_ffn_out", il);
+    }
+
+    cur = ggml_add(ctx0, cur, ffn_residual);
+    cb(cur, "mtp_post_ffn", il);
+
+    // final head norm: the MTP head has its own LayerNorm
+    GGML_ASSERT(layer.nextn.shared_head_norm && "NEMOTRON_H_MOE MTP: missing final head norm");
+    cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM, -1);
+
+    cb(cur, "h_nextn", -1);
+    res->t_h_nextn = cur;
+
+    if (!crop_before_ffn && inp_out_ids) {
+        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+    }
+
+    // LM head
+    ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;
+    ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;
+    GGML_ASSERT(head_w != nullptr && "NEMOTRON_H_MOE MTP requires an output projection");
+    cur = build_lora_mm(head_w, cur, head_s);
+    cb(cur, "result_output", -1);
+
+    res->t_logits = cur;
+    ggml_build_forward_expand(gf, cur);
+}

src/models/nemotron-h.cpp

@@ -7,13 +7,18 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) {
     ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank);
     ml.get_key(LLM_KV_SSM_GROUP_COUNT,    hparams.ssm_n_group);
 
+    // NextN/MTP: optional draft head appended as extra trailing block(s)
+    ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false);
+    GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all");
+
     // A layer is recurrent IFF the n_head_kv value is set to 0 and
-    // the n_ff value is set to 0
-    for (uint32_t i = 0; i < hparams.n_layer(); ++i) {
-        hparams.is_recr_impl[i] = (hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0);
+    // the n_ff value is set to 0. Appended MTP blocks are dense (non-recurrent)
+    for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
+        hparams.is_recr_impl[i] = i < hparams.n_layer() && hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0;
     }
 
     ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS,     hparams.f_norm_eps); // MTP head final_layernorm
 
     ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,        hparams.n_ff_exp,        false);
     ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp,      false);
@@ -30,9 +35,13 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) {
     }
 }
 
-void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) {
+void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) {
     LLAMA_LOAD_LOCALS;
 
+    const bool mtp_only    = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.attn_norm.weight") == nullptr;
+    const int  trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
+    const int  mtp_flags   = !ml.load_mtp ? TENSOR_SKIP : 0;
+
     // mamba2 Mixer SSM params
     // NOTE: int64_t for tensor dimensions
     const int64_t d_conv     = hparams.ssm_d_conv;
@@ -60,61 +69,94 @@ void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) {
         auto & layer = layers[i];
 
         // all blocks use the attn norm
-        layer.attn_norm  = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+        layer.attn_norm  = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, trunk_flags);
 
         if (hparams.is_recr(i)) {
             // ssm layers
-            layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, 0);
+            layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, trunk_flags);
 
-            layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, 0);
+            layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, trunk_flags);
             layer.ssm_conv1d_b = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "bias", i), {d_inner + 2*n_group*d_state}, TENSOR_NOT_REQUIRED);
 
-            layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, 0);
+            layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, trunk_flags);
 
             // no "weight" suffix for these
-            layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, 0);
-            layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, 0);
+            layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, trunk_flags);
+            layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, trunk_flags);
 
-            layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, 0);
+            layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, trunk_flags);
 
             // out_proj
-            layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, 0);
+            layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, trunk_flags);
         } else if (hparams.n_ff(i) == 0) {
             // attention layers (with optional bias)
             const int64_t n_head_i = hparams.n_head(i);
             const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i);
             const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i);
-            create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, 0);
-            layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, 0);
+            create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, trunk_flags);
+            layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, trunk_flags);
             layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);
         }  else {
             if (n_expert != 0) {
                 const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
                 const int64_t n_ff_shexp = hparams.n_ff_shexp;
 
-                layer.ffn_gate_inp    = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", i), { n_embd, n_expert}, 0);
-                layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert         }, 0);
+                layer.ffn_gate_inp    = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", i), { n_embd, n_expert}, trunk_flags);
+                layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert         }, trunk_flags);
 
                 // MoE branch
                 layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, TENSOR_NOT_REQUIRED);
                 layer.ffn_latent_up   = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP,   "weight", i), {moe_n_embd, n_embd}, TENSOR_NOT_REQUIRED);
 
-                layer.ffn_down_exps   = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp,   moe_n_embd, n_expert}, 0);
-                layer.ffn_up_exps     = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {moe_n_embd, n_ff_exp, n_expert}, 0);
+                layer.ffn_down_exps   = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp,   moe_n_embd, n_expert}, trunk_flags);
+                layer.ffn_up_exps     = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {moe_n_embd, n_ff_exp, n_expert}, trunk_flags);
 
                 // Shared expert branch
-                layer.ffn_down_shexp  = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0);
-                layer.ffn_up_shexp    = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd, n_ff_shexp}, 0);
+                layer.ffn_down_shexp  = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, trunk_flags);
+                layer.ffn_up_shexp    = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd, n_ff_shexp}, trunk_flags);
 
             } else {
                 // mlp layers
-                layer.ffn_down   = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  hparams.n_ff(i), n_embd}, 0);
-                layer.ffn_up     = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   hparams.n_ff(i)}, 0);
+                layer.ffn_down   = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  hparams.n_ff(i), n_embd}, trunk_flags);
+                layer.ffn_up     = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   hparams.n_ff(i)}, trunk_flags);
                 layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias",   i), {n_embd}, TENSOR_NOT_REQUIRED);
                 layer.ffn_up_b   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "bias",   i), {hparams.n_ff(i)}, TENSOR_NOT_REQUIRED);
             }
         }
     }
+
+    // NextN/MTP draft head: each predict layer folds an attention sub-layer and a MoE
+    // sub-layer into a single trailing block
+    for (int i = n_layer; i < n_layer_all; ++i) {
+        auto & layer = layers[i];
+
+        const int64_t n_head_i       = hparams.n_head(i);
+        const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i);
+        const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i);
+        const int64_t n_ff_exp       = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used;
+        const int64_t n_ff_shexp     = hparams.n_ff_shexp;
+
+        // NextN input-fusion tensors
+        layer.nextn.enorm            = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM,            "weight", i), {n_embd}, mtp_flags);
+        layer.nextn.hnorm            = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM,            "weight", i), {n_embd}, mtp_flags);
+        layer.nextn.eh_proj          = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ,          "weight", i), {2*n_embd, n_embd}, mtp_flags);
+        layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags);
+
+        // attention sub-layer
+        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, mtp_flags);
+        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, mtp_flags);
+        layer.wo   = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, mtp_flags);
+        layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias",   i), {n_embd}, mtp_flags | TENSOR_NOT_REQUIRED);
+
+        // MoE sub-layer
+        layer.attn_post_norm  = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM,  "weight", i), {n_embd}, mtp_flags);
+        layer.ffn_gate_inp    = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,    "weight", i), {n_embd, n_expert}, mtp_flags);
+        layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias",   i), {n_expert}, mtp_flags);
+        layer.ffn_down_exps   = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS,   "weight", i), {n_ff_exp,   moe_n_embd, n_expert}, mtp_flags);
+        layer.ffn_up_exps     = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,     "weight", i), {moe_n_embd, n_ff_exp,   n_expert}, mtp_flags);
+        layer.ffn_down_shexp  = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP,  "weight", i), {n_ff_shexp, n_embd}, mtp_flags);
+        layer.ffn_up_shexp    = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,    "weight", i), {n_embd, n_ff_shexp}, mtp_flags);
+    }
 }
 
 std::unique_ptr<llm_graph_context> llama_model_nemotron_h::build_arch_graph(const llm_graph_params & params) const {
@@ -153,7 +195,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
             cur = build_ffn_layer(cur, model, il);
         }
 
-        if (il == n_layer - 1 && inp_out_ids) {
+        if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
             cur   = ggml_get_rows(ctx0, cur, inp_out_ids);
             inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
         }
@@ -170,6 +212,14 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
 
     cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
 
+    // seed for the MTP/NextN draft head
+    cb(cur, "h_nextn", -1);
+    res->t_h_nextn = cur;
+
+    if (!cparams.embeddings_nextn_masked && inp_out_ids) {
+        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+    }
+
     cb(cur, "result_norm", -1);
     res->t_embd = cur;
 

tests/test-backend-ops.cpp

@@ -2584,6 +2584,7 @@ struct test_rms_norm_mul_rope : public test_case {
     const float eps;
     const bool multi_add; // test a sequence of adds feeding into rms_norm
     const bool set_rows;
+    const bool broadcast; // multiply by a 1D [ne0] weight, as model norm weights are
     int mode;
 
     std::string op_desc(ggml_tensor * t) override {
@@ -2594,12 +2595,12 @@ struct test_rms_norm_mul_rope : public test_case {
     bool run_whole_graph() override { return true; }
 
     std::string vars() override {
-        return VARS_TO_STR5(ne, eps, multi_add, set_rows, mode);
+        return VARS_TO_STR6(ne, eps, multi_add, set_rows, broadcast, mode);
     }
 
     test_rms_norm_mul_rope(std::array<int64_t, 4> ne, float eps = 1e-6f, bool multi_add = false,
-                           bool set_rows = false, int mode = GGML_ROPE_TYPE_NORMAL)
-        : ne(ne), eps(eps), multi_add(multi_add), set_rows(set_rows), mode(mode) {}
+                           bool set_rows = false, bool broadcast = false, int mode = GGML_ROPE_TYPE_NORMAL)
+        : ne(ne), eps(eps), multi_add(multi_add), set_rows(set_rows), broadcast(broadcast), mode(mode) {}
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], 1);
@@ -2610,7 +2611,9 @@ struct test_rms_norm_mul_rope : public test_case {
             a = ggml_add(ctx, ggml_add(ctx, a, b), c);
         }
 
-        a = ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), b);
+        ggml_tensor * w = broadcast ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[0]) : b;
+
+        a = ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), w);
 
         ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne[2]);
 
@@ -8756,16 +8759,18 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
 
     for (auto multi_add : {false, true}) {
         for (auto set_rows : {false, true}) {
-            for (auto rope : {GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX}) {
-                test_cases.emplace_back(new test_rms_norm_mul_rope({768, 1, 1, 1}, 1e-6f, multi_add, set_rows, rope));
-                test_cases.emplace_back(new test_rms_norm_mul_rope({768, 3, 1, 1}, 1e-6f, multi_add, set_rows, rope));
-                test_cases.emplace_back(new test_rms_norm_mul_rope({768, 3, 5, 1}, 1e-6f, multi_add, set_rows, rope));
-                test_cases.emplace_back(new test_rms_norm_mul_rope({128, 32, 2, 1}, 1e-6f, multi_add, set_rows, rope));
-                test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 2, 1}, 1e-6f, multi_add, set_rows, rope));
-                test_cases.emplace_back(new test_rms_norm_mul_rope({128, 32, 50, 1}, 1e-6f, multi_add, set_rows, rope));
-                test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 50, 1}, 1e-6f, multi_add, set_rows, rope));
-                test_cases.emplace_back(new test_rms_norm_mul_rope({8192, 2, 2, 1}, 1e-6f, multi_add, set_rows, rope));
-                test_cases.emplace_back(new test_rms_norm_mul_rope({8192, 2, 2, 1}, 1e-6f, multi_add, set_rows, rope));
+            for (auto broadcast : {false, true}) {
+                for (auto rope : {GGML_ROPE_TYPE_NORMAL, GGML_ROPE_TYPE_NEOX}) {
+                    test_cases.emplace_back(new test_rms_norm_mul_rope({768, 1, 1, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
+                    test_cases.emplace_back(new test_rms_norm_mul_rope({768, 3, 1, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
+                    test_cases.emplace_back(new test_rms_norm_mul_rope({768, 3, 5, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
+                    test_cases.emplace_back(new test_rms_norm_mul_rope({128, 32, 2, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
+                    test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 2, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
+                    test_cases.emplace_back(new test_rms_norm_mul_rope({128, 32, 50, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
+                    test_cases.emplace_back(new test_rms_norm_mul_rope({128, 4, 50, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
+                    test_cases.emplace_back(new test_rms_norm_mul_rope({8192, 2, 2, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
+                    test_cases.emplace_back(new test_rms_norm_mul_rope({8192, 2, 2, 1}, 1e-6f, multi_add, set_rows, broadcast, rope));
+                }
             }
         }
     }

tests/test-llama-archs.cpp

@@ -217,6 +217,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
 
     if (moe) {
         ms.add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, n_ff);
+        ms.add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, n_ff / 2);  // distinct from n_ff so a saver key-clobber surfaces on reload
         ms.add_kv(LLM_KV_INTERLEAVE_MOE_LAYER_STEP,  uint32_t(2));
         ms.add_kv(LLM_KV_EXPERT_COUNT,               uint32_t(2));
         ms.add_kv(LLM_KV_EXPERT_USED_COUNT,          uint32_t(1));
@@ -410,6 +411,9 @@ static bool arch_supported(const llm_arch arch) {
     if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) {
         return false; // FIXME @ngxson
     }
+    if (arch == LLM_ARCH_GRANITE_SWITCH) {
+        return false; // FIXME adapter fixture
+    }
     if (arch == LLM_ARCH_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) {
         return false; // FIXME Embedding (?) models produce inconsistent results.
     }

tools/server/server-tools.cpp

@@ -1669,8 +1669,13 @@ struct server_tool_get_info : server_tool {
 
         std::string cwd = json_value(params, "cwd", std::string());
         if (cwd.empty()) {
-            std::error_code ec;
-            cwd = path_to_utf8(fs::current_path(ec));
+            if (json_value(params, "runtime", std::string()).empty()) {
+                std::error_code ec;
+                cwd = path_to_utf8(fs::current_path(ec));
+            } else {
+                auto pwd = io->run({"pwd"}, SERVER_TOOL_GET_INFO_MAX_OUTPUT, SERVER_TOOL_GET_INFO_TIMEOUT);
+                cwd = pwd.exit_code == 0 && !pwd.timed_out ? string_strip(pwd.output) : "unknown";
+            }
         }
 
         return {

tools/server/tests/unit/test_tools_builtin.py

@@ -13,6 +13,9 @@ PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..
 # marker for the grep_search test to find in this file
 GREP_MARKER = "llama_cpp_test_tools_builtin_marker_grep_search"
 
+# image the container runtime tests run their shell in
+DOCKER_IMAGE = "busybox"
+
 
 @pytest.fixture(autouse=True)
 def create_server():
@@ -149,14 +152,16 @@ def test_tools_builtin_cwd_header():
 
 
 def _docker_unavailable_reason() -> str | None:
-    """None if docker can be used to run a container, otherwise the reason it can't."""
+    """None if docker can run the image these tests use, otherwise the reason it can't."""
     docker_bin = shutil.which("docker")
     if docker_bin is None:
         return "docker is not installed"
     try:
-        subprocess.run([docker_bin, "info"], capture_output=True, timeout=5, check=True)
+        # a daemon that answers `docker info` still cannot run a linux image when it serves
+        # windows containers, so probe the image itself, which also pulls it before the tests
+        subprocess.run([docker_bin, "run", "--rm", DOCKER_IMAGE, "true"], capture_output=True, timeout=60, check=True)
     except Exception as e:
-        return f"docker daemon is not usable: {e}"
+        return f"docker cannot run {DOCKER_IMAGE}: {e}"
     return None
 
 
@@ -167,7 +172,7 @@ def docker_container():
         pytest.skip(reason)  # ty: ignore[too-many-positional-arguments, invalid-argument-type]
 
     proc = subprocess.run(
-        ["docker", "run", "-d", "--rm", "busybox", "sleep", "300"],
+        ["docker", "run", "-d", "--rm", DOCKER_IMAGE, "sleep", "300"],
         capture_output=True, text=True,
     )
     if proc.returncode != 0:
@@ -214,7 +219,7 @@ def test_tools_builtin_docker_runtime_cleans_up_spawned_container():
         pytest.skip(reason)  # ty: ignore[too-many-positional-arguments, invalid-argument-type]
 
     global server
-    server.server_tools_runtime = "docker:busybox"
+    server.server_tools_runtime = f"docker:{DOCKER_IMAGE}"
     server.start()
 
     # exec_shell_command runs inside the container spawned for --tools-runtime; docker sets

tools/ui/eslint.config.js

@@ -1,14 +1,15 @@
 // For more info, see https://github.com/storybookjs/eslint-plugin-storybook#configuration-flat-config-format
-import storybook from 'eslint-plugin-storybook';
-
-import prettier from 'eslint-config-prettier';
+import svelteConfig from './svelte.config.js';
 import { includeIgnoreFile } from '@eslint/compat';
 import js from '@eslint/js';
+import prettier from 'eslint-config-prettier';
+import perfectionist from 'eslint-plugin-perfectionist';
+import simpleImportSort from 'eslint-plugin-simple-import-sort';
+import storybook from 'eslint-plugin-storybook';
 import svelte from 'eslint-plugin-svelte';
 import globals from 'globals';
 import { fileURLToPath } from 'node:url';
 import ts from 'typescript-eslint';
-import svelteConfig from './svelte.config.js';
 
 const gitignorePath = fileURLToPath(new URL('./.gitignore', import.meta.url));
 
@@ -21,32 +22,67 @@ export default ts.config(
 	...svelte.configs.prettier,
 	{
 		languageOptions: { globals: { ...globals.browser, ...globals.node } },
+		plugins: { perfectionist, 'simple-import-sort': simpleImportSort },
 		rules: {
-			// typescript-eslint strongly recommend that you do not use the no-undef lint rule on TypeScript projects.
-			// see: https://typescript-eslint.io/troubleshooting/faqs/eslint/#i-get-errors-from-the-no-undef-rule-about-global-variables-not-being-defined-even-though-there-are-no-typescript-errors
-			'no-undef': 'off',
-			'svelte/no-at-html-tags': 'off',
-			// This app uses hash-based routing (#/) where resolve() from $app/paths does not apply
-			'svelte/no-navigation-without-resolve': 'off',
-
 			// Snippet bodies often ignore one or more of the parent's params
 			// (e.g. `{#snippet children(_meta, ctx)}` when only ctx is read).
 			'@typescript-eslint/no-unused-vars': [
 				'error',
 				{ argsIgnorePattern: '^_', varsIgnorePattern: '^_' }
 			],
-
 			// Enforce empty line at end of file
-			'eol-last': 'error'
+			'eol-last': 'error',
+			// typescript-eslint strongly recommend that you do not use the no-undef lint rule on TypeScript projects.
+			// see: https://typescript-eslint.io/troubleshooting/faqs/eslint/#i-get-errors-from-the-no-undef-rule-about-global-variables-not-being-defined-even-though-there-are-no-typescript-errors
+			'no-undef': 'off',
+
+			'padding-line-between-statements': [
+				'error',
+				// Blank line between function/class declarations.
+				{ blankLine: 'always', next: ['function', 'class'], prev: ['function', 'class'] },
+				// Blank line around if blocks (if/else and else if stay one statement).
+				{ blankLine: 'always', next: '*', prev: 'if' },
+				{ blankLine: 'always', next: 'if', prev: '*' },
+				// Blank line after the last declaration in a group. Because the 'never'
+				// rules below are scoped per declaration kind, a const group and a let
+				// group get separated by a blank line, while same-kind declarations stay
+				// together.
+				{ blankLine: 'always', next: '*', prev: ['const', 'let', 'var'] },
+				// No blank line between consecutive declarations of the same kind (kept
+				// last so each takes precedence over the always rule above for matching
+				// declaration pairs).
+				{ blankLine: 'never', next: 'const', prev: 'const' },
+				{ blankLine: 'never', next: 'let', prev: 'let' },
+				{ blankLine: 'never', next: 'var', prev: 'var' },
+				// Blank line before a statement that follows another statement in the block
+				// (works for return/throw/break/continue). A blank line for a terminal
+				// statement that opens a block body can't be enforced here: Prettier removes
+				// the leading blank line of a block, so the two formatters would fight.
+				{ blankLine: 'always', next: ['return', 'throw', 'break', 'continue'], prev: '*' }
+			],
+
+			'perfectionist/sort-objects': ['error', { type: 'natural' }],
+
+			// Alphabetical order for variable declarations and object keys
+			'perfectionist/sort-variable-declarations': ['error', { type: 'natural' }],
+
+			// Sort imports alphabetically by module path, and sort named members within
+			// each statement. A single catch-all group keeps the list flat (no blank-line
+			// grouping); Prettier normalizes comma spacing afterwards.
+			'simple-import-sort/imports': ['error', { groups: [['.*']] }],
+			'svelte/no-at-html-tags': 'off',
+
+			// This app uses hash-based routing (#/) where resolve() from $app/paths does not apply
+			'svelte/no-navigation-without-resolve': 'off'
 		}
 	},
 	{
 		files: ['**/*.svelte', '**/*.svelte.ts', '**/*.svelte.js'],
 		languageOptions: {
 			parserOptions: {
-				projectService: true,
 				extraFileExtensions: ['.svelte'],
 				parser: ts.parser,
+				projectService: true,
 				svelteConfig
 			}
 		}

tools/ui/package-lock.json

@@ -39,6 +39,8 @@
 				"dompurify": "3.4.13",
 				"eslint": "9.39.4",
 				"eslint-config-prettier": "10.1.8",
+				"eslint-plugin-perfectionist": "^5.10.1",
+				"eslint-plugin-simple-import-sort": "^14.0.0",
 				"eslint-plugin-storybook": "10.5.6",
 				"eslint-plugin-svelte": "3.19.0",
 				"fflate": "0.8.3",
@@ -9281,6 +9283,226 @@
 				"eslint": ">=7.0.0"
 			}
 		},
+		"node_modules/eslint-plugin-perfectionist": {
+			"version": "5.10.1",
+			"resolved": "https://registry.npmjs.org/eslint-plugin-perfectionist/-/eslint-plugin-perfectionist-5.10.1.tgz",
+			"integrity": "sha512-Kprsp9Us0GqAesYaAIzUViw57xYp5WBqzXrcE0Mtww++E5fexWXYBipMuuD7yvyH4vvpBH0+oJ+OMAmZ0oYXkw==",
+			"dev": true,
+			"license": "MIT",
+			"dependencies": {
+				"@typescript-eslint/utils": "^8.65.0",
+				"natural-orderby": "^5.0.0"
+			},
+			"engines": {
+				"node": "^20.0.0 || >=22.0.0"
+			},
+			"peerDependencies": {
+				"eslint": "^8.45.0 || ^9.0.0 || ^10.0.0"
+			}
+		},
+		"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/project-service": {
+			"version": "8.66.0",
+			"resolved": "https://registry.npmjs.org/@typescript-eslint/project-service/-/project-service-8.66.0.tgz",
+			"integrity": "sha512-7MthGPTt4BP69lSryqpqq8HQqxuzynssckL/jyDyk3+TNMQ3y2jFWkptCrktWvBrP+EH787Nl5N5Qpw7WZg+5g==",
+			"dev": true,
+			"license": "MIT",
+			"dependencies": {
+				"@typescript-eslint/tsconfig-utils": "^8.66.0",
+				"@typescript-eslint/types": "^8.66.0",
+				"debug": "^4.4.3"
+			},
+			"engines": {
+				"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
+			},
+			"funding": {
+				"type": "opencollective",
+				"url": "https://opencollective.com/typescript-eslint"
+			},
+			"peerDependencies": {
+				"typescript": ">=4.8.4 <6.1.0"
+			}
+		},
+		"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/scope-manager": {
+			"version": "8.66.0",
+			"resolved": "https://registry.npmjs.org/@typescript-eslint/scope-manager/-/scope-manager-8.66.0.tgz",
+			"integrity": "sha512-8TGcH25j9zqJ/IULB/ppyhRvxA8QYfFEZ7nfbg6/BN9spDgb8fPWQXlE5l8TWBL50EtUx007uZ1o9VOwrq2/9g==",
+			"dev": true,
+			"license": "MIT",
+			"dependencies": {
+				"@typescript-eslint/types": "8.66.0",
+				"@typescript-eslint/visitor-keys": "8.66.0"
+			},
+			"engines": {
+				"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
+			},
+			"funding": {
+				"type": "opencollective",
+				"url": "https://opencollective.com/typescript-eslint"
+			}
+		},
+		"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/tsconfig-utils": {
+			"version": "8.66.0",
+			"resolved": "https://registry.npmjs.org/@typescript-eslint/tsconfig-utils/-/tsconfig-utils-8.66.0.tgz",
+			"integrity": "sha512-9D5gLYZG4rOjcoag8MQ/fWI8WqA9wcPDyOGyWtWFhvM1lHRbliqUSPIY5J3zqCU1tvSwzXxnnjhQhz5Ne7mJ4g==",
+			"dev": true,
+			"license": "MIT",
+			"engines": {
+				"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
+			},
+			"funding": {
+				"type": "opencollective",
+				"url": "https://opencollective.com/typescript-eslint"
+			},
+			"peerDependencies": {
+				"typescript": ">=4.8.4 <6.1.0"
+			}
+		},
+		"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/types": {
+			"version": "8.66.0",
+			"resolved": "https://registry.npmjs.org/@typescript-eslint/types/-/types-8.66.0.tgz",
+			"integrity": "sha512-H6gcYaSDOyvL3AD/jHUtUFo2jqGgn/F6nuyuZSu0QTesxL+cP4dQoIMrODRofuJC09g64+WgZ6tE19Y1N2YIFQ==",
+			"dev": true,
+			"license": "MIT",
+			"engines": {
+				"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
+			},
+			"funding": {
+				"type": "opencollective",
+				"url": "https://opencollective.com/typescript-eslint"
+			}
+		},
+		"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/typescript-estree": {
+			"version": "8.66.0",
+			"resolved": "https://registry.npmjs.org/@typescript-eslint/typescript-estree/-/typescript-estree-8.66.0.tgz",
+			"integrity": "sha512-8/x4INiiQb10jGgXYD7116/zQ+OL84ZIFn0za68wwFHCanT/VLbBEroWht8RV8fn0/ZCAoazHLQgwUC0UQcDfg==",
+			"dev": true,
+			"license": "MIT",
+			"dependencies": {
+				"@typescript-eslint/project-service": "8.66.0",
+				"@typescript-eslint/tsconfig-utils": "8.66.0",
+				"@typescript-eslint/types": "8.66.0",
+				"@typescript-eslint/visitor-keys": "8.66.0",
+				"debug": "^4.4.3",
+				"minimatch": "^10.2.2",
+				"semver": "^7.7.3",
+				"tinyglobby": "^0.2.15",
+				"ts-api-utils": "^2.5.0"
+			},
+			"engines": {
+				"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
+			},
+			"funding": {
+				"type": "opencollective",
+				"url": "https://opencollective.com/typescript-eslint"
+			},
+			"peerDependencies": {
+				"typescript": ">=4.8.4 <6.1.0"
+			}
+		},
+		"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/utils": {
+			"version": "8.66.0",
+			"resolved": "https://registry.npmjs.org/@typescript-eslint/utils/-/utils-8.66.0.tgz",
+			"integrity": "sha512-jasearZPolBw5NJNYGMwxzHMF83niVWmMU1VdHzG1CyfI2VS7f7nZltnKtHcg20hW+7Uo5GfK4MeDPoU3qI8EA==",
+			"dev": true,
+			"license": "MIT",
+			"dependencies": {
+				"@eslint-community/eslint-utils": "^4.9.1",
+				"@typescript-eslint/scope-manager": "8.66.0",
+				"@typescript-eslint/types": "8.66.0",
+				"@typescript-eslint/typescript-estree": "8.66.0"
+			},
+			"engines": {
+				"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
+			},
+			"funding": {
+				"type": "opencollective",
+				"url": "https://opencollective.com/typescript-eslint"
+			},
+			"peerDependencies": {
+				"eslint": "^8.57.0 || ^9.0.0 || ^10.0.0",
+				"typescript": ">=4.8.4 <6.1.0"
+			}
+		},
+		"node_modules/eslint-plugin-perfectionist/node_modules/@typescript-eslint/visitor-keys": {
+			"version": "8.66.0",
+			"resolved": "https://registry.npmjs.org/@typescript-eslint/visitor-keys/-/visitor-keys-8.66.0.tgz",
+			"integrity": "sha512-dkKR8q+lKciskj1Y3vthHktl+3cMLWGyVUP23bRiPZ5O9BRT++4EqDDV+TVeIKBL1VXVEqrJlz8MYbcnvJcAlg==",
+			"dev": true,
+			"license": "MIT",
+			"dependencies": {
+				"@typescript-eslint/types": "8.66.0",
+				"eslint-visitor-keys": "^5.0.0"
+			},
+			"engines": {
+				"node": "^18.18.0 || ^20.9.0 || >=21.1.0"
+			},
+			"funding": {
+				"type": "opencollective",
+				"url": "https://opencollective.com/typescript-eslint"
+			}
+		},
+		"node_modules/eslint-plugin-perfectionist/node_modules/balanced-match": {
+			"version": "4.0.4",
+			"resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-4.0.4.tgz",
+			"integrity": "sha512-BLrgEcRTwX2o6gGxGOCNyMvGSp35YofuYzw9h1IMTRmKqttAZZVU67bdb9Pr2vUHA8+j3i2tJfjO6C6+4myGTA==",
+			"dev": true,
+			"license": "MIT",
+			"engines": {
+				"node": "18 || 20 || >=22"
+			}
+		},
+		"node_modules/eslint-plugin-perfectionist/node_modules/brace-expansion": {
+			"version": "5.0.9",
+			"resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.9.tgz",
+			"integrity": "sha512-ScQ4IuvIEF1TMlP7Zt+vjJ//9zlPb2SDcxWxM3bk8s6t6GGdJ7KO1dCcTidOPJKePW30LE/2cT7wCyPho9/Wxg==",
+			"dev": true,
+			"license": "MIT",
+			"dependencies": {
+				"balanced-match": "^4.0.2"
+			},
+			"engines": {
+				"node": "20 || >=22"
+			}
+		},
+		"node_modules/eslint-plugin-perfectionist/node_modules/eslint-visitor-keys": {
+			"version": "5.0.1",
+			"resolved": "https://registry.npmjs.org/eslint-visitor-keys/-/eslint-visitor-keys-5.0.1.tgz",
+			"integrity": "sha512-tD40eHxA35h0PEIZNeIjkHoDR4YjjJp34biM0mDvplBe//mB+IHCqHDGV7pxF+7MklTvighcCPPZC7ynWyjdTA==",
+			"dev": true,
+			"license": "Apache-2.0",
+			"engines": {
+				"node": "^20.19.0 || ^22.13.0 || >=24"
+			},
+			"funding": {
+				"url": "https://opencollective.com/eslint"
+			}
+		},
+		"node_modules/eslint-plugin-perfectionist/node_modules/minimatch": {
+			"version": "10.2.6",
+			"resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.2.6.tgz",
+			"integrity": "sha512-vpLQEs+VLCr1nU0BXS07maYoFwlDAH0gngQuuttxIwutDFEMHq2blX+8vpgxDdK3J1PwjCJiep77OitTZ4Ll1A==",
+			"dev": true,
+			"license": "BlueOak-1.0.0",
+			"dependencies": {
+				"brace-expansion": "^5.0.8"
+			},
+			"engines": {
+				"node": "18 || 20 || >=22"
+			},
+			"funding": {
+				"url": "https://github.com/sponsors/isaacs"
+			}
+		},
+		"node_modules/eslint-plugin-simple-import-sort": {
+			"version": "14.0.0",
+			"resolved": "https://registry.npmjs.org/eslint-plugin-simple-import-sort/-/eslint-plugin-simple-import-sort-14.0.0.tgz",
+			"integrity": "sha512-NUJO0+XFCkk+o5EsAJruTgnfMEpeWrPWeJS15UVF60GgXmqz1BJ9/3hzlvG7lkL8Bubzos5cCLptThbFfPnSMQ==",
+			"dev": true,
+			"license": "MIT",
+			"peerDependencies": {
+				"eslint": ">=5.0.0"
+			}
+		},
 		"node_modules/eslint-plugin-storybook": {
 			"version": "10.5.6",
 			"resolved": "https://registry.npmjs.org/eslint-plugin-storybook/-/eslint-plugin-storybook-10.5.6.tgz",
@@ -13196,6 +13418,16 @@
 			"dev": true,
 			"license": "MIT"
 		},
+		"node_modules/natural-orderby": {
+			"version": "5.0.0",
+			"resolved": "https://registry.npmjs.org/natural-orderby/-/natural-orderby-5.0.0.tgz",
+			"integrity": "sha512-kKHJhxwpR/Okycz4HhQKKlhWe4ASEfPgkSWNmKFHd7+ezuQlxkA5cM3+XkBPvm1gmHen3w53qsYAv+8GwRrBlg==",
+			"dev": true,
+			"license": "MIT",
+			"engines": {
+				"node": ">=18"
+			}
+		},
 		"node_modules/negotiator": {
 			"version": "1.0.0",
 			"resolved": "https://registry.npmjs.org/negotiator/-/negotiator-1.0.0.tgz",

tools/ui/package.json

@@ -12,7 +12,7 @@
 		"check": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json",
 		"check:watch": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json --watch",
 		"reset": "rm -rf .svelte-kit node_modules",
-		"format": "prettier --write .",
+		"format": "eslint --fix . && prettier --write .",
 		"lint": "prettier --check . && eslint .",
 		"test": "npm run test:ui -- --run && npm run test:client -- --run && npm run test:unit -- --run && npm run test:e2e",
 		"test:e2e": "playwright test",
@@ -36,6 +36,7 @@
 		"@playwright/test": "1.56.1",
 		"@storybook/addon-a11y": "10.5.6",
 		"@storybook/addon-docs": "10.5.6",
+		"@storybook/addon-mcp": "0.7.0",
 		"@storybook/addon-svelte-csf": "5.1.2",
 		"@storybook/addon-vitest": "10.5.6",
 		"@storybook/sveltekit": "10.5.6",
@@ -57,6 +58,8 @@
 		"dompurify": "3.4.13",
 		"eslint": "9.39.4",
 		"eslint-config-prettier": "10.1.8",
+		"eslint-plugin-perfectionist": "^5.10.1",
+		"eslint-plugin-simple-import-sort": "^14.0.0",
 		"eslint-plugin-storybook": "10.5.6",
 		"eslint-plugin-svelte": "3.19.0",
 		"fflate": "0.8.3",
@@ -99,8 +102,7 @@
 		"vite-plugin-devtools-json": "0.2.1",
 		"vitest": "4.1.10",
 		"vitest-browser-svelte": "2.1.1",
-		"workbox-window": "7.4.1",
-		"@storybook/addon-mcp": "0.7.0"
+		"workbox-window": "7.4.1"
 	},
 	"overrides": {
 		"cookie": "1.1.1",

tools/ui/playwright.config.ts

@@ -1,31 +1,31 @@
 import { defineConfig, devices } from '@playwright/test';
 
 export default defineConfig({
-	testDir: 'tests/e2e',
-	testMatch: ['**/*.e2e.ts'],
-	timeout: 30000,
 	expect: {
 		timeout: 5000
 	},
-	fullyParallel: true,
 	forbidOnly: !!process.env.CI,
-	retries: process.env.CI ? 2 : 0,
-	workers: process.env.CI ? 1 : undefined,
-	reporter: 'line',
-	use: {
-		baseURL: 'http://localhost:8181',
-		trace: 'on-first-retry'
-	},
+	fullyParallel: true,
 	projects: [
 		{
 			name: 'chromium',
 			use: { ...devices['Desktop Chrome'] }
 		}
 	],
+	reporter: 'line',
+	retries: process.env.CI ? 2 : 0,
+	testDir: 'tests/e2e',
+	testMatch: ['**/*.e2e.ts'],
+	timeout: 30000,
+	use: {
+		baseURL: 'http://localhost:8181',
+		trace: 'on-first-retry'
+	},
 	webServer: {
 		command: 'npm run build && npx http-server ./dist -p 8181',
 		port: 8181,
-		timeout: 120000,
-		reuseExistingServer: !process.env.CI
-	}
+		reuseExistingServer: !process.env.CI,
+		timeout: 120000
+	},
+	workers: process.env.CI ? 1 : undefined
 });

tools/ui/pwa-assets-dark.config.ts

@@ -1,6 +1,6 @@
-import { defineConfig } from '@vite-pwa/assets-generator/config';
-import { FAVICON_COLORS, PWA_ASSET_GENERATOR } from './src/lib/constants/pwa';
 import { writeThemeFavicons } from './scripts/favicon-colorize';
+import { FAVICON_COLORS, PWA_ASSET_GENERATOR } from './src/lib/constants/pwa';
+import { defineConfig } from '@vite-pwa/assets-generator/config';
 
 writeThemeFavicons(FAVICON_COLORS.LIGHT, FAVICON_COLORS.DARK, {
 	padding: PWA_ASSET_GENERATOR.FAVICON_PADDING
@@ -10,18 +10,18 @@ export default defineConfig({
 	headLinkOptions: {
 		preset: '2023'
 	},
+	images: ['static/favicon-dark.svg'],
 	preset: {
-		transparent: {
-			sizes: [],
-			favicons: [[48, 'favicon-dark.ico']],
-			padding: PWA_ASSET_GENERATOR.FAVICON_PADDING
+		apple: {
+			sizes: []
 		},
 		maskable: {
 			sizes: []
 		},
-		apple: {
+		transparent: {
+			favicons: [[48, 'favicon-dark.ico']],
+			padding: PWA_ASSET_GENERATOR.FAVICON_PADDING,
 			sizes: []
 		}
-	},
-	images: ['static/favicon-dark.svg']
+	}
 });

tools/ui/pwa-assets.config.ts

@@ -1,3 +1,11 @@
+import { writeThemeFavicons } from './scripts/favicon-colorize';
+import {
+	FAVICON_COLORS,
+	PWA_ASSET_GENERATOR,
+	PWA_GENERATOR_DEVICES,
+	THEME_COLORS
+} from './src/lib/constants/pwa';
+import { SplashOrientation } from './src/lib/enums/splash.enums';
 import {
 	combinePresetAndAppleSplashScreens,
 	defineConfig,
@@ -5,14 +13,6 @@ import {
 } from '@vite-pwa/assets-generator/config';
 import { readFileSync } from 'node:fs';
 import { resolve } from 'node:path';
-import {
-	THEME_COLORS,
-	PWA_GENERATOR_DEVICES,
-	PWA_ASSET_GENERATOR,
-	FAVICON_COLORS
-} from './src/lib/constants/pwa';
-import { SplashOrientation } from './src/lib/enums/splash.enums';
-import { writeThemeFavicons } from './scripts/favicon-colorize';
 
 writeThemeFavicons(FAVICON_COLORS.LIGHT, FAVICON_COLORS.DARK, {
 	padding: PWA_ASSET_GENERATOR.FAVICON_PADDING
@@ -22,6 +22,7 @@ export default defineConfig({
 	headLinkOptions: {
 		preset: PWA_ASSET_GENERATOR.LINK_PRESET
 	},
+	images: ['static/favicon.svg'],
 	preset: combinePresetAndAppleSplashScreens(
 		{
 			...minimal2023Preset,
@@ -32,37 +33,37 @@ export default defineConfig({
 			}
 		},
 		{
-			padding: PWA_ASSET_GENERATOR.SPLASH_PADDING,
-			resizeOptions: {
-				background: THEME_COLORS.BACKGROUND_LIGHT,
-				fit: PWA_ASSET_GENERATOR.FIT_MODE
-			},
-			darkResizeOptions: {
-				background: THEME_COLORS.BACKGROUND_DARK,
-				fit: PWA_ASSET_GENERATOR.FIT_MODE
-			},
 			darkImageResolver: async (imageName: string) => {
 				if (imageName.endsWith('favicon.svg')) {
 					return readFileSync(resolve('static/favicon-dark.svg'));
 				}
 			},
+			darkResizeOptions: {
+				background: THEME_COLORS.BACKGROUND_DARK,
+				fit: PWA_ASSET_GENERATOR.FIT_MODE
+			},
 			linkMediaOptions: {
-				log: true,
 				addMediaScreen: PWA_ASSET_GENERATOR.ADD_MEDIA_SCREEN,
 				basePath: PWA_ASSET_GENERATOR.BASE_PATH,
+				log: true,
 				xhtml: PWA_ASSET_GENERATOR.XHTML
 			},
-			png: {
-				compressionLevel: PWA_ASSET_GENERATOR.PNG_COMPRESSION_LEVEL,
-				quality: PWA_ASSET_GENERATOR.PNG_QUALITY
-			},
 			name: (landscape, size, dark) => {
 				const orientation = landscape ? SplashOrientation.LANDSCAPE : SplashOrientation.PORTRAIT;
 				const darkPrefix = dark ? PWA_ASSET_GENERATOR.DARK_PREFIX : '';
+
 				return `apple-splash-${orientation}-${darkPrefix}${size.width}x${size.height}.png`;
+			},
+			padding: PWA_ASSET_GENERATOR.SPLASH_PADDING,
+			png: {
+				compressionLevel: PWA_ASSET_GENERATOR.PNG_COMPRESSION_LEVEL,
+				quality: PWA_ASSET_GENERATOR.PNG_QUALITY
+			},
+			resizeOptions: {
+				background: THEME_COLORS.BACKGROUND_LIGHT,
+				fit: PWA_ASSET_GENERATOR.FIT_MODE
 			}
 		},
 		PWA_GENERATOR_DEVICES
-	),
-	images: ['static/favicon.svg']
+	)
 });

tools/ui/scripts/favicon-colorize.ts

@@ -4,12 +4,10 @@ import { fileURLToPath } from 'node:url';
 
 const HERE = dirname(fileURLToPath(import.meta.url));
 const PROJECT_ROOT = resolve(HERE, '..');
-
 const DEFAULT_LOGO = resolve(PROJECT_ROOT, 'src/lib/assets/logo.svg');
 const DEFAULT_OUT_DIR = resolve(PROJECT_ROOT, 'static');
 const DEFAULT_OUT_LIGHT = resolve(DEFAULT_OUT_DIR, 'favicon.svg');
 const DEFAULT_OUT_DARK = resolve(DEFAULT_OUT_DIR, 'favicon-dark.svg');
-
 const CURRENT_COLOR = 'currentColor';
 
 export interface ColorizedFavicon {
@@ -39,8 +37,8 @@ export function colorizeFaviconSvg(
 	darkColor: string
 ): ColorizedFavicon {
 	return {
-		light: svg.replaceAll(CURRENT_COLOR, lightColor),
-		dark: svg.replaceAll(CURRENT_COLOR, darkColor)
+		dark: svg.replaceAll(CURRENT_COLOR, darkColor),
+		light: svg.replaceAll(CURRENT_COLOR, lightColor)
 	};
 }
 
@@ -54,33 +52,40 @@ export function padFaviconSvg(svg: string, padding: number): string {
 	if (!(padding > 0) || padding >= 1) return svg;
 
 	const viewBoxMatch = svg.match(/viewBox\s*=\s*["']([^"']+)["']/i);
+
 	if (!viewBoxMatch) return svg;
 
 	const parts = viewBoxMatch[1]
 		.trim()
 		.split(/[\s,]+/)
 		.map(Number);
+
 	if (parts.length !== 4 || parts.some((n) => !Number.isFinite(n))) return svg;
 
 	const [, , width, height] = parts;
+
 	if (width <= 0 || height <= 0) return svg;
 
 	const scale = 1 - padding;
 	const translateX = (padding * width) / 2;
 	const translateY = (padding * height) / 2;
-
 	const openTagStart = svg.search(/<svg\b/i);
+
 	if (openTagStart === -1) return svg;
+
 	const openTagEnd = svg.indexOf('>', openTagStart);
+
 	if (openTagEnd === -1) return svg;
+
 	const closeStart = svg.lastIndexOf('</svg');
+
 	if (closeStart === -1 || closeStart <= openTagEnd) return svg;
 
 	const openTag = svg.slice(0, openTagEnd + 1);
 	const inner = svg.slice(openTagEnd + 1, closeStart);
 	const closeTag = svg.slice(closeStart);
-
 	const group = `<g transform="translate(${translateX} ${translateY}) scale(${scale})">`;
+
 	return `${openTag}${group}${inner}</g>${closeTag}`;
 }
 
@@ -93,14 +98,15 @@ export function writeThemeFavicons(
 	lightColor: string,
 	darkColor: string,
 	{
-		sourcePath = DEFAULT_LOGO,
-		lightOutPath = DEFAULT_OUT_LIGHT,
 		darkOutPath = DEFAULT_OUT_DARK,
-		padding = 0
+		lightOutPath = DEFAULT_OUT_LIGHT,
+		padding = 0,
+		sourcePath = DEFAULT_LOGO
 	}: WriteThemeFaviconsOptions = {}
 ): void {
 	const source = readFileSync(sourcePath, 'utf-8');
-	const { light, dark } = colorizeFaviconSvg(source, lightColor, darkColor);
+	const { dark, light } = colorizeFaviconSvg(source, lightColor, darkColor);
+
 	mkdirSync(dirname(lightOutPath), { recursive: true });
 	writeFileSync(lightOutPath, padFaviconSvg(light, padding));
 	writeFileSync(darkOutPath, padFaviconSvg(dark, padding));

tools/ui/scripts/make-icons-circular.js

@@ -13,31 +13,28 @@
  * maskable-icon and apple-touch-icon are left untouched.
  */
 
-import sharp from 'sharp';
 import fs from 'fs';
 import path from 'path';
+import sharp from 'sharp';
 import { fileURLToPath } from 'url';
 
 const __filename = fileURLToPath(import.meta.url);
 const __dirname = path.dirname(__filename);
-
 const STATIC_DIR = path.resolve(__dirname, '..', 'static');
-
 const paddingPct = process.argv.reduce((acc, arg, i, args) => {
 	if (arg === '--padding-pct' && args[i + 1]) return parseFloat(args[i + 1]);
+
 	return acc;
 }, 0);
-
 // Scale down the source image before cropping to circle
 const scalePct = process.argv.reduce((acc, arg, i, args) => {
 	if (arg === '--scale-pct' && args[i + 1]) return parseFloat(args[i + 1]);
+
 	return acc;
 }, 85); // default 85% - icon fills 85% of the circular area
-
 // Source for circular icons: the maskable icon (white bg, full logo)
 const sourceIcon = 'maskable-icon-512x512.png';
 const targetIcons = ['pwa-64x64.png', 'pwa-192x192.png', 'pwa-512x512.png'];
-
 // maskable-icon and apple-touch-icon stay square
 const untouchedIcons = ['maskable-icon-512x512.png', 'apple-touch-icon-180x180.png'];
 
@@ -47,10 +44,13 @@ async function makeCircle(targetFilename) {
 
 	if (!fs.existsSync(sourcePath)) {
 		console.log(`⏭️  ${sourceIcon} not found, skipping`);
+
 		return;
 	}
+
 	if (!fs.existsSync(targetPath)) {
 		console.log(`⏭️  ${targetFilename} not found, skipping`);
+
 		return;
 	}
 
@@ -58,16 +58,18 @@ async function makeCircle(targetFilename) {
 	const size = Math.max(metadata.width, metadata.height);
 	const radius = Math.floor((size * (1 - paddingPct / 100)) / 2);
 	const center = Math.floor(size / 2);
-
 	// Build circular mask as RGBA buffer: white opaque circle on transparent bg
 	const maskBuf = Buffer.alloc(size * size * 4, 0);
+
 	for (let y = 0; y < size; y++) {
 		for (let x = 0; x < size; x++) {
 			const dx = x - center;
 			const dy = y - center;
 			const dist = Math.sqrt(dx * dx + dy * dy);
+
 			if (dist < radius) {
 				const i = (y * size + x) * 4;
+
 				maskBuf[i] = 255;
 				maskBuf[i + 1] = 255;
 				maskBuf[i + 2] = 255;
@@ -77,8 +79,9 @@ async function makeCircle(targetFilename) {
 	}
 
 	const tmpMask = path.join(STATIC_DIR, '.mask-tmp.png');
+
 	await sharp(maskBuf, {
-		raw: { width: size, height: size, channels: 4 }
+		raw: { channels: 4, height: size, width: size }
 	})
 		.png()
 		.toFile(tmpMask);
@@ -87,28 +90,26 @@ async function makeCircle(targetFilename) {
 	const circleDiameter = Math.floor(size * (1 - paddingPct / 100));
 	const scaledSize = Math.floor((circleDiameter * scalePct) / 100);
 	const offset = Math.floor((size - scaledSize) / 2);
-
 	const scaledBuf = await sharp(sourcePath)
 		.resize(scaledSize, scaledSize, {
-			fit: 'cover',
-			background: { r: 255, g: 255, b: 255, alpha: 1 }
+			background: { alpha: 1, b: 255, g: 255, r: 255 },
+			fit: 'cover'
 		})
 		.ensureAlpha()
 		.png()
 		.toBuffer();
-
 	// Step 2: Composite scaled image onto white background, then apply circular mask
 	const output = await sharp({
 		create: {
-			width: size,
-			height: size,
+			background: { alpha: 1, b: 255, g: 255, r: 255 },
 			channels: 4,
-			background: { r: 255, g: 255, b: 255, alpha: 1 }
+			height: size,
+			width: size
 		}
 	})
 		.composite([
-			{ input: scaledBuf, top: offset, left: offset },
-			{ input: tmpMask, top: 0, left: 0, blend: 'dest-in' }
+			{ input: scaledBuf, left: offset, top: offset },
+			{ blend: 'dest-in', input: tmpMask, left: 0, top: 0 }
 		])
 		.png()
 		.toBuffer();
@@ -130,6 +131,7 @@ async function main() {
 	console.log('\nUnchanged:');
 	for (const icon of untouchedIcons) {
 		const fp = path.join(STATIC_DIR, icon);
+
 		console.log(`  ${icon} (${fs.existsSync(fp) ? fs.statSync(fp).size + ' bytes' : 'missing'})`);
 	}
 }

tools/ui/scripts/vite-plugin-build-info.ts

@@ -1,7 +1,7 @@
-import { writeFileSync, existsSync } from 'node:fs';
+import { BUILD_CONFIG } from '../src/lib/constants/pwa';
+import { existsSync, writeFileSync } from 'node:fs';
 import { resolve } from 'path';
 import type { Plugin } from 'vite';
-import { BUILD_CONFIG } from '../src/lib/constants/pwa';
 
 let processed = false;
 
@@ -15,27 +15,29 @@ const OUTPUT_DIR = process.env.LLAMA_UI_OUT_DIR ?? BUILD_CONFIG.OUTPUT_DIR;
  */
 export function buildInfoPlugin(): Plugin {
 	return {
-		name: 'llamacpp:build-info',
 		apply: 'build',
 		closeBundle() {
 			setTimeout(() => {
 				try {
 					if (processed) return;
+
 					processed = true;
 
 					const buildNumber = process.env.LLAMA_BUILD_NUMBER || 'b0000';
-
 					const outDir = resolve(OUTPUT_DIR);
 					const indexPath = resolve(outDir, 'index.html');
+
 					if (!existsSync(indexPath)) return;
 
 					const buildJsonPath = resolve(outDir, 'build.json');
+
 					writeFileSync(buildJsonPath, JSON.stringify({ version: buildNumber }), 'utf-8');
 					console.log(`Created build.json (version: ${buildNumber})`);
 				} catch (error) {
 					console.error('Failed to write build.json:', error);
 				}
 			}, 100);
-		}
+		},
+		name: 'llamacpp:build-info'
 	};
 }

tools/ui/scripts/vite-plugin-nerdamer.ts

@@ -4,7 +4,6 @@ import { fileURLToPath } from 'url';
 import type { Plugin } from 'vite';
 
 const __dirname = dirname(fileURLToPath(import.meta.url));
-
 const VENDORS_DIR = resolve(__dirname, '../src/lib/vendors');
 const VIRTUAL_ID = 'virtual:nerdamer';
 const RESOLVED_ID = '\0' + VIRTUAL_ID;
@@ -21,29 +20,32 @@ export function nerdamerPlugin(): Plugin {
 	let bundled: string | null = null;
 
 	return {
-		name: 'llamacpp:nerdamer',
-		resolveId(id) {
-			return id === VIRTUAL_ID ? RESOLVED_ID : undefined;
-		},
 		async load(id) {
 			if (id !== RESOLVED_ID) return undefined;
+
 			if (bundled === null) {
 				const result = await build({
-					entryPoints: [resolve(VENDORS_DIR, 'nerdamer-prime/all.js')],
-					bundle: true,
-					minify: true,
-					format: 'iife',
-					globalName: 'nerdamer',
 					alias: {
 						'big-integer': resolve(VENDORS_DIR, 'big-integer/BigInteger.js'),
 						'decimal.js': resolve(VENDORS_DIR, 'decimal.js/decimal.js')
 					},
-					write: false,
-					logLevel: 'silent'
+					bundle: true,
+					entryPoints: [resolve(VENDORS_DIR, 'nerdamer-prime/all.js')],
+					format: 'iife',
+					globalName: 'nerdamer',
+					logLevel: 'silent',
+					minify: true,
+					write: false
 				});
+
 				bundled = result.outputFiles[0].text;
 			}
+
 			return `export default ${JSON.stringify(bundled)};`;
+		},
+		name: 'llamacpp:nerdamer',
+		resolveId(id) {
+			return id === VIRTUAL_ID ? RESOLVED_ID : undefined;
 		}
 	};
 }

tools/ui/scripts/vite-plugin-relativize-base.ts

@@ -1,7 +1,7 @@
-import { readFileSync, writeFileSync, existsSync } from 'node:fs';
+import { BUILD_CONFIG } from '../src/lib/constants/pwa';
+import { existsSync, readFileSync, writeFileSync } from 'node:fs';
 import { resolve } from 'path';
 import type { Plugin } from 'vite';
-import { BUILD_CONFIG } from '../src/lib/constants/pwa';
 
 let processed = false;
 
@@ -11,11 +11,15 @@ function rewrite(path: string, pairs: [string, string][]): void {
 	if (!existsSync(path)) {
 		return;
 	}
+
 	const text = readFileSync(path, 'utf-8');
+
 	let out = text;
+
 	for (const [from, to] of pairs) {
 		out = out.split(from).join(to);
 	}
+
 	if (out !== text) {
 		writeFileSync(path, out, 'utf-8');
 	}
@@ -32,12 +36,12 @@ function rewrite(path: string, pairs: [string, string][]): void {
  */
 export function relativizeBasePlugin(): Plugin {
 	return {
-		name: 'llamacpp:relativize-base',
 		apply: 'build',
 		closeBundle() {
 			setTimeout(() => {
 				try {
 					if (processed) return;
+
 					processed = true;
 
 					const outDir = resolve(OUTPUT_DIR);
@@ -56,6 +60,7 @@ export function relativizeBasePlugin(): Plugin {
 					console.error('Failed to relativize base refs:', error);
 				}
 			}, 100);
-		}
+		},
+		name: 'llamacpp:relativize-base'
 	};
 }

tools/ui/scripts/vite-plugin-splash-screen.ts

@@ -1,10 +1,10 @@
-import { readdirSync, readFileSync, writeFileSync, existsSync } from 'node:fs';
+import { NEWLINE, TAB } from '../src/lib/constants/code';
+import { APPLE_DEVICES, BUILD_CONFIG, REGEX_PATTERNS, SPLASH_LINK } from '../src/lib/constants/pwa';
+import { SplashOrientation } from '../src/lib/enums/splash.enums';
+import type { SplashDimensions } from '../src/lib/types';
+import { existsSync, readdirSync, readFileSync, writeFileSync } from 'node:fs';
 import { resolve } from 'path';
 import type { Plugin } from 'vite';
-import { TAB, NEWLINE } from '../src/lib/constants/code';
-import { APPLE_DEVICES, BUILD_CONFIG, REGEX_PATTERNS, SPLASH_LINK } from '../src/lib/constants/pwa';
-import type { SplashDimensions } from '../src/lib/types';
-import { SplashOrientation } from '../src/lib/enums/splash.enums';
 
 let processed = false;
 
@@ -16,23 +16,26 @@ const OUTPUT_DIR = process.env.LLAMA_UI_OUT_DIR ?? BUILD_CONFIG.OUTPUT_DIR;
  */
 export function generateSplashScreenLinks(outDir: string): string[] {
 	const files = readdirSync(outDir).filter((f) => f.match(REGEX_PATTERNS.SPLASH_FILE));
+
 	if (files.length === 0) return [];
 
 	const dimMap = new Map<string, SplashDimensions>();
+
 	for (const [dims, spec] of Object.entries(APPLE_DEVICES)) {
 		const [w, h] = dims.split('x').map(Number);
+
 		// logical-point dimensions
-		dimMap.set(`${w}x${h}`, { deviceW: spec.width, deviceH: spec.height, dpr: spec.dpr });
-		dimMap.set(`${h}x${w}`, { deviceW: spec.width, deviceH: spec.height, dpr: spec.dpr });
+		dimMap.set(`${w}x${h}`, { deviceH: spec.height, deviceW: spec.width, dpr: spec.dpr });
+		dimMap.set(`${h}x${w}`, { deviceH: spec.height, deviceW: spec.width, dpr: spec.dpr });
 		// pixel dimensions (used by actual generated splash files)
 		dimMap.set(`${w * spec.dpr}x${h * spec.dpr}`, {
-			deviceW: spec.width,
 			deviceH: spec.height,
+			deviceW: spec.width,
 			dpr: spec.dpr
 		});
 		dimMap.set(`${h * spec.dpr}x${w * spec.dpr}`, {
-			deviceW: spec.width,
 			deviceH: spec.height,
+			deviceW: spec.width,
 			dpr: spec.dpr
 		});
 	}
@@ -42,20 +45,23 @@ export function generateSplashScreenLinks(outDir: string): string[] {
 
 	for (const file of files) {
 		const match = file.match(REGEX_PATTERNS.SPLASH_FILE);
+
 		if (!match) continue;
+
 		const orientation = match[1] as SplashOrientation;
 		const isDark = !!match[2];
 		const pixelW = parseInt(match[3]);
 		const pixelH = parseInt(match[4]);
-
 		const key = `${pixelW}x${pixelH}`;
 		const spec = dimMap.get(key);
+
 		if (!spec) {
 			console.warn(`Unknown splash screen dimensions: ${key} (${file})`);
+
 			continue;
 		}
 
-		const { deviceW, deviceH, dpr } = spec;
+		const { deviceH, deviceW, dpr } = spec;
 		const media = `screen and (device-width: ${deviceW}px) and (device-height: ${deviceH}px) and (-webkit-device-pixel-ratio: ${dpr}) and (orientation: ${orientation})`;
 		const href = `./${file}`;
 
@@ -73,16 +79,17 @@ export function generateSplashScreenLinks(outDir: string): string[] {
 
 export function splashScreenPlugin(): Plugin {
 	return {
-		name: 'llamacpp:splash-screen',
 		apply: 'build',
 		closeBundle() {
 			setTimeout(() => {
 				try {
 					if (processed) return;
+
 					processed = true;
 
 					const outDir = resolve(OUTPUT_DIR);
 					const indexPath = resolve(outDir, 'index.html');
+
 					if (!existsSync(indexPath)) return;
 
 					let content = readFileSync(indexPath, 'utf-8');
@@ -91,9 +98,11 @@ export function splashScreenPlugin(): Plugin {
 					// The @vite-pwa/assets-generator generates apple-splash-*.png files;
 					// this scans them and creates the <link> tags SvelteKit needs.
 					const splashLinks = generateSplashScreenLinks(outDir);
+
 					if (splashLinks.length > 0) {
 						console.log(`Generated ${splashLinks.length} apple-splash link tags`);
 						const splashHtml = splashLinks.map((l) => TAB + TAB + l).join(NEWLINE);
+
 						content = content.replace(
 							REGEX_PATTERNS.HEAD_CLOSE,
 							splashHtml + NEWLINE + TAB + TAB + '</head>'
@@ -110,6 +119,7 @@ export function splashScreenPlugin(): Plugin {
 					console.error('Failed to process build output:', error);
 				}
 			}, 100);
-		}
+		},
+		name: 'llamacpp:splash-screen'
 	};
 }

tools/ui/src/app.d.ts

@@ -3,9 +3,8 @@
 
 import 'vite-plugin-pwa/pwa-assets';
 import 'vite-plugin-pwa/svelte';
-
+import { ModelModality, ServerModelStatus, ServerRole } from '$lib/enums';
 // Import chat types from dedicated module
-
 import type {
 	// API types
 	ApiChatCompletionRequest,
@@ -13,59 +12,57 @@ import type {
 	ApiChatCompletionStreamChunk,
 	ApiChatCompletionToolCall,
 	ApiChatCompletionToolCallDelta,
-	ApiChatMessageData,
 	ApiChatMessageContentPart,
+	ApiChatMessageData,
 	ApiContextSizeError,
 	ApiErrorResponse,
 	ApiLlamaCppServerProps,
 	ApiModelDataEntry,
+	ApiModelListResponse,
 	ApiModelLoadStage,
-	ApiModelsSseProgress,
 	ApiModelsSseData,
 	ApiModelsSseEvent,
-	ApiModelListResponse,
+	ApiModelsSseProgress,
 	ApiProcessingState,
 	ApiRouterModelMeta,
+	ApiRouterModelsListResponse,
 	ApiRouterModelsLoadRequest,
 	ApiRouterModelsLoadResponse,
 	ApiRouterModelsStatusRequest,
 	ApiRouterModelsStatusResponse,
-	ApiRouterModelsListResponse,
 	ApiRouterModelsUnloadRequest,
 	ApiRouterModelsUnloadResponse,
 	// Chat types
 	ChatAttachmentDisplayItem,
+	ChatMessagePromptProgress,
+	ChatMessageSiblingInfo,
+	ChatMessageTimings,
 	ChatMessageType,
 	ChatRole,
 	ChatUploadedFile,
-	ChatMessageSiblingInfo,
-	ChatMessagePromptProgress,
-	ChatMessageTimings,
 	// Database types
 	DatabaseConversation,
 	DatabaseMessage,
 	DatabaseMessageExtra,
 	DatabaseMessageExtraAudioFile,
-	DatabaseMessageExtraVideoFile,
 	DatabaseMessageExtraImageFile,
-	DatabaseMessageExtraTextFile,
-	DatabaseMessageExtraPdfFile,
 	DatabaseMessageExtraLegacyContext,
+	DatabaseMessageExtraPdfFile,
+	DatabaseMessageExtraTextFile,
+	DatabaseMessageExtraVideoFile,
 	ExportedConversation,
 	ExportedConversations,
+	ModelLoadProgress,
 	// Model types
 	ModelModalities,
 	ModelOption,
-	ModelLoadProgress,
 	// Settings types
 	SettingsChatServiceOptions,
+	SettingsConfigType,
 	SettingsConfigValue,
-	SettingsFieldConfig,
-	SettingsConfigType
+	SettingsFieldConfig
 } from '$lib/types';
 
-import { ServerRole, ServerModelStatus, ModelModality } from '$lib/enums';
-
 declare global {
 	// namespace App {
 	// interface Error {}

tools/ui/src/lib/components/app/actions/ActionIcon.svelte

@@ -1,8 +1,8 @@
 <script lang="ts">
-	import { Button, type ButtonVariant, type ButtonSize } from '$lib/components/ui/button';
+	import { Button, type ButtonSize, type ButtonVariant } from '$lib/components/ui/button';
 	import * as Tooltip from '$lib/components/ui/tooltip';
-	import type { Component } from 'svelte';
 	import { TooltipSide } from '$lib/enums';
+	import type { Component } from 'svelte';
 
 	interface Props {
 		ariaLabel?: string;
@@ -20,18 +20,18 @@
 	}
 
 	let {
-		icon,
-		tooltip,
-		variant = 'ghost',
-		href = '',
-		size = 'sm',
+		ariaLabel,
 		class: className = '',
 		disabled = false,
+		href = '',
+		icon,
 		iconSize = 'h-3 w-3',
-		tooltipSide = TooltipSide.TOP,
-		stopPropagationOnClick = false,
 		onclick,
-		ariaLabel
+		size = 'sm',
+		stopPropagationOnClick = false,
+		tooltip,
+		tooltipSide = TooltipSide.TOP,
+		variant = 'ghost'
 	}: Props = $props();
 
 	let innerWidth = $state(0);

tools/ui/src/lib/components/app/actions/ActionIconCopyToClipboard.svelte

@@ -1,8 +1,8 @@
 <script lang="ts">
-	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
-	import { Copy } from '@lucide/svelte';
-	import { copyToClipboard } from '$lib/utils';
 	import ActionIcon from './ActionIcon.svelte';
+	import { Copy } from '@lucide/svelte';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
+	import { copyToClipboard } from '$lib/utils';
 
 	export let ariaLabel: string = 'Copy to clipboard';
 	export let canCopy: boolean = true;

tools/ui/src/lib/components/app/badges/BadgesModality.svelte

@@ -7,7 +7,7 @@
 		class?: string;
 	}
 
-	let { modalities, class: className = '' }: Props = $props();
+	let { class: className = '', modalities }: Props = $props();
 </script>
 
 {#each modalities as modality (modality)}

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsList.svelte

@@ -28,18 +28,18 @@
 	}
 
 	let {
-		class: className = '',
-		style = '',
+		activeModelId,
 		attachments = [],
-		readonly = false,
-		onFileRemove,
-		uploadedFiles = $bindable([]),
+		class: className = '',
 		// Default to small size for form previews
 		imageClass = '',
 		imageHeight = 'h-24',
 		imageWidth = 'w-auto',
 		limitToSingleRow = false,
-		activeModelId
+		onFileRemove,
+		readonly = false,
+		style = '',
+		uploadedFiles = $bindable([])
 	}: Props = $props();
 
 	let carouselRef: HorizontalScrollCarousel | undefined = $state();
@@ -48,7 +48,7 @@
 	let previewFocusIndex = $state(0);
 	let viewAllDialogOpen = $state(false);
 
-	let displayItems = $derived(getAttachmentDisplayItems({ uploadedFiles, attachments }));
+	let displayItems = $derived(getAttachmentDisplayItems({ attachments, uploadedFiles }));
 
 	function openPreview(item: ChatAttachmentDisplayItem, event?: MouseEvent) {
 		event?.stopPropagation();

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItem.svelte

@@ -2,8 +2,8 @@
 	import {
 		ChatAttachmentsListItemMcpPrompt,
 		ChatAttachmentsListItemMcpResource,
-		ChatAttachmentsListItemThumbnailImage,
-		ChatAttachmentsListItemThumbnailFile
+		ChatAttachmentsListItemThumbnailFile,
+		ChatAttachmentsListItemThumbnailImage
 	} from '$lib/components/app';
 	import { AttachmentType } from '$lib/enums';
 	import type {
@@ -49,10 +49,10 @@
 		return {
 			id,
 			resource: {
-				uri: extra.uri,
 				name: extra.name,
+				serverName: extra.serverName,
 				title: extra.name,
-				serverName: extra.serverName
+				uri: extra.uri
 			}
 		};
 	}
@@ -64,12 +64,12 @@
 			? (item.attachment as DatabaseMessageExtraMcpPrompt)
 			: item.uploadedFile?.mcpPrompt
 				? {
-						type: AttachmentType.MCP_PROMPT as const,
-						name: item.name,
-						serverName: item.uploadedFile.mcpPrompt.serverName,
-						promptName: item.uploadedFile.mcpPrompt.promptName,
+						arguments: item.uploadedFile.mcpPrompt.arguments,
 						content: item.textContent ?? '',
-						arguments: item.uploadedFile.mcpPrompt.arguments
+						name: item.name,
+						promptName: item.uploadedFile.mcpPrompt.promptName,
+						serverName: item.uploadedFile.mcpPrompt.serverName,
+						type: AttachmentType.MCP_PROMPT as const
 					}
 				: null}
 	{#if mcpPrompt}

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemMcpPrompt.svelte

@@ -1,8 +1,8 @@
 <script lang="ts">
-	import { ChatMessageMcpPromptContent, ActionIcon } from '$lib/components/app';
 	import { X } from '@lucide/svelte';
-	import type { DatabaseMessageExtraMcpPrompt } from '$lib/types';
+	import { ActionIcon, ChatMessageMcpPromptContent } from '$lib/components/app';
 	import { McpPromptVariant } from '$lib/enums';
+	import type { DatabaseMessageExtraMcpPrompt } from '$lib/types';
 
 	interface Props {
 		class?: string;

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemMcpResource.svelte

@@ -1,11 +1,11 @@
 <script lang="ts">
-	import { Loader2, AlertCircle } from '@lucide/svelte';
+	import { AlertCircle, Loader2 } from '@lucide/svelte';
+	import { X } from '@lucide/svelte';
+	import { ActionIcon } from '$lib/components/app';
+	import * as Tooltip from '$lib/components/ui/tooltip';
 	import { mcpStore } from '$lib/stores/mcp.svelte';
 	import type { MCPResourceAttachment } from '$lib/types';
-	import * as Tooltip from '$lib/components/ui/tooltip';
-	import { ActionIcon } from '$lib/components/app';
-	import { X } from '@lucide/svelte';
-	import { getResourceIcon, getResourceDisplayName } from '$lib/utils';
+	import { getResourceDisplayName, getResourceIcon } from '$lib/utils';
 
 	interface Props {
 		attachment: MCPResourceAttachment;
@@ -24,6 +24,7 @@
 
 	function getStatusClass(attachment: MCPResourceAttachment): string {
 		if (attachment.error) return 'border-red-500/50 bg-red-500/10';
+
 		if (attachment.loading) return 'border-border/50 bg-muted/30';
 
 		return 'border-border/50 bg-muted/30';

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemThumbnailFile.svelte

@@ -1,17 +1,17 @@
 <script lang="ts">
+	import { Music, Video, X } from '@lucide/svelte';
+	import { ActionIcon } from '$lib/components/app';
 	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
-	import { X, Music, Video } from '@lucide/svelte';
+	import { AttachmentType } from '$lib/enums';
 	import {
 		formatFileSize,
 		getFileTypeLabel,
 		getPreviewText,
-		isPdfFile,
 		isAudioFile,
-		isVideoFile,
-		isTextFile
+		isPdfFile,
+		isTextFile,
+		isVideoFile
 	} from '$lib/utils';
-	import { ActionIcon } from '$lib/components/app';
-	import { AttachmentType } from '$lib/enums';
 
 	interface Props {
 		attachment?: DatabaseMessageExtra;
@@ -31,9 +31,9 @@
 		attachment,
 		class: className = '',
 		id,
+		name,
 		onclick,
 		onRemove,
-		name,
 		readonly = false,
 		size,
 		textContent,

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemThumbnailImage.svelte

@@ -1,6 +1,6 @@
 <script lang="ts">
-	import { ActionIcon } from '$lib/components/app';
 	import { X } from '@lucide/svelte';
+	import { ActionIcon } from '$lib/components/app';
 
 	interface Props {
 		class?: string;
@@ -20,9 +20,9 @@
 		height = 'h-16',
 		id,
 		imageClass = '',
+		name,
 		onclick,
 		onRemove,
-		name,
 		preview,
 		readonly = false,
 		width = 'w-auto'

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsPreview.svelte

@@ -12,12 +12,12 @@
 		getAttachmentDisplayItems,
 		getLanguageFromFilename,
 		isAudioFile,
-		isVideoFile,
 		isImageFile,
 		isMcpPrompt,
 		isMcpResource,
 		isPdfFile,
-		isTextFile
+		isTextFile,
+		isVideoFile
 	} from '$lib/utils';
 
 	interface PreviewItem {
@@ -42,21 +42,21 @@
 	}
 
 	let {
-		uploadedFiles = [],
-		attachments = [],
 		activeModelId,
+		attachments = [],
 		class: className = '',
-		previewFocusIndex = 0
+		previewFocusIndex = 0,
+		uploadedFiles = []
 	}: Props = $props();
 
 	let allItems = $derived(
-		getAttachmentDisplayItems({ uploadedFiles, attachments })
+		getAttachmentDisplayItems({ attachments, uploadedFiles })
 			.filter((item) => !isMcpPrompt(item) && !isMcpResource(item))
 			.map(
 				(item): PreviewItem => ({
 					...item,
-					isImage: isImageFile(item.attachment, item.uploadedFile),
 					isAudio: isAudioFile(item.attachment, item.uploadedFile),
+					isImage: isImageFile(item.attachment, item.uploadedFile),
 					isVideo: isVideoFile(item.attachment, item.uploadedFile)
 				})
 			)
@@ -88,10 +88,11 @@
 
 	$effect(() => {
 		const index = currentIndex;
+
 		setTimeout(() => {
 			const thumbnail = document.querySelector(`[data-thumbnail-index="${index}"]`);
 
-			thumbnail?.scrollIntoView({ behavior: 'smooth', inline: 'center', block: 'nearest' });
+			thumbnail?.scrollIntoView({ behavior: 'smooth', block: 'nearest', inline: 'center' });
 		}, 0);
 	});
 

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreviewCurrentItem/ChatAttachmentsPreviewCurrentItem.svelte

@@ -1,12 +1,12 @@
 <script lang="ts">
-	import type { ChatAttachmentDisplayItem } from '$lib/types';
-	import { Image, Music, Video, FileText, FileIcon } from '@lucide/svelte';
-	import ChatAttachmentsPreviewCurrentItemPdf from './ChatAttachmentsPreviewCurrentItemPdf.svelte';
-	import ChatAttachmentsPreviewCurrentItemImage from './ChatAttachmentsPreviewCurrentItemImage.svelte';
 	import ChatAttachmentsPreviewCurrentItemAudio from './ChatAttachmentsPreviewCurrentItemAudio.svelte';
-	import ChatAttachmentsPreviewCurrentItemVideo from './ChatAttachmentsPreviewCurrentItemVideo.svelte';
+	import ChatAttachmentsPreviewCurrentItemImage from './ChatAttachmentsPreviewCurrentItemImage.svelte';
+	import ChatAttachmentsPreviewCurrentItemPdf from './ChatAttachmentsPreviewCurrentItemPdf.svelte';
 	import ChatAttachmentsPreviewCurrentItemText from './ChatAttachmentsPreviewCurrentItemText.svelte';
 	import ChatAttachmentsPreviewCurrentItemUnavailable from './ChatAttachmentsPreviewCurrentItemUnavailable.svelte';
+	import ChatAttachmentsPreviewCurrentItemVideo from './ChatAttachmentsPreviewCurrentItemVideo.svelte';
+	import { FileIcon, FileText, Image, Music, Video } from '@lucide/svelte';
+	import type { ChatAttachmentDisplayItem } from '$lib/types';
 
 	interface Props {
 		currentItem: ChatAttachmentDisplayItem | null;
@@ -25,19 +25,19 @@
 	}
 
 	let {
+		activeModelId,
+		audioSrc,
 		currentItem,
-		isImage,
-		isAudio,
-		isVideo,
-		isPdf,
-		isText,
 		displayPreview,
 		displayTextContent,
-		audioSrc,
-		videoSrc,
-		language,
 		hasVisionModality,
-		activeModelId
+		isAudio,
+		isImage,
+		isPdf,
+		isText,
+		isVideo,
+		language,
+		videoSrc
 	}: Props = $props();
 
 	let IconComponent = $derived(

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreviewCurrentItem/ChatAttachmentsPreviewCurrentItemAudio.svelte

@@ -6,7 +6,7 @@
 		audioSrc: string | null;
 	}
 
-	let { currentItem, audioSrc }: Props = $props();
+	let { audioSrc, currentItem }: Props = $props();
 </script>
 
 <div class="flex flex-1 items-center justify-center p-8">

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreviewCurrentItem/ChatAttachmentsPreviewCurrentItemPdf.svelte

@@ -1,13 +1,13 @@
 <script lang="ts">
-	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
-	import type { ChatAttachmentDisplayItem } from '$lib/types';
-	import { FileText, Eye, Info } from '@lucide/svelte';
-	import { Button } from '$lib/components/ui/button';
-	import * as Alert from '$lib/components/ui/alert';
+	import { Eye, FileText, Info } from '@lucide/svelte';
 	import { SyntaxHighlightedCode } from '$lib/components/app';
+	import * as Alert from '$lib/components/ui/alert';
+	import { Button } from '$lib/components/ui/button';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
+	import { PdfViewMode } from '$lib/enums';
+	import type { ChatAttachmentDisplayItem } from '$lib/types';
 	import { getLanguageFromFilename } from '$lib/utils';
 	import { convertPDFToImage } from '$lib/utils/browser-only';
-	import { PdfViewMode } from '$lib/enums';
 
 	interface Props {
 		currentItem: ChatAttachmentDisplayItem | null;
@@ -17,7 +17,7 @@
 		activeModelId?: string;
 	}
 
-	let { currentItem, displayName, displayTextContent, hasVisionModality, activeModelId }: Props =
+	let { activeModelId, currentItem, displayName, displayTextContent, hasVisionModality }: Props =
 		$props();
 
 	let pdfViewMode = $state<PdfViewMode>(PdfViewMode.PAGES);
@@ -47,6 +47,7 @@
 					currentItem.attachment.images.length > 0
 				) {
 					pdfImages = currentItem.attachment.images;
+
 					return;
 				}
 
@@ -55,10 +56,12 @@
 					const base64Data = currentItem.attachment.base64Data;
 					const byteCharacters = atob(base64Data);
 					const byteNumbers = new Array(byteCharacters.length);
+
 					for (let i = 0; i < byteCharacters.length; i++) {
 						byteNumbers[i] = byteCharacters.charCodeAt(i);
 					}
 					const byteArray = new Uint8Array(byteNumbers);
+
 					file = new File([byteArray], displayName, { type: 'application/pdf' });
 				}
 			}

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreviewNavButtons.svelte

@@ -8,7 +8,7 @@
 		show: boolean;
 	}
 
-	let { onPrev, onNext, show }: Props = $props();
+	let { onNext, onPrev, show }: Props = $props();
 </script>
 
 {#if show}

tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsPreview/ChatAttachmentsPreviewThumbnailStrip.svelte

@@ -1,7 +1,7 @@
 <script lang="ts">
-	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
-	import { Music, Video, FileText } from '@lucide/svelte';
+	import { FileText, Music, Video } from '@lucide/svelte';
 	import { HorizontalScrollCarousel } from '$lib/components/app/misc';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 
 	interface PreviewItem {
 		id: string;
@@ -18,13 +18,15 @@
 		onNavigate: (index: number) => void;
 	}
 
-	let { items, currentIndex, onNavigate }: Props = $props();
+	let { currentIndex, items, onNavigate }: Props = $props();
 
 	function getFileExtension(name: string): string {
 		const parts = name.split('.');
+
 		if (parts.length > 1) {
 			return parts.pop()?.toUpperCase() ?? '';
 		}
+
 		return '';
 	}
 </script>

tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte

@@ -1,4 +1,5 @@
 <script lang="ts">
+	import ContextGaugePopup from './ChatFormContextGauge/ContextGaugePopup.svelte';
 	import {
 		ChatAttachmentsList,
 		ChatFormActions,
@@ -12,10 +13,10 @@
 	} from '$lib/components/app';
 	import {
 		CLIPBOARD_CONTENT_QUOTE_PREFIX,
-		INPUT_CLASSES,
-		SETTING_CONFIG_DEFAULT,
 		INITIAL_FILE_SIZE,
-		PROMPT_CONTENT_SEPARATOR
+		INPUT_CLASSES,
+		PROMPT_CONTENT_SEPARATOR,
+		SETTING_CONFIG_DEFAULT
 	} from '$lib/constants';
 	import {
 		ContentPartType,
@@ -24,20 +25,20 @@
 		MimeTypeText,
 		SpecialFileType
 	} from '$lib/enums';
-	import { config } from '$lib/stores/settings.svelte';
-	import ContextGaugePopup from './ChatFormContextGauge/ContextGaugePopup.svelte';
-	import { modelOptions, selectedModelId } from '$lib/stores/models.svelte';
-	import { isRouterMode } from '$lib/stores/server.svelte';
+	import { useChatFormPickers } from '$lib/hooks/use-chat-form-pickers.svelte';
 	import { chatStore } from '$lib/stores/chat.svelte';
-	import { mcpStore } from '$lib/stores/mcp.svelte';
-	import { mcpHasResourceAttachments } from '$lib/stores/mcp-resources.svelte';
-	import { toolsStore } from '$lib/stores/tools.svelte';
 	import {
-		conversationsStore,
-		activeMessages,
 		activeConversation,
+		activeMessages,
+		conversationsStore,
 		pendingCwd
 	} from '$lib/stores/conversations.svelte';
+	import { mcpStore } from '$lib/stores/mcp.svelte';
+	import { mcpHasResourceAttachments } from '$lib/stores/mcp-resources.svelte';
+	import { modelOptions, selectedModelId } from '$lib/stores/models.svelte';
+	import { isRouterMode } from '$lib/stores/server.svelte';
+	import { config } from '$lib/stores/settings.svelte';
+	import { toolsStore } from '$lib/stores/tools.svelte';
 	import type {
 		FileMentionEntry,
 		GetPromptResult,
@@ -56,7 +57,6 @@
 		parseClipboardContent,
 		uuid
 	} from '$lib/utils';
-	import { useChatFormPickers } from '$lib/hooks/use-chat-form-pickers.svelte';
 	import {
 		AudioRecorder,
 		convertToWav,
@@ -96,12 +96,6 @@
 		class: className = '',
 		disabled = false,
 		isLoading = false,
-		placeholder = 'Type a message...',
-		showMcpPromptButton = false,
-		showAddButton = true,
-		showModelSelector = true,
-		uploadedFiles = $bindable([]),
-		value = $bindable(''),
 		onAttachmentRemove,
 		onFilesAdd,
 		onStop,
@@ -109,7 +103,13 @@
 		onSystemPromptClick,
 		onUploadedFileRemove,
 		onUploadedFilesChange,
-		onValueChange
+		onValueChange,
+		placeholder = 'Type a message...',
+		showAddButton = true,
+		showMcpPromptButton = false,
+		showModelSelector = true,
+		uploadedFiles = $bindable([]),
+		value = $bindable('')
 	}: Props = $props();
 
 	// Component References
@@ -146,32 +146,35 @@
 	let cwd = $derived(activeConversation()?.cwd ?? pendingCwd());
 
 	const pickers = useChatFormPickers({
+		focusInput: refocusInput,
+		getCaretOffset: () => inputRef?.getCaretOffset(),
+		getCwd: () => cwd,
+		getPickersRef: () => pickersRef,
+		getServerHome: () => toolsStore.serverHome ?? null,
+		getShowModelSelector: () => showModelSelector,
 		getValue: () => value,
+		hasCwdTools: () => toolsStore.hasEnabledCwdTools,
+		hasPrompts: () => mcpStore.hasPromptsCapability(conversationsStore.getAllMcpServerOverrides()),
+		openModelSelector: () => chatFormActionsRef?.openModelSelector(),
+		setCaretOffset: (offset) => inputRef?.setCaretOffset(offset),
 		setValue: (v) => {
 			value = v;
 			onValueChange?.(v);
-		},
-		getCaretOffset: () => inputRef?.getCaretOffset(),
-		setCaretOffset: (offset) => inputRef?.setCaretOffset(offset),
-		focusInput: refocusInput,
-		getShowModelSelector: () => showModelSelector,
-		hasPrompts: () => mcpStore.hasPromptsCapability(conversationsStore.getAllMcpServerOverrides()),
-		hasCwdTools: () => toolsStore.hasEnabledCwdTools,
-		getCwd: () => cwd,
-		getServerHome: () => toolsStore.serverHome ?? null,
-		openModelSelector: () => chatFormActionsRef?.openModelSelector(),
-		getPickersRef: () => pickersRef
+		}
 	});
 
 	async function handleWorkingDirectoryChange(newDir: string | null) {
 		// Committing a directory consumes the `/cwd` token; the chip's
 		// clear-X path has no token to consume.
 		const token = findCommandToken(value);
+
 		if (token && token.name === 'cwd') {
 			value = '';
 			onValueChange?.('');
 		}
+
 		await conversationsStore.setCwd(newDir);
+
 		if (conversationsStore.activeConversation) {
 			await chatStore.recordCwdChange(newDir?.trim() || null);
 		}
@@ -185,6 +188,7 @@
 
 	let pasteLongTextToFileLength = $derived.by(() => {
 		const n = Number(currentConfig.pasteLongTextToFileLen);
+
 		return Number.isNaN(n) ? Number(SETTING_CONFIG_DEFAULT.pasteLongTextToFileLen) : n;
 	});
 
@@ -200,13 +204,16 @@
 		}
 
 		const selectedId = selectedModelId();
+
 		if (selectedId) {
 			const model = options.find((m) => m.id === selectedId);
+
 			if (model) return model.model;
 		}
 
 		if (conversationModel) {
 			const model = options.find((m) => m.model === conversationModel);
+
 			if (model) return model.model;
 		}
 
@@ -238,6 +245,7 @@
 	$effect(() => {
 		const wantContenteditable =
 			containsFileMentionLink(value ?? '') || containsCodeSpan(value ?? '');
+
 		if (useContenteditable === wantContenteditable) return;
 
 		if (!caretOffsetPinned) {
@@ -268,8 +276,10 @@
 	export function checkModelSelected(): boolean {
 		if (!hasModelSelected) {
 			chatFormActionsRef?.openModelSelector();
+
 			return false;
 		}
+
 		return true;
 	}
 
@@ -284,6 +294,7 @@
 	function handleFileRemove(fileId: string) {
 		if (fileId.startsWith('attachment-')) {
 			const index = parseInt(fileId.replace('attachment-', ''), 10);
+
 			if (!isNaN(index) && index >= 0 && index < attachments.length) {
 				onAttachmentRemove?.(index);
 			}
@@ -333,6 +344,7 @@
 		if (files.length > 0) {
 			event.preventDefault();
 			onFilesAdd?.(files);
+
 			return;
 		}
 
@@ -354,26 +366,27 @@
 								type: MimeTypeText.PLAIN
 							})
 					);
+
 					onFilesAdd?.(attachmentFiles);
 				}
 
 				// Handle MCP prompt attachments as ChatUploadedFile with mcpPrompt data
 				if (parsed.mcpPromptAttachments.length > 0) {
 					const mcpPromptFiles: ChatUploadedFile[] = parsed.mcpPromptAttachments.map((att) => ({
-						id: uuid(),
-						name: att.name,
-						size: att.content.length,
-						type: SpecialFileType.MCP_PROMPT,
 						file: new File([att.content], `${att.name}${FileExtensionText.TXT}`, {
 							type: MimeTypeText.PLAIN
 						}),
+						id: uuid(),
 						isLoading: false,
-						textContent: att.content,
 						mcpPrompt: {
-							serverName: att.serverName,
+							arguments: att.arguments,
 							promptName: att.promptName,
-							arguments: att.arguments
-						}
+							serverName: att.serverName
+						},
+						name: att.name,
+						size: att.content.length,
+						textContent: att.content,
+						type: SpecialFileType.MCP_PROMPT
 					}));
 
 					uploadedFiles = [...uploadedFiles, ...mcpPromptFiles];
@@ -412,17 +425,17 @@
 
 		const promptName = promptInfo.title || promptInfo.name;
 		const placeholder: ChatUploadedFile = {
-			id: placeholderId,
-			name: promptName,
-			size: INITIAL_FILE_SIZE,
-			type: SpecialFileType.MCP_PROMPT,
 			file: new File([], 'loading'),
+			id: placeholderId,
 			isLoading: true,
 			mcpPrompt: {
-				serverName: promptInfo.serverName,
+				arguments: args ? { ...args } : undefined,
 				promptName: promptInfo.name,
-				arguments: args ? { ...args } : undefined
-			}
+				serverName: promptInfo.serverName
+			},
+			name: promptName,
+			size: INITIAL_FILE_SIZE,
+			type: SpecialFileType.MCP_PROMPT
 		};
 
 		uploadedFiles = [...uploadedFiles, placeholder];
@@ -450,12 +463,12 @@
 			f.id === placeholderId
 				? {
 						...f,
-						isLoading: false,
-						textContent: promptText,
-						size: promptText.length,
 						file: new File([promptText], `${f.name}${FileExtensionText.TXT}`, {
 							type: MimeTypeText.PLAIN
-						})
+						}),
+						isLoading: false,
+						size: promptText.length,
+						textContent: promptText
 					}
 				: f
 		);
@@ -480,9 +493,11 @@
 	function handleMentionSelect(entry: FileMentionEntry) {
 		const cursor = inputRef?.getCaretOffset() ?? value.length;
 		const token = findMentionToken(value, cursor);
+
 		if (!token) return;
 
 		const built = buildMentionInsertion(entry, value, token);
+
 		if (!built) return;
 
 		// Pin the post-insertion caret BEFORE the swap effect runs;
@@ -504,6 +519,7 @@
 	async function handleMicClick() {
 		if (!audioRecorder || !recordingSupported) {
 			console.warn('Audio recording not supported');
+
 			return;
 		}
 
@@ -642,7 +658,7 @@
 				onFileUpload={handleFileUpload}
 				onMicClick={handleMicClick}
 				{onStop}
-				onSystemPromptClick={() => onSystemPromptClick?.({ message: value, files: uploadedFiles })}
+				onSystemPromptClick={() => onSystemPromptClick?.({ files: uploadedFiles, message: value })}
 				onMcpPromptClick={showMcpPromptButton ? () => pickers.openPromptPicker() : undefined}
 				onMcpResourcesClick={() => (isResourceDialogOpen = true)}
 			/>

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddButton.svelte

@@ -1,9 +1,9 @@
 <script lang="ts">
-	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { Plus } from '@lucide/svelte';
 	import { Button } from '$lib/components/ui/button';
 	import * as Tooltip from '$lib/components/ui/tooltip';
 	import { ATTACHMENT_TOOLTIP_TEXT } from '$lib/constants';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 
 	interface Props {
 		disabled?: boolean;

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddDropdown.svelte

@@ -1,20 +1,20 @@
 <script lang="ts">
-	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
-	import { Plus, File, MessageSquare, Zap, FolderOpen } from '@lucide/svelte';
+	import { File, FolderOpen, MessageSquare, Plus, Zap } from '@lucide/svelte';
+	import {
+		ChatFormActionAddMcpServersSubmenu,
+		ChatFormActionAddReasoningSubmenu,
+		ChatFormActionAddToolsSubmenu
+	} from '$lib/components/app';
+	import { buttonVariants } from '$lib/components/ui/button';
 	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
 	import * as Tooltip from '$lib/components/ui/tooltip';
-	import { buttonVariants } from '$lib/components/ui/button';
 	import { cn } from '$lib/components/ui/utils';
 	import {
 		ATTACHMENT_FILE_ITEMS,
 		ATTACHMENT_TOOLTIP_TEXT,
 		TOOLTIP_DELAY_DURATION
 	} from '$lib/constants';
-	import {
-		ChatFormActionAddToolsSubmenu,
-		ChatFormActionAddMcpServersSubmenu,
-		ChatFormActionAddReasoningSubmenu
-	} from '$lib/components/app';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte';
 
 	interface Props {
@@ -36,15 +36,15 @@
 		class: className = '',
 		disabled = false,
 		hasAudioModality = false,
-		hasVideoModality = false,
-		hasVisionModality = false,
 		hasMcpPromptsSupport = false,
 		hasMcpResourcesSupport = false,
+		hasVideoModality = false,
+		hasVisionModality = false,
 		onFileUpload,
-		onSystemPromptClick,
 		onMcpPromptClick,
+		onMcpResourcesClick,
 		onMcpSettingsClick,
-		onMcpResourcesClick
+		onSystemPromptClick
 	}: Props = $props();
 
 	let dropdownOpen = $state(false);
@@ -59,13 +59,13 @@
 
 	const attachmentMenu = useAttachmentMenu(
 		() => ({
-			hasVisionModality,
 			hasAudioModality,
-			hasVideoModality,
 			hasMcpPromptsSupport,
-			hasMcpResourcesSupport
+			hasMcpResourcesSupport,
+			hasVideoModality,
+			hasVisionModality
 		}),
-		() => ({ onFileUpload, onSystemPromptClick, onMcpPromptClick, onMcpResourcesClick }),
+		() => ({ onFileUpload, onMcpPromptClick, onMcpResourcesClick, onSystemPromptClick }),
 		() => {
 			dropdownOpen = false;
 		}

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpServersSubmenu.svelte

@@ -1,15 +1,15 @@
 <script lang="ts">
-	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
-	import { Settings, Plus } from '@lucide/svelte';
-	import { Switch } from '$lib/components/ui/switch';
+	import { Plus, Settings } from '@lucide/svelte';
+	import { goto } from '$app/navigation';
+	import { DropdownMenuSearchable, McpLogo, McpServerIdentity } from '$lib/components/app';
 	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
-	import { McpLogo, DropdownMenuSearchable, McpServerIdentity } from '$lib/components/app';
+	import { Switch } from '$lib/components/ui/switch';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
+	import { ROUTES } from '$lib/constants/routes';
+	import { HealthCheckStatus } from '$lib/enums';
 	import { conversationsStore } from '$lib/stores/conversations.svelte';
 	import { mcpStore } from '$lib/stores/mcp.svelte';
-	import { HealthCheckStatus } from '$lib/enums';
 	import type { MCPServerSettingsEntry } from '$lib/types';
-	import { goto } from '$app/navigation';
-	import { ROUTES } from '$lib/constants/routes';
 
 	interface Props {
 		onMcpSettingsClick?: () => void;
@@ -24,10 +24,13 @@
 	let hasMcpServers = $derived(mcpServers.length > 0);
 	let filteredMcpServers = $derived.by(() => {
 		const query = mcpSearchQuery.toLowerCase().trim();
+
 		if (!query) return mcpServers;
+
 		return mcpServers.filter((s) => {
 			const name = getServerLabel(s).toLowerCase();
 			const url = s.url.toLowerCase();
+
 			return name.includes(query) || url.includes(query);
 		});
 	});

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte

@@ -1,8 +1,8 @@
 <script lang="ts">
-	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
-	import { Lightbulb, LightbulbOff, Check, Info } from '@lucide/svelte';
+	import { Check, Info, Lightbulb, LightbulbOff } from '@lucide/svelte';
 	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
 	import * as Tooltip from '$lib/components/ui/tooltip';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { useReasoningMenu } from '$lib/hooks/use-reasoning-menu.svelte';
 
 	const reasoning = useReasoningMenu();

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte

@@ -1,30 +1,30 @@
 <script lang="ts">
-	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
-	import type { Snippet } from 'svelte';
-	import * as Tooltip from '$lib/components/ui/tooltip';
-	import * as Sheet from '$lib/components/ui/sheet';
-	import * as Collapsible from '$lib/components/ui/collapsible';
-	import { File, MessageSquare, Zap, FolderOpen } from '@lucide/svelte';
-	import { Switch } from '$lib/components/ui/switch';
-	import { Checkbox } from '$lib/components/ui/checkbox';
-	import { TOOLTIP_DELAY_DURATION } from '$lib/constants';
-	import { ATTACHMENT_FILE_ITEMS } from '$lib/constants/attachment-menu';
-	import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte';
-	import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte';
-	import { useReasoningMenu } from '$lib/hooks/use-reasoning-menu.svelte';
-	import { conversationsStore } from '$lib/stores/conversations.svelte';
-	import { mcpStore } from '$lib/stores/mcp.svelte';
-	import { McpLogo } from '$lib/components/app';
+	import { File, FolderOpen, MessageSquare, Zap } from '@lucide/svelte';
 	import {
-		PencilRuler,
+		Check,
 		ChevronDown,
 		ChevronRight,
 		Lightbulb,
 		LightbulbOff,
-		Check
+		PencilRuler
 	} from '@lucide/svelte';
+	import { McpLogo } from '$lib/components/app';
+	import { Checkbox } from '$lib/components/ui/checkbox';
+	import * as Collapsible from '$lib/components/ui/collapsible';
+	import * as Sheet from '$lib/components/ui/sheet';
+	import { Switch } from '$lib/components/ui/switch';
+	import * as Tooltip from '$lib/components/ui/tooltip';
+	import { TOOLTIP_DELAY_DURATION } from '$lib/constants';
+	import { ATTACHMENT_FILE_ITEMS } from '$lib/constants/attachment-menu';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { HealthCheckStatus } from '$lib/enums';
 	import { AttachmentAction } from '$lib/enums/attachment.enums';
+	import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte';
+	import { useReasoningMenu } from '$lib/hooks/use-reasoning-menu.svelte';
+	import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte';
+	import { conversationsStore } from '$lib/stores/conversations.svelte';
+	import { mcpStore } from '$lib/stores/mcp.svelte';
+	import type { Snippet } from 'svelte';
 
 	interface Props {
 		class?: string;
@@ -45,14 +45,14 @@
 		class: className = '',
 		disabled = false,
 		hasAudioModality = false,
-		hasVisionModality = false,
-		hasVideoModality = false,
 		hasMcpPromptsSupport = false,
 		hasMcpResourcesSupport = false,
+		hasVideoModality = false,
+		hasVisionModality = false,
 		onFileUpload,
-		onSystemPromptClick,
 		onMcpPromptClick,
 		onMcpResourcesClick,
+		onSystemPromptClick,
 		trigger
 	}: Props = $props();
 
@@ -64,13 +64,13 @@
 
 	const attachmentMenu = useAttachmentMenu(
 		() => ({
-			hasVisionModality,
 			hasAudioModality,
-			hasVideoModality,
 			hasMcpPromptsSupport,
-			hasMcpResourcesSupport
+			hasMcpResourcesSupport,
+			hasVideoModality,
+			hasVisionModality
 		}),
-		() => ({ onFileUpload, onSystemPromptClick, onMcpPromptClick, onMcpResourcesClick }),
+		() => ({ onFileUpload, onMcpPromptClick, onMcpResourcesClick, onSystemPromptClick }),
 		() => {
 			sheetOpen = false;
 		}

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte

@@ -1,14 +1,14 @@
 <script lang="ts">
-	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
-	import { PencilRuler, ChevronDown, ChevronRight, Loader2, Info, Check } from '@lucide/svelte';
+	import { Check, ChevronDown, ChevronRight, Info, Loader2, PencilRuler } from '@lucide/svelte';
 	import { Checkbox } from '$lib/components/ui/checkbox';
 	import * as Collapsible from '$lib/components/ui/collapsible';
 	import * as DropdownMenu from '$lib/components/ui/dropdown-menu';
 	import * as Tooltip from '$lib/components/ui/tooltip';
-	import { toolsStore } from '$lib/stores/tools.svelte';
 	import { CLI_FLAGS } from '$lib/constants';
-	import { mcpStore } from '$lib/stores/mcp.svelte';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte';
+	import { mcpStore } from '$lib/stores/mcp.svelte';
+	import { toolsStore } from '$lib/stores/tools.svelte';
 
 	const toolsPanel = useToolsPanel();
 	const hasMcpServersAvailable = $derived(mcpStore.getServers().length > 0);

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionsAdd.svelte

@@ -1,8 +1,8 @@
 <script lang="ts">
-	import { isMobile } from '$lib/stores/viewport.svelte';
+	import ChatFormActionAddButton from './ChatFormActionAddButton.svelte';
 	import ChatFormActionAddDropdown from './ChatFormActionAddDropdown.svelte';
 	import ChatFormActionAddSheet from './ChatFormActionAddSheet.svelte';
-	import ChatFormActionAddButton from './ChatFormActionAddButton.svelte';
+	import { isMobile } from '$lib/stores/viewport.svelte';
 
 	interface Props {
 		disabled?: boolean;
@@ -21,9 +21,9 @@
 	let {
 		disabled = false,
 		hasAudioModality = false,
-		hasVideoModality = false,
 		hasMcpPromptsSupport = false,
 		hasMcpResourcesSupport = false,
+		hasVideoModality = false,
 		hasVisionModality = false,
 		onFileUpload,
 		onMcpPromptClick,

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionModels.svelte

@@ -1,15 +1,15 @@
 <script lang="ts">
+	import { ModelsSelectorDropdown, ModelsSelectorSheet } from '$lib/components/app';
 	import { chatStore } from '$lib/stores/chat.svelte';
+	import { activeMessages } from '$lib/stores/conversations.svelte';
 	import {
-		modelsStore,
 		modelOptions,
+		modelsStore,
 		selectedModelId,
 		selectedModelName
 	} from '$lib/stores/models.svelte';
 	import { isRouterMode, serverError } from '$lib/stores/server.svelte';
-	import { ModelsSelectorDropdown, ModelsSelectorSheet } from '$lib/components/app';
 	import { isMobile } from '$lib/stores/viewport.svelte';
-	import { activeMessages } from '$lib/stores/conversations.svelte';
 
 	interface Props {
 		disabled?: boolean;
@@ -27,9 +27,9 @@
 		disabled = false,
 		forceForegroundText = false,
 		hasAudioModality = $bindable(false),
+		hasModelSelected = $bindable(false),
 		hasVideoModality = $bindable(false),
 		hasVisionModality = $bindable(false),
-		hasModelSelected = $bindable(false),
 		isSelectedModelInCache = $bindable(true),
 		submitTooltip = $bindable(''),
 		useGlobalSelection = false
@@ -46,6 +46,7 @@
 
 	let selectorModel = $derived.by(() => {
 		const storeModel = selectedModelName();
+
 		if (storeModel && storeModel !== conversationModel) {
 			return storeModel;
 		}
@@ -66,6 +67,7 @@
 				modelsStore.selectedModelName = null;
 				modelsStore.clearSelection();
 			}
+
 			lastSyncedConversationModel = conversationModel;
 		} else if (
 			isRouter &&
@@ -76,6 +78,7 @@
 		) {
 			lastSyncedConversationModel = null;
 			const first = modelOptions().find((m) => modelsStore.loadedModelIds.includes(m.model));
+
 			if (first) modelsStore.selectModelById(first.id);
 		}
 	});

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionRecord.svelte

@@ -1,8 +1,8 @@
 <script lang="ts">
-	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 	import { Mic, Square } from '@lucide/svelte';
 	import { Button } from '$lib/components/ui/button';
 	import * as Tooltip from '$lib/components/ui/tooltip';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
 
 	interface Props {
 		class?: string;

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte

@@ -1,28 +1,28 @@
 <script lang="ts">
-	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
-	import { Square, SkipForward } from '@lucide/svelte';
-	import { Button } from '$lib/components/ui/button';
-	import { ChatService } from '$lib/services';
+	import { SkipForward, Square } from '@lucide/svelte';
+	import { goto } from '$app/navigation';
+	import { page } from '$app/state';
 	import {
-		ChatFormActionsAdd,
 		ChatFormActionModels,
 		ChatFormActionRecord,
+		ChatFormActionsAdd,
 		ChatFormActionSubmit,
 		ChatFormContextGauge
 	} from '$lib/components/app';
+	import { Button } from '$lib/components/ui/button';
+	import { ICON_CLASS_DEFAULT } from '$lib/constants/css-classes';
+	import { ROUTES } from '$lib/constants/routes';
 	import { FileTypeCategory, MessageRole } from '$lib/enums';
-	import { mcpStore } from '$lib/stores/mcp.svelte';
-	import { config } from '$lib/stores/settings.svelte';
-	import { activeMessages, conversationsStore } from '$lib/stores/conversations.svelte';
+	import { ChatService } from '$lib/services';
 	import {
 		activeProcessingState,
 		isChatStreaming,
 		isLoading as chatIsLoading
 	} from '$lib/stores/chat.svelte';
+	import { activeMessages, conversationsStore } from '$lib/stores/conversations.svelte';
+	import { mcpStore } from '$lib/stores/mcp.svelte';
+	import { config } from '$lib/stores/settings.svelte';
 	import { getFileTypeCategory } from '$lib/utils';
-	import { goto } from '$app/navigation';
-	import { page } from '$app/state';
-	import { ROUTES } from '$lib/constants/routes';
 
 	interface Props {
 		canSend?: boolean;
@@ -51,15 +51,15 @@
 		isLoading = false,
 		isReasoning = false,
 		isRecording = false,
-		showAddButton = true,
-		showModelSelector = true,
-		uploadedFiles = [],
 		onFileUpload,
+		onMcpPromptClick,
+		onMcpResourcesClick,
 		onMicClick,
 		onStop,
 		onSystemPromptClick,
-		onMcpPromptClick,
-		onMcpResourcesClick
+		showAddButton = true,
+		showModelSelector = true,
+		uploadedFiles = []
 	}: Props = $props();
 
 	let currentConfig = $derived(config());
@@ -105,29 +105,39 @@
 		if (!page.params.id) return false;
 
 		const messages = activeMessages() as DatabaseMessage[];
+
 		let totalHistoricalTokens = 0;
+
 		for (const m of messages) {
 			if (m.role !== MessageRole.ASSISTANT) continue;
+
 			const timings = m.timings;
+
 			if (!timings) continue;
+
 			const agenticLlm = timings.agentic?.llm;
+
 			if (agenticLlm?.prompt_n != null || agenticLlm?.predicted_n != null) {
 				totalHistoricalTokens += (agenticLlm?.prompt_n ?? 0) + (agenticLlm?.predicted_n ?? 0);
 			} else {
 				totalHistoricalTokens += (timings.prompt_n ?? 0) + (timings.predicted_n ?? 0);
 			}
 		}
+
 		if (totalHistoricalTokens > 0) return true;
 
 		if (!chatIsLoading() && !isChatStreaming()) return false;
 
 		const processingState = activeProcessingState();
+
 		if (!processingState) return false;
+
 		const livePromptTokens = Math.max(
 			processingState.promptTokens ?? 0,
 			processingState.promptProgress?.processed ?? 0
 		);
 		const liveOutputTokens = processingState.outputTokensUsed ?? 0;
+
 		return livePromptTokens > 0 || liveOutputTokens > 0;
 	});
 </script>

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContenteditable.svelte

@@ -1,11 +1,8 @@
 <script lang="ts">
-	import { onDestroy, onMount, untrack } from 'svelte';
-	import { mode } from 'mode-watcher';
-	import githubDarkCss from 'highlight.js/styles/github-dark.css?inline';
-	import githubLightCss from 'highlight.js/styles/github.css?inline';
-	import { isMobile } from '$lib/stores/viewport.svelte';
-	import { ColorMode } from '$lib/enums';
 	import { TRIM_LEADING_PADDING_REGEX, TRIM_TRAILING_PADDING_REGEX } from '$lib/constants';
+	import { ColorMode } from '$lib/enums';
+	import { isMobile } from '$lib/stores/viewport.svelte';
+	import type { ContentToken, SourceHistoryEntry } from '$lib/utils';
 	import {
 		badgeAwareWordJump,
 		buildFragment,
@@ -19,10 +16,13 @@
 		SourceHistory,
 		stripBlockBoundaryLineBreaks,
 		syncCodeBlockHatches,
-		tokenizeContent,
-		textOffsetToRange
+		textOffsetToRange,
+		tokenizeContent
 	} from '$lib/utils';
-	import type { ContentToken, SourceHistoryEntry } from '$lib/utils';
+	import githubLightCss from 'highlight.js/styles/github.css?inline';
+	import githubDarkCss from 'highlight.js/styles/github-dark.css?inline';
+	import { mode } from 'mode-watcher';
+	import { onDestroy, onMount, untrack } from 'svelte';
 
 	interface Props {
 		class?: string;
@@ -57,7 +57,9 @@
 	// serialized source, not the DOM shape.
 	function syncEmptyState(serialized?: string) {
 		if (!rootElement) return;
+
 		const source = serialized ?? serializeContent(rootElement);
+
 		rootElement.dataset.empty = source.length === 0 ? 'true' : 'false';
 	}
 
@@ -93,23 +95,24 @@
 	 */
 	function highlightCodeBlockElement(el: HTMLElement): boolean {
 		const segment = el.textContent ?? '';
+
 		if (highlightedSegments.get(el) === segment) return false;
 
 		const open = CODE_BLOCK_OPEN_RE.exec(segment);
+
 		if (!open) return false;
 
 		const prefix = open[0];
 		const language = open[1].trim().split(/\s+/)[0] ?? '';
 		const content = segment.slice(prefix.length, -3);
-
 		const leading = content.match(TRIM_LEADING_PADDING_REGEX)?.[0] ?? '';
 		const trailing = content.match(TRIM_TRAILING_PADDING_REGEX)?.[0] ?? '';
 		const core = content.slice(leading.length, content.length - trailing.length);
-
 		// autoDetect off: re-guessing the language on every keystroke
 		// costs ~38ms a call and flickers while typing
 		const html = core ? highlightCode(core, language || 'text', false) : '';
 		const tpl = document.createElement('template');
+
 		tpl.innerHTML = html;
 
 		el.replaceChildren(
@@ -118,6 +121,7 @@
 			document.createTextNode(trailing + '```')
 		);
 		highlightedSegments.set(el, segment);
+
 		return true;
 	}
 
@@ -136,9 +140,11 @@
 		if (!rootElement) return;
 
 		const range = safeRange();
+
 		if (!range) return;
 
 		let node: Node | null = range.startContainer;
+
 		if (node === rootElement) {
 			node = rootElement.childNodes[range.startOffset - 1] ?? null;
 		}
@@ -146,11 +152,14 @@
 		while (node && node !== rootElement) {
 			if (node instanceof HTMLElement && node.dataset.codeToken === 'block') {
 				const caret = rangeToTextOffset(rootElement, range);
+
 				if (highlightCodeBlockElement(node)) {
 					restoreCaret(caret);
 				}
+
 				return;
 			}
+
 			node = node.parentNode;
 		}
 	}
@@ -182,6 +191,7 @@
 		document.querySelectorAll('style[data-highlight-theme-preview]').forEach((s) => s.remove());
 
 		const style = document.createElement('style');
+
 		style.setAttribute('data-highlight-theme-preview', 'true');
 		style.textContent = isDark ? githubDarkCss : githubLightCss;
 
@@ -196,6 +206,7 @@
 		if (!rootElement) return null;
 
 		const selection = window.getSelection();
+
 		if (!selection || selection.rangeCount === 0) return null;
 
 		const range = selection.getRangeAt(0);
@@ -212,6 +223,7 @@
 
 		const target = textOffsetToRange(rootElement, offset);
 		const selection = window.getSelection();
+
 		if (!selection) return;
 
 		if (extend && selection.anchorNode) {
@@ -221,6 +233,7 @@
 				target.startContainer,
 				target.startOffset
 			);
+
 			return;
 		}
 
@@ -230,14 +243,16 @@
 
 	function resizeHeight() {
 		if (!rootElement) return;
+
 		rootElement.style.height = 'auto';
 		rootElement.style.height = `${rootElement.scrollHeight}px`;
 	}
 
 	function recordHistory(newGroup: boolean) {
 		if (!rootElement) return;
+
 		history.push(
-			{ value: lastEmittedValue, caret: rangeToTextOffset(rootElement, safeRange()) },
+			{ caret: rangeToTextOffset(rootElement, safeRange()), value: lastEmittedValue },
 			Date.now(),
 			newGroup
 		);
@@ -262,10 +277,12 @@
 		// lands on the line directly below the block.
 		if (inputType === 'insertLineBreak' || inputType === 'insertParagraph') {
 			const caret = rangeToTextOffset(rootElement, safeRange());
+
 			if (stripBlockBoundaryLineBreaks(rootElement)) {
 				restoreCaret(caret);
 			} else {
 				const source = serializeContent(rootElement);
+
 				let end = caret;
 
 				// the caret must end up after the inserted \n; some browsers
@@ -284,7 +301,9 @@
 				// \n doubles as a block's separator line (source ends with
 				// \n\n) or sits inside a block element.
 				let last = rootElement.lastChild;
+
 				while (last && last.nodeName === 'BR') last = last.previousSibling;
+
 				if (
 					end === source.length &&
 					source.endsWith('\n') &&
@@ -302,7 +321,9 @@
 		syncCodeBlockHatches(rootElement);
 
 		const serialized = serializeContent(rootElement);
+
 		syncEmptyState(serialized);
+
 		if (serialized === lastEmittedValue) return;
 
 		// Plain typing/deletes coalesce per time window; structural edits
@@ -316,6 +337,7 @@
 		// completed or broken) - the browser-owned text nodes cannot
 		// restyle themselves across element boundaries.
 		const tokens = tokenizeContent(serialized);
+
 		if (!domMatchesTokens(rootElement, tokens)) {
 			renderTokens(tokens);
 
@@ -324,6 +346,7 @@
 			// block element and the rebuild splits it back out, which
 			// synthesizes the separator newline) - keep value in sync.
 			const reserialized = serializeContent(rootElement);
+
 			if (reserialized !== serialized) {
 				lastEmittedValue = reserialized;
 				value = reserialized;
@@ -361,6 +384,7 @@
 		if (!rootElement) return;
 
 		const range = safeRange();
+
 		if (!range) return;
 
 		if (!range.collapsed) {
@@ -374,16 +398,22 @@
 		// a break at the very end of a code block exits the block (the
 		// new line belongs below it, not inside)
 		let exitBlock: HTMLElement | null = null;
+
 		if (container.nodeType === Node.TEXT_NODE) {
 			let node: Node | null = container.parentNode;
+
 			while (node && node !== rootElement) {
 				if (node instanceof HTMLElement && node.dataset.codeToken === 'block') {
 					const tail = document.createRange();
+
 					tail.setStart(container, offset);
 					tail.setEnd(node, node.childNodes.length);
+
 					if (tail.toString().length === 0) exitBlock = node;
+
 					break;
 				}
+
 				node = node.parentNode;
 			}
 		}
@@ -392,6 +422,7 @@
 			exitBlock.after(nl);
 		} else if (container.nodeType === Node.TEXT_NODE) {
 			const text = container as Text;
+
 			if (offset === 0) {
 				text.before(nl);
 			} else if (offset === text.length) {
@@ -405,6 +436,7 @@
 
 		const selection = window.getSelection();
 		const after = document.createRange();
+
 		after.setStartAfter(nl);
 		after.collapse(true);
 		selection?.removeAllRanges();
@@ -428,9 +460,11 @@
 		if (rootElement.firstChild?.nodeName === 'BR') return false;
 
 		const first = rootElement.firstChild;
+
 		if (!(first instanceof HTMLElement) || first.dataset.codeToken !== 'block') return false;
 
 		const range = safeRange();
+
 		if (!range || !range.collapsed) return false;
 
 		// the caret must sit inside the block: on its very first
@@ -439,16 +473,19 @@
 		if (!first.contains(range.startContainer)) return false;
 
 		const caret = rangeToTextOffset(rootElement, range);
+
 		if (key === 'ArrowLeft') {
 			if (caret !== 0) return false;
 		} else {
 			const firstLineEnd = (first.textContent ?? '').indexOf('\n');
+
 			if (firstLineEnd !== -1 && caret > firstLineEnd) return false;
 		}
 
 		// eslint-disable-next-line svelte/no-dom-manipulating -- the token layer is owned imperatively; Svelte renders only the contenteditable host, never its children
 		rootElement.prepend(document.createElement('br'));
 		restoreCaret(0, extend);
+
 		return true;
 	}
 
@@ -464,14 +501,17 @@
 		if (!rootElement) return;
 
 		const first = rootElement.firstChild;
+
 		if (first?.nodeName !== 'BR') return;
 
 		const second = first.nextSibling;
+
 		if (!(second instanceof HTMLElement) || second.dataset.codeToken !== 'block') return;
 
 		const range = safeRange();
 		const onHatch =
 			range !== null && range.startContainer === rootElement && range.startOffset === 0;
+
 		if (!onHatch) {
 			first.remove();
 		}
@@ -491,6 +531,7 @@
 	 */
 	function handleKeydown(event: KeyboardEvent) {
 		const mod = event.ctrlKey || event.metaKey;
+
 		if (mod && !event.altKey && !isComposing && rootElement) {
 			const key = event.key.toLowerCase();
 			const isUndo = key === 'z' && !event.shiftKey;
@@ -499,11 +540,13 @@
 			if (isUndo || isRedo) {
 				event.preventDefault();
 				const current = {
-					value: lastEmittedValue,
-					caret: rangeToTextOffset(rootElement, safeRange())
+					caret: rangeToTextOffset(rootElement, safeRange()),
+					value: lastEmittedValue
 				};
 				const entry = isUndo ? history.undo(current) : history.redo(current);
+
 				if (entry) applyHistoryEntry(entry);
+
 				return;
 			}
 		}
@@ -524,6 +567,7 @@
 			// stuck on the old line (see insertLineBreak).
 			event.preventDefault();
 			insertLineBreak();
+
 			return;
 		}
 
@@ -543,6 +587,7 @@
 			// re-tokenize/re-highlight follows.
 			event.preventDefault();
 			document.execCommand('insertLineBreak');
+
 			return;
 		}
 
@@ -555,6 +600,7 @@
 		) {
 			if (moveCaretBeforeLeadingCodeBlock(event.key, event.shiftKey)) {
 				event.preventDefault();
+
 				return;
 			}
 		}
@@ -574,6 +620,7 @@
 				if (target !== null) {
 					event.preventDefault();
 					restoreCaret(target, event.shiftKey);
+
 					return;
 				}
 			}
@@ -586,6 +633,7 @@
 	// change as our own and does not re-render.
 	function applyHistoryEntry(entry: SourceHistoryEntry) {
 		if (!rootElement) return;
+
 		renderTokens(tokenizeContent(entry.value));
 		lastEmittedValue = entry.value;
 		value = entry.value;
@@ -602,6 +650,7 @@
 	 */
 	function handlePasteEvent(event: ClipboardEvent) {
 		const pasted = event.clipboardData?.getData('text/plain');
+
 		if (pasted && pasted.length > 0) {
 			event.preventDefault();
 
@@ -609,6 +658,7 @@
 			// element-boundary carets (e.g. right before a badge) Chromium's
 			// insertText can drop the preceding text node's trailing whitespace.
 			const range = safeRange();
+
 			if (rootElement && range && range.collapsed) {
 				restoreCaret(rangeToTextOffset(rootElement, range));
 			}
@@ -621,6 +671,7 @@
 	// consumes the event (files, quoted prompts, long text).
 	function handlePaste(event: ClipboardEvent) {
 		onPaste?.(event);
+
 		if (!event.defaultPrevented) {
 			handlePasteEvent(event);
 		}
@@ -634,20 +685,23 @@
 		if (!rootElement) return null;
 
 		const range = safeRange();
+
 		if (!range || range.collapsed) return null;
 
 		const startRange = range.cloneRange();
+
 		startRange.collapse(true);
 
 		const source = serializeContent(rootElement);
 		const start = rangeToTextOffset(rootElement, startRange);
 		const end = rangeToTextOffset(rootElement, range);
 
-		return { text: source.slice(start, end), range };
+		return { range, text: source.slice(start, end) };
 	}
 
 	function handleCopy(event: ClipboardEvent) {
 		const slice = selectionSourceSlice();
+
 		if (!slice) return;
 
 		event.clipboardData?.setData('text/plain', slice.text);
@@ -656,6 +710,7 @@
 
 	function handleCut(event: ClipboardEvent) {
 		const slice = selectionSourceSlice();
+
 		if (!slice) return;
 
 		event.clipboardData?.setData('text/plain', slice.text);
@@ -675,6 +730,7 @@
 		resizeHeight();
 		syncEmptyState();
 		document.addEventListener('selectionchange', handleSelectionChange);
+
 		if (!isMobile.current) {
 			rootElement?.focus({ preventScroll: true });
 		}
@@ -689,6 +745,7 @@
 	// browser already owns the right shape.
 	$effect(() => {
 		const incoming = value ?? '';
+
 		if (incoming === lastEmittedValue) return;
 
 		recordHistory(true); // external edit (mention insert, clear, ...): own undo step
@@ -702,6 +759,7 @@
 
 	export function getCaretOffset(): number {
 		if (!rootElement) return 0;
+
 		return rangeToTextOffset(rootElement, safeRange());
 	}
 
@@ -710,11 +768,13 @@
 		if (rootElement && rootElement !== document.activeElement) {
 			rootElement.focus({ preventScroll: true });
 		}
+
 		restoreCaret(offset);
 	}
 
 	export function focus() {
 		if (isMobile.current) return;
+
 		rootElement?.focus({ preventScroll: true });
 	}
 

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ChatFormContextGauge.svelte

@@ -1,9 +1,7 @@
 <script lang="ts">
-	import { untrack } from 'svelte';
-	import { activeConversation, activeMessages } from '$lib/stores/conversations.svelte';
-	import { chatStore, isChatStreaming, isLoading } from '$lib/stores/chat.svelte';
-	import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte';
 	import ContextGaugeDial from './ContextGaugeDial.svelte';
+	import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte';
+	import { chatStore, isChatStreaming, isLoading } from '$lib/stores/chat.svelte';
 	import {
 		gaugeTriggerClick,
 		gaugeTriggerEnter,
@@ -11,22 +9,28 @@
 		gaugeTriggerLeave,
 		gaugeTriggerPointerDown
 	} from '$lib/stores/context-gauge-popup.svelte';
+	import { activeConversation, activeMessages } from '$lib/stores/conversations.svelte';
+	import { untrack } from 'svelte';
 
 	const gauge = useContextGauge();
 
 	$effect(() => {
 		const conv = activeConversation();
+
 		untrack(() => chatStore.setActiveProcessingConversation(conv?.id ?? null));
 	});
 
 	$effect(() => {
 		const conv = activeConversation();
 		const messages = activeMessages() as DatabaseMessage[];
+
 		if (!conv) return;
+
 		if (isLoading() || isChatStreaming()) return;
 
 		if (messages.length === 0) {
 			untrack(() => chatStore.clearProcessingState(conv.id));
+
 			return;
 		}
 

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetailRow.svelte

@@ -5,7 +5,7 @@
 		subtitle?: string;
 	}
 
-	let { label, value, subtitle }: Props = $props();
+	let { label, subtitle, value }: Props = $props();
 </script>
 
 <div class="grid gap-1.5">

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte

@@ -1,8 +1,8 @@
 <script lang="ts">
+	import ContextGaugeDetailRow from './ContextGaugeDetailRow.svelte';
 	import { ChevronDown } from '@lucide/svelte';
 	import * as Collapsible from '$lib/components/ui/collapsible';
 	import { STATS_UNITS } from '$lib/constants';
-	import ContextGaugeDetailRow from './ContextGaugeDetailRow.svelte';
 
 	interface Props {
 		currentRead: number;
@@ -18,15 +18,15 @@
 	}
 
 	let {
-		currentRead,
-		currentFresh,
-		currentCache,
-		currentOutput,
-		kvTotal,
-		cumulativeRead,
-		cumulativeOutput,
-		cumulativeCacheTotal,
 		averageTokensPerSecond,
+		cumulativeCacheTotal,
+		cumulativeOutput,
+		cumulativeRead,
+		currentCache,
+		currentFresh,
+		currentOutput,
+		currentRead,
+		kvTotal,
 		transientDetails
 	}: Props = $props();
 

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDial.svelte

@@ -8,7 +8,7 @@
 		size?: 'sm' | 'md';
 	}
 
-	let { percent, level, size = 'sm' }: Props = $props();
+	let { level, percent, size = 'sm' }: Props = $props();
 
 	const RADIUS = 11;
 	const CIRCUMFERENCE = 2 * Math.PI * RADIUS;

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeLoadModel.svelte

@@ -8,7 +8,7 @@
 		onLoad: () => void;
 	}
 
-	let { modelId, isLoading, onLoad }: Props = $props();
+	let { isLoading, modelId, onLoad }: Props = $props();
 </script>
 
 {#if modelId !== null && !isLoading}

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugePopup.svelte

@@ -1,15 +1,15 @@
 <script lang="ts">
-	import { formatParameters } from '$lib/utils/formatters';
-	import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte';
+	import { colorLevelBgClass, colorLevelTextClass } from './context-gauge';
 	import ContextGaugeDetails from './ContextGaugeDetails.svelte';
 	import ContextGaugeLoadModel from './ContextGaugeLoadModel.svelte';
-	import { colorLevelBgClass, colorLevelTextClass } from './context-gauge';
+	import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte';
 	import {
-		gaugePopup,
 		gaugeCardEnter,
 		gaugeCardLeave,
+		gaugePopup,
 		gaugePopupClose
 	} from '$lib/stores/context-gauge-popup.svelte';
+	import { formatParameters } from '$lib/utils/formatters';
 
 	const gauge = useContextGauge();
 
@@ -30,13 +30,18 @@
 
 		const onPointerDown = (event: PointerEvent) => {
 			const target = event.target;
+
 			if (!(target instanceof Node)) return;
+
 			if (cardEl?.contains(target)) return;
+
 			if (target instanceof Element && target.closest('[data-context-gauge-trigger]')) return;
+
 			gaugePopupClose();
 		};
 
 		document.addEventListener('pointerdown', onPointerDown, true);
+
 		return () => document.removeEventListener('pointerdown', onPointerDown, true);
 	});
 

tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/context-gauge.ts

@@ -5,8 +5,11 @@ const CRITICAL_THRESHOLD = 95;
 
 export function colorLevelFromPercent(percent: number | null): ColorLevel {
 	if (percent === null) return 'neutral';
+
 	if (percent >= CRITICAL_THRESHOLD) return 'critical';
+
 	if (percent >= WARNING_THRESHOLD) return 'warning';
+
 	return 'ok';
 }