model : Granite-Switch Architecture (#25107)

* granite-switch: add llama.cpp backend (POC, CPU)

New "granite-switch" architecture: a dense, all-attention Granite-4.1
model with N embedded LoRA adapters selected per-token by control tokens.

- gguf-py schema (arch, KV keys, stacked LoRA tensor names) + writer helpers
- conversion/granite.py: GraniteSwitchModel converter (stacks N adapters +
  zero base slot into per-projection A/B tensors; emits switch metadata)
- C++ arch registration (llama-arch.{h,cpp}, llama-model.{h,cpp})
- src/models/granite_switch.cpp: load + per-token switched-LoRA graph via
  ggml_mul_mat_id over stacked tensors; sticky per-token index + control-token
  substitution in llm_graph_input_switch::set_input
- llm_graph_input_switch in src/models/models.h

Runs end-to-end on CPU: convert 3b checkpoint (842 tensors, stacked dim 13)
and generate on both base and control-token paths. Sticky switch state is
single-sequence (POC); full multi-sequence machinery is a follow-up.

* granite-switch: add Mac (Metal) build + mid-sequence switch demo script

Self-contained script to build llama.cpp on Apple Silicon (Metal),
convert the composed 3b checkpoint, and run the crisp mid-sequence
adapter-switch demos verified on Vela:
  - answerability: <|answerability|> mid-seq -> "unanswerable"
  - query_rewrite: <|query_rewrite|> mid-seq -> {"rewritten_question": ...}
Each demo runs the same prompt twice, differing only by a control token
placed before the assistant turn, so the per-token switch is visible.

* granite-switch mac demo: add -no-cnv so each run is one-shot

The composed model ships a chat template, so llama-completion auto-enables
interactive conversation mode and halts at a `>` prompt after generating,
stalling the script. -no-cnv disables conversation mode: generate once from
the raw prompt and exit (also prints special tokens, making the switch visible).

* granite-switch: replace global sticky index with in-graph router attention

The POC computed the per-token adapter index on the CPU and carried it
across ubatches in ONE global `mutable int32_t poc_sticky_index`, reset
only when a ubatch contained sequence position 0. That global had two
problems:

  1. Concurrency: with multiple sequences in a batch it was last-writer-
     wins — one sequence's adapter leaked into the others.
  2. Multi-turn: an interactive `ollama run` chat continues one KV cache,
     so turn 2 never saw position 0 and the index never reset — the
     adapter stayed stuck on across turns.

Port the vLLM/HF backend mechanism faithfully: a single-head causal
"router" attention recovers the adapter index in-graph. Per token, only
dim 0 carries signal — Q[0]=1, K[0]=+gain for a control token / -gain
otherwise, V[0]=adapter slot / 0 — and the causal softmax over the single
visible control token recovers that adapter's slot (readback =
clamp(round(V[0]), 0, n_adapters)). gain=15 matches config.py and is
F16-safe (no F32 cache).

The router's K/V live in the model KV cache at an extra layer
R == hparams.router_layer (== n_layer). We bump n_layer_all to n_real+1
so the cache allocator gives the router its own per-sequence slot, and
set n_layer_nextn=1 so n_layer() stays n_real — the decoder loop and
tensor loading are untouched and never reference layer R. The router K is
exempted from the k-shift RoPE loop (its dim-0 value is a literal
magnitude, not a rotation).

Because the selection now lives in the per-sequence KV cache, CONCURRENT
requests are isolated for free (problem 1 fixed; verified by
scratch/concurrent_switch_test.cpp). set_input becomes stateless pure
per-token maps; the global is gone.

Single-switch contract / known limitation, identical to vLLM & HF: the
gain is flat (no recency), so within one sequence there is no mechanism to
revert to base mid-sequence — once an adapter fires it stays on until that
sequence ends (problem 2 is therefore NOT fixed by a faithful copy; vLLM/HF
avoid it only because each served request is a fresh sequence). A client
continuing one KV cache across turns must start a fresh sequence per turn,
or opt into a recency-biased router (a deliberate divergence, not done
here). Documented in granite_switch.cpp and asserted by
scratch/multiturn_leak_test.cpp.

Verified (CPU): both demos unchanged (answerability -> "unanswerable",
query_rewrite -> rewritten query); concurrent two-sequence isolation
passes; multi-turn carry-over matches the vLLM/HF contract.

* granite-switch: drop scratch tests and mac demo for upstream PR

Remove the local-only development artifacts that should not ship in the
upstream PR:
  - granite-switch-mac-demo.sh (local Metal build + demo driver)
  - scratch/concurrent_switch_test.cpp
  - scratch/multiturn_leak_test.cpp

Also drop the now-dangling reference to the scratch tests from the
granite_switch.cpp header comment. Leaves only the core architecture
support (conversion, gguf constants, llama-arch/model/kv-cache, and the
granite_switch graph).

* granite-switch: trim comments to match native llama.cpp style

* granite-switch: trim conversion comments to match native style

* granite-switch: drop unused adapter_ranks metadata

* granite-switch: rename arch to graniteswitch and drop obid alias

* granite-switch: fix non-ASCII comments and document router gain assumption

* granite-switch: drop section comments from constants.py to match native style

* granite-switch: add functional tensor block comments matching Granite4 Vision style

* granite-switch: clarify n_expert_used comment

State the actual constraint: mul_mat_id needs n_expert_used == 1, and
since the GGUF carries expert_count = 0 the generic loader's
n_expert == 0 => n_expert_used == 0 assertion has already passed by the
time load_arch_hparams runs, so it is forced to 1 here.

* granite-switch: note n_layer_nextn reuse has no MTP

The router carving reuses n_layer_nextn, normally the MTP/next-token
count. Clarify in the comment that it is borrowed here purely as the
trailing-layers lever and that there is no MTP head, to spare readers
the double-take.

* granite-switch: rename source file and apply review nits

* granite-switch: don't force LoRA tensors to F16, follow --outtype instead

* granite-switch: drop redundant _permute_qk wrapper, call LlamaModel.permute directly

* granite-switch: read router gain from GGUF (control_token_gain) instead of hardcoding 15.0

* granite-switch: derive n_slots()

* granite-switch: move llm_graph_input_switch into granite-switch.cpp

* granite-switch: cut AI-style narration comments

* granite-switch: collapse multi-line comments

* granite-switch: rename control_token_* maps to adapter_token_*

* granite-switch: cut noise comments

* granite-switch: rename embedded LoRA tensors to <base>.lora_a/lora_b

* granite-switch: GGML_ASSERT token input to avoid UB on embeddings

* granite-switch: TODO for raw embedding input support

* granite-switch: collapse LoRA tensor constants to .lora_a/.lora_b suffix

* granite-switch: drop n_expert_used hack, guard mul_mat_id buft probe

* granite-switch: stop forcing dense expert counts, read from config

* granite-switch: renamed control_token_gain metadata key to router_gain

* granite-switch: trim header comments to match native style

* granite-switch: collapse LoRA tensors to base name + suffix

* granite-switch: inline suffix checks in tensor op resolution

* granite-switch: drop switch-lora struct comment

* granite-switch: guard router layer index and inline n_slots

* granite-switch: group adapter metadata under {arch}.adapters.* namespace

* granite-switch: add hparams.has_rope(il) for KV-shift rope skipping

* granite-switch: skip arch in test-llama-archs (adapter fixture missing, TODO)

* granite-switch: Keys.Adapters namespace + simplify n_slots

* granite-switch: validate substitute token ids against n_vocab

* granite-switch: bound adapter count and lora rank from GGUF

* granite-switch: reject MTP context type when router_layer is set

* granite-switch: throw on bad adapter metadata instead of GGML_ASSERT

* granite-switch: use ASCII +/- in router K signal comment

* granite-switch: document n_layer_nextn repurpose and its leak points

* granite-switch: gate lora_a/lora_b op mapping on router_layer

* granite-switch: label all three preview model sizes

.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

common/arg.cpp

@@ -3308,6 +3308,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.server_tools = parse_csv_row(value);
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS"));
+    add_opt(common_arg(
+        {"--tools-runtime"}, "OPTION",
+        "experimental: run tools in a separate runtime environment (default: none, use host environment)\n"
+        "available options:\n"
+        "  'docker:<image>': spin up a new Docker container and reuse it for all invocations, clean up on server exit\n"
+        "  'docker-container:<id>': use an existing Docker container by ID, won't stop on server exit\n",
+        [](common_params & params, const std::string & value) {
+            params.server_tools_runtime = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_TOOLS_RUNTIME"));
     add_opt(common_arg(
         {"--mcp-servers-config"}, "PATH",
         "experimental: path to JSON file with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)\n"

common/common.h

@@ -655,6 +655,7 @@ struct common_params {
 
     // enable built-in tools
     std::vector<std::string> server_tools;
+    std::string server_tools_runtime;
 
     // MCP server configs (Cursor-compatible JSON)
     std::string mcp_servers_config;   // path to JSON file with MCP server definitions

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

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/cpy.cu

@@ -253,9 +253,9 @@ static void ggml_cpy_f32_q8_0_cuda(
     const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
 
     GGML_ASSERT(ne % QK8_0 == 0);
-    const int64_t num_blocks = ne / QK8_0;
+    const int64_t num_blocks = (ne/QK8_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
     GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_q8_0, QK8_0><<<num_blocks, 1, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_q8_0, QK8_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
         (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
 
@@ -264,9 +264,9 @@ static void ggml_cpy_q8_0_f32_cuda(
     const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t nb00, const int64_t nb01, const int64_t nb02,
     const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
 
-    const int64_t num_blocks = ne;
+    const int64_t num_blocks = (ne/QK8_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
     GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_q_f32<cpy_blck_q8_0_f32, QK8_0><<<num_blocks, 1, 0, stream>>>
+    cpy_q_f32<cpy_blck_q8_0_f32, QK8_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
         (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
 
@@ -276,9 +276,9 @@ static void ggml_cpy_f32_q4_0_cuda(
     const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
 
     GGML_ASSERT(ne % QK4_0 == 0);
-    const int64_t num_blocks = ne / QK4_0;
+    const int64_t num_blocks = (ne/QK4_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
     GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_q4_0, QK4_0><<<num_blocks, 1, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_q4_0, QK4_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
         (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
 
@@ -289,9 +289,9 @@ static void ggml_cpy_q4_0_f32_cuda(
     const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
     const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
     cudaStream_t stream) {
-    const int64_t num_blocks = ne;
+    const int64_t num_blocks = (ne/QK4_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
     GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_q_f32<cpy_blck_q_f32<dequantize_q4_0, QK4_0>, QK4_0><<<num_blocks, 1, 0, stream>>>(
+    cpy_q_f32<cpy_blck_q_f32<dequantize_q4_0, QK4_0>, QK4_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
         cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
          ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
@@ -302,9 +302,9 @@ static void ggml_cpy_f32_q4_1_cuda(
     const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
 
     GGML_ASSERT(ne % QK4_1 == 0);
-    const int64_t num_blocks = ne / QK4_1;
+    const int64_t num_blocks = (ne/QK4_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
     GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_q4_1, QK4_1><<<num_blocks, 1, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_q4_1, QK4_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
         (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
 
@@ -315,9 +315,9 @@ static void ggml_cpy_q4_1_f32_cuda(
     const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
     const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
     cudaStream_t stream) {
-    const int64_t num_blocks = ne;
+    const int64_t num_blocks = (ne/QK4_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
     GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_q_f32<cpy_blck_q_f32<dequantize_q4_1, QK4_1>, QK4_1><<<num_blocks, 1, 0, stream>>>(
+    cpy_q_f32<cpy_blck_q_f32<dequantize_q4_1, QK4_1>, QK4_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
         cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
          ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
@@ -328,9 +328,9 @@ static void ggml_cpy_f32_q5_0_cuda(
     const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
 
     GGML_ASSERT(ne % QK5_0 == 0);
-    const int64_t num_blocks = ne / QK5_0;
+    const int64_t num_blocks = (ne/QK5_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
     GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_q5_0, QK5_0><<<num_blocks, 1, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_q5_0, QK5_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
         (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
 
@@ -341,9 +341,9 @@ static void ggml_cpy_q5_0_f32_cuda(
     const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
     const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
     cudaStream_t stream) {
-    const int64_t num_blocks = ne;
+    const int64_t num_blocks = (ne/QK5_0 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
     GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_q_f32<cpy_blck_q_f32<dequantize_q5_0, QK5_0>, QK5_0><<<num_blocks, 1, 0, stream>>>(
+    cpy_q_f32<cpy_blck_q_f32<dequantize_q5_0, QK5_0>, QK5_0><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
         cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
         ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
@@ -354,9 +354,9 @@ static void ggml_cpy_f32_q5_1_cuda(
     const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
 
     GGML_ASSERT(ne % QK5_1 == 0);
-    const int64_t num_blocks = ne / QK5_1;
+    const int64_t num_blocks = (ne/QK5_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
     GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_q5_1, QK5_1><<<num_blocks, 1, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_q5_1, QK5_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
         (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
 
@@ -367,9 +367,9 @@ static void ggml_cpy_q5_1_f32_cuda(
     const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12,
     const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13,
     cudaStream_t stream) {
-    const int64_t num_blocks = ne;
+    const int64_t num_blocks = (ne/QK5_1 + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
     GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_q_f32<cpy_blck_q_f32<dequantize_q5_1, QK5_1>, QK5_1><<<num_blocks, 1, 0, stream>>>(
+    cpy_q_f32<cpy_blck_q_f32<dequantize_q5_1, QK5_1>, QK5_1><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>(
         cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03,
         ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
@@ -380,9 +380,9 @@ static void ggml_cpy_f32_iq4_nl_cuda(
     const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, cudaStream_t stream) {
 
     GGML_ASSERT(ne % QK4_NL == 0);
-    const int64_t num_blocks = ne / QK4_NL;
+    const int64_t num_blocks = (ne/QK4_NL + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE;
     GGML_ASSERT(num_blocks <= INT_MAX);
-    cpy_f32_q<cpy_blck_f32_iq4_nl, QK4_NL><<<num_blocks, 1, 0, stream>>>
+    cpy_f32_q<cpy_blck_f32_iq4_nl, QK4_NL><<<num_blocks, CUDA_CPY_BLOCK_SIZE, 0, stream>>>
         (cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13);
 }
 

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",
@@ -3972,6 +3981,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)
 

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);
@@ -2596,6 +2599,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

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

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]);
 
@@ -8576,6 +8579,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
             test_cases.emplace_back(new test_cpy(type_src, type_dst, {256, 2, 3, 4}, {-1,-1,-1,-1}, {1, 0, 2, 3})); // cpy not-contiguous
         }
     }
+    // quant block count not a multiple of the kernel block size
+    test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_Q4_0, {96, 1, 1, 1}));
+    test_cases.emplace_back(new test_cpy(GGML_TYPE_Q4_0, GGML_TYPE_F32, {96, 1, 1, 1}));
     test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_I32, {256, 2, 3, 4}));
     test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_I32, {256, 2, 3, 4}, {-1,-1,-1,-1}, {1, 0, 2, 3}));
     test_cases.emplace_back(new test_cpy(GGML_TYPE_I32, GGML_TYPE_F32, {256, 2, 3, 4}));
@@ -8753,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/cli/README.md

@@ -54,7 +54,6 @@
 | `-ctv, --cache-type-v TYPE` | KV cache data type for V<br/>allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1<br/>(default: f16)<br/>(env: LLAMA_ARG_CACHE_TYPE_V) |
 | `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)<br/>(env: LLAMA_ARG_DEFRAG_THOLD) |
 | `-np, --parallel N` | number of parallel sequences to decode (default: 1)<br/>(env: LLAMA_ARG_N_PARALLEL) |
-| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)<br/>(env: LLAMA_ARG_RPC) |
 | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
 | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>(env: LLAMA_ARG_MMAP) |
 | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available<br/>(env: LLAMA_ARG_DIO) |

tools/completion/README.md

@@ -137,7 +137,6 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1
 | `-ctv, --cache-type-v TYPE` | KV cache data type for V<br/>allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1<br/>(default: f16)<br/>(env: LLAMA_ARG_CACHE_TYPE_V) |
 | `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)<br/>(env: LLAMA_ARG_DEFRAG_THOLD) |
 | `-np, --parallel N` | number of parallel sequences to decode (default: 1)<br/>(env: LLAMA_ARG_N_PARALLEL) |
-| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)<br/>(env: LLAMA_ARG_RPC) |
 | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
 | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>(env: LLAMA_ARG_MMAP) |
 | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available<br/>(env: LLAMA_ARG_DIO) |

tools/server/README-dev.md

@@ -201,6 +201,7 @@ Invoke a tool call, request body is a JSON object with:
 
 Headers:
 - `x-tool-cwd`: optional; if set, use as the CWD for tool; this is not part of tool's params because it's meant to be set by the runtime, not the LLM itself
+- `x-tool-runtime`: optional; if set, run the tool inside this isolate instead of on the host. Only `docker-container:<id>` is supported for now, using an already-running container
 
 Returns JSON object. There are two response formats (MCP tools use the same two formats: their result content is concatenated into `plain_text_response`, and RPC or tool errors are surfaced as the `error` string):
 

tools/server/README.md

@@ -71,7 +71,6 @@ For the full list of features, please refer to [server's changelog](https://gith
 | `-ctk, --cache-type-k TYPE` | KV cache data type for K<br/>allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1<br/>(default: f16)<br/>(env: LLAMA_ARG_CACHE_TYPE_K) |
 | `-ctv, --cache-type-v TYPE` | KV cache data type for V<br/>allowed values: f32, f16, bf16, q8_0, q4_0, q4_1, iq4_nl, q5_0, q5_1<br/>(default: f16)<br/>(env: LLAMA_ARG_CACHE_TYPE_V) |
 | `-dt, --defrag-thold N` | KV cache defragmentation threshold (DEPRECATED)<br/>(env: LLAMA_ARG_DEFRAG_THOLD) |
-| `--rpc SERVERS` | comma-separated list of RPC servers (host:port)<br/>(env: LLAMA_ARG_RPC) |
 | `--mlock` | DEPRECATED in favor of `--load-mode`: force system to keep model in RAM rather than swapping or compressing<br/>(env: LLAMA_ARG_MLOCK) |
 | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)<br/>(env: LLAMA_ARG_MMAP) |
 | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available<br/>(env: LLAMA_ARG_DIO) |
@@ -198,6 +197,8 @@ For the full list of features, please refer to [server's changelog](https://gith
 | `--ui-config, --webui-config JSON` | JSON that provides default UI settings (overrides UI defaults)<br/>(env: LLAMA_ARG_UI_CONFIG) |
 | `--ui-config-file, --webui-config-file PATH` | JSON file that provides default UI settings (overrides UI defaults)<br/>(env: LLAMA_ARG_UI_CONFIG_FILE) |
 | `--ui-mcp-proxy, --webui-mcp-proxy, --no-ui-mcp-proxy, --no-webui-mcp-proxy` | experimental: whether to enable MCP CORS proxy - do not enable in untrusted environments (default: disabled)<br/>(env: LLAMA_ARG_UI_MCP_PROXY) |
+| `--tools TOOL1,TOOL2,...` | experimental: whether to enable built-in tools for AI agents - do not enable in untrusted environments (default: no tools)<br/>specify "all" to enable all tools<br/>available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, get_datetime<br/>note: for security reasons, this will limit --cors-origins to localhost by default<br/>(env: LLAMA_ARG_TOOLS) |
+| `--tools-runtime OPTION` | experimental: run tools in a separate runtime environment (default: none, use host environment)<br/>available options:<br/>  'docker:<image>': spin up a new Docker container and reuse it for all invocations, clean up on server exit<br/>  'docker-container:<id>': use an existing Docker container by ID, won't stop on server exit<br/><br/>(env: LLAMA_ARG_TOOLS_RUNTIME) |
 | `--tools TOOL1,TOOL2,...` | experimental: whether to enable built-in tools for AI agents - do not enable in untrusted environments (default: no tools)<br/>specify "all" to enable all tools<br/>available tools: read_file, file_glob_search, grep_search, exec_shell_command, write_file, edit_file, get_datetime, get_info<br/>note: for security reasons, this will limit --cors-origins to localhost by default<br/>(env: LLAMA_ARG_TOOLS) |
 | `--mcp-servers-config PATH` | experimental: path to JSON file with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)<br/>note: for security reasons, this will limit --cors-origins to localhost by default<br/>(env: LLAMA_ARG_MCP_SERVERS_CONFIG) |
 | `--mcp-servers-json JSON` | experimental: inline JSON with MCP server definitions (Cursor-compatible format) - do not enable in untrusted environments (default: none)<br/>note: for security reasons, this will limit --cors-origins to localhost by default<br/>(env: LLAMA_ARG_MCP_SERVERS_JSON) |

tools/server/server-tools.cpp

@@ -10,12 +10,15 @@
 #include <ctime>
 #include <atomic>
 #include <cstring>
+#include <cstdint>
 #include <cstdlib>
 #include <algorithm>
+#include <iterator>
 #include <unordered_set>
 #include <tuple>
 #include <functional>
 #include <memory>
+#include <mutex>
 
 #if defined(_WIN32)
 #   ifndef NOMINMAX
@@ -71,6 +74,7 @@ json server_tool::to_json() const {
         {"permissions", json{
             {"write", permission_write}
         }},
+        {"uses_cwd", uses_cwd},
         {"definition", get_definition()},
     };
 }
@@ -127,6 +131,13 @@ static int entry_depth(const std::string & rel) {
     return 1 + (int) std::count(rel.begin(), rel.end(), '/');
 }
 
+// directories that a listing reports but never descends into: they can be enormous
+// lowercase only, the local walker case-folds a name before the lookup
+static const char * const SERVER_TOOL_JUNK_DIR_NAMES[] = {
+    ".git", ".svn", ".hg", "node_modules", "__pycache__",
+    ".venv", "venv", "dist", "build", "target", ".cache", ".idea", ".vscode",
+};
+
 class tools_io {
 public:
     struct exec_result {
@@ -165,6 +176,85 @@ public:
             const std::function<bool(const std::string &)> & on_chunk = nullptr) const = 0;
 };
 
+// shared subprocess execution helper, used by both the local and the docker-backed tools_io implementations.
+// combine_stderr=false when the raw stdout bytes must not be tainted by stderr, e.g. reading file contents.
+static tools_io::exec_result run_subprocess(
+        const std::vector<std::string> & args,
+        size_t max_output,
+        int timeout_secs,
+        const std::function<bool(const std::string &)> & on_chunk,
+        bool combine_stderr,
+        const std::string & cwd = "") {
+    tools_io::exec_result res;
+
+    common_subproc proc;
+
+    int options = subprocess_option_no_window
+                | subprocess_option_inherit_environment
+                | subprocess_option_search_user_path;
+    if (combine_stderr) {
+        options |= subprocess_option_combined_stdout_stderr;
+    }
+
+    if (!proc.create(args, options, {}, cwd.empty() ? nullptr : cwd.c_str())) {
+        res.output = "failed to spawn process";
+        return res;
+    }
+
+    std::atomic<bool> done{false};
+    std::atomic<bool> timed_out{false};
+
+    std::thread timeout_thread([&]() {
+        auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout_secs);
+        while (!done.load()) {
+            if (std::chrono::steady_clock::now() >= deadline) {
+                timed_out.store(true);
+                proc.terminate();
+                return;
+            }
+            std::this_thread::sleep_for(std::chrono::milliseconds(100));
+        }
+    });
+
+    FILE * f = proc.stdout_file();
+    std::string output;
+    bool truncated = false;
+    if (f) {
+        char buf[4096];
+        while (fgets(buf, sizeof(buf), f) != nullptr) {
+            if (!truncated) {
+                size_t len = strlen(buf);
+                if (output.size() + len <= max_output) {
+                    output.append(buf, len);
+                    if (on_chunk && !on_chunk(console_output_to_utf8(std::string(buf, len)))) {
+                        proc.terminate();
+                        break;
+                    }
+                } else {
+                    size_t remaining = max_output - output.size();
+                    output.append(buf, remaining);
+                    if (on_chunk && remaining > 0) on_chunk(console_output_to_utf8(std::string(buf, remaining)));
+                    truncated = true;
+                }
+            }
+        }
+    }
+
+    done.store(true);
+    if (timeout_thread.joinable()) {
+        timeout_thread.join();
+    }
+
+    res.exit_code = proc.join();
+
+    res.output    = console_output_to_utf8(output);
+    res.timed_out = timed_out.load();
+    if (truncated) {
+        res.output += "\n[output truncated]";
+    }
+    return res;
+}
+
 class tools_io_basic : public tools_io {
 public:
     // cwd, if non-empty, is used to resolve relative paths and as the working directory for run()
@@ -276,72 +366,7 @@ public:
             size_t max_output,
             int timeout_secs,
             const std::function<bool(const std::string &)> & on_chunk = nullptr) const override {
-        exec_result res;
-
-        common_subproc proc;
-
-        int options = subprocess_option_no_window
-                    | subprocess_option_combined_stdout_stderr
-                    | subprocess_option_inherit_environment
-                    | subprocess_option_search_user_path;
-
-        if (!proc.create(args, options, {}, cwd.empty() ? nullptr : cwd.c_str())) {
-            res.output = "failed to spawn process";
-            return res;
-        }
-
-        std::atomic<bool> done{false};
-        std::atomic<bool> timed_out{false};
-
-        std::thread timeout_thread([&]() {
-            auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(timeout_secs);
-            while (!done.load()) {
-                if (std::chrono::steady_clock::now() >= deadline) {
-                    timed_out.store(true);
-                    proc.terminate();
-                    return;
-                }
-                std::this_thread::sleep_for(std::chrono::milliseconds(100));
-            }
-        });
-
-        FILE * f = proc.stdout_file();
-        std::string output;
-        bool truncated = false;
-        if (f) {
-            char buf[4096];
-            while (fgets(buf, sizeof(buf), f) != nullptr) {
-                if (!truncated) {
-                    size_t len = strlen(buf);
-                    if (output.size() + len <= max_output) {
-                        output.append(buf, len);
-                        if (on_chunk && !on_chunk(console_output_to_utf8(std::string(buf, len)))) {
-                            proc.terminate();
-                            break;
-                        }
-                    } else {
-                        size_t remaining = max_output - output.size();
-                        output.append(buf, remaining);
-                        if (on_chunk && remaining > 0) on_chunk(console_output_to_utf8(std::string(buf, remaining)));
-                        truncated = true;
-                    }
-                }
-            }
-        }
-
-        done.store(true);
-        if (timeout_thread.joinable()) {
-            timeout_thread.join();
-        }
-
-        res.exit_code = proc.join();
-
-        res.output    = console_output_to_utf8(output);
-        res.timed_out = timed_out.load();
-        if (truncated) {
-            res.output += "\n[output truncated]";
-        }
-        return res;
+        return run_subprocess(args, max_output, timeout_secs, on_chunk, /*combine_stderr=*/true, cwd);
     }
 
 private:
@@ -384,10 +409,8 @@ private:
     }
 
     static const std::unordered_set<std::string> & junk_dir_names() {
-        static const std::unordered_set<std::string> names = {
-            ".git", ".svn", ".hg", "node_modules", "__pycache__",
-            ".venv", "venv", "dist", "build", "target", ".cache", ".idea", ".vscode",
-        };
+        static const std::unordered_set<std::string> names(
+            std::begin(SERVER_TOOL_JUNK_DIR_NAMES), std::end(SERVER_TOOL_JUNK_DIR_NAMES));
         return names;
     }
 
@@ -450,9 +473,274 @@ private:
     }
 };
 
+// timeout for auxiliary isolate calls (stat/mkdir/ls/cp helpers); exec_shell_command uses its own
+// caller-controlled timeout instead, enforced separately in run()
+static constexpr int SERVER_TOOL_ISOLATE_EXEC_TIMEOUT = 15; // seconds
+static constexpr size_t SERVER_TOOL_ISOLATE_READ_FILE_MAX_SIZE = 64 * 1024 * 1024; // 64 MB
+
+// runs every tools_io operation as a command inside an isolate: a container, a remote host, ...
+// the isolate is created, mounted, and torn down externally by the caller
+// it must provide a POSIX environment: sh, cat, wc, mkdir, dirname, find, timeout
+class tools_io_isolate : public tools_io {
+public:
+    // cwd, if non-empty, is used to resolve relative paths and as the working directory for run()
+    explicit tools_io_isolate(std::string cwd = "") : cwd(std::move(cwd)) {}
+
+    // resolves `path` against `cwd` if `path` is relative and `cwd` is set; otherwise returns `path` unchanged.
+    // isolate paths are always POSIX-style ('/'), regardless of host OS.
+    std::string resolve(const std::string & path) const override {
+        if (cwd.empty() || (!path.empty() && path[0] == '/')) {
+            return path;
+        }
+        return cwd + "/" + path;
+    }
+
+    bool is_directory(const std::string & path) const override {
+        return shell_test("-d", resolve(path));
+    }
+
+    bool is_regular_file(const std::string & path) const override {
+        return shell_test("-f", resolve(path));
+    }
+
+    bool file_size(const std::string & path, uintmax_t & out_size) const override {
+        auto res = exec({"sh", "-c", "wc -c < \"$1\"", "_", resolve(path)}, 64, true);
+        if (res.exit_code != 0 || res.timed_out) return false;
+        try {
+            size_t pos;
+            out_size = (uintmax_t) std::stoull(res.output, &pos);
+        } catch (...) {
+            return false;
+        }
+        return true;
+    }
+
+    bool read_file(const std::string & path, std::string & out) const override {
+        // combine_stderr=false: stderr must not be spliced into raw file bytes
+        auto res = exec({"cat", "--", resolve(path)}, SERVER_TOOL_ISOLATE_READ_FILE_MAX_SIZE, false);
+        if (res.exit_code != 0 || res.timed_out) return false;
+        out = res.output;
+        return true;
+    }
+
+    bool write_file(const std::string & path, const std::string & content) const override {
+        std::string abs_path = resolve(path);
+
+        std::error_code ec;
+        fs::path tmp_dir = fs::temp_directory_path(ec);
+        if (ec) return false;
+
+        static std::atomic<uint64_t> tmp_counter{0};
+        fs::path tmp = tmp_dir / string_format(
+            "llama-tools-io-isolate-%zu-%llu.tmp",
+            std::hash<std::thread::id>{}(std::this_thread::get_id()),
+            (unsigned long long) tmp_counter.fetch_add(1));
+
+        {
+            std::ofstream f(tmp, std::ios::binary);
+            if (!f) return false;
+            f << content;
+            if (!f) return false;
+        }
+
+        bool ok = shell_run({"sh", "-c", "mkdir -p \"$(dirname \"$1\")\"", "_", abs_path});
+        if (ok) {
+            ok = upload(tmp.string(), abs_path);
+        }
+
+        std::error_code rm_ec;
+        fs::remove(tmp, rm_ec);
+        return ok;
+    }
+
+    list_result list_entries(const std::string & base, int max_depth, list_kind kind) const override {
+        list_result out;
+
+        const std::string abs_base = resolve(base);
+        if (!is_directory(base)) {
+            out.err = "path does not exist or is not a directory";
+            return out;
+        }
+
+        // git ls-files cannot list directories; use the walker when they are requested
+        if (kind == list_kind::files) {
+            auto res = exec(
+                {"sh", "-c", "cd \"$1\" && git ls-files --cached --others --exclude-standard", "_", abs_base},
+                SERVER_TOOL_GIT_LS_FILES_MAX_OUTPUT, true);
+
+            if (res.exit_code == 0 && !res.timed_out) {
+                for (const auto & rel : split_lines(res.output, /*strip_dot_slash=*/false)) {
+                    if (max_depth > 0 && entry_depth(rel) > max_depth) continue;
+                    out.entries.push_back({rel, false});
+                }
+                return out;
+            }
+        }
+
+        if (kind == list_kind::dirs || kind == list_kind::all) {
+            for (auto & rel : find_entries(abs_base, max_depth, /*dirs=*/true, out.truncated)) {
+                out.entries.push_back({std::move(rel), true});
+            }
+        }
+        if (kind == list_kind::files || kind == list_kind::all) {
+            for (auto & rel : find_entries(abs_base, max_depth, /*dirs=*/false, out.truncated)) {
+                out.entries.push_back({std::move(rel), false});
+            }
+        }
+
+        return out;
+    }
+
+    // wraps the command with an in-isolate `timeout`, since killing the host-side client
+    // does not kill the process tree running inside the isolate
+    exec_result run(
+            const std::vector<std::string> & args,
+            size_t max_output,
+            int timeout_secs,
+            const std::function<bool(const std::string &)> & on_chunk = nullptr) const override {
+        std::vector<std::string> inner = {"timeout", std::to_string(timeout_secs) + "s"};
+        inner.insert(inner.end(), args.begin(), args.end());
+        // small buffer over timeout_secs so the in-isolate `timeout` has a chance to exit cleanly
+        // before the host-side supervisory timeout forcibly kills the client
+        return run_subprocess(
+            build_argv(with_cwd(inner), /*needs_stdin=*/true),
+            max_output, timeout_secs + 5, on_chunk, true);
+    }
+
+protected:
+    // wrap `inner` (a complete POSIX argv) into the host-side argv that runs it in the isolate
+    // a transport that re-parses its args in a remote shell (ssh) must join `inner` with shell_quote_join()
+    virtual std::vector<std::string> build_argv(const std::vector<std::string> & inner, bool needs_stdin) const = 0;
+
+    // copy a host file into the isolate, `isolate_path` is absolute and its parent already exists
+    virtual bool upload(const std::string & host_path, const std::string & isolate_path) const = 0;
+
+    // quote `argv` into a single string that a POSIX shell re-parses into exactly `argv`
+    static std::string shell_quote_join(const std::vector<std::string> & argv) {
+        std::string out;
+        for (const auto & arg : argv) {
+            if (!out.empty()) out += ' ';
+            out += '\'';
+            for (const char c : arg) {
+                // a single quote cannot be escaped inside single quotes: close, escape, reopen
+                if (c == '\'') out += "'\\''";
+                else           out += c;
+            }
+            out += '\'';
+        }
+        return out;
+    }
+
+private:
+    std::string cwd;
+
+    // set the working directory in the command itself, docker's `-w` has no equivalent on every transport
+    // auxiliary calls do not need this, they use the absolute paths from resolve()
+    std::vector<std::string> with_cwd(const std::vector<std::string> & inner) const {
+        if (cwd.empty()) {
+            return inner;
+        }
+        // 127 is what a shell reports for a command it could not run
+        std::vector<std::string> out = {"sh", "-c", "cd \"$1\" || exit 127; shift; exec \"$@\"", "_", cwd};
+        out.insert(out.end(), inner.begin(), inner.end());
+        return out;
+    }
+
+    exec_result exec(const std::vector<std::string> & inner, size_t max_output, bool combine_stderr) const {
+        return run_subprocess(
+            build_argv(inner, /*needs_stdin=*/false),
+            max_output, SERVER_TOOL_ISOLATE_EXEC_TIMEOUT, nullptr, combine_stderr);
+    }
+
+    bool shell_run(const std::vector<std::string> & inner) const {
+        auto res = exec(inner, 4096, true);
+        return res.exit_code == 0 && !res.timed_out;
+    }
+
+    bool shell_test(const char * flag, const std::string & path) const {
+        return shell_run({"sh", "-c", std::string("[ ") + flag + " \"$1\" ]", "_", path});
+    }
+
+    static std::vector<std::string> split_lines(const std::string & text, bool strip_dot_slash) {
+        std::vector<std::string> result;
+        std::istringstream iss(text);
+        std::string line;
+        while (std::getline(iss, line)) {
+            if (!line.empty() && line.back() == '\r') line.pop_back();
+            if (line.empty()) continue;
+            if (strip_dot_slash && line.rfind("./", 0) == 0) line = line.substr(2);
+            std::replace(line.begin(), line.end(), '\\', '/');
+            result.push_back(line);
+        }
+        return result;
+    }
+
+    // one `find` pass in the isolate. junk directories stay selectable but are never descended into,
+    // and -mindepth/-maxdepth keep a busybox image working as well as a GNU one
+    std::vector<std::string> find_entries(const std::string & abs_base, int max_depth, bool dirs, bool & truncated) const {
+        std::string prune_expr;
+        for (const char * n : SERVER_TOOL_JUNK_DIR_NAMES) {
+            if (!prune_expr.empty()) prune_expr += " -o ";
+            prune_expr += std::string("-name ") + n;
+        }
+
+        std::string cmd = "cd \"$1\" && find . -mindepth 1";
+        if (max_depth > 0) {
+            cmd += " -maxdepth " + std::to_string(max_depth);
+        }
+        cmd += " \\( " + prune_expr + " \\) -prune";
+        cmd += dirs ? " -print -o -type d -print" : " -o -type f -print";
+
+        auto res = exec({"sh", "-c", cmd, "_", abs_base}, SERVER_TOOL_GIT_LS_FILES_MAX_OUTPUT, true);
+        truncated = truncated || res.timed_out;
+        return split_lines(res.output, /*strip_dot_slash=*/true);
+    }
+};
+
+// an already-running docker container, driven through `docker exec` and `docker cp`
+class tools_io_docker : public tools_io_isolate {
+public:
+    tools_io_docker(std::string container_id, std::string cwd = "")
+        : tools_io_isolate(std::move(cwd)), container_id(std::move(container_id)) {}
+
+protected:
+    std::vector<std::string> build_argv(const std::vector<std::string> & inner, bool needs_stdin) const override {
+        std::vector<std::string> argv = {"docker", "exec"};
+        if (needs_stdin) {
+            argv.push_back("-i");
+        }
+        argv.push_back(container_id);
+        argv.insert(argv.end(), inner.begin(), inner.end());
+        return argv;
+    }
+
+    bool upload(const std::string & host_path, const std::string & isolate_path) const override {
+        auto res = run_subprocess(
+            {"docker", "cp", host_path, container_id + ":" + isolate_path},
+            4096, SERVER_TOOL_ISOLATE_EXEC_TIMEOUT, nullptr, true);
+        return res.exit_code == 0 && !res.timed_out;
+    }
+
+private:
+    std::string container_id;
+};
+
+// runtime spec used by --tools-runtime and the x-tool-runtime header
+// this is the only scheme for now, ssh: and podman: can be added next to it
+static const std::string SERVER_TOOL_RUNTIME_DOCKER_CONTAINER = "docker-container:";
+
+// an empty runtime runs the tools on the host
 static std::unique_ptr<tools_io> make_tools_io(const json & params) {
-    std::string cwd = json_value(params, "cwd", std::string());
-    return std::make_unique<tools_io_basic>(cwd);
+    std::string cwd     = json_value(params, "cwd", std::string());
+    std::string runtime = json_value(params, "runtime", std::string());
+    if (runtime.empty()) {
+        return std::make_unique<tools_io_basic>(cwd);
+    }
+    if (runtime.rfind(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER, 0) == 0) {
+        return std::make_unique<tools_io_docker>(runtime.substr(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER.size()), cwd);
+    }
+    // do not fall back to the host, the caller asked for an isolate
+    throw std::runtime_error("unknown tool runtime: " + runtime);
 }
 
 // no '/' in pattern -> match basename at any depth; else match full relative path
@@ -476,6 +764,7 @@ struct server_tool_read_file : server_tool {
     server_tool_read_file() {
         name = "read_file";
         display_name = "Read file";
+        uses_cwd = true;
         permission_write = false;
     }
 
@@ -564,6 +853,7 @@ struct server_tool_file_glob_search : server_tool {
     server_tool_file_glob_search() {
         name = "file_glob_search";
         display_name = "File search";
+        uses_cwd = true;
         permission_write = false;
     }
 
@@ -678,6 +968,7 @@ struct server_tool_grep_search : server_tool {
     server_tool_grep_search() {
         name = "grep_search";
         display_name = "Grep search";
+        uses_cwd = true;
         permission_write = false;
     }
 
@@ -830,6 +1121,7 @@ struct server_tool_exec_shell_command : server_tool {
     server_tool_exec_shell_command() {
         name = "exec_shell_command";
         display_name = "Execute shell command";
+        uses_cwd = true;
         permission_write = true;
         support_stream = true;
     }
@@ -861,8 +1153,11 @@ struct server_tool_exec_shell_command : server_tool {
         timeout    = std::min(timeout,    SERVER_TOOL_EXEC_SHELL_COMMAND_MAX_TIMEOUT);
         max_output = std::min(max_output, SERVER_TOOL_EXEC_SHELL_COMMAND_MAX_OUTPUT_SIZE);
 
+        // an isolate is always POSIX regardless of host OS, so it always gets `sh -c`
 #ifdef _WIN32
-        std::vector<std::string> args = {"cmd", "/c", command};
+        std::vector<std::string> args = !json_value(params, "runtime", std::string()).empty()
+            ? std::vector<std::string>{"sh", "-c", command}
+            : std::vector<std::string>{"cmd", "/c", command};
 #else
         std::vector<std::string> args = {"sh", "-c", command};
 #endif
@@ -905,6 +1200,7 @@ struct server_tool_write_file : server_tool {
     server_tool_write_file() {
         name = "write_file";
         display_name = "Write file";
+        uses_cwd = true;
         permission_write = true;
     }
 
@@ -947,6 +1243,7 @@ struct server_tool_edit_file : server_tool {
     server_tool_edit_file() {
         name = "edit_file";
         display_name = "Edit file";
+        uses_cwd = true;
         permission_write = true;
     }
 
@@ -1335,6 +1632,7 @@ struct server_tool_get_info : server_tool {
     server_tool_get_info() {
         name = "get_info";
         display_name = "Get Runtime Info";
+        uses_cwd = true;
         permission_write = false;
     }
 
@@ -1355,19 +1653,29 @@ struct server_tool_get_info : server_tool {
     json invoke(json params, server_tool::stream *) const override {
         auto io = make_tools_io(params);
 
+        // inside an isolate, we always use the linux command
 #ifdef _WIN32
-        auto res = io->run({"cmd", "/c", "ver"}, SERVER_TOOL_GET_INFO_MAX_OUTPUT, SERVER_TOOL_GET_INFO_TIMEOUT);
+        std::vector<std::string> args = !json_value(params, "runtime", std::string()).empty()
+            ? std::vector<std::string>{"uname", "-a"}
+            : std::vector<std::string>{"cmd", "/c", "ver"};
 #else
-        auto res = io->run({"uname", "-a"}, SERVER_TOOL_GET_INFO_MAX_OUTPUT, SERVER_TOOL_GET_INFO_TIMEOUT);
+        std::vector<std::string> args = {"uname", "-a"};
 #endif
+
+        auto res = io->run(args, SERVER_TOOL_GET_INFO_MAX_OUTPUT, SERVER_TOOL_GET_INFO_TIMEOUT);
         // "ver" prints a blank line before the version, so the output is stripped on both ends;
         // a failed spawn or a timeout leaves a diagnostic in res.output, which is not an OS name
         std::string os_info = res.exit_code == 0 && !res.timed_out ? string_strip(res.output) : "unknown";
 
         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 {
@@ -1461,6 +1769,103 @@ struct server_mcp_tool : server_tool {
     }
 };
 
+// owns the docker container used as the sandboxed runtime for tool invocations, as configured by
+// --tools-runtime. "spawned" mode starts and stops the container itself; "existing" mode just reuses
+// a container id the user already has running and never stops it.
+struct server_tools_docker_runtime {
+    server_tools_docker_runtime(const server_tools_docker_runtime &) = delete;
+
+    explicit server_tools_docker_runtime(const std::string & spec) {
+        static const std::string docker_prefix = "docker:";
+        if (spec.rfind(docker_prefix, 0) == 0) {
+            spawned = true;
+            image   = spec.substr(docker_prefix.size());
+            if (image.empty()) {
+                throw std::runtime_error("--tools-runtime docker:<image> requires an image name");
+            }
+            spawn();
+        } else if (spec.rfind(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER, 0) == 0) {
+            spawned      = false;
+            container_id = spec.substr(SERVER_TOOL_RUNTIME_DOCKER_CONTAINER.size());
+            if (container_id.empty()) {
+                throw std::runtime_error("--tools-runtime docker-container:<id> requires a container id");
+            }
+        } else {
+            throw std::runtime_error("unknown --tools-runtime option: " + spec);
+        }
+    }
+
+    ~server_tools_docker_runtime() {
+        if (spawned && !container_id.empty()) {
+            // closing stdin signals the container's shell (its pid 1) to exit; --rm then removes it
+            proc.close_stdin();
+            proc.join();
+        }
+    }
+
+    // container id to use for the next tool call; respawns a spawned container that died on its own,
+    // or throws if an externally-managed one is no longer reachable
+    std::string get_container_id() {
+        std::lock_guard<std::mutex> lock(mutex);
+        if (!spawned) {
+            if (!is_running(container_id)) {
+                throw std::runtime_error(string_format(
+                    "docker container \"%s\" is no longer running, restart it to keep using tools",
+                    container_id.c_str()));
+            }
+            return container_id;
+        }
+
+        if (!proc.alive()) {
+            SRV_WRN("docker tools runtime container \"%s\" died, respawning\n", container_id.c_str());
+            spawn();
+        }
+        return container_id;
+    }
+
+private:
+    bool spawned = false;
+    std::string image; // spawned mode only
+    std::string container_id;
+    common_subproc proc; // spawned mode only: `docker run` client that keeps the container alive
+    std::mutex mutex;
+
+    // spawns "docker run --rm -i <image> sh" and keeps its stdin open; the shell blocks reading stdin,
+    // so the container stays alive until we close it (see destructor) or it is killed from the outside
+    void spawn() {
+        std::error_code ec;
+        fs::path cidfile = fs::temp_directory_path(ec) / string_format(
+            "llama-tools-runtime-cid-%zu.tmp", std::hash<std::thread::id>{}(std::this_thread::get_id()));
+        fs::remove(cidfile, ec);
+
+        std::vector<std::string> args = {"docker", "run", "--rm", "-i", "--cidfile", cidfile.string(), image, "sh"};
+        int options = subprocess_option_no_window
+                    | subprocess_option_inherit_environment
+                    | subprocess_option_search_user_path;
+        if (!proc.create(args, options)) {
+            throw std::runtime_error("failed to spawn docker container for tools runtime (image: " + image + ")");
+        }
+
+        std::string cid;
+        for (int i = 0; i < 100 && cid.empty(); i++) {
+            std::ifstream f(cidfile);
+            if (f) std::getline(f, cid);
+            if (cid.empty()) std::this_thread::sleep_for(std::chrono::milliseconds(100));
+        }
+        fs::remove(cidfile, ec);
+        if (cid.empty()) {
+            proc.terminate();
+            throw std::runtime_error("timed out waiting for docker container to start (image: " + image + ")");
+        }
+        container_id = cid;
+    }
+
+    static bool is_running(const std::string & id) {
+        auto res = run_subprocess({"docker", "inspect", "-f", "{{.State.Running}}", id}, 16, 5, nullptr, true);
+        return res.exit_code == 0 && !res.timed_out && res.output.rfind("true", 0) == 0;
+    }
+};
+
 static server_tool & find_tool(std::vector<std::unique_ptr<server_tool>> & tools, const std::string & name, bool require_stream) {
     for (auto & t : tools) {
         if (t->name == name) {
@@ -1506,8 +1911,16 @@ static std::string get_header(const std::map<std::string, std::string> & headers
     return default_value;
 }
 
+server_tools::server_tools() = default;
+server_tools::~server_tools() = default;
+
 void server_tools::setup(const std::vector<std::string> & enabled_tools,
-                         server_mcp & mcp_mgr) {
+                         server_mcp & mcp_mgr,
+                         const std::string & tools_runtime) {
+    if (!tools_runtime.empty()) {
+        docker_runtime = std::make_unique<server_tools_docker_runtime>(tools_runtime);
+    }
+
     if (!enabled_tools.empty()) {
         if (!common_subproc::is_supported()) {
             throw std::runtime_error("subprocess is not enabled on this build");
@@ -1590,11 +2003,26 @@ void server_tools::setup(const std::vector<std::string> & enabled_tools,
             bool stream = body.value("stream", false);
 
             // accept x-tool-cwd header to override of the process
+            if (params.contains("cwd")) {
+                params.erase("cwd");
+            }
             auto cwd = get_header(req.headers, "x-tool-cwd");
             if (!cwd.empty()) {
                 params["cwd"] = cwd;
             }
 
+            // accept x-tool-runtime header to route tool I/O through an isolate, e.g. "docker-container:<id>";
+            // falls back to the --tools-runtime isolate, if configured
+            if (params.contains("runtime")) {
+                params.erase("runtime");
+            }
+            auto runtime = get_header(req.headers, "x-tool-runtime");
+            if (!runtime.empty()) {
+                params["runtime"] = runtime;
+            } else if (docker_runtime) {
+                params["runtime"] = SERVER_TOOL_RUNTIME_DOCKER_CONTAINER + docker_runtime->get_container_id();
+            }
+
             server_tool & tool = find_tool(tools, tool_name, stream);
 
             if (stream) {

tools/server/server-tools.h

@@ -14,6 +14,7 @@ struct server_tool {
     std::string display_name;
     bool permission_write = false;
     bool support_stream = false; // if true, output can be streamed
+    bool uses_cwd = false;       // if true, the tool resolves paths and runs against the working directory
 
     virtual ~server_tool() = default;
     virtual json get_definition() const = 0;
@@ -30,6 +31,8 @@ struct server_tool {
     json to_json() const;
 };
 
+struct server_tools_docker_runtime; // impl detail, defined in server-tools.cpp
+
 struct server_tools {
     std::vector<std::unique_ptr<server_tool>> tools;
 
@@ -37,9 +40,16 @@ struct server_tools {
     server_response queue_res;
     std::atomic<int> res_id{0};
 
+    // set when --tools-runtime is configured; owns the docker container used to run tools, if any
+    std::unique_ptr<server_tools_docker_runtime> docker_runtime;
+
     void setup(const std::vector<std::string> & enabled_tools,
-               server_mcp & mcp_mgr);
+               server_mcp & mcp_mgr,
+               const std::string & tools_runtime);
 
     server_http_context::handler_t handle_get;
     server_http_context::handler_t handle_post;
+
+    server_tools();
+    ~server_tools();
 };

tools/server/server.cpp

@@ -338,7 +338,7 @@ int llama_server(common_params & params, int argc, char ** argv) {
 
     if (!params.server_tools.empty() || !mcp_mgr.empty()) {
         try {
-            tools.setup(params.server_tools, mcp_mgr);
+            tools.setup(params.server_tools, mcp_mgr, params.server_tools_runtime);
         } catch (const std::exception & e) {
             SRV_ERR("tools setup failed: %s\n", e.what());
             return 1;
@@ -348,6 +348,9 @@ int llama_server(common_params & params, int argc, char ** argv) {
         if (!params.server_tools.empty()) {
             warn_names.push_back("built-in tools (experimental)");
         }
+        if (!params.server_tools_runtime.empty()) {
+            warn_names.push_back("tools runtime (experimental)");
+        }
         if (!mcp_mgr.empty()) {
             warn_names.push_back("MCP servers (experimental)");
         }

tools/server/tests/unit/test_tools_builtin.py

@@ -1,4 +1,6 @@
 import os
+import shutil
+import subprocess
 
 import pytest
 from utils import *
@@ -11,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():
@@ -146,6 +151,97 @@ def test_tools_builtin_cwd_header():
             os.remove(marker_path)
 
 
+def _docker_unavailable_reason() -> str | None:
+    """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:
+        # 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 cannot run {DOCKER_IMAGE}: {e}"
+    return None
+
+
+@pytest.fixture
+def docker_container():
+    reason = _docker_unavailable_reason()
+    if reason is not None:
+        pytest.skip(reason)  # ty: ignore[too-many-positional-arguments, invalid-argument-type]
+
+    proc = subprocess.run(
+        ["docker", "run", "-d", "--rm", DOCKER_IMAGE, "sleep", "300"],
+        capture_output=True, text=True,
+    )
+    if proc.returncode != 0:
+        pytest.skip(f"failed to start docker container: {proc.stderr.strip()}")  # ty: ignore[too-many-positional-arguments, invalid-argument-type]
+
+    container_id = proc.stdout.strip()
+    try:
+        yield container_id
+    finally:
+        subprocess.run(["docker", "rm", "-f", container_id], capture_output=True)
+
+
+def test_tools_builtin_runtime_header(docker_container: str):
+    global server
+    server.start()
+
+    headers = {"x-tool-runtime": f"docker-container:{docker_container}", "x-tool-cwd": "/tmp"}
+
+    write_res = call_tool("write_file", {"path": "test.log", "content": "hello docker\n"}, headers=headers)
+    assert write_res["result"] == "file written successfully"
+
+    read_res = call_tool("read_file", {"path": "test.log"}, headers=headers)
+    assert read_res["plain_text_response"] == "hello docker\n"
+
+    exec_res = call_tool("exec_shell_command", {"command": "cat test.log"}, headers=headers)
+    assert "hello docker" in exec_res["plain_text_response"]
+
+
+def test_tools_builtin_runtime_header_unknown_scheme():
+    global server
+    server.start()
+
+    # an unknown runtime must fail, never silently fall back to running on the host
+    res = server.make_request("POST", "/tools",
+                              data={"tool": "exec_shell_command", "params": {"command": "echo hi"}},
+                              headers={"x-tool-runtime": "ssh:example.com"})
+    assert res.status_code == 500, res.body
+    assert "unknown tool runtime" in str(res.body)
+
+
+def test_tools_builtin_docker_runtime_cleans_up_spawned_container():
+    reason = _docker_unavailable_reason()
+    if reason is not None:
+        pytest.skip(reason)  # ty: ignore[too-many-positional-arguments, invalid-argument-type]
+
+    global server
+    server.server_tools_runtime = f"docker:{DOCKER_IMAGE}"
+    server.start()
+
+    # exec_shell_command runs inside the container spawned for --tools-runtime; docker sets
+    # the container's hostname to its own short id, so this also tells us which one to check
+    res = call_tool("exec_shell_command", {"command": "hostname"})
+    container_id = res["plain_text_response"].splitlines()[0].strip()
+    assert len(container_id) >= 8, res
+
+    running = subprocess.run(
+        ["docker", "inspect", "-f", "{{.State.Running}}", container_id],
+        capture_output=True, text=True,
+    )
+    assert running.returncode == 0 and running.stdout.strip() == "true", running.stderr
+
+    server.stop()
+
+    # a clean server shutdown must stop and remove the container it spawned (it runs with --rm),
+    # not leave it behind as an abandoned child
+    leftover = subprocess.run(["docker", "inspect", container_id], capture_output=True, text=True)
+    assert leftover.returncode != 0, f"container {container_id} was not cleaned up after server exit"
+
+
 def test_tools_builtin_edit_file_rejects_overlapping_edits():
     global server
     server.start()

tools/server/tests/utils.py

@@ -115,6 +115,7 @@ class ServerProcess:
     backend_sampling: bool = False
     gcp_compat: bool = False
     server_tools: str | None = None
+    server_tools_runtime: str | None = None
     mcp_servers_config: str | None = None
     mcp_servers_json: str | None = None
     cors_origins: str | None = None
@@ -270,6 +271,8 @@ class ServerProcess:
             server_args.append("--ui-mcp-proxy")
         if self.server_tools:
             server_args.extend(["--tools", self.server_tools])
+        if self.server_tools_runtime:
+            server_args.extend(["--tools-runtime", self.server_tools_runtime])
         if self.mcp_servers_config:
             server_args.extend(["--mcp-servers-config", self.mcp_servers_config])
         if self.mcp_servers_json:

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()),
-		hasBuiltinTools: () => toolsStore.builtinTools.length > 0,
-		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)}
 			/>
@@ -651,7 +667,7 @@
 
 	<ContextGaugePopup />
 
-	{#if toolsStore.builtinTools.length > 0}
+	{#if toolsStore.hasEnabledCwdTools}
 		<ChatFormWorkingDirectory
 			directory={cwd}
 			isOpen={pickers.isWorkingDirectoryPickerOpen}

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