server: on-demand mmproj - free encoder VRAM on the text path

The vision encoder sits idle in VRAM on every text request. With
LLAMA_MMPROJ_ONDEMAND=1 the server keeps the encoder in RAM (released
after load) and brings it into VRAM just-in-time before mtmd_batch_encode,
releasing it again afterwards - so the encoder's VRAM is free for KV /
expert cache on the common text-only path.

When the encoder does not fit (that VRAM has been claimed), evict the LLM
backbone weights first: they are not needed while the encoder runs (the
decode that uses them happens afterwards), their host shadow is read-only
(cheap free, no D2H), and the KV / prompt cache is left untouched. Order
on the way back matters: free the encoder before re-uploading the weights
so the peak stays within VRAM.

Also wire the encoder release/restore into the VRAM arbiter: vram_go_cold
releases it; the just-in-time encode-path restore replaces the previous
restore in vram_ensure_warm.

Assisted-by: Claude

.github/workflows/release.yml

@@ -1109,6 +1109,8 @@ jobs:
             -DGGML_SYCL=ON \
             -DCMAKE_C_COMPILER=icx \
             -DCMAKE_CXX_COMPILER=icpx \
+            -DCMAKE_INSTALL_RPATH='$ORIGIN' \
+            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
             -DLLAMA_OPENSSL=OFF \
             -DGGML_NATIVE=OFF \
             -DGGML_SYCL_F16=${{ matrix.fp16 }}

README.md

@@ -12,6 +12,77 @@
 
 LLM inference in C/C++
 
+## ⚡ This fork — Fable's MoE-offload prefill optimizations
+
+Two **opt-in** optimizations for large MoE models whose experts are offloaded to system RAM
+(`--n-cpu-moe`), found and implemented by Fable. Both are **off by default**, toggled via
+environment variables, and produce **token-identical** output to mainline.
+
+| Env var | What it does |
+| --- | --- |
+| `GGML_CUDA_REGISTER_HOST=1` | Page-locks (pins) the mmap'd CPU expert weights so host->device copies go straight over DMA instead of through the driver's hidden bounce buffer (~6-7 -> ~20 GB/s). Works on CUDA and Vulkan (also honored as `GGML_VK_REGISTER_HOST`). Note: it is a presence check, so `=0` still enables it. |
+| `GGML_SCHED_PREFETCH_EXPERTS=1` | Prefetches each layer's experts on a second stream, so the weight uploads overlap compute instead of stalling the GPU. **CUDA only** - on the Vulkan backend the second backend instance shares one device queue, giving no overlap, so this regresses (see Vulkan note below). Leave it off on Vulkan. |
+
+### Benchmark
+
+Measured on an **RTX 3060 12GB** with **Qwen3.6-35B-A3B** (`--n-cpu-moe 26`), prompt-processing at 2048 (`MODEL` = path to your `.gguf`):
+
+```bash
+# baseline (patches off):
+./build/bin/llama-bench -m MODEL -ngl 99 -ncmoe 26 -p 2048 -n 0 -r 5 -b 2048 -ub 2048
+
+# patched (both optimizations on):
+GGML_CUDA_REGISTER_HOST=1 GGML_SCHED_PREFETCH_EXPERTS=1 \
+./build/bin/llama-bench -m MODEL -ngl 99 -ncmoe 26 -p 2048 -n 0 -r 5 -b 2048 -ub 2048
+```
+
+Result: **~1143 → ~1880 t/s** prefill (**+64%**) — same GPU, same settings, token-identical.
+
+Branches: [`fable5/host-register`](https://github.com/thecodacus/llama.cpp/tree/fable5/host-register) (pinning only) · [`fable5/prefetch-experts`](https://github.com/thecodacus/llama.cpp/tree/fable5/prefetch-experts) (both — this branch).
+
+### Vulkan (older AMD, e.g. RX 580 / Polaris)
+
+On the Vulkan backend the CUDA-oriented flags above behave differently, and this fork adds a
+Polaris-specific flash-attention fix. Findings on an **RX 580 8GB** (Polaris / GCN, PCIe 3.0 x16,
+no fp16, no matrix cores) with **Qwen3.5-35B-A3B Q4_K_M**, `-b 2048 -ub 2048`:
+
+- **Flash-attention `mask_opt` is now enabled for GCN large head sizes (this fork's own change).**
+  Upstream disables it on GCN; it is a **lossless** win in high-context prefill - it skips
+  fully-masked causal blocks and the per-block mask add on fully-visible ones, which is real work on
+  a card whose attention is compute-bound (no matrix cores). Auto-on, no flag. On Qwen3.5-35B
+  (head_dim 256): pp2048 **+8% @ 16k, +12% @ 32k**, growing with depth; perplexity bit-identical.
+- **For a long-running server, load with `--no-mmap`, not pinning.** `GGML_CUDA_REGISTER_HOST=1`
+  (pinning) gives ~+17% in an isolated `llama-bench` run, but in a server the RADV host-pointer
+  import fails and the fallback pre-stage buffer allocation fails for large / co-resident models, so
+  it silently reverts to slow staging (and can trip warnings/OOM). `--no-mmap` (weights in RAM) is
+  both faster and clean there. Pinning is still fine for one-off `llama-bench` numbers.
+- **`GGML_SCHED_PREFETCH_EXPERTS=1` regresses - do not use it** on Vulkan (its second backend shares
+  one device queue, so uploads never overlap compute).
+- **`-b 2048 -ub 2048` is the biggest prefill lever** (the default `-ub 512` roughly halves pp).
+- **Tune `--n-cpu-moe` to context length.** Keep some expert layers resident in spare VRAM for short
+  prompts (e.g. `ncmoe 28` on the 35B, ~+5% over all-host); at long context the KV cache needs that
+  VRAM, so raise it (`ncmoe 40`, all experts on host). Keep flash attention on (`-fa 1`).
+
+Prefill throughput (isolated `llama-bench`, pinned unless noted):
+
+| Config | pp2048 (t/s) |
+| --- | ---: |
+| baseline, unpinned, `ncmoe 40` | ~252 |
+| pinned, `ncmoe 40` | ~294 |
+| pinned, `ncmoe 28` | ~308 |
+| server default (`--no-mmap`, `ncmoe 40`) | ~285 |
+| + `mask_opt`, @ 32k context | **+12%** |
+
+**Recommended RX 580 / Polaris serving command** (per model):
+
+```bash
+llama-server -hf <repo>:<quant> --no-mmap -ngl 99 --n-cpu-moe 40 -b 2048 -ub 2048 -fa 1
+```
+
+Lower `--n-cpu-moe` (e.g. 28) if the model plus your context budget leave spare VRAM; keep it high
+for long-context / agentic use. At long context the bottleneck is attention compute (GPU-bound), so
+`mask_opt` (above) is where the remaining prefill wins come from, not the MoE-transfer path.
+
 ## Recent API changes
 
 - [Changelog for `libllama` API](https://github.com/ggml-org/llama.cpp/issues/9289)

common/arg.cpp

@@ -351,6 +351,10 @@ static std::string get_default_local_path(const std::string & url) {
     return fs_get_cache_file(string_split<std::string>(f, '/').back());
 }
 
+static bool spec_types_is_default(const common_params & params) {
+    return params.speculative.types == std::vector<enum common_speculative_type>{COMMON_SPECULATIVE_TYPE_NONE};
+}
+
 common_models_handler common_models_handler_init(const common_params & params, llama_example curr_ex) {
     common_download_hf_plan plan;
     common_download_hf_plan plan_spec;
@@ -391,7 +395,14 @@ common_models_handler common_models_handler_init(const common_params & params, l
     }
 
     if (!params.speculative.draft.mparams.hf_repo.empty()) {
-        plan_spec = common_download_get_hf_plan(params.speculative.draft.mparams, opts);
+        // without a requested type, discover every sidecar the draft repo ships to infer the type later
+        auto opts_spec = opts;
+        if (spec_types_is_default(params)) {
+            opts_spec.download_mtp    = true;
+            opts_spec.download_dflash = true;
+            opts_spec.download_eagle3 = true;
+        }
+        plan_spec = common_download_get_hf_plan(params.speculative.draft.mparams, opts_spec);
     }
 
     if (!params.vocoder.model.hf_repo.empty()) {
@@ -527,8 +538,57 @@ void common_models_handler_apply(common_models_handler & handler, common_params
         }
     };
 
+    // infer the speculative type from the sidecar shipped by the draft repo when none is requested
+    if (spec_types_is_default(params)) {
+        if (!plan_spec.mtp.local_path.empty()) {
+            params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_MTP };
+            plan_spec.dflash = {};
+            plan_spec.eagle3 = {};
+        } else if (!plan_spec.dflash.local_path.empty()) {
+            params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH };
+            plan_spec.eagle3 = {};
+        } else if (!plan_spec.eagle3.local_path.empty()) {
+            params.speculative.types = { COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 };
+        }
+    }
+
+    // when a sidecar type is requested, the draft repo resolves to its sidecar instead of a full model
+    const bool spec_sidecar_found = !plan_spec.mtp.local_path.empty() ||
+                                    !plan_spec.dflash.local_path.empty() ||
+                                    !plan_spec.eagle3.local_path.empty();
+    if (!plan_spec.mtp.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan_spec.mtp, opts, [&]() {
+            // only use the discovered MTP head when no draft path is set yet
+            if (params.speculative.draft.mparams.path.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.mtp);
+            } else {
+                hf_cache::finalize_file(plan_spec.mtp);
+            }
+        });
+    }
+    if (!plan_spec.dflash.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan_spec.dflash, opts, [&]() {
+            // only use the discovered DFlash sidecar when no draft path is set yet
+            if (params.speculative.draft.mparams.path.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.dflash);
+            } else {
+                hf_cache::finalize_file(plan_spec.dflash);
+            }
+        });
+    }
+    if (!plan_spec.eagle3.local_path.empty() && !had_spec_url) {
+        tasks.emplace_back(plan_spec.eagle3, opts, [&]() {
+            // only use the discovered Eagle3 sidecar when no draft path is set yet
+            if (params.speculative.draft.mparams.path.empty()) {
+                params.speculative.draft.mparams.path = hf_cache::finalize_file(plan_spec.eagle3);
+            } else {
+                hf_cache::finalize_file(plan_spec.eagle3);
+            }
+        });
+    }
+
     // handle plan_spec (e.g. --spec-draft-hf)
-    if (!plan_spec.model_files.empty() && !had_spec_url) {
+    if (!plan_spec.model_files.empty() && !had_spec_url && !spec_sidecar_found) {
         add_tasks(plan_spec.model_files, plan_spec.primary, params.speculative.draft.mparams);
         had_spec_url = true;
     }
@@ -2515,15 +2575,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     ).set_env("LLAMA_ARG_CPU_MOE"));
     add_opt(common_arg(
         {"-ncmoe", "--n-cpu-moe"}, "N",
-        "keep the Mixture of Experts (MoE) weights of the first N layers in the CPU",
+        "keep the Mixture of Experts (MoE) weights of the first N layers in the CPU; "
-        [](common_params & params, int value) {
+        "fractional N offloads part of the boundary layer at tensor granularity",
-            if (value < 0) {
+        [](common_params & params, const std::string & value) {
+            const double n = std::stod(value);
+            if (n < 0) {
                 throw std::invalid_argument("invalid value");
             }
-            for (int i = 0; i < value; ++i) {
+            for (const std::string & re : llm_ffn_exps_cpu_block_regexes(n)) {
                 // keep strings alive and avoid leaking memory by storing them in a static vector
                 static std::list<std::string> buft_overrides;
-                buft_overrides.push_back(llm_ffn_exps_block_regex(i));
+                buft_overrides.push_back(re);
                 params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()});
             }
         }

common/chat-auto-parser-generator.cpp

@@ -47,6 +47,8 @@ common_chat_params peg_generator::generate_parser(const common_chat_template &
     data.generation_prompt = common_chat_template_generation_prompt(tmpl, inputs);
     data.format            = COMMON_CHAT_FORMAT_PEG_NATIVE;
     data.preserved_tokens  = autoparser.preserved_tokens;
+    data.additional_stops.insert(data.additional_stops.end(),
+        autoparser.additional_stops.begin(), autoparser.additional_stops.end());
 
     std::string parser_generation_prompt = data.generation_prompt;
 
@@ -286,7 +288,13 @@ common_peg_parser analyze_tools::build_func_parser(common_chat_peg_builder & p,
         // we only emit tool_close when we can actually see the closing marker. This prevents
         // premature closing during partial parsing when we've seen e.g. "</" which could be
         // either "</tool_call>" (end) or "<arg_key>" prefix that failed to match.
-        func_parser = func_parser + p.tool_close(p.peek(p.literal(format.per_call_end)));
+        // Laguna (v4): the model may emit whitespace between the last </arg_value> and
+        // </tool_call> even though the template renders them tight. Tolerate optional
+        // leading space in the close lookahead so the tool call still closes.
+        auto close_peek = arguments.tolerate_intertag_whitespace
+                              ? p.peek(p.space() + p.literal(format.per_call_end))
+                              : p.peek(p.literal(format.per_call_end));
+        func_parser = func_parser + p.tool_close(close_peek);
     } else {
         func_parser = func_parser + p.tool_close(p.space());  // force this to process tool closing callbacks in mapper
     }

common/chat-auto-parser.h

@@ -206,6 +206,7 @@ struct tool_arguments_analysis {
     std::string value_prefix;  // e.g., "", "<arg_value>", ""
     std::string value_suffix;  // e.g., "</param>", "</arg_value>", ""
     std::string separator;     // e.g., "", "\n", ","
+    bool tolerate_intertag_whitespace = false; // Laguna: accept optional whitespace between arg tags
 };
 
 struct tool_id_analysis {
@@ -388,6 +389,7 @@ struct autoparser {
 
     // Preserved tokens for tokenizer (union of all non-empty markers)
     std::vector<std::string> preserved_tokens;
+    std::vector<std::string> additional_stops;  // literal stop strings (e.g. Laguna </assistant>) caught however tokenized
 
     autoparser() = default;
 

common/chat-diff-analyzer.cpp

@@ -173,6 +173,26 @@ static std::vector<std::function<void(const common_chat_template & tmpl, autopar
               LOG_DBG(ANSI_ORANGE "[Patch: JSON name/parameters tool instruction]\n" ANSI_RESET);
           }
       },
+      // Laguna (poolside) - the v4 chat template renders reasoning and tool-arg
+      // delimiters with formatting whitespace ("<think>\n", "</arg_value>\n") that
+      // the model does not emit, so the inferred delimiters carry a spurious
+      // newline and never match the model output. Trim to the bare tag. (v8
+      // renders without the whitespace, so this is a no-op there.)
+      [](const common_chat_template & tmpl, autoparser & analysis) -> void {
+          if (tmpl.src.find("laguna_glm_thinking") != std::string::npos) {
+              analysis.reasoning.start              = trim_whitespace(analysis.reasoning.start);
+              analysis.reasoning.end                = trim_whitespace(analysis.reasoning.end);
+              analysis.tools.arguments.value_prefix = trim_whitespace(analysis.tools.arguments.value_prefix);
+              analysis.tools.arguments.value_suffix = trim_whitespace(analysis.tools.arguments.value_suffix);
+              analysis.tools.arguments.separator    = trim_whitespace(analysis.tools.arguments.separator);
+              analysis.tools.arguments.tolerate_intertag_whitespace = true;
+              // The CONTROL/eot </assistant> token only halts generation when emitted as the
+              // single token; after tool calls the model can spell it out as text tokens.
+              // A literal stop string catches it either way.
+              analysis.additional_stops.push_back("</assistant>");
+              LOG_DBG(ANSI_ORANGE "[Patch: Laguna]\n" ANSI_RESET);
+          }
+      },
 
     });
 

common/chat.cpp

@@ -15,11 +15,13 @@
 
 #include "nlohmann/json.hpp"
 
+#include <algorithm>
 #include <cstdio>
 #include <cstdlib>
 #include <ctime>
 #include <exception>
 #include <functional>
+#include <map>
 
 #include <optional>
 #include <sstream>
@@ -1855,12 +1857,89 @@ static common_chat_params common_chat_params_init_gigachat_v3(
     return data;
 }
 
+// The DeepSeek V4 reference implementation renders consecutive tool results into a single
+// user block, ordered by the tool call order of the preceding assistant message (matched
+// by tool call id) rather than by the order they appear in the conversation.
+static json deepseek_v4_sort_tool_results(const json & messages) {
+    json adjusted = messages;
+    std::map<std::string, size_t> call_order;
+
+    for (size_t i = 0; i < adjusted.size();) {
+        const auto & msg  = adjusted[i];
+        const auto   role = msg.value("role", "");
+
+        if (role == "assistant" && msg.contains("tool_calls") &&
+                msg.at("tool_calls").is_array() && !msg.at("tool_calls").empty()) {
+            call_order.clear();
+            const auto & tool_calls = msg.at("tool_calls");
+            for (size_t idx = 0; idx < tool_calls.size(); idx++) {
+                auto id = tool_calls[idx].value("id", "");
+                if (!id.empty()) {
+                    call_order[id] = idx;
+                }
+            }
+            i++;
+            continue;
+        }
+
+        if (role != "user" && role != "tool") {
+            i++;
+            continue;
+        }
+
+        // collect a maximal run of user/tool messages - they render into one user block
+        std::vector<size_t> tool_positions;
+        size_t run_end = i;
+        for (; run_end < adjusted.size(); run_end++) {
+            const auto r = adjusted[run_end].value("role", "");
+            if (r == "tool") {
+                tool_positions.push_back(run_end);
+            } else if (r != "user") {
+                break;
+            }
+        }
+
+        if (tool_positions.size() > 1 && !call_order.empty()) {
+            std::vector<json> results;
+            results.reserve(tool_positions.size());
+            for (auto pos : tool_positions) {
+                results.push_back(adjusted[pos]);
+            }
+            std::stable_sort(results.begin(), results.end(), [&](const json & a, const json & b) {
+                const auto order = [&](const json & m) {
+                    auto it = call_order.find(m.value("tool_call_id", ""));
+                    return it == call_order.end() ? (size_t) 0 : it->second;
+                };
+                return order(a) < order(b);
+            });
+            for (size_t k = 0; k < tool_positions.size(); k++) {
+                adjusted[tool_positions[k]] = std::move(results[k]);
+            }
+        }
+
+        i = run_end;
+    }
+
+    return adjusted;
+}
+
 static common_chat_params common_chat_params_init_deepseek_v3_2(const common_chat_template &    tmpl,
                                                                  const autoparser::generation_params & inputs) {
     common_chat_params data;
 
-    data.prompt             = common_chat_template_direct_apply_impl(tmpl, inputs);
+    // V4 uses the same DSML markup as V3.2, but names the tool call block "tool_calls"
-    data.generation_prompt  = common_chat_template_generation_prompt_impl(tmpl, inputs);
+    // instead of "function_calls", renders tool results in tool call order and its
+    // non-thinking generation prompt ends with a bare </think> instead of an empty
+    // <think></think> pair.
+    const bool is_v4 = tmpl.source().find("function_calls") == std::string::npos;
+
+    std::optional<json> adjusted_messages;
+    if (is_v4) {
+        adjusted_messages = deepseek_v4_sort_tool_results(inputs.messages);
+    }
+
+    data.prompt             = common_chat_template_direct_apply_impl(tmpl, inputs, adjusted_messages);
+    data.generation_prompt  = common_chat_template_generation_prompt_impl(tmpl, inputs, adjusted_messages);
     data.format             = COMMON_CHAT_FORMAT_PEG_NATIVE;
     data.supports_thinking  = true;
     data.thinking_start_tag = "<think>";
@@ -1879,8 +1958,9 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
     const std::string DSML         = "｜DSML｜";
     const std::string THINK_START  = "<think>";
     const std::string THINK_END    = "</think>";
-    const std::string FC_START     = "<" + DSML + "function_calls>";
+    const std::string TC_BLOCK     = is_v4 ? "tool_calls" : "function_calls";
-    const std::string FC_END       = "</" + DSML + "function_calls>";
+    const std::string FC_START     = "<" + DSML + TC_BLOCK + ">";
+    const std::string FC_END       = "</" + DSML + TC_BLOCK + ">";
     const std::string INVOKE_START = "<" + DSML + "invoke";
     const std::string INVOKE_END   = "</" + DSML + "invoke>";
     const std::string PARAM_START  = "<" + DSML + "parameter";
@@ -1907,8 +1987,11 @@ static common_chat_params common_chat_params_init_deepseek_v3_2(const common_cha
             reasoning = p.optional(THINK_START + p.reasoning(p.until(THINK_END)) + THINK_END);
         } else if (extract_reasoning) {
             // Thinking disabled but reasoning extraction requested: the generation prompt
-            // contains an empty <think></think> pair that must still be consumed.
+            // contains an empty <think></think> pair (V3.2) or a bare </think> (V4) that
-            reasoning = p.optional(p.literal(THINK_START) + p.until(THINK_END) + p.literal(THINK_END));
+            // must still be consumed.
+            reasoning = is_v4
+                ? p.optional(p.literal(THINK_END))
+                : p.optional(p.literal(THINK_START) + p.until(THINK_END) + p.literal(THINK_END));
         }
 
         if (has_response_format) {
@@ -2612,12 +2695,14 @@ std::optional<common_chat_params> common_chat_try_specialized_template(
         return common_chat_params_init_gigachat_v3(tmpl, params);
     }
 
-    // DeepSeek V3.2 format detection: template defines dsml_token and uses it for tool calls.
+    // DeepSeek V3.2/V4 format detection: template defines dsml_token and uses it for tool calls.
     // The template source contains the token as a variable assignment, not as a literal in markup.
+    // V3.2 names the tool call block "function_calls", V4 names it "tool_calls".
     if (src.find("dsml_token") != std::string::npos &&
-        src.find("function_calls") != std::string::npos &&
+        src.find("DSML") != std::string::npos &&
-        src.find("DSML") != std::string::npos) {
+        (src.find("function_calls") != std::string::npos ||
-        LOG_DBG("Using specialized template: DeepSeek V3.2\n");
+         src.find("tool_calls") != std::string::npos)) {
+        LOG_DBG("Using specialized template: DeepSeek V3.2/V4\n");
         return common_chat_params_init_deepseek_v3_2(tmpl, params);
     }
 

common/common.h

@@ -1080,6 +1080,28 @@ inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() {
     return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() };
 }
 
+// ATSInfer-style tensor-granularity static placement of MoE expert weights.
+// Offloads the expert weights of the first floor(n) layers to the CPU, plus a
+// subset of the boundary layer's three expert tensors for the fractional part.
+// Expert tensors are dropped from the GPU in ascending performance-density
+// order (gate, then up), keeping the higher-value down_proj resident longest.
+inline std::vector<std::string> llm_ffn_exps_cpu_block_regexes(double n_cpu_moe) {
+    std::vector<std::string> regexes;
+    const int n_full = n_cpu_moe > 0 ? (int) n_cpu_moe : 0;
+    for (int i = 0; i < n_full; ++i) {
+        regexes.push_back(llm_ffn_exps_block_regex(i));
+    }
+    const int k = (int) ((n_cpu_moe - n_full) * 3.0 + 0.5);
+    if (k >= 3) {
+        regexes.push_back(llm_ffn_exps_block_regex(n_full));
+    } else if (k == 2) {
+        regexes.push_back(string_format("blk\\.%d\\.ffn_(gate|up)_(ch|)exps", n_full));
+    } else if (k == 1) {
+        regexes.push_back(string_format("blk\\.%d\\.ffn_gate_(ch|)exps", n_full));
+    }
+    return regexes;
+}
+
 //
 // training utils
 //

common/jinja/caps.cpp

@@ -23,6 +23,7 @@ void caps_apply_preserve_reasoning(jinja::context & ctx, bool enabled) {
     ctx.set_val("preserve_thinking",         mk_val<value_bool>(enabled));
     ctx.set_val("clear_thinking",            mk_val<value_bool>(!enabled));
     ctx.set_val("truncate_history_thinking", mk_val<value_bool>(!enabled));
+    ctx.set_val("drop_thinking",            mk_val<value_bool>(!enabled));
 }
 
 static void caps_try_execute(jinja::program & prog,

conversion/__init__.py

@@ -18,6 +18,7 @@ __all__ = [
 
 TEXT_MODEL_MAP: dict[str, str] = {
     "AfmoeForCausalLM": "afmoe",
+    "LagunaForCausalLM": "laguna",
     "ApertusForCausalLM": "llama",
     "ArceeForCausalLM": "llama",
     "ArcticForCausalLM": "arctic",

conversion/base.py

@@ -1682,6 +1682,9 @@ class TextModel(ModelBase):
         if chkhsh == "9dcf830ee9990cdbf78cc523a5f7bd9ad8f3f9890c2d3581d2785ad10f07049d":
             # ref: https://huggingface.co/JetBrains/Mellum2-12B-A2.5B-Base
             res = "mellum2"
+        if chkhsh == "972da7b59cec44d1f0a490a86c96df53859e486e481563e5dddac155013d87ac":
+            # ref: https://huggingface.co/poolside/Laguna-XS.2
+            res = "laguna"
 
         if res is None:
             logger.warning("\n")

conversion/hunyuan.py

@@ -338,6 +338,12 @@ class HunyuanVLTextModel(HunYuanModel):
 
     def __init__(self, dir_model: Path, *args, **kwargs):
         super().__init__(dir_model, *args, **kwargs)
+        # transformers 5.13.0 encodes HunyuanVL XD-RoPE as dynamic + mrope_section.
+        # Normalize it to avoid the HunYuan dynamic-RoPE context assertion.
+        if self.rope_parameters.get("rope_type") == "dynamic" and "mrope_section" in self.rope_parameters:
+            self.rope_parameters["rope_type"] = "xdrope"
+            self.rope_parameters["type"] = "xdrope"
+            self.rope_parameters["xdrope_section"] = list(self.rope_parameters["mrope_section"])
 
     def set_gguf_parameters(self):
         super().set_gguf_parameters()

conversion/laguna.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from torch import Tensor
+
+from .base import ModelBase, TextModel, gguf, logger
+
+
+@ModelBase.register("LagunaForCausalLM")
+class LagunaModel(TextModel):
+    model_arch = gguf.MODEL_ARCH.LAGUNA
+    _experts: list[dict] | None = None
+    _gate_types: list[str] | None = None
+
+    # --- vocab ---------------------------------------------------------------
+
+    def set_vocab(self) -> None:
+        self._set_vocab_gpt2()
+
+        # Some Laguna releases wrap the chat template in tokenizer_config.json as
+        # "{% include 'chat_template.jinja' %}", which SpecialVocab embeds verbatim
+        # and llama.cpp's jinja engine cannot process. Prefer the resolved template
+        # from the chat_template.jinja file so the GGUF is self-contained.
+        tmpl_file = self.dir_model / "chat_template.jinja"
+        if tmpl_file.is_file():
+            self.gguf_writer.add_chat_template(tmpl_file.read_text(encoding="utf-8"))
+            logger.info("gguf: embedded resolved chat_template.jinja (overriding include directive)")
+
+        # eos_token_id is a list [2, 24]: token 2 (EOS, also BOS) and token 24
+        # (</assistant>, the turn-end). _set_vocab_gpt2 only records the scalar
+        # eos, so register the extra id as eot; llama.cpp folds eot into its EOG
+        # set, so the model halts on </assistant> natively.
+        eos_ids = self.hparams.get("eos_token_id")
+        if isinstance(eos_ids, list):
+            bos_id = self.hparams.get("bos_token_id")
+            extra = [e for e in eos_ids if e != bos_id]
+            if extra:
+                self.gguf_writer.add_eot_token_id(extra[0])
+                logger.info(f"gguf: registered eot_token_id={extra[0]} from eos list {eos_ids}")
+
+    def get_vocab_base(self) -> tuple[list[str], list[int], str]:
+        # </assistant> is the assistant turn-end (registered as eot below). The
+        # HF tokenizer flags it special=false, so the base classifies it as
+        # USER_DEFINED and llama.cpp renders its text into generated content,
+        # leaking "</assistant>" and breaking response parsing. It is a control
+        # marker, so promote it to CONTROL: llama.cpp then treats it as
+        # end-of-generation and suppresses its text.
+        tokens, toktypes, tokpre = super().get_vocab_base()
+        for i, tok in enumerate(tokens):
+            if tok == "</assistant>":
+                toktypes[i] = gguf.TokenType.CONTROL
+                logger.info(f"gguf: marked </assistant> (id {i}) as CONTROL token")
+        return tokens, toktypes, tokpre
+
+    # --- hparams -------------------------------------------------------------
+
+    def set_gguf_parameters(self) -> None:
+        super().set_gguf_parameters()
+        hparams = self.hparams
+
+        # super() does not emit vocab_size for the gpt2 vocab path; head_count is
+        # overridden with a per-layer array (XS.2 varies heads per layer via
+        # num_attention_heads_per_layer; M.1 is uniform and omits it).
+        self.gguf_writer.add_vocab_size(hparams["vocab_size"])
+
+        per_layer_heads = hparams.get("num_attention_heads_per_layer")
+        if not per_layer_heads:
+            per_layer_heads = [hparams["num_attention_heads"]] * hparams["num_hidden_layers"]
+        assert len(per_layer_heads) == hparams["num_hidden_layers"], (
+            f"num_attention_heads_per_layer length {len(per_layer_heads)} != "
+            f"num_hidden_layers {hparams['num_hidden_layers']}"
+        )
+        self.gguf_writer.add_head_count(per_layer_heads)
+
+        # Resolve + validate the attention gate type now so an inconsistent
+        # `gating` field fails at conversion time. See _attn_gate_types.
+        self._attn_gate_types()
+
+        # SWA window size (M.1 has none -> key omitted, swa_type stays NONE).
+        sliding_window = hparams.get("sliding_window") or 0
+        if sliding_window > 0:
+            self.gguf_writer.add_sliding_window(sliding_window)
+
+        # MoE (expert_count / expert_used_count come from super().set_gguf_parameters())
+        self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"])
+        self.gguf_writer.add_expert_shared_feed_forward_length(hparams["shared_expert_intermediate_size"])
+        self.gguf_writer.add_expert_weights_norm(True)  # HF reference always sum-normalises after top-k
+        self.gguf_writer.add_expert_weights_scale(float(hparams["moe_routed_scaling_factor"]))
+        self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID)
+
+        # Leading dense layers (XS.2 has 1, M.1 has 3) before the MoE layers.
+        mlp_layer_types: list[str] = hparams["mlp_layer_types"]
+        leading_dense = 0
+        for t in mlp_layer_types:
+            if t == "dense":
+                leading_dense += 1
+            else:
+                break
+        self.gguf_writer.add_leading_dense_block_count(leading_dense)
+
+        # Per-layer-type RoPE dimension count (partial rotary). base emits
+        # rope_freq_base(_swa) and the YaRN params from self.rope_parameters.
+        head_dim = hparams["head_dim"]
+        full_rope = self.rope_parameters["full_attention"]
+        self.gguf_writer.add_rope_dimension_count(
+            int(head_dim * float(full_rope.get("partial_rotary_factor", 1.0))))
+        swa_rope = self.rope_parameters.get("sliding_attention")
+        if swa_rope is not None:
+            self.gguf_writer.add_rope_dimension_count_swa(
+                int(head_dim * float(swa_rope.get("partial_rotary_factor", 1.0))))
+
+    def _attn_gate_types(self) -> list[str]:
+        """Per-layer attention output gate type: "per_head" or "per_element".
+
+        `gating_types` (per layer) is authoritative when present; otherwise the
+        scalar `gating` field is used (the "per-element"/"per-head" string, or
+        the legacy boolean True == per-head, as in Laguna-XS.2).
+
+        Fails loudly when the model is per-element but the `gating` field does
+        not declare that as a string: runtimes that key off `gating` (vLLM,
+        transformers) ignore gating_types and read a bare boolean True as
+        per-head, silently corrupting the model. Surfacing it here keeps a
+        broken checkpoint from being packaged as if it were fine.
+        """
+        if self._gate_types is not None:
+            return self._gate_types
+        hparams = self.hparams
+        n_layer = hparams["num_hidden_layers"]
+        gating = hparams.get("gating")
+        gating_types = hparams.get("gating_types")
+
+        def _norm(t: object) -> str:
+            sval = str(t).replace("-", "_")
+            if sval in ("per_element", "per_head"):
+                return sval
+            raise ValueError(f"Laguna: unrecognised attention gate type {t!r}")
+
+        if gating_types:
+            assert len(gating_types) == n_layer, (
+                f"gating_types length {len(gating_types)} != num_hidden_layers {n_layer}")
+            types = [_norm(t) for t in gating_types]
+        elif isinstance(gating, str):
+            types = [_norm(gating)] * n_layer
+        elif gating is True:
+            types = ["per_head"] * n_layer
+        else:
+            raise ValueError(
+                f"Laguna: cannot determine attention gate type "
+                f"(gating={gating!r}, gating_types={gating_types!r})")
+
+        if any(t == "per_element" for t in types) and not (
+                isinstance(gating, str) and _norm(gating) == "per_element"):
+            raise ValueError(
+                f"Laguna config declares a per-element attention gate but "
+                f"`gating`={gating!r} is not the string \"per-element\". Runtimes that "
+                f"read `gating` (vLLM, transformers) will mis-handle this checkpoint as "
+                f"per-head. Set gating=\"per-element\" in the source config.")
+
+        self._gate_types = types
+        return types
+
+    # --- tensor handling -----------------------------------------------------
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        # Per-expert MoE weights: model.layers.{bid}.mlp.experts.{xid}.{w}.weight.
+        # Only the NUMBERED per-expert weights are stacked; the router bias
+        # (mlp.experts.e_score_correction_bias) takes the normal mapping path.
+        if re.search(r"mlp\.experts\.\d+\.", name):
+            n_experts = self.find_hparam(["num_local_experts", "num_experts"])
+            assert bid is not None
+            if self._experts is None:
+                self._experts = [{} for _ in range(self.block_count)]
+            self._experts[bid][name] = data_torch
+            needed = [f"model.layers.{bid}.mlp.experts.{x}.{w}.weight"
+                      for x in range(n_experts) for w in ("gate_proj", "up_proj", "down_proj")]
+            if all(e in self._experts[bid] for e in needed):
+                for w_name in ["gate_proj", "up_proj", "down_proj"]:
+                    datas = [self._experts[bid][f"model.layers.{bid}.mlp.experts.{x}.{w_name}.weight"]
+                             for x in range(n_experts)]
+                    stacked = torch.stack(datas, dim=0)
+                    merged = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
+                    yield from TextModel.modify_tensors(self, stacked, merged, bid)
+                self._experts[bid].clear()
+                return
+            return
+        # Cross-check the gate projection width against the declared gate type;
+        # a mismatch means the weights and config disagree -> fail, do not guess.
+        if bid is not None and name.endswith("self_attn.g_proj.weight"):
+            heads = (self.hparams.get("num_attention_heads_per_layer")
+                     or [self.hparams["num_attention_heads"]] * self.hparams["num_hidden_layers"])
+            n_head = heads[bid]
+            head_dim = self.hparams["head_dim"]
+            gate_type = self._attn_gate_types()[bid]
+            expected = n_head * head_dim if gate_type == "per_element" else n_head
+            out_features = int(data_torch.shape[0])
+            if out_features != expected:
+                raise ValueError(
+                    f"Laguna layer {bid}: g_proj output width {out_features} contradicts the "
+                    f"declared {gate_type} gate (expected {expected}); weights and config disagree.")
+
+        yield from TextModel.modify_tensors(self, data_torch, name, bid)

convert_hf_to_gguf_update.py

@@ -162,6 +162,7 @@ models = [
     {"name": "granite-embed-multi-97m", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/ibm-granite/granite-embedding-97m-multilingual-r2", },
     {"name": "granite-embed-multi-311m", "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/ibm-granite/granite-embedding-311m-multilingual-r2", },
     {"name": "mellum2",          "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/JetBrains/Mellum2-12B-A2.5B-Base"},
+    {"name": "laguna",           "tokt": TOKENIZER_TYPE.BPE, "repo": "https://huggingface.co/poolside/Laguna-XS.2", },
 ]
 
 # some models are known to be broken upstream, so we will skip them as exceptions

docs/ops.md

@@ -25,10 +25,10 @@ Legend:
 |                             CEIL | ❌ | ❌ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                            CLAMP | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | 🟡 | ✅ | 🟡 | ✅ | ❌ | ❌ |
 |                        COL2IM_1D | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
-|                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           CONCAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                          CONV_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          CONV_3D | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_1D | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
@@ -41,6 +41,9 @@ Legend:
 |                             DIAG | ❌ | ❌ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    DIAG_MASK_INF | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              DIV | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                     DSV4_HC_COMB | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                     DSV4_HC_POST | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
+|                      DSV4_HC_PRE | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                              DUP | ❌ | ✅ | ✅ | 🟡 | ❌ | 🟡 | 🟡 | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              ELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                              EXP | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
@@ -63,16 +66,17 @@ Legend:
 |                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                           IM2COL | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                        IM2COL_3D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
-|                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                          L2_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ❌ | 🟡 | ❌ | ✅ | 🟡 | ❌ | ❌ | ❌ |
+|                LIGHTNING_INDEXER | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             MEAN | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                              MUL | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                          MUL_MAT | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 | 🟡 |
-|                 MUL_MAT_HADAMARD | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
+|                 MUL_MAT_HADAMARD | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ✅ | ✅ | 🟡 | 🟡 | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | 🟡 | ✅ | ❌ | ❌ |
+|                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | 🟡 | 🟡 | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
 |                     OPT_STEP_SGD | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
 |                         OUT_PROD | 🟡 | 🟡 | 🟡 | 🟡 | ❌ | ❌ | ❌ | 🟡 | ❌ | ❌ | ❌ | 🟡 |
@@ -82,7 +86,7 @@ Legend:
 |                          POOL_2D | ❌ | 🟡 | ✅ | ✅ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                            REGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                             RELU | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
-|                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ | ❌ |
+|                           REPEAT | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ | ❌ |
 |                      REPEAT_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
 |                         RMS_NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 |                    RMS_NORM_BACK | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |

ggml/include/ggml-backend.h

@@ -104,6 +104,11 @@ extern "C" {
     GGML_API enum ggml_status ggml_backend_graph_compute      (ggml_backend_t backend, struct ggml_cgraph * cgraph);
     GGML_API enum ggml_status ggml_backend_graph_compute_async(ggml_backend_t backend, struct ggml_cgraph * cgraph);
 
+    // Free transient/scratch device memory the backend holds outside of any allocated buffer
+    // (compute preallocations, staging buffers). No-op if the backend does not implement it.
+    // The backend remains usable; scratch is reallocated lazily on the next compute.
+    GGML_API void ggml_backend_free_scratch(ggml_backend_t backend);
+
     // NOTE: will be removed, use device version instead
     GGML_API bool ggml_backend_supports_op(ggml_backend_t backend, const struct ggml_tensor * op);
     GGML_API bool ggml_backend_supports_buft(ggml_backend_t backend, ggml_backend_buffer_type_t buft);

ggml/src/CMakeLists.txt

@@ -430,7 +430,7 @@ if (GGML_CPU_ALL_VARIANTS)
             message(FATAL_ERROR "Unsupported ARM target OS: ${CMAKE_SYSTEM_NAME}")
         endif()
     elseif (GGML_SYSTEM_ARCH STREQUAL "PowerPC")
-        if (CMAKE_SYSTEM_NAME MATCHES "Linux")
+        if (CMAKE_SYSTEM_NAME MATCHES "Linux|AIX")
             ggml_add_cpu_backend_variant(power0)
             ggml_add_cpu_backend_variant(power7_1       POWER7)
             ggml_add_cpu_backend_variant(power7_2       POWER7  VSX)

ggml/src/ggml-backend-impl.h

@@ -137,6 +137,11 @@ extern "C" {
 
         // (optional) sort/optimize the nodes in the graph
         void                      (*graph_optimize)    (ggml_backend_t backend, struct ggml_cgraph * cgraph);
+
+        // (optional) free transient/scratch device memory the backend holds outside of any buffer
+        // (e.g. compute preallocations and staging buffers). The backend stays usable; the scratch
+        // is reallocated lazily on the next compute. Used to shrink an idle model's device footprint.
+        void                      (*free_scratch)       (ggml_backend_t backend);
     };
 
     struct ggml_backend {

ggml/src/ggml-backend.cpp

@@ -420,6 +420,15 @@ void ggml_backend_synchronize(ggml_backend_t backend) {
     backend->iface.synchronize(backend);
 }
 
+void ggml_backend_free_scratch(ggml_backend_t backend) {
+    GGML_ASSERT(backend);
+    if (backend->iface.free_scratch == NULL) {
+        return;
+    }
+
+    backend->iface.free_scratch(backend);
+}
+
 ggml_backend_graph_plan_t ggml_backend_graph_plan_create(ggml_backend_t backend, struct ggml_cgraph * cgraph) {
     GGML_ASSERT(backend);
     GGML_ASSERT(backend->iface.graph_plan_create != NULL);
@@ -761,6 +770,10 @@ static bool ggml_is_view_op(enum ggml_op op) {
 #define GGML_SCHED_MAX_COPIES 4
 #endif
 
+#ifndef GGML_SCHED_MAX_PREFETCH_SLOTS
+#define GGML_SCHED_MAX_PREFETCH_SLOTS 8
+#endif
+
 struct ggml_backend_sched_split {
     int backend_id;
     int i_start;
@@ -818,6 +831,19 @@ struct ggml_backend_sched {
 
     bool op_offload;
 
+    // full-tensor prefetch of offloaded MUL_MAT_ID weights (GGML_SCHED_PREFETCH_EXPERTS)
+    // with a large batch virtually every expert is used, so the routing ids are not worth
+    // waiting for; uploads run through a second backend instance on the same device so
+    // they overlap compute, alternating between two staging slots
+    bool prefetch_experts;
+    ggml_backend_t prefetch_backend;
+    int prefetch_n_slots;
+    ggml_backend_buffer_t prefetch_slots[GGML_SCHED_MAX_PREFETCH_SLOTS];
+    ggml_backend_event_t prefetch_ready[GGML_SCHED_MAX_PREFETCH_SLOTS];
+    ggml_backend_event_t prefetch_free[GGML_SCHED_MAX_PREFETCH_SLOTS];
+    bool prefetch_used[GGML_SCHED_MAX_PREFETCH_SLOTS];
+    int prefetch_cur;
+
     int debug;
 
     // used for debugging graph reallocations [GGML_SCHED_DEBUG_REALLOC]
@@ -1538,6 +1564,94 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
     return true;
 }
 
+static void ggml_backend_sched_prefetch_disable(ggml_backend_sched_t sched, ggml_backend_t split_backend) {
+    sched->prefetch_experts = false;
+    if (sched->prefetch_backend) {
+        ggml_backend_synchronize(split_backend);
+        ggml_backend_synchronize(sched->prefetch_backend);
+    }
+    for (int i = 0; i < sched->prefetch_n_slots; i++) {
+        ggml_backend_buffer_free(sched->prefetch_slots[i]);
+        sched->prefetch_slots[i] = NULL;
+        sched->prefetch_used[i] = false;
+    }
+}
+
+// slots are sized once for the largest offloaded expert tensor in the current graph so
+// that they never need to grow mid-eval
+static size_t ggml_backend_sched_prefetch_max_size(ggml_backend_sched_t sched) {
+    size_t max_size = 0;
+    for (int split_id = 0; split_id < sched->n_splits; split_id++) {
+        struct ggml_backend_sched_split * split = &sched->splits[split_id];
+        if (split->graph.n_nodes == 0 || split->graph.nodes[0]->op != GGML_OP_MUL_MAT_ID) {
+            continue;
+        }
+        for (int input_id = 0; input_id < split->n_inputs; input_id++) {
+            const ggml_tensor * input = split->inputs[input_id];
+            if (input->buffer &&
+                ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
+                ggml_backend_buffer_is_host(input->buffer)) {
+                max_size = std::max(max_size, ggml_nbytes(input));
+            }
+        }
+    }
+    return max_size;
+}
+
+static bool ggml_backend_sched_prefetch_init(ggml_backend_sched_t sched, ggml_backend_t split_backend, size_t size) {
+    if (sched->prefetch_backend == NULL) {
+        ggml_backend_dev_t dev = split_backend->device;
+        ggml_backend_dev_props props;
+        ggml_backend_dev_get_props(dev, &props);
+        if (!props.caps.async || !props.caps.events) {
+            sched->prefetch_experts = false;
+            return false;
+        }
+        sched->prefetch_backend = ggml_backend_dev_init(dev, NULL);
+        if (sched->prefetch_backend == NULL) {
+            sched->prefetch_experts = false;
+            return false;
+        }
+        for (int i = 0; i < sched->prefetch_n_slots; i++) {
+            sched->prefetch_ready[i] = ggml_backend_event_new(dev);
+            sched->prefetch_free[i]  = ggml_backend_event_new(dev);
+            if (sched->prefetch_ready[i] == NULL || sched->prefetch_free[i] == NULL) {
+                sched->prefetch_experts = false;
+                return false;
+            }
+        }
+    }
+
+    size = std::max(size, ggml_backend_sched_prefetch_max_size(sched));
+
+    ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(split_backend);
+    for (int i = 0; i < sched->prefetch_n_slots; i++) {
+        if (sched->prefetch_slots[i] == NULL || ggml_backend_buffer_get_size(sched->prefetch_slots[i]) < size) {
+            // allocate before freeing so a failure leaves the old slot intact
+            ggml_backend_buffer_t new_buf = ggml_backend_buft_alloc_buffer(buft, size);
+            if (new_buf == NULL) {
+                // overlap needs at least 2 slots, otherwise run with what fits
+                if (i >= 2 && sched->prefetch_slots[0] != NULL &&
+                    ggml_backend_buffer_get_size(sched->prefetch_slots[0]) >= size) {
+                    sched->prefetch_n_slots = i;
+                    sched->prefetch_cur = 0;
+                    return true;
+                }
+                ggml_backend_sched_prefetch_disable(sched, split_backend);
+                return false;
+            }
+            if (sched->prefetch_slots[i] != NULL) {
+                ggml_backend_synchronize(split_backend);
+                ggml_backend_synchronize(sched->prefetch_backend);
+                ggml_backend_buffer_free(sched->prefetch_slots[i]);
+            }
+            sched->prefetch_slots[i] = new_buf;
+            sched->prefetch_used[i] = false;
+        }
+    }
+    return true;
+}
+
 static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
     GGML_ASSERT(sched);
     struct ggml_backend_sched_split * splits = sched->splits;
@@ -1550,6 +1664,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
         struct ggml_backend_sched_split * split = &splits[split_id];
         int split_backend_id = split->backend_id;
         ggml_backend_t split_backend = sched->backends[split_backend_id];
+        int split_prefetch_slot = -1;
+        ggml_tensor * prefetch_input_cpy = NULL;
+        ggml_backend_buffer_t prefetch_saved_buffer = NULL;
+        void * prefetch_saved_data = NULL;
 
         // copy the input tensors to the split backend
         for (int input_id = 0; input_id < split->n_inputs; input_id++) {
@@ -1566,6 +1684,41 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                 }
                 ggml_backend_tensor_copy(input, input_cpy);
             } else {
+                // with a large batch virtually every expert is used, so instead of waiting
+                // for the routing ids, upload the full tensor through the prefetch backend
+                // and let the copy overlap compute of the previous split
+                if (sched->prefetch_experts && !sched->callback_eval && split_prefetch_slot == -1 && split->graph.n_nodes > 0) {
+                    ggml_tensor * node = split->graph.nodes[0];
+                    if (ggml_backend_buffer_get_usage(input->buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS &&
+                        ggml_backend_buffer_is_host(input->buffer) &&
+                        node->op == GGML_OP_MUL_MAT_ID && node->src[0] == input_cpy) {
+                        const ggml_tensor * ids = node->src[2];
+                        const int64_t n_expert = input->ne[2];
+                        if (ids->ne[0]*ids->ne[1] >= 2*n_expert &&
+                            ggml_backend_sched_prefetch_init(sched, split_backend, ggml_nbytes(input))) {
+                            const int slot = sched->prefetch_cur;
+                            sched->prefetch_cur = (sched->prefetch_cur + 1) % sched->prefetch_n_slots;
+                            // wait for the previous user of this slot to finish computing
+                            if (sched->prefetch_used[slot]) {
+                                ggml_backend_event_wait(sched->prefetch_backend, sched->prefetch_free[slot]);
+                            }
+                            // point the staging copy at the slot only for the duration of
+                            // this split, so a fallback to the regular path on a later
+                            // eval can never see a dangling slot pointer
+                            prefetch_input_cpy    = input_cpy;
+                            prefetch_saved_buffer = input_cpy->buffer;
+                            prefetch_saved_data   = input_cpy->data;
+                            input_cpy->buffer = sched->prefetch_slots[slot];
+                            input_cpy->data   = ggml_backend_buffer_get_base(sched->prefetch_slots[slot]);
+                            ggml_backend_tensor_set_async(sched->prefetch_backend, input_cpy, input->data, 0, ggml_nbytes(input));
+                            ggml_backend_event_record(sched->prefetch_ready[slot], sched->prefetch_backend);
+                            ggml_backend_event_wait(split_backend, sched->prefetch_ready[slot]);
+                            split_prefetch_slot = slot;
+                            continue;
+                        }
+                    }
+                }
+
                 // wait for the split backend to finish using the input before overwriting it
                 if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
                     ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]);
@@ -1676,6 +1829,13 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
 
         if (!sched->callback_eval) {
             enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph);
+            if (split_prefetch_slot != -1) {
+                // the kernels have captured the slot address at launch, safe to restore
+                ggml_backend_event_record(sched->prefetch_free[split_prefetch_slot], split_backend);
+                sched->prefetch_used[split_prefetch_slot] = true;
+                prefetch_input_cpy->buffer = prefetch_saved_buffer;
+                prefetch_input_cpy->data   = prefetch_saved_data;
+            }
             if (ec != GGML_STATUS_SUCCESS) {
                 return ec;
             }
@@ -1788,6 +1948,15 @@ ggml_backend_sched_t ggml_backend_sched_new(
     sched->galloc = ggml_gallocr_new_n(sched->bufts, n_backends);
     sched->op_offload = op_offload;
 
+    // GGML_SCHED_PREFETCH_EXPERTS=1 enables the default slot count, higher values set it
+    // directly; more slots let uploads run further ahead of compute at the cost of one
+    // max-sized expert tensor of device memory per slot
+    const char * GGML_SCHED_PREFETCH_EXPERTS = getenv("GGML_SCHED_PREFETCH_EXPERTS");
+    const int prefetch_n_slots = GGML_SCHED_PREFETCH_EXPERTS ? atoi(GGML_SCHED_PREFETCH_EXPERTS) : 0;
+    sched->prefetch_experts = op_offload && prefetch_n_slots > 0;
+    // default of 3 covers the gate/up/down expert tensors of one MoE layer
+    sched->prefetch_n_slots = prefetch_n_slots <= 1 ? 3 : std::min(prefetch_n_slots, GGML_SCHED_MAX_PREFETCH_SLOTS);
+
     ggml_backend_sched_reset(sched);
 
     return sched;
@@ -1802,6 +1971,16 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) {
             ggml_backend_event_free(sched->events[b][c]);
         }
     }
+    if (sched->prefetch_backend) {
+        ggml_backend_synchronize(sched->prefetch_backend);
+        // the slot count may have been reduced after a failed allocation, free everything
+        for (int i = 0; i < GGML_SCHED_MAX_PREFETCH_SLOTS; i++) {
+            ggml_backend_event_free(sched->prefetch_ready[i]);
+            ggml_backend_event_free(sched->prefetch_free[i]);
+            ggml_backend_buffer_free(sched->prefetch_slots[i]);
+        }
+        ggml_backend_free(sched->prefetch_backend);
+    }
     ggml_gallocr_free(sched->galloc);
     ggml_free(sched->ctx);
     ggml_hash_set_free(&sched->hash_set);
@@ -1906,6 +2085,9 @@ void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) {
     for (int i = 0; i < sched->n_backends; i++) {
         ggml_backend_synchronize(sched->backends[i]);
     }
+    if (sched->prefetch_backend) {
+        ggml_backend_synchronize(sched->prefetch_backend);
+    }
     if (!sched->is_alloc) {
         // if the graph is not already allocated, always use copy 0 after a synchronization
         // this ensures that during generation the same copy is used every time,

ggml/src/ggml-cpu/kleidiai/kleidiai.cpp

@@ -1719,6 +1719,7 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
 
             return true;
         }
+
         return false;
     }
 
@@ -1727,6 +1728,20 @@ class extra_buffer_type : ggml::cpu::extra_buffer_type {
             if (op->src[0]->buffer && op->src[0]->buffer->buft == ggml_backend_cpu_kleidiai_buffer_type()) {
                 return (ggml::cpu::tensor_traits *) op->src[0]->extra;
             } else {
+                // KleidiAI only has kernels for Q4_0 and Q8_0. For a quantized weight of any
+                // other type (K-quants, IQ) it declines the op and returns nullptr below, so
+                // KleidiAI does not accelerate it. Another CPU backend may still take the op,
+                // and this can run during graph planning, so the message says what KleidiAI
+                // did rather than what ends up executing. Warn once per process.
+                if (ggml_is_quantized(op->src[0]->type) &&
+                    op->src[0]->type != GGML_TYPE_Q4_0 && op->src[0]->type != GGML_TYPE_Q8_0) {
+                    static std::atomic<bool> warned(false);
+                    if (!warned.exchange(true)) {
+                        GGML_LOG_WARN("kleidiai: no kernel for tensor type %s, not accelerated by KleidiAI "
+                                      "(kernels available for Q4_0 and Q8_0)\n",
+                                      ggml_type_name(op->src[0]->type));
+                    }
+                }
                 if (op->src[0]->type != GGML_TYPE_F16) {
                     return nullptr;
                 }

ggml/src/ggml-cpu/llamafile/sgemm.cpp

@@ -2329,7 +2329,7 @@ class tinyBLAS_Q0_PPC {
         mc = 32;
         nc = 32;
         kc = 32;
-        n_chunk = 32
+        n_chunk = 32;
     #endif
         int64_t n_aligned = 0;
         if (n % n_chunk == 0) {

ggml/src/ggml-cuda/common.cuh

@@ -362,6 +362,15 @@ static bool blackwell_mma_available(const int cc) {
            ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_RUBIN;
 }
 
+// Checks whether the tensor's base data pointer and higher-dimensional strides are byte-aligned to `alignment` bytes.
+static bool ggml_cuda_is_aligned(const ggml_tensor * tensor, const size_t alignment) {
+    GGML_ASSERT(tensor != nullptr);
+    return (reinterpret_cast<uintptr_t>(tensor->data) % alignment) == 0 &&
+           tensor->nb[1] % alignment == 0 &&
+           tensor->nb[2] % alignment == 0 &&
+           tensor->nb[3] % alignment == 0;
+}
+
 static constexpr __device__ int ggml_cuda_get_physical_warp_size() {
 #if defined(GGML_USE_HIP) && (defined(__GFX9__) || defined(__GFX8__))
     return 64;
@@ -937,6 +946,9 @@ static __device__ __forceinline__ uint2 fast_div_modulo(uint32_t n, const uint3
 
 typedef void (*dequantize_kernel_t)(const void * vx, const int64_t ib, const int iqs, float2 & v);
 
+template<typename dst_t>
+using dequantize_kq_t = void (*)(const void * vx, const int64_t ib, dst_t * y, const int tid);
+
 static __device__ __forceinline__ float get_alibi_slope(
     const float max_bias, const uint32_t h, const uint32_t n_head_log2, const float m0, const float m1
 ) {

ggml/src/ggml-cuda/convert.cu

@@ -140,358 +140,107 @@ static __global__ void dequantize_block_q4_1(const void * __restrict__ vx, dst_t
 
 template<typename dst_t>
 static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
     const int64_t i = blockIdx.x;
-    const block_q2_K * x = (const block_q2_K *) vx;
 
-    const int64_t tid = threadIdx.x;
+    dequantize_q2_K(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t n   = tid/32;
-    const int64_t l   = tid - 32*n;
-    const int64_t is  = 8*n + l/16;
-
-    const uint8_t q = x[i].qs[32*n + l];
-    dst_t * y = yy + i*QK_K + 128*n;
-
-    float dall = __low2half(x[i].dm);
-    float dmin = __high2half(x[i].dm);
-    y[l+ 0] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[i].scales[is+0] >> 4));
-    y[l+32] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[i].scales[is+2] >> 4));
-    y[l+64] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[i].scales[is+4] >> 4));
-    y[l+96] = ggml_cuda_cast<dst_t>(dall * (x[i].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[i].scales[is+6] >> 4));
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
     const int64_t i = blockIdx.x;
-    const block_q3_K * x = (const block_q3_K *) vx;
 
-    const int64_t r = threadIdx.x/4;
+    dequantize_q3_K(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t tid = r/2;
-    const int64_t is0 = r%2;
-    const int64_t l0 = 16*is0 + 4*(threadIdx.x%4);
-    const int64_t n = tid / 4;
-    const int64_t j = tid - 4*n;
-
-    uint8_t m = 1 << (4*n + j);
-    int64_t is = 8*n + 2*j + is0;
-    int shift = 2*j;
-
-    int8_t us = is <  4 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+8] >> 0) & 3) << 4) :
-                is <  8 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+4] >> 2) & 3) << 4) :
-                is < 12 ? (x[i].scales[is-8] >>  4) | (((x[i].scales[is+0] >> 4) & 3) << 4) :
-                          (x[i].scales[is-8] >>  4) | (((x[i].scales[is-4] >> 6) & 3) << 4);
-    float d_all = x[i].d;
-    float dl = d_all * (us - 32);
-
-    dst_t * y = yy + i*QK_K + 128*n + 32*j;
-    const uint8_t * q = x[i].qs + 32*n;
-    const uint8_t * hm = x[i].hmask;
-
-    for (int l = l0; l < l0+4; ++l) {
-        y[l] = ggml_cuda_cast<dst_t>(dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)));
-    }
-}
-
-static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) {
-    if (j < 4) {
-        d = q[j] & 63; m = q[j + 4] & 63;
-    } else {
-        d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4);
-        m = (q[j+4] >>  4) | ((q[j-0] >> 6) << 4);
-    }
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-    const block_q4_K * x = (const block_q4_K *) vx;
-
     const int64_t i = blockIdx.x;
 
-    // assume 32 threads
+    dequantize_q4_K(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t tid = threadIdx.x;
-    const int64_t il  = tid/8;
-    const int64_t ir  = tid%8;
-    const int64_t is  = 2*il;
-    const int64_t n   = 4;
-
-    dst_t * y = yy + i*QK_K + 64*il + n*ir;
-
-    const float dall = __low2half(x[i].dm);
-    const float dmin = __high2half(x[i].dm);
-
-    const uint8_t * q = x[i].qs + 32*il + n*ir;
-
-    uint8_t sc, m;
-    get_scale_min_k4(is + 0, x[i].scales, sc, m);
-    const float d1 = dall * sc; const float m1 = dmin * m;
-    get_scale_min_k4(is + 1, x[i].scales, sc, m);
-    const float d2 = dall * sc; const float m2 = dmin * m;
-    for (int l = 0; l < n; ++l) {
-        y[l + 0] = ggml_cuda_cast<dst_t>(d1 * (q[l] & 0xF) - m1);
-        y[l +32] = ggml_cuda_cast<dst_t>(d2 * (q[l] >>  4) - m2);
-    }
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-    const block_q5_K * x = (const block_q5_K *) vx;
-
     const int64_t i = blockIdx.x;
 
-    // assume 64 threads - this is very slightly better than the one below
+    dequantize_q5_K(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t tid = threadIdx.x;
-    const int64_t il  = tid/16;   // il is in 0...3
-    const int64_t ir  = tid%16;   // ir is in 0...15
-    const int64_t is  = 2*il;     // is is in 0...6
-
-    dst_t * y = yy + i*QK_K + 64*il + 2*ir;
-
-    const float dall = __low2half(x[i].dm);
-    const float dmin = __high2half(x[i].dm);
-
-    const uint8_t * ql = x[i].qs + 32*il + 2*ir;
-    const uint8_t * qh = x[i].qh + 2*ir;
-
-    uint8_t sc, m;
-    get_scale_min_k4(is + 0, x[i].scales, sc, m);
-    const float d1 = dall * sc; const float m1 = dmin * m;
-    get_scale_min_k4(is + 1, x[i].scales, sc, m);
-    const float d2 = dall * sc; const float m2 = dmin * m;
-
-    uint8_t   hm  = 1 << (2*il);
-    y[ 0] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1);
-    y[ 1] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1);
-    hm <<= 1;
-    y[32] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 0] >>  4) + (qh[ 0] & hm ? 16 : 0)) - m2);
-    y[33] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 1] >>  4) + (qh[ 1] & hm ? 16 : 0)) - m2);
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-    const block_q6_K * x = (const block_q6_K *) vx;
-
     const int64_t i = blockIdx.x;
 
-    // assume 64 threads - this is very slightly better than the one below
+    dequantize_q6_K(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t tid = threadIdx.x;
-    const int64_t ip  = tid/32;   // ip is 0 or 1
-    const int64_t il  = tid - 32*ip; // 0...32
-    const int64_t is  = 8*ip + il/16;
-
-    dst_t * y = yy + i*QK_K + 128*ip + il;
-
-    const float d = x[i].d;
-
-    const uint8_t * ql = x[i].ql + 64*ip + il;
-    const uint8_t   qh = x[i].qh[32*ip + il];
-    const int8_t  * sc = x[i].scales + is;
-
-    y[ 0] = ggml_cuda_cast<dst_t>(d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32));
-    y[32] = ggml_cuda_cast<dst_t>(d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32));
-    y[64] = ggml_cuda_cast<dst_t>(d * sc[4] * ((int8_t)((ql[ 0]  >> 4) | (((qh >> 4) & 3) << 4)) - 32));
-    y[96] = ggml_cuda_cast<dst_t>(d * sc[6] * ((int8_t)((ql[32]  >> 4) | (((qh >> 6) & 3) << 4)) - 32));
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
     const int64_t i = blockIdx.x;
-    const block_iq2_xxs * x = (const block_iq2_xxs  *) vx;
 
-    const int64_t tid = threadIdx.x;
+    dequantize_iq2_xxs(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint16_t * q2 = x[i].qs + 4*ib;
-    const uint8_t  * aux8 = (const uint8_t *)q2;
-    const uint8_t  * grid = (const uint8_t *)(iq2xxs_grid + aux8[il]);
-    const uint32_t aux32 = q2[2] | (q2[3] << 16);
-    const float d = (float)x[i].d * (0.5f + (aux32 >> 28)) * 0.25f;
-    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
-    for (int j = 0; j < 8; ++j) {
-        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
-    }
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
     const int64_t i = blockIdx.x;
-    const block_iq2_xs * x = (const block_iq2_xs *) vx;
 
-    const int64_t tid = threadIdx.x;
+    dequantize_iq2_xs(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint16_t * q2 = x[i].qs + 4*ib;
-    const uint8_t  * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511));
-    const float d = (float)x[i].d * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
-    const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
-    for (int j = 0; j < 8; ++j) {
-        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
-    }
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
     const int64_t i = blockIdx.x;
-    const block_iq2_s * x = (const block_iq2_s *) vx;
 
-    const int64_t tid = threadIdx.x;
+    dequantize_iq2_s(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[i].qs[4*ib+il] | ((x[i].qh[ib] << (8-2*il)) & 0x300)));
-    const float d = (float)x[i].d * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
-    const uint8_t signs = x[i].qs[QK_K/8+4*ib+il];
-    for (int j = 0; j < 8; ++j) {
-        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
-    }
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
     const int64_t i = blockIdx.x;
-    const block_iq3_xxs * x = (const block_iq3_xxs  *) vx;
 
-    const int64_t tid = threadIdx.x;
+    dequantize_iq3_xxs(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint8_t  * q3 = x[i].qs + 8*ib;
-    const uint16_t * gas = (const uint16_t *)(x[i].qs + QK_K/4) + 2*ib;
-    const uint8_t  * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*il+0]);
-    const uint8_t  * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*il+1]);
-    const uint32_t aux32 = gas[0] | (gas[1] << 16);
-    const float d = (float)x[i].d * (0.5f + (aux32 >> 28)) * 0.5f;
-    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
-    for (int j = 0; j < 4; ++j) {
-        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
-        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
-    }
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_iq3_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
     const int64_t i = blockIdx.x;
-    const block_iq3_s * x = (const block_iq3_s *) vx;
 
-    const int64_t tid = threadIdx.x;
+    dequantize_iq3_s(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint8_t * qs = x[i].qs + 8*ib;
-    const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*il+0] | ((x[i].qh[ib] << (8-2*il)) & 256)));
-    const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*il+1] | ((x[i].qh[ib] << (7-2*il)) & 256)));
-    const float d = (float)x[i].d * (1 + 2*((x[i].scales[ib/2] >> 4*(ib%2)) & 0xf));
-    const uint8_t signs = x[i].signs[4*ib + il];
-    for (int j = 0; j < 4; ++j) {
-        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
-        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
-    }
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_iq1_s(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
     const int64_t i = blockIdx.x;
-    const block_iq1_s * x = (const block_iq1_s  *) vx;
 
-    const int64_t tid = threadIdx.x;
+    dequantize_iq1_s(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const float delta = x[i].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA;
-    const float d = (float)x[i].d * (2*((x[i].qh[ib] >> 12) & 7) + 1);
-    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
-    grid32[0] = iq1s_grid_gpu[x[i].qs[4*ib+il] | (((x[i].qh[ib] >> 3*il) & 7) << 8)];
-    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
-    grid32[0] &= 0x0f0f0f0f;
-    for (int j = 0; j < 8; ++j) {
-        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
-    }
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
     const int64_t i = blockIdx.x;
-    const block_iq1_m * x = (const block_iq1_m  *) vx;
 
-    const int64_t tid = threadIdx.x;
+    dequantize_iq1_m(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 8*il;
-    const uint16_t * sc = (const uint16_t *)x[i].scales;
-    iq1m_scale_t scale;
-    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
-    const int64_t ib16 = 2*ib + il/2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4);
-    const float d = (float)scale.f16 * (2*((sc[ib16/4] >> 3*(ib16%4)) & 0x7) + 1);
-    const float delta = x[i].qh[2*ib+il/2] & (0x08 << 4*(il%2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA;
-    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
-    grid32[0] = iq1s_grid_gpu[x[i].qs[4*ib+il] | (((x[i].qh[2*ib+il/2] >> 4*(il%2)) & 7) << 8)];
-    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
-    grid32[0] &= 0x0f0f0f0f;
-    for (int j = 0; j < 8; ++j) {
-        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
-    }
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
     const int64_t i = blockIdx.x;
-    const block_iq4_nl * x = (const block_iq4_nl *) vx + i*(QK_K/QK4_NL);
 
-    const int64_t tid = threadIdx.x;
+    dequantize_iq4_nl(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 4*il;
-    const uint8_t  * q4 = x[ib].qs + 4*il;
-    const float d = (float)x[ib].d;
-    for (int j = 0; j < 4; ++j) {
-        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
-        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
-    }
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst_t * __restrict__ yy) {
     const int64_t i = blockIdx.x;
-    const block_iq4_xs * x = (const block_iq4_xs *)vx;
 
-    const int64_t tid = threadIdx.x;
+    dequantize_iq4_xs(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 4*il;
-    const uint8_t  * q4 = x[i].qs + 16*ib + 4*il;
-    const float d = (float)x[i].d * ((((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4)) - 32);
-    for (int j = 0; j < 4; ++j) {
-        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
-        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
-    }
 }
 
 template<typename dst_t>
 static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_t * __restrict__ yy) {
-
     const int64_t i = blockIdx.x;
-    const block_mxfp4 * x = (const block_mxfp4 *) vx + i*(QK_K/QK_MXFP4);
 
-    const int64_t tid = threadIdx.x;
+    dequantize_mxfp4(vx, i, yy + i*QK_K, threadIdx.x);
-    const int64_t il = tid/8; // 0...3
-    const int64_t ib = tid%8; // 0...7
-    dst_t * y = yy + i*QK_K + 32*ib + 4*il;
-    const uint8_t  * q4 = x[ib].qs + 4*il;
-    const float d = ggml_cuda_e8m0_to_fp32(x[ib].e);
-    for (int j = 0; j < 4; ++j) {
-        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] & 0xf]*0.5f);
-        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] >>  4]*0.5f);
-    }
 }
 
 template <int qk, int qr, dequantize_kernel_t dequantize_kernel, typename dst_t>

ggml/src/ggml-cuda/dequantize.cuh

@@ -1,4 +1,5 @@
 #include "common.cuh"
+#include "convert.cuh"
 
 static __device__ __forceinline__ void dequantize_q1_0(const void * vx, const int64_t ib, const int iqs, float2 & v){
     const block_q1_0 * x = (const block_q1_0 *) vx;
@@ -97,3 +98,335 @@ static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const in
     v.x *= d;
     v.y *= d;
 }
+
+//================================== k-quants
+
+// Each call dequantizes one super-block of QK_K values into y using the
+// thread layout of the caller: 32 threads for q4_K, 64 threads otherwise.
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q2_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q2_K * x = (const block_q2_K *) vx;
+
+    const int64_t n   = tid/32;
+    const int64_t l   = tid - 32*n;
+    const int64_t is  = 8*n + l/16;
+
+    const uint8_t q = x[ib].qs[32*n + l];
+    dst_t * y = yy + 128*n;
+
+    float dall = __low2half(x[ib].dm);
+    float dmin = __high2half(x[ib].dm);
+    y[l+ 0] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[ib].scales[is+0] >> 4));
+    y[l+32] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[ib].scales[is+2] >> 4));
+    y[l+64] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[ib].scales[is+4] >> 4));
+    y[l+96] = ggml_cuda_cast<dst_t>(dall * (x[ib].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[ib].scales[is+6] >> 4));
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q3_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q3_K * x = (const block_q3_K *) vx;
+
+    const int64_t r = tid/4;
+    const int64_t t = r/2;
+    const int64_t is0 = r%2;
+    const int64_t l0 = 16*is0 + 4*(tid%4);
+    const int64_t n = t / 4;
+    const int64_t j = t - 4*n;
+
+    uint8_t m = 1 << (4*n + j);
+    int64_t is = 8*n + 2*j + is0;
+    int shift = 2*j;
+
+    int8_t us = is <  4 ? (x[ib].scales[is-0] & 0xF) | (((x[ib].scales[is+8] >> 0) & 3) << 4) :
+                is <  8 ? (x[ib].scales[is-0] & 0xF) | (((x[ib].scales[is+4] >> 2) & 3) << 4) :
+                is < 12 ? (x[ib].scales[is-8] >>  4) | (((x[ib].scales[is+0] >> 4) & 3) << 4) :
+                          (x[ib].scales[is-8] >>  4) | (((x[ib].scales[is-4] >> 6) & 3) << 4);
+    float d_all = x[ib].d;
+    float dl = d_all * (us - 32);
+
+    dst_t * y = yy + 128*n + 32*j;
+    const uint8_t * q = x[ib].qs + 32*n;
+    const uint8_t * hm = x[ib].hmask;
+
+    for (int l = l0; l < l0+4; ++l) {
+        y[l] = ggml_cuda_cast<dst_t>(dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)));
+    }
+}
+
+static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) {
+    if (j < 4) {
+        d = q[j] & 63; m = q[j + 4] & 63;
+    } else {
+        d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4);
+        m = (q[j+4] >>  4) | ((q[j-0] >> 6) << 4);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q4_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q4_K * x = (const block_q4_K *) vx;
+
+    // assume 32 threads
+    const int64_t il  = tid/8;
+    const int64_t ir  = tid%8;
+    const int64_t is  = 2*il;
+    const int64_t n   = 4;
+
+    dst_t * y = yy + 64*il + n*ir;
+
+    const float dall = __low2half(x[ib].dm);
+    const float dmin = __high2half(x[ib].dm);
+
+    const uint8_t * q = x[ib].qs + 32*il + n*ir;
+
+    uint8_t sc, m;
+    get_scale_min_k4(is + 0, x[ib].scales, sc, m);
+    const float d1 = dall * sc; const float m1 = dmin * m;
+    get_scale_min_k4(is + 1, x[ib].scales, sc, m);
+    const float d2 = dall * sc; const float m2 = dmin * m;
+    for (int l = 0; l < n; ++l) {
+        y[l + 0] = ggml_cuda_cast<dst_t>(d1 * (q[l] & 0xF) - m1);
+        y[l +32] = ggml_cuda_cast<dst_t>(d2 * (q[l] >>  4) - m2);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q5_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q5_K * x = (const block_q5_K *) vx;
+
+    // assume 64 threads - this is very slightly better than the one below
+    const int64_t il  = tid/16;   // il is in 0...3
+    const int64_t ir  = tid%16;   // ir is in 0...15
+    const int64_t is  = 2*il;     // is is in 0...6
+
+    dst_t * y = yy + 64*il + 2*ir;
+
+    const float dall = __low2half(x[ib].dm);
+    const float dmin = __high2half(x[ib].dm);
+
+    const uint8_t * ql = x[ib].qs + 32*il + 2*ir;
+    const uint8_t * qh = x[ib].qh + 2*ir;
+
+    uint8_t sc, m;
+    get_scale_min_k4(is + 0, x[ib].scales, sc, m);
+    const float d1 = dall * sc; const float m1 = dmin * m;
+    get_scale_min_k4(is + 1, x[ib].scales, sc, m);
+    const float d2 = dall * sc; const float m2 = dmin * m;
+
+    uint8_t   hm  = 1 << (2*il);
+    y[ 0] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1);
+    y[ 1] = ggml_cuda_cast<dst_t>(d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1);
+    hm <<= 1;
+    y[32] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 0] >>  4) + (qh[ 0] & hm ? 16 : 0)) - m2);
+    y[33] = ggml_cuda_cast<dst_t>(d2 * ((ql[ 1] >>  4) + (qh[ 1] & hm ? 16 : 0)) - m2);
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_q6_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) {
+    const block_q6_K * x = (const block_q6_K *) vx;
+
+    // assume 64 threads - this is very slightly better than the one below
+    const int64_t ip  = tid/32;   // ip is 0 or 1
+    const int64_t il  = tid - 32*ip; // 0...32
+    const int64_t is  = 8*ip + il/16;
+
+    dst_t * y = yy + 128*ip + il;
+
+    const float d = x[ib].d;
+
+    const uint8_t * ql = x[ib].ql + 64*ip + il;
+    const uint8_t   qh = x[ib].qh[32*ip + il];
+    const int8_t  * sc = x[ib].scales + is;
+
+    y[ 0] = ggml_cuda_cast<dst_t>(d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32));
+    y[32] = ggml_cuda_cast<dst_t>(d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32));
+    y[64] = ggml_cuda_cast<dst_t>(d * sc[4] * ((int8_t)((ql[ 0]  >> 4) | (((qh >> 4) & 3) << 4)) - 32));
+    y[96] = ggml_cuda_cast<dst_t>(d * sc[6] * ((int8_t)((ql[32]  >> 4) | (((qh >> 6) & 3) << 4)) - 32));
+}
+
+//================================== i-quants
+
+// Each call dequantizes one super-block of QK_K values into y with 32
+// threads; iq4_nl packs QK_K/QK4_NL sub-blocks per super-block.
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq2_xxs * x = (const block_iq2_xxs  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint16_t * q2 = x[ibs].qs + 4*ib;
+    const uint8_t  * aux8 = (const uint8_t *)q2;
+    const uint8_t  * grid = (const uint8_t *)(iq2xxs_grid + aux8[il]);
+    const uint32_t aux32 = q2[2] | (q2[3] << 16);
+    const float d = (float)x[ibs].d * (0.5f + (aux32 >> 28)) * 0.25f;
+    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq2_xs * x = (const block_iq2_xs *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint16_t * q2 = x[ibs].qs + 4*ib;
+    const uint8_t  * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511));
+    const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
+    const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq2_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq2_s * x = (const block_iq2_s *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[ibs].qs[4*ib+il] | ((x[ibs].qh[ib] << (8-2*il)) & 0x300)));
+    const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
+    const uint8_t signs = x[ibs].qs[QK_K/8+4*ib+il];
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq3_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq3_xxs * x = (const block_iq3_xxs  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint8_t  * q3 = x[ibs].qs + 8*ib;
+    const uint16_t * gas = (const uint16_t *)(x[ibs].qs + QK_K/4) + 2*ib;
+    const uint8_t  * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*il+0]);
+    const uint8_t  * grid2 = (const uint8_t *)(iq3xxs_grid + q3[2*il+1]);
+    const uint32_t aux32 = gas[0] | (gas[1] << 16);
+    const float d = (float)x[ibs].d * (0.5f + (aux32 >> 28)) * 0.5f;
+    const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
+    for (int j = 0; j < 4; ++j) {
+        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
+        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq3_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq3_s * x = (const block_iq3_s *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint8_t * qs = x[ibs].qs + 8*ib;
+    const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*il+0] | ((x[ibs].qh[ib] << (8-2*il)) & 256)));
+    const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*il+1] | ((x[ibs].qh[ib] << (7-2*il)) & 256)));
+    const float d = (float)x[ibs].d * (1 + 2*((x[ibs].scales[ib/2] >> 4*(ib%2)) & 0xf));
+    const uint8_t signs = x[ibs].signs[4*ib + il];
+    for (int j = 0; j < 4; ++j) {
+        y[j+0] = ggml_cuda_cast<dst_t>(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
+        y[j+4] = ggml_cuda_cast<dst_t>(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq1_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq1_s * x = (const block_iq1_s  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const float delta = x[ibs].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA;
+    const float d = (float)x[ibs].d * (2*((x[ibs].qh[ib] >> 12) & 7) + 1);
+    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
+    grid32[0] = iq1s_grid_gpu[x[ibs].qs[4*ib+il] | (((x[ibs].qh[ib] >> 3*il) & 7) << 8)];
+    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
+    grid32[0] &= 0x0f0f0f0f;
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq1_m(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq1_m * x = (const block_iq1_m  *) vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 8*il;
+    const uint16_t * sc = (const uint16_t *)x[ibs].scales;
+    iq1m_scale_t scale;
+    scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
+    const int64_t ib16 = 2*ib + il/2; // sc[ib16/4] >> 3*(ib16%4) -> sc[ib/2] >> 3*((2*ib+il/2)%4);
+    const float d = (float)scale.f16 * (2*((sc[ib16/4] >> 3*(ib16%4)) & 0x7) + 1);
+    const float delta = x[ibs].qh[2*ib+il/2] & (0x08 << 4*(il%2)) ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA;
+    uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32;
+    grid32[0] = iq1s_grid_gpu[x[ibs].qs[4*ib+il] | (((x[ibs].qh[2*ib+il/2] >> 4*(il%2)) & 7) << 8)];
+    grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
+    grid32[0] &= 0x0f0f0f0f;
+    for (int j = 0; j < 8; ++j) {
+        y[j] = ggml_cuda_cast<dst_t>(d * (q[j] + delta));
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq4_nl(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_iq4_nl * x = (const block_iq4_nl *) vx + ibs*(QK_K/QK4_NL);
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 4*il;
+    const uint8_t  * q4 = x[ib].qs + 4*il;
+    const float d = (float)x[ib].d;
+    for (int j = 0; j < 4; ++j) {
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_iq4_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+    const block_iq4_xs * x = (const block_iq4_xs *)vx;
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 4*il;
+    const uint8_t  * q4 = x[ibs].qs + 16*ib + 4*il;
+    const float d = (float)x[ibs].d * ((((x[ibs].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[ibs].scales_h >> 2*ib) & 3) << 4)) - 32);
+    for (int j = 0; j < 4; ++j) {
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] & 0xf]);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_iq4nl[q4[j] >>  4]);
+    }
+}
+
+template<typename dst_t>
+static __device__ __forceinline__ void dequantize_mxfp4(const void * vx, const int64_t ibs, dst_t * yy, const int tid) {
+
+    const block_mxfp4 * x = (const block_mxfp4 *) vx + ibs*(QK_K/QK_MXFP4);
+
+    const int64_t il = tid/8; // 0...3
+    const int64_t ib = tid%8; // 0...7
+    dst_t * y = yy + 32*ib + 4*il;
+    const uint8_t  * q4 = x[ib].qs + 4*il;
+    const float d = ggml_cuda_e8m0_to_fp32(x[ib].e);
+    for (int j = 0; j < 4; ++j) {
+        y[j+ 0] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] & 0xf]*0.5f);
+        y[j+16] = ggml_cuda_cast<dst_t>(d * kvalues_mxfp4[q4[j] >>  4]*0.5f);
+    }
+}

ggml/src/ggml-cuda/getrows.cu

@@ -40,6 +40,35 @@ static __global__ void k_get_rows(
     }
 }
 
+template<typename dst_t, dequantize_kq_t<dst_t> dequantize_kq>
+static __global__ void k_get_rows_kq(
+        const void * __restrict__ src0, const int32_t * __restrict__ src1, dst_t * __restrict__ dst,
+        const int64_t ne00, /*const int64_t ne01, const int64_t ne02, const int64_t ne03,*/
+        /*const int64_t ne10,*/ const int64_t ne11, const uint3 ne12_fdv, /*const int64_t ne13,*/
+        /*const size_t s0,*/ const size_t s1, const size_t s2, const size_t s3,
+        /*const size_t nb00,*/ const size_t nb01, const size_t nb02, const size_t nb03,
+        const size_t s10, const size_t s11, const size_t s12/*, const size_t s13*/) {
+
+    ggml_cuda_pdl_sync();
+    const int64_t nsb = ne00/QK_K; // super-blocks per row
+    for (int64_t z = blockIdx.z; z < ne11*(int64_t)ne12_fdv.z; z += gridDim.z) {
+        // The x and y dimensions of the grid are swapped because the maximum allowed grid size for x is higher.
+        const int i10 = blockIdx.x;
+        const uint2 dm  = fast_div_modulo((uint32_t)z, ne12_fdv);
+        const int i11 = dm.x;
+        const int i12 = dm.y;
+
+        const int i01 = src1[i10*s10 + i11*s11 + i12*s12];
+
+        dst_t * dst_row = dst + i10*s1 + i11*s2 + i12*s3;
+        const void * src0_row = (const char *) src0 + i01*nb01 + i11*nb02 + i12*nb03;
+
+        for (int64_t ib = blockIdx.y; ib < nsb; ib += gridDim.y) {
+            dequantize_kq(src0_row, ib, dst_row + ib*QK_K, threadIdx.x);
+        }
+    }
+}
+
 template<typename src0_t, typename dst_t>
 static __global__ void k_get_rows_float(
         const src0_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr,
@@ -164,6 +193,43 @@ static void get_rows_cuda_q(
         s10, s11, s12/*, s13*/);
 }
 
+template<int block_dim, typename dst_t, dequantize_kq_t<dst_t> dequantize_kq>
+static void get_rows_cuda_kq(
+        const void * src0_d, const int32_t * src1_d, dst_t * dst_d,
+        const int64_t ne00, const size_t nb01, const size_t nb02, const size_t nb03,
+        const int64_t ne10, const int64_t ne11, const int64_t ne12, const size_t nb10, const size_t nb11, const size_t nb12,
+        const size_t nb1, const size_t nb2, const size_t nb3,
+        cudaStream_t stream) {
+    GGML_ASSERT(ne00 % QK_K == 0);
+    const int64_t nsb = ne00/QK_K;
+
+    const dim3 block_dims(block_dim, 1, 1);
+    const dim3 block_nums(ne10, MIN(nsb, UINT16_MAX), MIN(ne11*ne12, UINT16_MAX));
+
+    // strides in elements
+    // const size_t s0 = nb0 / sizeof(dst_t);
+    const size_t s1 = nb1 / sizeof(dst_t);
+    const size_t s2 = nb2 / sizeof(dst_t);
+    const size_t s3 = nb3 / sizeof(dst_t);
+
+    const size_t s10 = nb10 / sizeof(int32_t);
+    const size_t s11 = nb11 / sizeof(int32_t);
+    const size_t s12 = nb12 / sizeof(int32_t);
+    // const size_t s13 = nb13 / sizeof(int32_t);
+
+    GGML_ASSERT(ne12 > 0);
+    GGML_ASSERT(ne11 <= std::numeric_limits<uint32_t>::max() / ne12);
+    const uint3 ne12_fdv = init_fastdiv_values(ne12);
+
+    k_get_rows_kq<dst_t, dequantize_kq><<<block_nums, block_dims, 0, stream>>>(
+        src0_d, src1_d, dst_d,
+        ne00, /*ne01, ne02, ne03,*/
+        /*ne10,*/ ne11, ne12_fdv, /*ne13,*/
+        /* s0,*/ s1, s2, s3,
+        /* nb00,*/ nb01, nb02, nb03,
+        s10, s11, s12/*, s13*/);
+}
+
 template<typename src0_t, typename dst_t>
 static void get_rows_cuda_float(
         const src0_t * src0_d, const int32_t * src1_d, dst_t * dst_d,
@@ -274,8 +340,67 @@ static void ggml_cuda_get_rows_switch_src0_type(
             get_rows_cuda_q<QK8_0, QR8_0, dequantize_q8_0>(src0_d, src1_d, dst_d,
                 ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
             break;
+        case GGML_TYPE_Q2_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q2_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q3_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q3_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q4_K:
+            get_rows_cuda_kq<32, dst_t, dequantize_q4_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q5_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q5_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_Q6_K:
+            get_rows_cuda_kq<64, dst_t, dequantize_q6_K<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ2_XXS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq2_xxs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ2_XS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq2_xs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ2_S:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq2_s<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ3_XXS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq3_xxs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ3_S:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq3_s<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ1_S:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq1_s<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ1_M:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq1_m<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ4_NL:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq4_nl<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_IQ4_XS:
+            get_rows_cuda_kq<32, dst_t, dequantize_iq4_xs<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
+        case GGML_TYPE_MXFP4:
+            get_rows_cuda_kq<32, dst_t, dequantize_mxfp4<dst_t>>(src0_d, src1_d, dst_d,
+                ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream);
+            break;
         default:
-            // TODO: k-quants
             GGML_ABORT("%s: unsupported src0 type: %s\n", __func__, ggml_type_name(src0_type));
             break;
     }

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

@@ -2703,6 +2703,7 @@ static int ggml_cuda_try_gdn_cache_fusion(
 
 static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) {
     args.sigmoid         = false;
+    args.sqrt_softplus   = false;
     args.softmax         = false;
     args.delayed_softmax = false;
     args.prob_bias       = false;
@@ -2716,10 +2717,17 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod
     }
 
     if (nodes[node_idx]->op == GGML_OP_UNARY) {
-        if (ggml_get_unary_op(nodes[node_idx]) != GGML_UNARY_OP_SIGMOID) {
+        const ggml_unary_op unary_op = ggml_get_unary_op(nodes[node_idx]);
+        if (unary_op == GGML_UNARY_OP_SIGMOID) {
+            args.sigmoid = true;
+        } else if (unary_op == GGML_UNARY_OP_SOFTPLUS && node_idx + 1 < n_nodes &&
+                   nodes[node_idx + 1]->op == GGML_OP_SQRT && nodes[node_idx + 1]->src[0] == nodes[node_idx]) {
+            // sqrt(softplus(x)) scoring (DeepSeek-V4)
+            args.sqrt_softplus = true;
+            node_idx++;
+        } else {
             return false;
         }
-        args.sigmoid = true;
     }
 
     if (nodes[node_idx]->op == GGML_OP_ARGSORT) {
@@ -2728,7 +2736,7 @@ static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int nod
 
     node_idx++;
 
-    if (args.sigmoid || args.softmax) {
+    if (args.sigmoid || args.sqrt_softplus || args.softmax) {
         // SOFTMAX -> RESHAPE
         if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE ||
                 nodes[node_idx]->src[0] != nodes[node_idx - 1]) {
@@ -3172,21 +3180,27 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
             const ggml_tensor * scale   = nullptr;
 
             if (!args.delayed_softmax) {
-                ggml_op gating_op = args.sigmoid ? GGML_OP_UNARY : GGML_OP_SOFT_MAX;
                 int out_nodes[2];  // nodes which can't be elided
 
-                if (args.prob_bias) {
+                if (args.sigmoid) {
-                    bias = cgraph->nodes[i + 2]->src[1];
+                    ops.insert(ops.end(), { GGML_OP_UNARY });
-                    ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
+                } else if (args.sqrt_softplus) {
-                                            GGML_OP_GET_ROWS });
+                    ops.insert(ops.end(), { GGML_OP_UNARY, GGML_OP_SQRT });
-                    out_nodes[0] = i + 4;
-                    ids          = cgraph->nodes[i + 4];
                 } else {
-                    ops.insert(ops.end(),
+                    ops.insert(ops.end(), { GGML_OP_SOFT_MAX });
-                               { gating_op, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
-                    out_nodes[0] = i + 3;
-                    ids          = cgraph->nodes[i + 3];
                 }
+                const int i_probs = i + (int) ops.size() - 1;  // last node of the gating activation
+
+                if (args.prob_bias) {
+                    bias = cgraph->nodes[i_probs + 2]->src[1];
+                    ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, GGML_OP_VIEW,
+                                            GGML_OP_GET_ROWS });
+                    out_nodes[0] = i_probs + 4;
+                } else {
+                    ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS });
+                    out_nodes[0] = i_probs + 3;
+                }
+                ids = cgraph->nodes[out_nodes[0]];
 
                 if (args.norm) {
                     ops.insert(ops.end(),
@@ -4831,7 +4845,25 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                     case GGML_TYPE_Q5_0:
                     case GGML_TYPE_Q5_1:
                     case GGML_TYPE_Q8_0:
+                    case GGML_TYPE_Q2_K:
+                    case GGML_TYPE_Q3_K:
+                    case GGML_TYPE_Q4_K:
+                    case GGML_TYPE_Q5_K:
+                    case GGML_TYPE_Q6_K:
+                    case GGML_TYPE_IQ2_XXS:
+                    case GGML_TYPE_IQ2_XS:
+                    case GGML_TYPE_IQ2_S:
+                    case GGML_TYPE_IQ3_XXS:
+                    case GGML_TYPE_IQ3_S:
+                    case GGML_TYPE_IQ1_S:
+                    case GGML_TYPE_IQ1_M:
+                    case GGML_TYPE_IQ4_XS:
                         return true;
+                    case GGML_TYPE_IQ4_NL:
+                    case GGML_TYPE_MXFP4:
+                        // 32-value sub-blocks, the row size does not guarantee
+                        // the QK_K super-blocks the get_rows kernel iterates on
+                        return op->src[0]->ne[0] % QK_K == 0;
                     default:
                         return false;
                 }

ggml/src/ggml-cuda/mmq.cu

@@ -130,14 +130,20 @@ void ggml_cuda_mul_mat_q(
         const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * y_block_size/y_values_per_block +
             ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
         ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
+        ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
+        if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
+            src1_scale.alloc(ne13*ne12*ne11);
+        }
 
         {
             const int64_t s11 = src1->nb[1] / ts_src1;
             const int64_t s12 = src1->nb[2] / ts_src1;
             const int64_t s13 = src1->nb[3] / ts_src1;
             if (use_native_fp4) {
+                static constexpr size_t align_float8 = 32;
+                const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
                 static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1));
-                quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded,
+                quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded,
                                         ne11, ne12, ne13, stream);
 
             } else {
@@ -155,6 +161,7 @@ void ggml_cuda_mul_mat_q(
 
         const mmq_args args = {
             src0_d, src0->type, (const int *) src1_q8_1.ptr, nullptr, nullptr, dst_d,
+            src0->type == GGML_TYPE_NVFP4 && use_native_fp4 ? src1_scale.ptr : nullptr,
             ne00, ne01, ne1, s01, ne11, s1,
             ne02, ne12, s02, s12, s2,
             ne03, ne13, s03, s13, s3,
@@ -192,6 +199,10 @@ void ggml_cuda_mul_mat_q(
     const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * y_block_size/y_values_per_block +
         ggml_cuda_mmq_get_J_max(src0->type, fallback, cc, ne11) * sizeof(block_q8_1_mmq);
     ggml_cuda_pool_alloc<char> src1_q8_1(ctx.pool(), nbytes_src1_q8_1);
+    ggml_cuda_pool_alloc<float> src1_scale(ctx.pool());
+    if (src0->type == GGML_TYPE_NVFP4 && use_native_fp4) {
+        src1_scale.alloc(ne12*n_expert_used);
+    }
 
     const int64_t ne11_flat = ne12*n_expert_used;
     const int64_t ne12_flat = 1;
@@ -202,18 +213,19 @@ void ggml_cuda_mul_mat_q(
         const int64_t s12 = src1->nb[2] / ts_src1;
         const int64_t s13 = src1->nb[3] / ts_src1;
 
-        if (dedup_bcast) {
-            // quantize each token once, scatter its block to all n_expert_used slots
         if (use_native_fp4) {
-                quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
+            static constexpr size_t align_float8 = 32;
+            const bool use_aligned_float8 = ggml_cuda_is_aligned(src1, align_float8);
+            if (dedup_bcast) {
+                quantize_scatter_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10,
                                         /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
             } else {
+                quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13,
+                                        ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
+            }
+        } else if (dedup_bcast) {
             quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10,
                                     /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream);
-            }
-        } else if (use_native_fp4) {
-            quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
-                                    ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
         } else {
             quantize_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, s11, s12, s13,
                                    ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream);
@@ -229,6 +241,7 @@ void ggml_cuda_mul_mat_q(
     // Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid.
     const mmq_args args = {
         src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d,
+        src1_scale.ptr,
         ne00, ne01, ne_get_rows, s01, ne_get_rows, s1,
         ne02, ne02, s02, s12, s2,
         ne03, ne13, s03, s13, s3,

ggml/src/ggml-cuda/mmq.cuh

@@ -13,7 +13,7 @@
 typedef void (*ggml_cuda_mmq_load_tiles_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride);
 typedef void (*ggml_cuda_mmq_vec_dot_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00);
 typedef void (*ggml_cuda_mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted,
-    float * __restrict__ dst, const int stride, const int i_max, const int j_max);
+    float * __restrict__ dst, const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max);
 
 enum mmq_q8_1_ds_layout {
     MMQ_Q8_1_DS_LAYOUT_D4,
@@ -413,11 +413,13 @@ static __host__ int ggml_cuda_mmq_get_nbytes_shared_x(const ggml_cuda_mmq_config
 
 template <ggml_type type, int J, bool fallback> static __device__ __forceinline__ void ggml_cuda_mmq_write_back_dp4a(
         const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst,
-        const int stride, const int i_max, const int j_max) {
+        const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
     constexpr int warp_size = ggml_cuda_get_physical_warp_size();
     constexpr int nwarps    = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size;
     constexpr int I         = ggml_cuda_mmq_get_I(type, J, fallback);
 
+    const bool y_scale_used = y_scale != nullptr;
+
 #pragma unroll
     for (int j0 = 0; j0 < J; j0 += nwarps) {
         const int j = j0 + threadIdx.y;
@@ -434,15 +436,25 @@ template <ggml_type type, int J, bool fallback> static __device__ __forceinline_
                 continue;
             }
 
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                if (y_scale_used) {
+                    dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+                } else {
                     dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
                 }
+            } else {
+                dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (I/warp_size) + i0/warp_size];
+                GGML_UNUSED(y_scale_used);
+            }
+        }
     }
 }
 
 template<ggml_type type, int J, bool fallback>
 static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
             const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst,
-            const int stride, const int i_max, const int j_max) {
+            const float * __restrict__ y_scale, const int stride, const int i_max, const int j_max) {
+
 #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
     typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C;
 #else
@@ -457,6 +469,8 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
 
     const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I);
 
+    const bool y_scale_used = y_scale != nullptr;
+
 #pragma unroll
     for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) {
 #pragma unroll
@@ -475,8 +489,17 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma(
                     continue;
                 }
 
+                if constexpr (type == GGML_TYPE_NVFP4) {
+                    if (y_scale_used) {
+                        dst[ids_dst[j]*stride + i] = y_scale[j] * sum[(j0/tile_C::J + n)*tile_C::ne + l];
+                    } else {
                         dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
                     }
+                } else {
+                    dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l];
+                    GGML_UNUSED(y_scale_used);
+                }
+            }
         }
     }
 }
@@ -819,6 +842,7 @@ template <ggml_type type, int J, bool fallback, bool fixup>
 static __device__ __forceinline__ void mul_mat_q_process_tile(
         const char * __restrict__ x, const int offset_x, const int * __restrict__ y,
         const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup,
+        const float * __restrict__ y_scale,
         const int stride_row_x, const int ncols_y, const int stride_col_dst,
         const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) {
 
@@ -884,9 +908,9 @@ static __device__ __forceinline__ void mul_mat_q_process_tile(
     }
 
     if (fixup) {
-        write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), I, I, J);
+        write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J);
     } else {
-        write_back(sum, ids_dst, dst, stride_col_dst, tile_x_max_i, tile_y_max_j);
+        write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j);
     }
 }
 
@@ -898,6 +922,7 @@ __launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_g
 static __global__ void mul_mat_q(
         const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst,
         const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup,
+        const float * __restrict__ y_scale,
         const uint3 blocks_per_ne00, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst,
         const uint3 channel_ratio, const uint3 nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst,
         const uint3 sample_ratio, const uint3 nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst,
@@ -945,6 +970,12 @@ static __global__ void mul_mat_q(
         int col_diff   = ncols_dst;
         int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
         int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+        int offset_y_scale;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+        } else {
+            GGML_UNUSED(offset_y_scale);
+        }
 
         if (ids_dst) {
             col_low  = expert_bounds[zt + 0];
@@ -953,6 +984,9 @@ static __global__ void mul_mat_q(
 
             offset_y   = 0;
             offset_dst = 0;
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                offset_y_scale = 0;
+            }
 
             if (jt*J >= col_diff) {
                 return;
@@ -974,6 +1008,11 @@ static __global__ void mul_mat_q(
 
         offset_y   += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int));
         offset_dst += it*I;
+        const float * y_scale_tile = nullptr;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale += col_low + jt*J;
+            y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+        }
 
         const int tile_x_max_i = nrows_x  - it*I - 1;
         const int tile_y_max_j = col_diff - jt*J - 1;
@@ -982,7 +1021,8 @@ static __global__ void mul_mat_q(
 
         constexpr bool fixup = false;
         mul_mat_q_process_tile<type, J, fallback, fixup>
-            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
+            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+             stride_row_x, ncols_y, stride_col_dst,
              tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z);
         return;
     }
@@ -1018,6 +1058,12 @@ static __global__ void mul_mat_q(
         int col_diff   = ncols_dst;
         int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
         int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+        int offset_y_scale;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+        } else {
+            GGML_UNUSED(offset_y_scale);
+        }
 
         if (ids_dst) {
             col_low  = expert_bounds[zt + 0];
@@ -1026,6 +1072,9 @@ static __global__ void mul_mat_q(
 
             offset_y   = 0;
             offset_dst = 0;
+            if constexpr (type == GGML_TYPE_NVFP4) {
+                offset_y_scale = 0;
+            }
 
             if (jt*J >= col_diff) {
                 kbc += blocks_per_ne00.z;
@@ -1053,6 +1102,11 @@ static __global__ void mul_mat_q(
 
         offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
         offset_dst += it*I;
+        const float * y_scale_tile = nullptr;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale += col_low + jt * J;
+            y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+        }
 
         const int tile_x_max_i = nrows_x  - it*I - 1;
         const int tile_y_max_j = col_diff - jt*J - 1;
@@ -1061,7 +1115,8 @@ static __global__ void mul_mat_q(
 
         constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer.
         mul_mat_q_process_tile<type, J, fallback, fixup>
-            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
+            (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+             stride_row_x, ncols_y, stride_col_dst,
              tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
 
         kbc += blocks_per_ne00.z;
@@ -1092,6 +1147,12 @@ static __global__ void mul_mat_q(
     int col_diff   = ncols_dst;
     int offset_y       = wt*stride_sample_y   + zt*stride_channel_y;
     int offset_dst     = wt*stride_sample_dst + zt*stride_channel_dst + jt*J*stride_col_dst;
+    int offset_y_scale;
+    if constexpr (type == GGML_TYPE_NVFP4) {
+        offset_y_scale = wt*nchannels_y.z*ncols_y + zt*ncols_y;
+    } else {
+        GGML_UNUSED(offset_y_scale);
+    }
 
     if (ids_dst) {
         col_low  = expert_bounds[zt + 0];
@@ -1100,6 +1161,9 @@ static __global__ void mul_mat_q(
 
         offset_y   = 0;
         offset_dst = 0;
+        if constexpr (type == GGML_TYPE_NVFP4) {
+            offset_y_scale = 0;
+        }
 
         if (jt*J >= col_diff) {
             return;
@@ -1122,6 +1186,11 @@ static __global__ void mul_mat_q(
 
     offset_y += (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int));
     offset_dst += it*I;
+    const float * y_scale_tile = nullptr;
+    if constexpr (type == GGML_TYPE_NVFP4) {
+        offset_y_scale += col_low + jt * J;
+        y_scale_tile = y_scale ? y_scale + offset_y_scale : nullptr;
+    }
 
     const int tile_x_max_i = nrows_x  - it*I - 1;
     const int tile_y_max_j = col_diff - jt*J - 1;
@@ -1130,7 +1199,8 @@ static __global__ void mul_mat_q(
 
     constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks.
     mul_mat_q_process_tile<type, J, fallback, fixup>
-        (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst,
+        (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, y_scale_tile,
+         stride_row_x, ncols_y, stride_col_dst,
          tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop);
 }
 
@@ -1274,6 +1344,7 @@ static __global__ void mul_mat_q_stream_k_fixup(
 
 struct mmq_args {
     const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; float * dst;
+    const float * y_scale;
     int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst;
     int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst;
     int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst;
@@ -1323,7 +1394,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
 
     if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) {
         mul_mat_q<type, J, fallback><<<block_nums_xy_tiling, block_dims, nbytes_shared, stream>>>
-            (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr,
+            (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale,
              blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
              channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
              sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,
@@ -1352,7 +1423,7 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
     const dim3 block_dims_fixup(block_dims.x, block_dims.y/2, block_dims.z);
 
     mul_mat_q<type, J, fallback><<<block_nums_stream_k, block_dims, nbytes_shared, stream>>>
-        (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr,
+        (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup.ptr, args.y_scale,
          blocks_per_ne00_fd, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst,
          channel_ratio_fd, nchannels_y_fd, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst,
          sample_ratio_fd, nsamples_y_fd, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst,

ggml/src/ggml-cuda/quantize.cu

@@ -1,6 +1,55 @@
 #include "quantize.cuh"
 #include <cstdint>
 
+#if defined(BLACKWELL_MMA_AVAILABLE)
+// this maps to 256-bit loads in PTX on supported devices,
+// and otherwise falls back to 2 128-bit loads
+struct __builtin_align__(32) float8 {
+    float x; float y; float z; float w;
+    float p; float q; float r; float s;
+};
+#endif
+
+#if CUDART_VERSION >= 12080
+static __device__ __forceinline__ float nvfp4_native_scale_error(
+        const float vals[QK_NVFP4_SUB], const float inv_col_scale, const float inv_scale, const float scale) {
+    const float scale_dequant = 2.0f * scale;
+    float err = 0.0f;
+
+#pragma unroll
+    for (int k = 0; k < QK_NVFP4_SUB; k += 4) {
+        const float v0 = vals[k + 0] * inv_col_scale;
+        const float v1 = vals[k + 1] * inv_col_scale;
+        const float v2 = vals[k + 2] * inv_col_scale;
+        const float v3 = vals[k + 3] * inv_col_scale;
+
+        const __nv_fp4x4_e2m1 q(make_float4(v0 * inv_scale, v1 * inv_scale, v2 * inv_scale, v3 * inv_scale));
+        const __nv_fp4x4_storage_t q_storage = q.__x;
+        const __nv_fp4x2_storage_t q_lo = static_cast<__nv_fp4x2_storage_t>(q_storage);
+        const __nv_fp4x2_storage_t q_hi = static_cast<__nv_fp4x2_storage_t>(q_storage >> 8U);
+
+        const __half2_raw hraw2_lo = __nv_cvt_fp4x2_to_halfraw2(q_lo, __NV_E2M1);
+        const __half2_raw hraw2_hi = __nv_cvt_fp4x2_to_halfraw2(q_hi, __NV_E2M1);
+        const __half2 h2_lo = static_cast<__half2>(hraw2_lo);
+        const __half2 h2_hi = static_cast<__half2>(hraw2_hi);
+        const float2 dq_lo = __half22float2(h2_lo);
+        const float2 dq_hi = __half22float2(h2_hi);
+
+        const float err0 = fabsf(v0) - fabsf(dq_lo.x) * scale_dequant;
+        const float err1 = fabsf(v1) - fabsf(dq_lo.y) * scale_dequant;
+        const float err2 = fabsf(v2) - fabsf(dq_hi.x) * scale_dequant;
+        const float err3 = fabsf(v3) - fabsf(dq_hi.y) * scale_dequant;
+
+        err = fmaf(err0, err0, err);
+        err = fmaf(err1, err1, err);
+        err = fmaf(err2, err2, err);
+        err = fmaf(err3, err3, err);
+    }
+
+    return err;
+}
+#endif // CUDART_VERSION >= 12080
+
 __launch_bounds__(CUDA_QUANTIZE_BLOCK_SIZE, 1)
 static __global__ void quantize_q8_1(
         const float * x_ptr, void * vy_ptr,
@@ -74,95 +123,189 @@ __device__ __forceinline__ uint8_t compute_e8m0_scale(float amax) {
     return static_cast<uint8_t>(biased);
 }
 
-
 // scatter: grid over tokens, quantize once, write to all the token's compact rows
-template <bool scatter>
+template <bool scatter, bool use_aligned_float8>
 static __global__ void quantize_mmq_nvfp4(
-        const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy,
+        const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, float * __restrict__ scale,
         const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
         const int64_t ne0, const int64_t ne1, const int64_t ne2, const int n_expert_used) {
 #if defined(BLACKWELL_MMA_AVAILABLE)
 
-    const int64_t i0_base = ((int64_t) blockDim.x * blockIdx.y + threadIdx.x) * QK_NVFP4_SUB;
-    if (i0_base >= ne0) {
-        return;
-    }
-
-    const int64_t k_block = i0_base / QK_FP4_MMQ;
     const int64_t blocks_per_col = (ne0 + QK_FP4_MMQ - 1) / QK_FP4_MMQ;
-    if (k_block >= blocks_per_col) {
-        return;
-    }
-    const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB;
 
     int64_t base_idx;
     if constexpr (scatter) {
         base_idx = (int64_t) blockIdx.x * s02; // one physical row per token
     } else {
-        const int64_t i2  = blockIdx.z % ne2;
+        const int64_t i2  = blockIdx.y % ne2;
-        const int64_t i3  = blockIdx.z / ne2;
+        const int64_t i3  = blockIdx.y / ne2;
         const int64_t i01 = ids ? ids[blockIdx.x] : blockIdx.x;
         base_idx = i3 * s03 + i2 * s02 + i01 * s01;
     }
+    const float * __restrict__  x_row = x + base_idx;
 
-    float vals_raw[QK_NVFP4_SUB];
+    float amax = 0.0f;
-    float amax_raw = 0.0f;
+    if constexpr (use_aligned_float8) {
-#pragma unroll
+        for (int64_t i0 = 8 * threadIdx.x; i0 < ne00; i0 += 8 * blockDim.x) {
-    for (int k = 0; k < QK_NVFP4_SUB; k++) {
+            const float * x_base = x_row + i0;
-        const int64_t i00 = i0_base + k;
+            const float8 v = reinterpret_cast<const float8 *>(x_base)[0];
-        if (i00 < ne00) {
+            amax = fmaxf(amax, fabsf(v.x));
-            const float v = x[base_idx + i00];
+            amax = fmaxf(amax, fabsf(v.y));
-            vals_raw[k] = v;
+            amax = fmaxf(amax, fabsf(v.z));
-            amax_raw = fmaxf(amax_raw, fabsf(v));
+            amax = fmaxf(amax, fabsf(v.w));
-        } else {
+            amax = fmaxf(amax, fabsf(v.p));
-            vals_raw[k] = 0.0f;
+            amax = fmaxf(amax, fabsf(v.q));
+            amax = fmaxf(amax, fabsf(v.r));
+            amax = fmaxf(amax, fabsf(v.s));
         }
+    } else {
+        for (int64_t i0 = threadIdx.x; i0 < ne00; i0 += blockDim.x) {
+            amax = fmaxf(amax, fabsf(x_row[i0]));
+        }
+    }
+
+    amax = warp_reduce_max<WARP_SIZE>(amax);
+
+    __shared__ float warp_amax[CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE];
+    const int lane = threadIdx.x % WARP_SIZE;
+    const int warp = threadIdx.x / WARP_SIZE;
+
+    if (lane == 0) {
+        warp_amax[warp] = amax;
+    }
+    __syncthreads();
+
+    if (warp == 0) {
+        amax = threadIdx.x < int(CUDA_QUANTIZE_BLOCK_SIZE_MMQ / WARP_SIZE) ? warp_amax[lane] : 0.0f;
+        amax = warp_reduce_max<WARP_SIZE>(amax);
+        if (lane == 0) {
+            warp_amax[0] = amax / (6.0f * 448.0f);
+            if constexpr (scatter) {
+#pragma unroll
+                for (int slot = 0; slot < n_expert_used; ++slot) {
+                    const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot];
+                    scale[i] = warp_amax[0];
+                }
+            } else {
+                scale[blockIdx.y * ne1 + blockIdx.x] = warp_amax[0];
+            }
+        }
+    }
+    __syncthreads();
+
+    block_fp4_mmq * y = (block_fp4_mmq *) vy;
+    const int64_t n_subblocks = (ne0 + QK_NVFP4_SUB - 1) / QK_NVFP4_SUB;
+
+    for (int64_t isb = threadIdx.x; isb < n_subblocks; isb += blockDim.x) {
+        const int64_t i0_base = isb * QK_NVFP4_SUB;
+        const int64_t k_block = i0_base / QK_FP4_MMQ;
+        const int sub = (i0_base % QK_FP4_MMQ) / QK_NVFP4_SUB;
+
+        const float row_scale = warp_amax[0];
+        const float inv_col_scale = row_scale > 0.0f ? 1.0f / row_scale : 0.0f;
+
+        float vals[QK_NVFP4_SUB];
+        if constexpr (use_aligned_float8) {
+            const float * x_base = x_row + i0_base;
+            const float8 v0 = i0_base +  7 < ne00 ? reinterpret_cast<const float8 *>(x_base)[0]     : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+            const float8 v1 = i0_base + 15 < ne00 ? reinterpret_cast<const float8 *>(x_base + 8)[0] : float8{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
+            vals[0] = v0.x; vals[1] = v0.y; vals[2] = v0.z; vals[3] = v0.w;
+            vals[4] = v0.p; vals[5] = v0.q; vals[6] = v0.r; vals[7] = v0.s;
+            vals[8] = v1.x; vals[9] = v1.y; vals[10] = v1.z; vals[11] = v1.w;
+            vals[12] = v1.p; vals[13] = v1.q; vals[14] = v1.r; vals[15] = v1.s;
+        } else {
+#pragma unroll
+            for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+                const int64_t i00 = i0_base + k;
+                vals[k] = i00 < ne00 ? x_row[i00] : 0.0f;
+            }
+        }
+
+        uint32_t q0 = 0;
+        uint32_t q1 = 0;
+
+        float amax_sub = 0.0f;
+#pragma unroll
+        for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+            amax_sub = fmaxf(amax_sub, fabsf(vals[k] * inv_col_scale));
         }
 
         static constexpr int test_offsets[5] = { 0, -1, 1, -2, 2 };
-    const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_raw / 6.0f);
+        const int first_fp8_code = (int) ggml_cuda_fp32_to_ue4m3(amax_sub / 6.0f);
 
-    float best_err = FLT_MAX;
+        uint8_t fp8_code = (uint8_t) first_fp8_code;
-    uint8_t fp8_code = 0;
+        float subblock_scale = ggml_cuda_ue4m3_to_fp32(fp8_code);
-    float subblock_scale = 0.0f;
+        float inv_scale_err = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+#if CUDART_VERSION >= 12080
+        float best_err = nvfp4_native_scale_error(vals, inv_col_scale, inv_scale_err, subblock_scale);
+#else
+        float best_err = 0.0f;
+#pragma unroll
+        for (int k = 0; k < QK_NVFP4_SUB; ++k) {
+            const float v = vals[k] * inv_col_scale;
+            const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, inv_scale_err);
+            const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * subblock_scale;
+            best_err = fmaf(err_diff, err_diff, best_err);
+        }
+#endif // CUDART_VERSION >= 12080
 
-#pragma unroll // Check +/- 2 to find best code to reduce NVFP4 activation loss. Negligible overhead on Blackwell.
+#pragma unroll
-    for (int i = 0; i < 5; i++) {
+        for (int i = 1; i < 5; ++i) {
             const int test_code = first_fp8_code + test_offsets[i];
             if (test_code < 0 || test_code > 0x7e) {
                 continue;
             }
-        const uint8_t code = (uint8_t) test_code;
+
-        const float test_scale = ggml_cuda_ue4m3_to_fp32(code);
+            const float test_scale = ggml_cuda_ue4m3_to_fp32((uint8_t) test_code);
             const float test_inv_scale = test_scale > 0.0f ? 0.5f / test_scale : 0.0f;
+#if CUDART_VERSION >= 12080
+            const float cur_err = nvfp4_native_scale_error(vals, inv_col_scale, test_inv_scale, test_scale);
+#else
             float cur_err = 0.0f;
 #pragma unroll
             for (int k = 0; k < QK_NVFP4_SUB; ++k) {
-            const float v = vals_raw[k];
+                const float v = vals[k] * inv_col_scale;
                 const uint8_t q = ggml_cuda_float_to_fp4_e2m1(v, test_inv_scale);
-            const float err_diff = fabsf(v) - fabsf(kvalues_mxfp4[q & 0x7]) * test_scale;
+                const float err_diff = fabsf(v) - fabsf(kvalues_fp4[q & 0x7]) * test_scale;
                 cur_err = fmaf(err_diff, err_diff, cur_err);
             }
+#endif // CUDART_VERSION >= 12080
 
             if (cur_err < best_err) {
                 best_err = cur_err;
-            fp8_code = test_code;
+                fp8_code = (uint8_t) test_code;
                 subblock_scale = test_scale;
             }
         }
-
+#if CUDART_VERSION >= 12080
         const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
-    uint32_t q0 = 0;
+        const float s = inv_col_scale * inv_scale;
-    uint32_t q1 = 0;
+
-#pragma unroll // this is faster than the previous __nv_fp4x4_e2m1
+        __nv_fp4x4_e2m1 q0_lo(make_float4(vals[0] * s, vals[8]  * s, vals[1] * s, vals[9]  * s));
+        __nv_fp4x4_e2m1 q0_hi(make_float4(vals[2] * s, vals[10] * s, vals[3] * s, vals[11] * s));
+        __nv_fp4x4_e2m1 q1_lo(make_float4(vals[4] * s, vals[12] * s, vals[5] * s, vals[13] * s));
+        __nv_fp4x4_e2m1 q1_hi(make_float4(vals[6] * s, vals[14] * s, vals[7] * s, vals[15] * s));
+
+        const char2 q0_lo_c = *reinterpret_cast<char2 *>(&q0_lo);
+        const char2 q0_hi_c = *reinterpret_cast<char2 *>(&q0_hi);
+        const char2 q1_lo_c = *reinterpret_cast<char2 *>(&q1_lo);
+        const char2 q1_hi_c = *reinterpret_cast<char2 *>(&q1_hi);
+
+        q0 = uint32_t(uint8_t(q0_lo_c.x)) | (uint32_t(uint8_t(q0_lo_c.y)) <<  8) |
+            (uint32_t(uint8_t(q0_hi_c.x)) << 16) | (uint32_t(uint8_t(q0_hi_c.y)) << 24);
+        q1 = uint32_t(uint8_t(q1_lo_c.x)) | (uint32_t(uint8_t(q1_lo_c.y)) <<  8) |
+            (uint32_t(uint8_t(q1_hi_c.x)) << 16) | (uint32_t(uint8_t(q1_hi_c.y)) << 24);
+#else
+        const float inv_scale = subblock_scale > 0.0f ? 0.5f / subblock_scale : 0.0f;
+#pragma unroll
         for (int k = 0; k < QK_NVFP4_SUB / 4; ++k) {
-        q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k +  0], inv_scale) << (8 * k);
+            q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 0] * inv_col_scale, inv_scale)) << (8 * k);
-        q0 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k +  8], inv_scale) << (8 * k + 4);
+            q0 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 8] * inv_col_scale, inv_scale)) << (8 * k + 4);
-        q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k +  4], inv_scale) << (8 * k);
+            q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 4] * inv_col_scale, inv_scale)) << (8 * k);
-        q1 |= (uint32_t) ggml_cuda_float_to_fp4_e2m1(vals_raw[k + 12], inv_scale) << (8 * k + 4);
+            q1 |= uint32_t(ggml_cuda_float_to_fp4_e2m1(vals[k + 12] * inv_col_scale, inv_scale)) << (8 * k + 4);
         }
+#endif // CUDART_VERSION >= 12080
 
-    block_fp4_mmq * y = (block_fp4_mmq *) vy;
         if constexpr (scatter) {
 #pragma unroll
             for (int slot = 0; slot < n_expert_used; ++slot) {
@@ -174,15 +317,15 @@ static __global__ void quantize_mmq_nvfp4(
                 reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
             }
         } else {
-        block_fp4_mmq * yb = y + (blockIdx.z * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x);
+            block_fp4_mmq * yb = y + (blockIdx.y * ((int64_t) blocks_per_col * ne1) + k_block * ne1 + blockIdx.x);
             uint32_t * yqs = reinterpret_cast<uint32_t *>(yb->qs);
             yqs[2 * sub + 0] = q0;
             yqs[2 * sub + 1] = q1;
             reinterpret_cast<uint8_t *>(yb->d4)[sub] = fp8_code;
         }
-    GGML_UNUSED(n_expert_used);
+    }
 #else
-    GGML_UNUSED(n_expert_used);
+    GGML_UNUSED_VARS(x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, n_expert_used);
     NO_DEVICE_CODE; // This is for Blackwell NVFP4 activations only.
 #endif // defined(BLACKWELL_MMA_AVAILABLE)
 
@@ -491,18 +634,22 @@ void quantize_scatter_mmq_q8_1_cuda(
 
 // scatter=true reuses the quant kernels: grid over tokens, ids = inverse map (token slot -> compact row)
 void quantize_scatter_mmq_fp4_cuda(
-        const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0,
+        const float * x, const int32_t * ids_src1_inv, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8,
         const int64_t ne00, const int64_t stride_token, const int64_t ne0,
         const int64_t n_tokens, const int64_t nrows_dst, const int n_expert_used, cudaStream_t stream) {
     GGML_ASSERT(ne0 > 0);
     if (type_src0 == GGML_TYPE_NVFP4) {
+        GGML_ASSERT(scale);
         GGML_ASSERT(ne00 % QK_NVFP4 == 0);
-        constexpr int nvfp4_block_size = 128;
+        const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
-        const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size);
+        const dim3 num_blocks(n_tokens, 1, 1);
-        const dim3 block_size(nvfp4_block_size, 1, 1);
+        if (use_aligned_float8) {
-        const dim3 num_blocks(n_tokens, block_num_y, 1);
+            quantize_mmq_nvfp4<true, true><<<num_blocks, block_size, 0, stream>>>(
-        quantize_mmq_nvfp4<true><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
-            x, ids_src1_inv, vy, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+        } else {
+            quantize_mmq_nvfp4<true, false><<<num_blocks, block_size, 0, stream>>>(
+                x, ids_src1_inv, vy, scale, ne00, /*s01=*/0, /*s02=*/stride_token, /*s03=*/0, ne0, /*ne1=*/nrows_dst, /*ne2=*/1, n_expert_used);
+        }
     } else {
         GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4);
         constexpr int nwarps = 8;
@@ -516,20 +663,24 @@ void quantize_scatter_mmq_fp4_cuda(
 }
 
 void quantize_mmq_fp4_cuda(
-        const float * x, const int32_t * ids, void * vy, const ggml_type type_src0,
+        const float * x, const int32_t * ids, void * vy, float * scale, const ggml_type type_src0, const bool use_aligned_float8,
         const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03,
         const int64_t ne0, const int64_t ne1, const int64_t ne2, const int64_t ne3, cudaStream_t stream) {
     GGML_ASSERT(type_src0 == GGML_TYPE_MXFP4 || type_src0 == GGML_TYPE_NVFP4);
     GGML_ASSERT(ne0 > 0);
 
     if (type_src0 == GGML_TYPE_NVFP4) {
+        GGML_ASSERT(scale);
         GGML_ASSERT(ne00 % QK_NVFP4 == 0);
-        constexpr int nvfp4_block_size = 128;
+        const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1);
-        const int64_t block_num_y = (ne0 + QK_NVFP4_SUB * nvfp4_block_size - 1) / (QK_NVFP4_SUB * nvfp4_block_size);
+        const dim3 num_blocks(ne1, ne2 * ne3, 1);
-        const dim3 block_size(nvfp4_block_size, 1, 1);
+        if (use_aligned_float8) {
-        const dim3 num_blocks(ne1, block_num_y, ne2 * ne3);
+            quantize_mmq_nvfp4<false, true><<<num_blocks, block_size, 0, stream>>>(
-        quantize_mmq_nvfp4<false><<<num_blocks, block_size, 0, stream>>>(
+                x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
-            x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+        } else {
+            quantize_mmq_nvfp4<false, false><<<num_blocks, block_size, 0, stream>>>(
+                x, ids, vy, scale, ne00, s01, s02, s03, ne0, ne1, ne2, /*n_expert_used=*/0);
+        }
     } else {
         GGML_ASSERT(ne0 % (2 * QK_MXFP4) == 0);
 

ggml/src/ggml-cuda/quantize.cuh

@@ -29,7 +29,9 @@ void quantize_mmq_q8_1_cuda(
 void quantize_mmq_fp4_cuda(const float *   x,
                              const int32_t * ids,
                              void *          vy,
+                             float *         scale,
                              ggml_type       type_src0,
+                             bool            use_aligned_float8,
                              int64_t         ne00,
                              int64_t         s01,
                              int64_t         s02,
@@ -44,7 +46,9 @@ void quantize_mmq_fp4_cuda(const float *   x,
 void quantize_scatter_mmq_fp4_cuda(const float *   x,
                                    const int32_t * ids_src1_inv,
                                    void *          vy,
+                                   float *         scale,
                                    ggml_type       type_src0,
+                                   bool            use_aligned_float8,
                                    int64_t         ne00,
                                    int64_t         stride_token,
                                    int64_t         ne0,

ggml/src/ggml-cuda/topk-moe.cu

@@ -8,6 +8,7 @@
 // Kernel config struct - passed by value to CUDA kernel
 struct topk_moe_config {
     bool use_sigmoid;
+    bool use_sqrt_softplus;
     bool with_norm;
     bool delayed_softmax;
 };
@@ -67,6 +68,16 @@ __device__ void sigmoid_warp_inplace(float (&vals)[experts_per_thread], const in
     }
 }
 
+template <int experts_per_thread, bool use_limit>
+__device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], const int limit, const int lane) {
+#pragma unroll
+    for (int i = 0; i < experts_per_thread; i++) {
+        const int  idx    = lane + i * WARP_SIZE;
+        const bool active = !use_limit || (idx < limit);
+        vals[i]           = active ? sqrtf(vals[i] > 20.0f ? vals[i] : logf(1.0f + expf(vals[i]))) : -INFINITY;
+    }
+}
+
 /*
     This kernel does the following:
     1. optionally softmax over the logits per token [n_experts, n_tokens]
@@ -115,6 +126,8 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float *
     if (!config.delayed_softmax) {
         if (config.use_sigmoid) {
            sigmoid_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
+        } else if (config.use_sqrt_softplus) {
+           sqrt_softplus_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
         } else {
            softmax_warp_inplace<experts_per_thread, false>(wt, n_experts, threadIdx.x);
         }
@@ -365,6 +378,7 @@ void ggml_cuda_op_topk_moe(ggml_backend_cuda_context &     ctx,
 
     topk_moe_config config;
     config.use_sigmoid       = args.sigmoid;
+    config.use_sqrt_softplus = args.sqrt_softplus;
     config.with_norm         = with_norm;
     config.delayed_softmax   = args.delayed_softmax;
 
@@ -415,7 +429,7 @@ bool ggml_cuda_should_use_topk_moe(const ggml_tensor * gating_op,
     } else if (gating_op->op == GGML_OP_UNARY) {
         ggml_unary_op op = ggml_get_unary_op(gating_op);
 
-        if (op != GGML_UNARY_OP_SIGMOID) {
+        if (op != GGML_UNARY_OP_SIGMOID && op != GGML_UNARY_OP_SOFTPLUS) {
             return false;
         }
     }

ggml/src/ggml-cuda/topk-moe.cuh

@@ -5,6 +5,7 @@
 
 struct ggml_cuda_topk_moe_args {
     bool sigmoid{};
+    bool sqrt_softplus{};
     bool softmax{};
     bool delayed_softmax{};
     bool prob_bias{};

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

@@ -1286,7 +1286,8 @@ struct ggml_hexagon_opbatch {
         int64_t nb2 = is_repack ? nb1 * ne1 : t->nb[2];
         int64_t nb3 = is_repack ? nb2 * t->ne[2] : t->nb[3];
 
-        return (h->ne[0] == ne0) && (h->ne[1] == ne1) && (h->ne[2] == t->ne[2]) && (h->ne[3] == t->ne[3]) &&
+        return (h->type == t->type) &&
+               (h->ne[0] == ne0) && (h->ne[1] == ne1) && (h->ne[2] == t->ne[2]) && (h->ne[3] == t->ne[3]) &&
                (h->nb[0] == t->nb[0]) && (h->nb[1] == nb1) && (h->nb[2] == nb2) && (h->nb[3] == nb3);
     }
 
@@ -3083,7 +3084,10 @@ static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session
         return false;
     }
 
-    if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(dst)) {
+    if (!ggml_is_contiguous_1(src0)) {
+        return false;
+    }
+    if (!ggml_is_contiguous(dst)) {
         return false;
     }
 
@@ -3094,7 +3098,7 @@ static bool ggml_hexagon_supported_activations(const struct ggml_hexagon_session
         if (!ggml_are_same_shape(src0, src1)) {
             return false;
         }
-        if (!ggml_is_contiguous(src1)) {
+        if (!ggml_is_contiguous_1(src1)) {
             return false;
         }
     }
@@ -4151,12 +4155,10 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
                 case GGML_UNARY_OP_SIGMOID:
                 case GGML_UNARY_OP_SOFTPLUS:
                 case GGML_UNARY_OP_TANH:
-                    supp = ggml_hexagon_supported_unary(sess, op);
-                    break;
                 case GGML_UNARY_OP_SILU:
                 case GGML_UNARY_OP_GELU:
                 case GGML_UNARY_OP_GELU_QUICK:
-                    supp = ggml_hexagon_supported_activations(sess, op);
+                    supp = ggml_hexagon_supported_unary(sess, op);
                     break;
                 default:
                     break;

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

@@ -723,14 +723,14 @@ static int execute_op(struct htp_ops_context * octx) {
         case HTP_OP_SQRT:
         case HTP_OP_UNARY_SOFTPLUS:
         case HTP_OP_UNARY_SIGMOID:
+        case HTP_OP_UNARY_SILU:
+        case HTP_OP_UNARY_GELU:
         case HTP_OP_UNARY_NEG:
         case HTP_OP_UNARY_EXP:
         case HTP_OP_UNARY_TANH:
         case HTP_OP_L2_NORM:
             return op_unary(octx);
 
-        case HTP_OP_UNARY_SILU:
-        case HTP_OP_UNARY_GELU:
         case HTP_OP_GLU_SWIGLU:
         case HTP_OP_GLU_SWIGLU_OAI:
         case HTP_OP_GLU_GEGLU:

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

@@ -276,6 +276,39 @@ static void sigmoid_f32(const float * restrict src,
     }
 }
 
+// silu(x) = x * sigmoid(x)
+static void silu_f32(const float * restrict src,
+                     float * restrict dst,
+                     const uint32_t num_rows,
+                     const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_sigmoid_f32_aa(dst_local, src_local, ne0);
+        hvx_mul_f32_aaa(dst_local, src_local, dst_local, ne0);
+    }
+}
+
+// gelu(x) = x * sigmoid(1.702 * x)  (quick/sigmoid approximation, matches CPU GELU_QUICK reference)
+static void gelu_f32(const float * restrict src,
+                     float * restrict dst,
+                     const uint32_t num_rows,
+                     const struct htp_unary_context * uctx) {
+    htp_unary_op_preamble;
+
+    for (uint32_t ir = 0; ir < num_rows; ir++) {
+        const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned);
+        uint8_t * restrict dst_local       = (uint8_t *)dst + (ir * dst_row_size_aligned);
+
+        hvx_mul_scalar_f32(dst_local, src_local, 1.702f, ne0);
+        hvx_sigmoid_f32_aa(dst_local, dst_local, ne0);
+        hvx_mul_f32_aaa(dst_local, src_local, dst_local, ne0);
+    }
+}
+
 static void tri_f32(const float * restrict src,
                     float * restrict dst,
                     const uint32_t num_rows,
@@ -566,6 +599,8 @@ DEFINE_UNARY_TASK(sqrt,           false, false, sqrt_f32(src0_vtcm, dst_vtcm, bl
 DEFINE_UNARY_TASK(unary_neg,      false, false, neg_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_exp,      false, false, exp_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_sigmoid,  false, false, sigmoid_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_silu,     false, false, silu_f32(src0_vtcm, dst_vtcm, block_size, uctx))
+DEFINE_UNARY_TASK(unary_gelu,     false, false, gelu_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_softplus, false, false, softplus_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(unary_tanh,     false, false, tanh_f32(src0_vtcm, dst_vtcm, block_size, uctx))
 DEFINE_UNARY_TASK(l2_norm,        false, false, l2_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx))
@@ -717,6 +752,19 @@ static inline void tile_unary_softplus_f32(uint8_t * dst_vtcm, const uint8_t * s
     }
 }
 
+// silu(x) = x * sigmoid(x)
+static inline void tile_silu_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) {
+    hvx_sigmoid_f32_aa(dst_vtcm, src_vtcm, tw);
+    hvx_mul_f32_aaa(dst_vtcm, src_vtcm, dst_vtcm, tw);
+}
+
+// gelu(x) = x * sigmoid(1.702 * x)  (quick/sigmoid approximation, matches CPU GELU_QUICK reference)
+static inline void tile_gelu_f32(uint8_t * dst_vtcm, const uint8_t * src_vtcm, uint32_t tw) {
+    hvx_mul_scalar_f32(dst_vtcm, src_vtcm, 1.702f, tw);
+    hvx_sigmoid_f32_aa(dst_vtcm, dst_vtcm, tw);
+    hvx_mul_f32_aaa(dst_vtcm, src_vtcm, dst_vtcm, tw);
+}
+
 // Triangular mask applied to one column tile. Boundary is an absolute column index, so
 // each vector compares against its absolute column position (col_start + i*VLEN_FP32).
 static inline void tri_apply_tile_f32(const uint8_t * restrict src, uint8_t * restrict dst,
@@ -798,6 +846,8 @@ DEFINE_UNARY_TILED_TASK(sqrt,           false, hvx_sqrt_f32_aa(dst_vtcm, src_vtc
 DEFINE_UNARY_TILED_TASK(unary_neg,      false, hvx_scale_f32_aa(dst_vtcm, src_vtcm, tw, -1.0f))
 DEFINE_UNARY_TILED_TASK(unary_exp,      false, hvx_exp_f32(dst_vtcm, src_vtcm, tw, false))
 DEFINE_UNARY_TILED_TASK(unary_sigmoid,  false, hvx_sigmoid_f32_aa(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(unary_silu,     false, tile_silu_f32(dst_vtcm, src_vtcm, tw))
+DEFINE_UNARY_TILED_TASK(unary_gelu,     false, tile_gelu_f32(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(unary_softplus, false, tile_unary_softplus_f32(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(unary_tanh,     false, hvx_tanh_f32_aa(dst_vtcm, src_vtcm, tw))
 DEFINE_UNARY_TILED_TASK(tri,            true,  tri_apply_tile_f32(src_vtcm, dst_vtcm, tw, col, i01, ne0, tri_ttype))
@@ -821,6 +871,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
         case HTP_OP_UNARY_NEG:       op_type = "neg-f32";          break;
         case HTP_OP_UNARY_EXP:       op_type = "exp-f32";          break;
         case HTP_OP_UNARY_SIGMOID:   op_type = "sigmoid-f32";      break;
+        case HTP_OP_UNARY_SILU:      op_type = "silu-f32";         break;
+        case HTP_OP_UNARY_GELU:      op_type = "gelu-f32";         break;
         case HTP_OP_UNARY_SOFTPLUS:  op_type = "softplus-f32";     break;
         case HTP_OP_UNARY_TANH:      op_type = "tanh-f32";         break;
         case HTP_OP_L2_NORM:         op_type = "l2norm-f32";       break;
@@ -917,6 +969,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
                 case HTP_OP_UNARY_NEG:       task_func = unary_task_f32_tiled_unary_neg;      break;
                 case HTP_OP_UNARY_EXP:       task_func = unary_task_f32_tiled_unary_exp;      break;
                 case HTP_OP_UNARY_SIGMOID:   task_func = unary_task_f32_tiled_unary_sigmoid;  break;
+                case HTP_OP_UNARY_SILU:      task_func = unary_task_f32_tiled_unary_silu;     break;
+                case HTP_OP_UNARY_GELU:      task_func = unary_task_f32_tiled_unary_gelu;     break;
                 case HTP_OP_UNARY_SOFTPLUS:  task_func = unary_task_f32_tiled_unary_softplus; break;
                 case HTP_OP_UNARY_TANH:      task_func = unary_task_f32_tiled_unary_tanh;     break;
                 case HTP_OP_TRI:             task_func = unary_task_f32_tiled_tri;            break;
@@ -934,6 +988,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) {
                 case HTP_OP_UNARY_NEG:       task_func = unary_task_f32_unary_neg;            break;
                 case HTP_OP_UNARY_EXP:       task_func = unary_task_f32_unary_exp;            break;
                 case HTP_OP_UNARY_SIGMOID:   task_func = unary_task_f32_unary_sigmoid;        break;
+                case HTP_OP_UNARY_SILU:      task_func = unary_task_f32_unary_silu;           break;
+                case HTP_OP_UNARY_GELU:      task_func = unary_task_f32_unary_gelu;           break;
                 case HTP_OP_UNARY_SOFTPLUS:  task_func = unary_task_f32_unary_softplus;       break;
                 case HTP_OP_UNARY_TANH:      task_func = unary_task_f32_unary_tanh;           break;
                 case HTP_OP_L2_NORM:         task_func = unary_task_f32_l2_norm;              break;

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

@@ -51,6 +51,8 @@ static inline bool htp_op_is_unary(uint32_t opcode) {
         case HTP_OP_UNARY_NEG:
         case HTP_OP_UNARY_EXP:
         case HTP_OP_UNARY_SIGMOID:
+        case HTP_OP_UNARY_SILU:
+        case HTP_OP_UNARY_GELU:
         case HTP_OP_UNARY_SOFTPLUS:
         case HTP_OP_UNARY_TANH:
         case HTP_OP_L2_NORM:

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

@@ -156,6 +156,24 @@ typedef struct VkPhysicalDeviceShaderFloat8FeaturesEXT {
 } VkPhysicalDeviceShaderFloat8FeaturesEXT;
 #endif
 
+#ifndef VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME
+#define VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME "VK_KHR_internally_synchronized_queues"
+#define VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_INTERNALLY_SYNCHRONIZED_QUEUES_FEATURES_KHR ((VkStructureType)1000504000)
+#define VK_DEVICE_QUEUE_CREATE_INTERNALLY_SYNCHRONIZED_BIT_KHR ((VkDeviceQueueCreateFlagBits)0x00000004)
+
+// Compile-time constant guaranteed; no runtime initialization overhead
+static constexpr vk::DeviceQueueCreateFlagBits eInternallySynchronizedKHR =
+    static_cast<vk::DeviceQueueCreateFlagBits>(0x00000004);
+
+typedef struct VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR {
+    VkStructureType    sType;
+    void* pNext;
+    VkBool32           internallySynchronizedQueues;
+} VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR;
+#else
+static constexpr vk::DeviceQueueCreateFlagBits eInternallySynchronizedKHR = vk::DeviceQueueCreateFlagBits::eInternallySynchronizedKHR;
+#endif
+
 #define ROUNDUP_POW2(M, N) (((M) + (N) - 1) & ~((N) - 1))
 #define CEIL_DIV(M, N) (((M) / (N)) + (((M) % (N)) != 0))
 static bool is_pow2(uint32_t x) { return x > 1 && (x & (x-1)) == 0; }
@@ -285,27 +303,41 @@ struct vk_command_pool {
 };
 
 // Prevent simultaneous submissions to the same queue.
-// This could be per vk_queue if we stopped having two vk_queue structures
+struct vk_queue_handle {
-// sharing the same vk::Queue.
+    vk::Queue queue;
-static std::mutex queue_mutex;
+    virtual void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) = 0;
+    virtual void lock()   {}   // no-op by default (internally synchronized case)
+    virtual void unlock() {}
+    virtual ~vk_queue_handle() = default;
+};
+
+struct vk_queue_handle_synchronized : vk_queue_handle {
+    std::mutex mutex;
+    void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) override {
+        std::lock_guard<std::mutex> guard(mutex);
+        queue.submit(submits, fence);
+    }
+    void lock()   override { mutex.lock(); }
+    void unlock() override { mutex.unlock(); }
+};
+
+struct vk_queue_handle_unsynchronized : vk_queue_handle {
+    void submit(vk::ArrayProxy<const vk::SubmitInfo> submits, vk::Fence fence) override {
+        // Driver guarantees internal synchronization via VK_KHR_internally_synchronized_queues
+        queue.submit(submits, fence);
+    }
+    // lock()/unlock() inherited no-ops
+};
 
 struct vk_queue {
     uint32_t queue_family_index;
-    vk::Queue queue;
+    std::shared_ptr<vk_queue_handle> handle;
 
     vk_command_pool cmd_pool;
 
     vk::PipelineStageFlags stage_flags;
 
     bool transfer_only;
-
-    // copy everything except the cmd_pool
-    void copyFrom(vk_queue &other) {
-        queue_family_index = other.queue_family_index;
-        queue = other.queue;
-        stage_flags = other.stage_flags;
-        transfer_only = other.transfer_only;
-    }
 };
 
 static const char * ggml_backend_vk_buffer_type_name(ggml_backend_buffer_type_t buft);
@@ -712,11 +744,12 @@ struct vk_device_struct {
     uint32_t vendor_id;
     vk::DriverId driver_id;
     vk_device_architecture architecture;
-    vk_queue compute_queue;
+    std::unique_ptr<vk_queue> compute_queue;
-    vk_queue transfer_queue;
+    std::unique_ptr<vk_queue> transfer_queue;
     bool single_queue;
     bool support_async;
     bool async_use_transfer_queue;
+    bool has_internally_synchronized_queues = false;
     uint32_t subgroup_size;
     uint32_t subgroup_size_log2;
     uint32_t shader_core_count;
@@ -1019,8 +1052,13 @@ struct vk_device_struct {
 
         ggml_vk_destroy_buffer(sync_staging);
 
-        compute_queue.cmd_pool.destroy(device);
+        if (compute_queue) compute_queue->cmd_pool.destroy(device);
-        transfer_queue.cmd_pool.destroy(device);
+        if (transfer_queue) transfer_queue->cmd_pool.destroy(device);
+
+        // Explicitly clear to ensure queues drop their shared_ptrs to handles
+        // before the Vulkan logical device instance is destroyed
+        compute_queue.reset();
+        transfer_queue.reset();
 
         for (auto& pipeline : all_pipelines) {
             if (pipeline.expired()) {
@@ -2909,8 +2947,7 @@ static vk_command_buffer* ggml_vk_create_cmd_buffer(vk_device& device, vk_comman
 static void ggml_vk_submit(vk_context& ctx, vk::Fence fence) {
     if (ctx->seqs.empty()) {
         if (fence) {
-            std::lock_guard<std::mutex> guard(queue_mutex);
+            ctx->p->q->handle->submit({}, fence);
-            ctx->p->q->queue.submit({}, fence);
         }
         return;
     }
@@ -2979,8 +3016,7 @@ static void ggml_vk_submit(vk_context& ctx, vk::Fence fence) {
         }
     }
 
-    std::lock_guard<std::mutex> guard(queue_mutex);
+    ctx->p->q->handle->submit(submit_infos, fence);
-    ctx->p->q->queue.submit(submit_infos, fence);
 
     ctx->seqs.clear();
 }
@@ -3031,18 +3067,44 @@ static uint32_t ggml_vk_find_queue_family_index(std::vector<vk::QueueFamilyPrope
     abort();
 }
 
-static void ggml_vk_create_queue(vk_device& device, vk_queue& q, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) {
+static std::unique_ptr<vk_queue> ggml_vk_create_queue(vk_device& device, uint32_t queue_family_index, uint32_t queue_index, vk::PipelineStageFlags&& stage_flags, bool transfer_only) {
     VK_LOG_DEBUG("ggml_vk_create_queue()");
     std::lock_guard<std::recursive_mutex> guard(device->mutex);
 
-    q.queue_family_index = queue_family_index;
+    auto q = std::make_unique<vk_queue>();
-    q.transfer_only = transfer_only;
+    q->queue_family_index = queue_family_index;
+    q->transfer_only = transfer_only;
 
-    q.cmd_pool.init(device, &q);
+    std::shared_ptr<vk_queue_handle> h;
+    vk::DeviceQueueInfo2 queue_info2{};
+    queue_info2.queueFamilyIndex = queue_family_index;
+    queue_info2.queueIndex = queue_index;
 
-    q.queue = device->device.getQueue(queue_family_index, queue_index);
+    if (device->has_internally_synchronized_queues) {
+        h = std::make_shared<vk_queue_handle_unsynchronized>();
+        queue_info2.flags = eInternallySynchronizedKHR;
+    } else {
+        h = std::make_shared<vk_queue_handle_synchronized>();
+    }
 
-    q.stage_flags = stage_flags;
+    h->queue = device->device.getQueue2(queue_info2);
+    q->handle = h;
+
+    q->cmd_pool.init(device, q.get());
+
+    q->stage_flags = stage_flags;
+    return q;
+}
+
+static std::unique_ptr<vk_queue> ggml_vk_create_aliased_queue(vk_device& device, const std::unique_ptr<vk_queue>& source) {
+    std::lock_guard<std::recursive_mutex> guard(device->mutex);
+    auto q = std::make_unique<vk_queue>();
+    q->handle = source->handle;
+    q->queue_family_index = source->queue_family_index;
+    q->stage_flags = source->stage_flags;
+    q->transfer_only = source->transfer_only;
+    q->cmd_pool.init(device, q.get());
+    return q;
 }
 
 static vk_context ggml_vk_create_context(ggml_backend_vk_context * ctx, vk_command_pool& p) {
@@ -3107,11 +3169,11 @@ static void ggml_vk_queue_command_pools_cleanup(vk_device& device) {
     // Arbitrary frequency to cleanup/reuse command buffers
     static constexpr uint32_t cleanup_frequency = 10;
 
-    if (device->compute_queue.cmd_pool.buffers_in_use() >= cleanup_frequency) {
+    if (device->compute_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
-        ggml_vk_command_pool_cleanup(device, device->compute_queue.cmd_pool);
+        ggml_vk_command_pool_cleanup(device, device->compute_queue->cmd_pool);
     }
-    if (device->transfer_queue.cmd_pool.buffers_in_use() >= cleanup_frequency) {
+    if (device->transfer_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) {
-        ggml_vk_command_pool_cleanup(device, device->transfer_queue.cmd_pool);
+        ggml_vk_command_pool_cleanup(device, device->transfer_queue->cmd_pool);
     }
 }
 
@@ -3221,6 +3283,7 @@ static vk_buffer ggml_vk_create_buffer(vk_device& device, size_t size, const std
             import_info.setPNext(&mem_flags_info);
             buf->device_memory = device->device.allocateMemory({ size, memory_type_idx, &import_info });
         } catch (const vk::SystemError& e) {
+            GGML_LOG_WARN("ggml_vulkan: host pointer memory import failed (%s)\n", e.what());
         }
     } else {
         for (auto it = req_flags_list.begin(); it != req_flags_list.end(); it++) {
@@ -5886,6 +5949,7 @@ static vk_device ggml_vk_get_device(size_t idx) {
         bool coopmat2_support = false;
         bool coopmat2_decode_vector_support = false;
         bool pipeline_executable_properties_support = false;
+        bool internally_sync_support = false;
         device->coopmat_support = false;
         device->integer_dot_product = false;
         device->shader_64b_indexing = false;
@@ -5957,6 +6021,8 @@ static vk_device ggml_vk_get_device(size_t idx) {
             } else if (strcmp("VK_EXT_shader_64bit_indexing", properties.extensionName) == 0) {
                 device->shader_64b_indexing = true;
 #endif
+            } else if (strcmp(VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME, properties.extensionName) == 0) {
+                internally_sync_support = true;
             }
         }
 
@@ -6143,14 +6209,6 @@ static vk_device ggml_vk_get_device(size_t idx) {
         device->single_queue = compute_queue_family_index == transfer_queue_family_index && queue_family_props[compute_queue_family_index].queueCount == 1;
 
         std::vector<vk::DeviceQueueCreateInfo> device_queue_create_infos;
-        if (compute_queue_family_index != transfer_queue_family_index) {
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities});
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), transfer_queue_family_index, 1, priorities + 1});
-        } else if(!device->single_queue) {
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 2, priorities});
-        } else {
-            device_queue_create_infos.push_back({vk::DeviceQueueCreateFlags(), compute_queue_family_index, 1, priorities});
-        }
         vk::DeviceCreateInfo device_create_info{};
         std::vector<const char *> device_extensions;
         vk::PhysicalDeviceFeatures device_features = device->physical_device.getFeatures();
@@ -6172,6 +6230,17 @@ static vk_device ggml_vk_get_device(size_t idx) {
 
         last_struct = (VkBaseOutStructure *)&vk12_features;
 
+        VkPhysicalDeviceInternallySynchronizedQueuesFeaturesKHR internally_synchronized_queues_features{};
+        internally_synchronized_queues_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_INTERNALLY_SYNCHRONIZED_QUEUES_FEATURES_KHR;
+        internally_synchronized_queues_features.pNext = nullptr;
+        internally_synchronized_queues_features.internallySynchronizedQueues = VK_FALSE;
+
+        if (internally_sync_support) {
+            last_struct->pNext = (VkBaseOutStructure *)&internally_synchronized_queues_features;
+            last_struct = (VkBaseOutStructure *)&internally_synchronized_queues_features;
+            device_extensions.push_back(VK_KHR_INTERNALLY_SYNCHRONIZED_QUEUES_EXTENSION_NAME);
+        }
+
         VkPhysicalDevicePipelineRobustnessFeaturesEXT pl_robustness_features;
         pl_robustness_features.pNext = nullptr;
         pl_robustness_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PIPELINE_ROBUSTNESS_FEATURES_EXT;
@@ -6310,6 +6379,23 @@ static vk_device ggml_vk_get_device(size_t idx) {
 
         vkGetPhysicalDeviceFeatures2(device->physical_device, &device_features2);
 
+        device->has_internally_synchronized_queues = internally_synchronized_queues_features.internallySynchronizedQueues;
+
+        // Build queue create infos only after querying whether internally synchronized queues are enabled.
+        // getQueue2() later uses the same flag, so creation/retrieval must stay consistent.
+        vk::DeviceQueueCreateFlags queue_flags = device->has_internally_synchronized_queues ?
+                                                eInternallySynchronizedKHR :
+                                                vk::DeviceQueueCreateFlags();
+
+        if (compute_queue_family_index != transfer_queue_family_index) {
+            device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 1, priorities});
+            device_queue_create_infos.push_back({queue_flags, transfer_queue_family_index, 1, priorities + 1});
+        } else if(!device->single_queue) {
+            device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 2, priorities});
+        } else {
+            device_queue_create_infos.push_back({queue_flags, compute_queue_family_index, 1, priorities});
+        }
+
         device->pipeline_executable_properties_support = pipeline_executable_properties_support;
 
         device->fp16 = device->fp16 && vk12_features.shaderFloat16;
@@ -6592,7 +6678,7 @@ static vk_device ggml_vk_get_device(size_t idx) {
         device->device = device->physical_device.createDevice(device_create_info);
 
         // Queues
-        ggml_vk_create_queue(device, device->compute_queue, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false);
+        device->compute_queue = ggml_vk_create_queue(device, compute_queue_family_index, 0, { vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer }, false);
 
         // Shaders
         // Disable matmul tile sizes early if performance low or not supported
@@ -6694,13 +6780,11 @@ static vk_device ggml_vk_get_device(size_t idx) {
 
         if (!device->single_queue) {
             const uint32_t transfer_queue_index = compute_queue_family_index == transfer_queue_family_index ? 1 : 0;
-            ggml_vk_create_queue(device, device->transfer_queue, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true);
+            device->transfer_queue = ggml_vk_create_queue(device, transfer_queue_family_index, transfer_queue_index, { vk::PipelineStageFlagBits::eTransfer }, true);
 
             device->async_use_transfer_queue = prefers_transfer_queue || (getenv("GGML_VK_ASYNC_USE_TRANSFER_QUEUE") != nullptr);
         } else {
-            // TODO: Use pointer or reference to avoid copy
+            device->transfer_queue = ggml_vk_create_aliased_queue(device, device->compute_queue);
-            device->transfer_queue.copyFrom(device->compute_queue);
-            device->transfer_queue.cmd_pool.init(device, &device->transfer_queue);
 
             device->async_use_transfer_queue = false;
         }
@@ -7263,7 +7347,7 @@ static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) {
     ctx->fence = ctx->device->device.createFence({});
     ctx->almost_ready_fence = ctx->device->device.createFence({});
 
-    ctx->compute_cmd_pool.init(ctx->device, &ctx->device->compute_queue);
+    ctx->compute_cmd_pool.init(ctx->device, ctx->device->compute_queue.get());
     if (ctx->device->async_use_transfer_queue) {
         vk::SemaphoreTypeCreateInfo tci{ vk::SemaphoreType::eTimeline, 0 };
         vk::SemaphoreCreateInfo ci{};
@@ -7271,7 +7355,7 @@ static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) {
         ctx->transfer_semaphore.s = ctx->device->device.createSemaphore(ci);
         ctx->transfer_semaphore.value = 0;
 
-        ctx->transfer_cmd_pool.init(ctx->device, &ctx->device->transfer_queue);
+        ctx->transfer_cmd_pool.init(ctx->device, ctx->device->transfer_queue.get());
     }
 
     if (vk_perf_logger_enabled) {
@@ -8045,6 +8129,11 @@ static bool ggml_vk_buffer_write_2d_async(vk_context subctx, vk_buffer& dst, siz
     ggml_vk_host_get(dst->device, src, buf, buf_offset);
 
     if (buf != nullptr) {
+        // extent of the read in pinned source memory; guard against tensors that
+        // straddle a pinned-chunk boundary (they fall back to staging below)
+        size_t src_extent = (width == spitch) ? (size_t) width * height
+                                              : (height > 0 ? (height - 1) * spitch + width : 0);
+        if (buf_offset + src_extent <= buf->size) {
             // Memory is pinned, use as staging buffer
             std::vector<vk::BufferCopy> slices(1);
             if (width == spitch && width == dpitch) {
@@ -8065,6 +8154,8 @@ static bool ggml_vk_buffer_write_2d_async(vk_context subctx, vk_buffer& dst, siz
             subctx->s->buffer->buf.copyBuffer(buf->buffer, dst->buffer, slices);
             return true;
         }
+        // straddles a chunk boundary: fall through to staging
+    }
     VK_LOG_DEBUG("STAGING");
 
     if (!sync_staging) {
@@ -8126,7 +8217,7 @@ static void ggml_vk_buffer_write_2d(vk_buffer& dst, size_t offset, const void *
     } else {
         std::lock_guard<std::recursive_mutex> guard(dst->device->mutex);
 
-        vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue->cmd_pool);
         ggml_vk_ctx_begin(dst->device, subctx);
         bool ret = ggml_vk_buffer_write_2d_async(subctx, dst, offset, src, spitch, dpitch, width, height, true);
         GGML_ASSERT(ret);
@@ -8241,7 +8332,7 @@ static void ggml_vk_buffer_read_2d(vk_buffer& src, size_t offset, void * dst, si
         GGML_ASSERT(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent);
 
         std::lock_guard<std::recursive_mutex> guard(src->device->mutex);
-        vk_context subctx = ggml_vk_create_temporary_context(src->device->compute_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(src->device->compute_queue->cmd_pool);
         ggml_vk_ctx_begin(src->device, subctx);
         subctx->s->buffer->buf.pipelineBarrier(
             vk::PipelineStageFlagBits::eComputeShader | vk::PipelineStageFlagBits::eTransfer,
@@ -8267,7 +8358,7 @@ static void ggml_vk_buffer_read_2d(vk_buffer& src, size_t offset, void * dst, si
     } else {
         std::lock_guard<std::recursive_mutex> guard(src->device->mutex);
 
-        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue->cmd_pool);
         ggml_vk_ctx_begin(src->device, subctx);
         bool ret = ggml_vk_buffer_read_2d_async(subctx, src, offset, dst, spitch, dpitch, width, height, true);
         GGML_ASSERT(ret);
@@ -8304,7 +8395,7 @@ static void ggml_vk_buffer_copy(vk_buffer& dst, size_t dst_offset, vk_buffer& sr
         std::lock_guard<std::recursive_mutex> guard(src->device->mutex);
         VK_LOG_DEBUG("ggml_vk_buffer_copy(SINGLE_DEVICE, " << size << ")");
         // Copy within the device
-        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue.cmd_pool);
+        vk_context subctx = ggml_vk_create_temporary_context(src->device->transfer_queue->cmd_pool);
         ggml_vk_ctx_begin(src->device, subctx);
         ggml_vk_buffer_copy_async(subctx, dst, dst_offset, src, src_offset, size);
         ggml_vk_ctx_end(subctx);
@@ -8347,7 +8438,7 @@ static void ggml_vk_buffer_memset(vk_buffer& dst, size_t offset, uint32_t c, siz
     }
 
     std::lock_guard<std::recursive_mutex> guard(dst->device->mutex);
-    vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue.cmd_pool);
+    vk_context subctx = ggml_vk_create_temporary_context(dst->device->transfer_queue->cmd_pool);
     ggml_vk_ctx_begin(dst->device, subctx);
     subctx->s->buffer->buf.fillBuffer(dst->buffer, offset, size, c);
     ggml_vk_ctx_end(subctx);
@@ -10501,8 +10592,10 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx
     }
 
     // Only use mask opt when the mask is fairly large. This hasn't been tuned extensively.
+    // GCN with large head size (>= 256) benefits in high-context prefill (skipping the
+    // per-block mask add on fully-visible blocks dominates), so enable it there too.
     bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16
-                        && (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK > 256 || HSV > 256);
+                        && (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK >= 256 || HSV >= 256);
     vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc,
                                                                    mask != nullptr, use_mask_opt, logit_softcap != 0, k->type, v->type);
 
@@ -12429,9 +12522,108 @@ static void ggml_vk_opt_step_sgd(ggml_backend_vk_context * ctx, vk_context& subc
     ggml_vk_op_f32<vk_op_push_constants>(ctx, subctx, src0, src1, src2, nullptr, dst, GGML_OP_OPT_STEP_SGD, { (uint32_t)n, 0, 0.0f, 0.0f, 0.0f, 0.0f });
 }
 
+// Fast path for concat along dim 0 where one source is stored "transposed"
+// (nb[1] == type_size, dim1 innermost) and the other source + dst are
+// contiguous along dim 0. The generic concat shader reads the transposed
+// source with a catastrophic uncoalesced stride; here we instead copy the
+// contiguous source with copy.comp and transpose the other source into the
+// matching dst sub-region with the tiled copy_transpose shader (shared-memory
+// transpose, coalesced read+write). Mirrors the precedent in
+// ggml_vk_cpy_to_contiguous: direct dispatch with custom push constants.
+static void ggml_vk_concat_transpose_fastpath(ggml_backend_vk_context * ctx, vk_context& subctx,
+                                              const ggml_tensor * ctg, const ggml_tensor * trp,
+                                              ggml_tensor * dst, uint32_t off_ctg, uint32_t off_trp) {
+    const uint32_t ts = ggml_type_size(dst->type);
+
+    vk_pipeline pipeline_cpy = (ts == 4) ? ctx->device->pipeline_cpy_f32_f32
+                                          : ctx->device->pipeline_cpy_f16_f16;
+    vk_pipeline pipeline_trp = (ts == 4) ? ctx->device->pipeline_cpy_transpose_32
+                                          : ctx->device->pipeline_cpy_transpose_16;
+
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline_cpy, 1);
+    ggml_pipeline_request_descriptor_sets(ctx, pipeline_trp, 1);
+
+    vk_subbuffer ctg_buf = ggml_vk_tensor_subbuffer(ctx, ctg, true);
+    vk_subbuffer trp_buf = ggml_vk_tensor_subbuffer(ctx, trp, true);
+    vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true);
+
+    const uint32_t a_misalign_ctg = get_misalign_bytes(ctx, ctg) / ts;
+    const uint32_t a_misalign_trp = get_misalign_bytes(ctx, trp) / ts;
+    const uint32_t d_misalign     = get_misalign_bytes(ctx, dst) / ts;
+
+    // Dispatch A: contiguous copy of `ctg` into dst[off_ctg : off_ctg + ctg->ne[0], :]
+    if (ctg->ne[0] > 0) {
+        const uint32_t ne_ctg = (uint32_t) ggml_nelements(ctg);
+        vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(ctg, dst, ne_ctg);
+        pc.ne10 = (uint32_t) ctg->ne[0];                 // only ctg's columns in dst
+        pc.misalign_offsets = (a_misalign_ctg << 16) | (d_misalign + off_ctg);
+        init_pushconst_fastdiv(pc);
+        std::array<uint32_t, 3> el = ne_ctg > 262144 ? std::array<uint32_t,3>{512, 512, CEIL_DIV(ne_ctg, 262144)}
+                                      : ne_ctg > 512  ? std::array<uint32_t,3>{512, CEIL_DIV(ne_ctg, 512), 1}
+                                                      : std::array<uint32_t,3>{ne_ctg, 1, 1};
+        ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_cpy, { ctg_buf, dst_buf }, pc, el);
+    }
+
+    // Dispatch B: tiled transpose of `trp` into dst[off_trp : off_trp + trp->ne[0], :]
+    if (trp->ne[0] > 0) {
+        vk_op_unary_push_constants pc = vk_op_unary_push_constants_init(trp, dst, ggml_nelements(trp));
+        pc.ne10 = (uint32_t) trp->ne[0];                 // dst bound = trp columns (NOT dst->ne[0])
+        pc.misalign_offsets = (a_misalign_trp << 16) | (d_misalign + off_trp);
+        init_pushconst_fastdiv(pc);
+        std::array<uint32_t, 3> el = {
+            (uint32_t) CEIL_DIV(trp->ne[0], 32),
+            (uint32_t) CEIL_DIV(trp->ne[1], 32),
+            (uint32_t) (trp->ne[2] * trp->ne[3]),
+        };
+        el[0] = std::min(el[0], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[0]);
+        el[1] = std::min(el[1], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[1]);
+        el[2] = std::min(el[2], (uint32_t) ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
+        ggml_vk_dispatch_pipeline(ctx, subctx, pipeline_trp, { trp_buf, dst_buf }, pc, el);
+    }
+
+    ggml_vk_sync_buffers(ctx, subctx);
+}
+
 static void ggml_vk_concat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
     int * op_params = (int *)dst->op_params;
 
+    // Fast path: concat along dim 0 with one source "transposed" (nb[1]==type_size,
+    // dim1 innermost) and the other source + dst contiguous along dim 0. The generic
+    // shader reads the transposed source uncoalesced (~5.6ms on RX 580 vs ~130us for
+    // a coalesced copy); here we instead use a tiled shared-memory transpose.
+    auto src_transposed_2d = [&](const ggml_tensor * s) {
+        const uint32_t ts = ggml_type_size(s->type);
+        return s->nb[1] == ts                                              // dim1 innermost
+            && s->nb[0] == s->ne[1] * ts                                   // consistent 2D transpose
+            && s->ne[2] == 1 && s->ne[3] == 1;
+    };
+    const uint32_t dst_ts = ggml_type_size(dst->type);
+    const bool dim0_ok   = (op_params[0] == 0) && (dst->nb[0] == dst_ts);
+    const bool types_ok  = (dst_ts == 4 || dst_ts == 2)
+                        && (src0->type == src1->type && src0->type == dst->type)
+                        && (ggml_blck_size(dst->type) == 1);
+    const bool shapes_ok = (src0->ne[1] == src1->ne[1] && src1->ne[1] == dst->ne[1])
+                        && (src0->ne[2] == src1->ne[2] && src0->ne[2] == dst->ne[2])
+                        && (src0->ne[3] == src1->ne[3] && src0->ne[3] == dst->ne[3])
+                        && (src0->ne[0] + src1->ne[0] == dst->ne[0]);
+    // doffset is 16-bit for unary shaders; guard against overflow
+    const bool off_fits  = ((get_misalign_bytes(ctx, dst)/dst_ts + dst->ne[0])) < 0xFFFFu;
+
+    const bool s1_trans = dim0_ok && types_ok && shapes_ok && off_fits
+                        && src_transposed_2d(src1) && ggml_is_contiguous(src0);
+    const bool s0_trans = dim0_ok && types_ok && shapes_ok && off_fits
+                        && src_transposed_2d(src0) && ggml_is_contiguous(src1);
+
+    if (s1_trans || s0_trans) {
+        const ggml_tensor * ctg = s1_trans ? src0 : src1;     // contiguous source
+        const ggml_tensor * trp = s1_trans ? src1 : src0;     // transposed source
+        // dst offsets (in elements) where each source region begins
+        const uint32_t off_ctg = s1_trans ? 0u                       : (uint32_t) src1->ne[0];
+        const uint32_t off_trp = s1_trans ? (uint32_t) src0->ne[0]   : 0u;
+        ggml_vk_concat_transpose_fastpath(ctx, subctx, ctg, trp, dst, off_ctg, off_trp);
+        return;
+    }
+
     const uint32_t src0_type_size = ggml_type_size(src0->type);
     const uint32_t src1_type_size = ggml_type_size(src1->type);
     const uint32_t dst_type_size = ggml_type_size(dst->type);
@@ -15651,6 +15843,35 @@ static const char * ggml_backend_vk_name(ggml_backend_t backend) {
     return ctx->name.c_str();
 }
 
+// Free the compute-scratch device buffers (prealloc_* and the transfer staging buffer) without
+// tearing down the backend (pipelines, command pools, fences and device stay alive). These buffers
+// are reallocated lazily by ggml_vk_preallocate_buffers() on the next compute, so this just shrinks
+// an idle model's device footprint. Mirrors the buffer-freeing subset of ggml_vk_cleanup().
+static void ggml_backend_vk_free_scratch(ggml_backend_t backend) {
+    ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context;
+    VK_LOG_DEBUG("ggml_backend_vk_free_scratch(" << ctx->name << ")");
+
+    // discard any unsubmitted command buffer and wait for in-flight work before freeing
+    ctx->compute_ctx.reset();
+    ggml_vk_synchronize(ctx);
+
+    ggml_vk_destroy_buffer(ctx->prealloc_x);
+    ggml_vk_destroy_buffer(ctx->prealloc_y);
+    ggml_vk_destroy_buffer(ctx->prealloc_split_k);
+    ggml_vk_destroy_buffer(ctx->prealloc_add_rms_partials);
+    ggml_vk_destroy_buffer(ctx->sync_staging);
+
+    ctx->prealloc_y_last_pipeline_used = nullptr;
+    ctx->prealloc_y_last_tensor_used = nullptr;
+    ctx->prealloc_y_last_decode_vector_staging = false;
+
+    ctx->prealloc_size_x = 0;
+    ctx->prealloc_size_y = 0;
+    ctx->prealloc_size_split_k = 0;
+    ctx->prealloc_size_add_rms_partials = 0;
+    ctx->prealloc_size_add_rms_partials_offset = 0;
+}
+
 static void ggml_backend_vk_free(ggml_backend_t backend) {
     ggml_backend_vk_context * ctx = (ggml_backend_vk_context *)backend->context;
     VK_LOG_DEBUG("ggml_backend_vk_free(" << ctx->name << ")");
@@ -15875,19 +16096,17 @@ static void ggml_vk_synchronize(ggml_backend_vk_context * ctx) {
                 1, &ctx->transfer_semaphore.value,
                 0, nullptr,
             };
-            vk::PipelineStageFlags stage = ctx->device->transfer_queue.stage_flags;
+            vk::PipelineStageFlags stage = ctx->device->transfer_queue->stage_flags;
             vk::SubmitInfo si{
                 1, &ctx->transfer_semaphore.s, &stage,
                 0, nullptr,
                 0, nullptr,
             };
             si.setPNext(&tl_info);
-            std::lock_guard<std::mutex> guard(queue_mutex);
+            ctx->device->compute_queue->handle->submit({ si }, ctx->fence);
-            ctx->device->compute_queue.queue.submit({ si }, ctx->fence);
             ctx->transfer_semaphore_last_submitted = ctx->transfer_semaphore.value;
         } else {
-            std::lock_guard<std::mutex> guard(queue_mutex);
+            ctx->device->compute_queue->handle->submit({}, ctx->fence);
-            ctx->device->compute_queue.queue.submit({}, ctx->fence);
         }
         ggml_vk_wait_for_fence(ctx);
         ctx->submit_pending = false;
@@ -16449,7 +16668,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
         vk::DebugUtilsLabelEXT dul = {};
         dul.pLabelName = "ggml_backend_vk_graph_compute";
         dul.color = std::array<float,4>{1.0f, 1.0f, 1.0f, 1.0f};
-        vk_instance.pfn_vkQueueBeginDebugUtilsLabelEXT(ctx->device->compute_queue.queue, reinterpret_cast<VkDebugUtilsLabelEXT*>(&dul));
+
+        std::lock_guard<vk_queue_handle> guard(*ctx->device->compute_queue->handle);
+        vk_instance.pfn_vkQueueBeginDebugUtilsLabelEXT(ctx->device->compute_queue->handle->queue, reinterpret_cast<VkDebugUtilsLabelEXT*>(&dul));
     }
 
     ctx->prealloc_size_add_rms_partials_offset = 0;
@@ -16495,7 +16716,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg
         std::fill(ctx->query_nodes.begin(), ctx->query_nodes.end(), nullptr);
         std::fill(ctx->query_node_idx.begin(), ctx->query_node_idx.end(), 0);
 
-        GGML_ASSERT(ctx->compute_ctx.expired());
+        // compute_ctx may hold async host uploads recorded before the graph; append the timestamp after them
         compute_ctx = ggml_vk_get_compute_ctx(ctx);
         ctx->query_idx = 0;
         compute_ctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++);
@@ -17186,6 +17407,7 @@ static ggml_backend_i ggml_backend_vk_interface = {
     /* .event_record            = */ ggml_backend_vk_event_record,
     /* .event_wait              = */ ggml_backend_vk_event_wait,
     /* .graph_optimize          = */ ggml_vk_graph_optimize,
+    /* .free_scratch            = */ ggml_backend_vk_free_scratch,
 };
 
 static ggml_guid_t ggml_backend_vk_guid() {
@@ -18160,11 +18382,112 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg,
     return devices[device];
 }
 
+// Import an mmap-backed host region as a Vulkan pinned buffer via
+// VK_EXT_external_memory_host so H2D uploads DMA straight from system RAM
+// instead of bouncing through the staging buffer + host memcpy. Mirrors the
+// GGML_CUDA_REGISTER_HOST path; populates device->pinned_memory, which the
+// existing pinned fast path in ggml_vk_buffer_write_2d_async looks up.
+//
+// A single Vulkan buffer cannot cover a whole multi-GB mmap (it is capped at
+// device->max_buffer_size), so the region is imported in page-aligned chunks.
+// ggml_vk_host_get resolves a tensor pointer to the chunk that contains it,
+// and ggml_vk_buffer_write_2d_async falls back to staging for any tensor that
+// straddles a chunk boundary, so correctness is preserved.
+static bool ggml_backend_vk_register_host_buffer(void * buffer, size_t size) {
+    if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr && getenv("GGML_VK_REGISTER_HOST") == nullptr) {
+        return false;
+    }
+    if (size == 0) {
+        return false;
+    }
+
+    bool success = false;
+    for (size_t i = 0; i < GGML_VK_MAX_DEVICES; i++) {
+        vk_device& device = vk_instance.devices[i];
+        if (!device || !device->external_memory_host || device->max_buffer_size == 0) {
+            continue;
+        }
+
+        const size_t align = device->min_imported_host_pointer_alignment;
+        size_t chunk = device->max_buffer_size & ~(align - 1);
+        if (chunk == 0) {
+            continue;
+        }
+
+        uint8_t * p = static_cast<uint8_t *>(buffer);
+        size_t remaining = size;
+        bool dev_success = false;
+        bool import_ok = true;   // flips to false once VK_EXT_external_memory_host import fails
+        while (remaining > 0) {
+            size_t cur = std::min(remaining, chunk);
+            vk_buffer buf;
+            if (import_ok) {
+                buf = ggml_vk_buffer_from_host_ptr(device, p, cur);
+            }
+            if (!buf || !buf->buffer) {
+                // Fallback for drivers that can't import file-backed mmap pages (e.g. RADV):
+                // make a one-time copy of the region into a host-visible Vulkan buffer. The GPU
+                // then DMAs straight from it every eval at full PCIe bandwidth, instead of paying
+                // a slow single-threaded pageable memcpy into the staging buffer each time.
+                import_ok = false;
+                buf = ggml_vk_create_buffer_check(device, cur,
+                    vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached,
+                    vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent);
+                if (buf && buf->buffer && buf->ptr) {
+                    memcpy(buf->ptr, p, cur);
+                } else {
+                    break;
+                }
+            }
+            {
+                std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
+                device->pinned_memory.emplace_back(p, cur, buf);
+            }
+            dev_success = true;
+            p += cur;
+            remaining -= cur;
+        }
+        if (dev_success) {
+            success = true;
+        }
+    }
+    return success;
+}
+
+static void ggml_backend_vk_unregister_host_buffer(void * buffer) {
+    for (size_t i = 0; i < GGML_VK_MAX_DEVICES; i++) {
+        vk_device& device = vk_instance.devices[i];
+        if (!device) {
+            continue;
+        }
+        std::lock_guard<std::shared_mutex> guard(device->pinned_memory_mutex);
+        for (auto it = device->pinned_memory.begin(); it != device->pinned_memory.end(); ++it) {
+            if (std::get<0>(*it) == buffer) {
+                vk_buffer buf = std::get<2>(*it);
+                device->pinned_memory.erase(it);
+                ggml_vk_destroy_buffer(buf);
+                break;
+            }
+        }
+    }
+}
+
+static void * ggml_backend_vk_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) {
+    GGML_UNUSED(reg);
+    if (strcmp(name, "ggml_backend_register_host_buffer") == 0) {
+        return (void *) ggml_backend_vk_register_host_buffer;
+    }
+    if (strcmp(name, "ggml_backend_unregister_host_buffer") == 0) {
+        return (void *) ggml_backend_vk_unregister_host_buffer;
+    }
+    return nullptr;
+}
+
 static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = {
     /* .get_name         = */ ggml_backend_vk_reg_get_name,
     /* .get_device_count = */ ggml_backend_vk_reg_get_device_count,
     /* .get_device       = */ ggml_backend_vk_reg_get_device,
-    /* .get_proc_address = */ NULL,
+    /* .get_proc_address = */ ggml_backend_vk_reg_get_proc_address,
 };
 
 ggml_backend_reg_t ggml_backend_vk_reg() {

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

@@ -355,6 +355,30 @@ struct ggml_webgpu_conv2d_pipeline_key_hash {
     }
 };
 
+// Same type fields as conv2d plus the input layout (WHCN vs CWHN).
+struct ggml_webgpu_conv2d_dw_pipeline_key {
+    ggml_type weight_type;
+    ggml_type input_type;
+    ggml_type output_type;
+    bool      whcn;
+
+    bool operator==(const ggml_webgpu_conv2d_dw_pipeline_key & other) const {
+        return weight_type == other.weight_type && input_type == other.input_type && output_type == other.output_type &&
+               whcn == other.whcn;
+    }
+};
+
+struct ggml_webgpu_conv2d_dw_pipeline_key_hash {
+    size_t operator()(const ggml_webgpu_conv2d_dw_pipeline_key & key) const {
+        size_t seed = 0;
+        ggml_webgpu_hash_combine(seed, key.weight_type);
+        ggml_webgpu_hash_combine(seed, key.input_type);
+        ggml_webgpu_hash_combine(seed, key.output_type);
+        ggml_webgpu_hash_combine(seed, key.whcn);
+        return seed;
+    }
+};
+
 /** Im2Col **/
 struct ggml_webgpu_im2col_pipeline_key {
     ggml_type input_type;
@@ -1210,6 +1234,8 @@ class ggml_webgpu_shader_lib {
         soft_max_pipelines;
     std::unordered_map<ggml_webgpu_conv2d_pipeline_key, webgpu_pipeline, ggml_webgpu_conv2d_pipeline_key_hash>
         conv2d_pipelines;
+    std::unordered_map<ggml_webgpu_conv2d_dw_pipeline_key, webgpu_pipeline, ggml_webgpu_conv2d_dw_pipeline_key_hash>
+        conv2d_dw_pipelines;
     std::unordered_map<ggml_webgpu_im2col_pipeline_key, webgpu_pipeline, ggml_webgpu_im2col_pipeline_key_hash>
         im2col_pipelines;
 
@@ -3172,6 +3198,50 @@ class ggml_webgpu_shader_lib {
         return conv2d_pipelines[key];
     }
 
+    // whcn selects the input layout: contiguous WHCN vs contiguous-channels CWHN
+    webgpu_pipeline get_conv2d_dw_pipeline(const ggml_webgpu_shader_lib_context & context, bool whcn) {
+        ggml_webgpu_conv2d_dw_pipeline_key key = {};
+        key.weight_type                        = context.src0->type;
+        key.input_type                         = context.src1->type;
+        key.output_type                        = context.dst->type;
+        key.whcn                               = whcn;
+
+        auto it = conv2d_dw_pipelines.find(key);
+        if (it != conv2d_dw_pipelines.end()) {
+            return it->second;
+        }
+
+        std::vector<std::string> defines;
+        std::string              variant = whcn ? "conv_2d_dw_whcn" : "conv_2d_dw_cwhn";
+
+        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");
+            } else if (type == GGML_TYPE_F16) {
+                defines.push_back(s_prefix + "_F16");
+            } else {
+                GGML_ABORT("Unsupported type for CONV_2D_DW shader");
+            }
+        };
+
+        push_type_defines("WEIGHT", key.weight_type);
+        push_type_defines("INPUT", key.input_type);
+        push_type_defines("OUTPUT", key.output_type);
+        if (whcn) {
+            defines.push_back("WHCN");
+        }
+        defines.push_back(std::string("WG_SIZE=") + std::to_string(context.max_wg_size));
+
+        auto processed           = preprocessor.preprocess(wgsl_conv2d_dw, defines);
+        auto decisions           = std::make_shared<ggml_webgpu_generic_shader_decisions>();
+        decisions->wg_size       = context.max_wg_size;
+        webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
+        pipeline.context         = decisions;
+        conv2d_dw_pipelines[key] = pipeline;
+        return conv2d_dw_pipelines[key];
+    }
+
     webgpu_pipeline get_im2col_pipeline(const ggml_webgpu_shader_lib_context & context) {
         ggml_webgpu_im2col_pipeline_key key = {};
         key.input_type                      = context.src1->type;

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

@@ -978,6 +978,67 @@ static webgpu_encoded_op ggml_webgpu_conv_2d(webgpu_context & ctx,
     return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
 }
 
+// Same param/binding layout as conv_2d; the shader differs
+static webgpu_encoded_op ggml_webgpu_conv_2d_dw(webgpu_context & ctx,
+                                                ggml_tensor *    src0,
+                                                ggml_tensor *    src1,
+                                                ggml_tensor *    dst) {
+    const int32_t s0 = ggml_get_op_params_i32(dst, 0);
+    const int32_t s1 = ggml_get_op_params_i32(dst, 1);
+    const int32_t p0 = ggml_get_op_params_i32(dst, 2);
+    const int32_t p1 = ggml_get_op_params_i32(dst, 3);
+    const int32_t d0 = ggml_get_op_params_i32(dst, 4);
+    const int32_t d1 = ggml_get_op_params_i32(dst, 5);
+
+    // Scalar params matching conv2d_dw.wgsl (weight src0 [KW,KH,1,C], input src1, output dst).
+    std::vector<uint32_t> params = {
+        (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)),
+        (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)),
+        (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
+
+        (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],
+        (uint32_t) src1->ne[1],
+        (uint32_t) src0->ne[0],
+        (uint32_t) src0->ne[1],
+
+        (uint32_t) s0,
+        (uint32_t) s1,
+        (uint32_t) p0,
+        (uint32_t) p1,
+        (uint32_t) d0,
+        (uint32_t) d1,
+    };
+
+    std::vector<wgpu::BindGroupEntry> entries = {
+        ggml_webgpu_make_tensor_bind_group_entry(ctx, 0, src0),
+        ggml_webgpu_make_tensor_bind_group_entry(ctx, 1, src1),
+        ggml_webgpu_make_tensor_bind_group_entry(ctx, 2, dst),
+    };
+
+    ggml_webgpu_shader_lib_context shader_lib_ctx = {};
+    shader_lib_ctx.src0                           = src0;
+    shader_lib_ctx.src1                           = src1;
+    shader_lib_ctx.dst                            = dst;
+    shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
+
+    // Input layout: contiguous -> WHCN, contiguous-channels -> CWHN
+    const bool      whcn      = ggml_is_contiguous(src1);
+    webgpu_pipeline pipeline  = ctx->shader_lib->get_conv2d_dw_pipeline(shader_lib_ctx, whcn);
+    auto *          decisions = static_cast<ggml_webgpu_generic_shader_decisions *>(pipeline.context.get());
+
+    uint32_t wg_x;
+    uint32_t wg_y;
+    uint32_t total_wg = CEIL_DIV((uint32_t) ggml_nelements(dst), decisions->wg_size);
+    compute_2d_workgroups(total_wg, ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension, wg_x, wg_y);
+
+    return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
+}
+
 static webgpu_encoded_op ggml_webgpu_im2col(webgpu_context & ctx,
                                             ggml_tensor *    src0,
                                             ggml_tensor *    src1,
@@ -3164,6 +3225,8 @@ static std::optional<webgpu_encoded_op> ggml_webgpu_encode(webgpu_context ctx,
             return ggml_webgpu_sum_rows(ctx, src0, node);
         case GGML_OP_CONV_2D:
             return ggml_webgpu_conv_2d(ctx, src0, src1, node);
+        case GGML_OP_CONV_2D_DW:
+            return ggml_webgpu_conv_2d_dw(ctx, src0, src1, node);
         case GGML_OP_IM2COL:
             return ggml_webgpu_im2col(ctx, src0, src1, node);
         case GGML_OP_UPSCALE:
@@ -4349,6 +4412,12 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
                           (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16) &&
                           (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16);
             break;
+        case GGML_OP_CONV_2D_DW:
+            supports_op = (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) &&
+                          (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16) &&
+                          (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16) &&
+                          (ggml_is_contiguous(src1) || ggml_is_contiguous_channels(src1));
+            break;
         case GGML_OP_IM2COL:
             supports_op = (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) &&
                           (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);

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

@@ -0,0 +1,137 @@
+#include "common_decls.tmpl"
+enable f16;
+
+// Ported from the Vulkan backend's conv2d_dw.comp. Two variants (based on WHCN)
+// selected by the input (src1) layout: contiguous -> WHCN, else CWHN.
+// 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
+
+@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
+
+@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
+
+struct Params {
+    offset_w: u32,
+    offset_i: u32,
+    offset_o: u32,
+
+    ne: u32,
+    channels: u32,
+    batches: u32,
+    dst_w: u32, dst_h: u32,
+    src_w: u32, src_h: u32,
+    knl_w: u32, knl_h: u32,
+
+    stride_x: i32, stride_y: i32,
+    pad_x: i32, pad_y: i32,
+    dilation_x: i32, dilation_y: i32,
+};
+
+@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 {
+    let i0    = idx / params.dst_w;
+    let dst_x = idx - i0 * params.dst_w;
+    let i1    = i0 / params.dst_h;
+    let dst_y = i0 - i1 * params.dst_h;
+    let n     = i1 / params.channels;
+    let c     = i1 - n * params.channels;
+
+    let src_i = params.offset_i + n * params.channels * params.src_h * params.src_w
+                                + c * params.src_h * params.src_w;
+    let knl_i = params.offset_w + c * params.knl_h * params.knl_w;
+
+    var sum: f32 = 0.0;
+    for (var ky: u32 = 0u; ky < params.knl_h; ky += 1u) {
+        let src_y = i32(dst_y) * params.stride_y + i32(ky) * params.dilation_y - params.pad_y;
+        if (src_y < 0 || src_y >= i32(params.src_h)) { continue; }
+        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);
+            sum += v * k;
+        }
+    }
+    return sum;
+}
+#else
+// Channels contiguous (CWHN): channel is the innermost axis.
+fn conv_2d_dw(idx: u32) -> f32 {
+    let i0    = idx / params.channels;
+    let c     = idx - i0 * params.channels;
+    let i1    = i0 / params.dst_w;
+    let dst_x = i0 - i1 * params.dst_w;
+    let n     = i1 / params.dst_h;
+    let dst_y = i1 - n * params.dst_h;
+
+    let src_i   = params.offset_i + n * params.channels * params.src_h * params.src_w;
+    let src_row = params.src_w * params.channels;
+    let knl_row = params.knl_w * params.channels;
+
+    var sum: f32 = 0.0;
+    for (var ky: u32 = 0u; ky < params.knl_h; ky += 1u) {
+        let src_y = i32(dst_y) * params.stride_y + i32(ky) * params.dilation_y - params.pad_y;
+        if (src_y < 0 || src_y >= i32(params.src_h)) { continue; }
+        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);
+            sum += v * k;
+        }
+    }
+    return sum;
+}
+#endif
+
+@compute @workgroup_size(WG_SIZE)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_wg: vec3<u32>
+) {
+    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));
+}

gguf-py/gguf/constants.py

@@ -507,6 +507,7 @@ class MODEL_ARCH(IntEnum):
     DOTS1            = auto()
     ARCEE            = auto()
     AFMOE            = auto()
+    LAGUNA           = auto()
     ERNIE4_5         = auto()
     ERNIE4_5_MOE     = auto()
     HUNYUAN_MOE      = auto()
@@ -1088,6 +1089,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
     MODEL_ARCH.DOTS1:            "dots1",
     MODEL_ARCH.ARCEE:            "arcee",
     MODEL_ARCH.AFMOE:            "afmoe",
+    MODEL_ARCH.LAGUNA:           "laguna",
     MODEL_ARCH.ERNIE4_5:         "ernie4_5",
     MODEL_ARCH.ERNIE4_5_MOE:     "ernie4_5-moe",
     MODEL_ARCH.FALCON_H1:        "falcon-h1",
@@ -3823,6 +3825,31 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_POST_NORM,
         MODEL_TENSOR.FFN_EXP_PROBS_B,
     ],
+    MODEL_ARCH.LAGUNA: [
+        MODEL_TENSOR.TOKEN_EMBD,
+        MODEL_TENSOR.OUTPUT_NORM,
+        MODEL_TENSOR.OUTPUT,
+        MODEL_TENSOR.ATTN_NORM,
+        MODEL_TENSOR.ATTN_Q,
+        MODEL_TENSOR.ATTN_K,
+        MODEL_TENSOR.ATTN_V,
+        MODEL_TENSOR.ATTN_OUT,
+        MODEL_TENSOR.ATTN_Q_NORM,
+        MODEL_TENSOR.ATTN_K_NORM,
+        MODEL_TENSOR.ATTN_GATE,
+        MODEL_TENSOR.FFN_NORM,
+        MODEL_TENSOR.FFN_GATE,
+        MODEL_TENSOR.FFN_DOWN,
+        MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.FFN_GATE_INP,
+        MODEL_TENSOR.FFN_EXP_PROBS_B,
+        MODEL_TENSOR.FFN_GATE_EXP,
+        MODEL_TENSOR.FFN_DOWN_EXP,
+        MODEL_TENSOR.FFN_UP_EXP,
+        MODEL_TENSOR.FFN_GATE_SHEXP,
+        MODEL_TENSOR.FFN_UP_SHEXP,
+        MODEL_TENSOR.FFN_DOWN_SHEXP,
+    ],
     MODEL_ARCH.ERNIE4_5: [
         MODEL_TENSOR.TOKEN_EMBD,
         MODEL_TENSOR.OUTPUT_NORM,

gguf-py/gguf/tensor_mapping.py

@@ -479,6 +479,7 @@ class TensorNameMap:
             "model.layers.{bid}.mlp.e_score_correction",                    # exaone-moe
             "model.layers.{bid}.block_sparse_moe.gate.e_score_correction",  # kimi
             "model.layers.{bid}.moe.router_bias",                           # step3.5 expert selection bias
+            "model.layers.{bid}.mlp.experts.e_score_correction",       # laguna
         ),
 
         # Feed-forward up

include/llama.h

@@ -557,6 +557,16 @@ extern "C" {
 
     LLAMA_API const struct llama_model * llama_get_model   (const struct llama_context * ctx);
     LLAMA_API           llama_memory_t   llama_get_memory  (const struct llama_context * ctx);
+
+    // On-demand device (VRAM) residency: free the model's GPU weight buffers (keeping a host
+    // shadow so no reload is needed) to hand VRAM to another model, and rebuild them on demand.
+    // Any subsequent llama_decode auto-restores. If evict_kv is true, the KV cache device buffers
+    // are also evicted to a host shadow (for when the KV would not fit alongside the other model),
+    // at the cost of a D2H/H2D copy of the live KV each cycle; otherwise the KV stays resident.
+    // Intended for time-sharing a single GPU between several always-loaded models.
+    LLAMA_API void llama_context_release_device(struct llama_context * ctx, bool evict_kv);
+    LLAMA_API void llama_context_restore_device(struct llama_context * ctx);
+    LLAMA_API bool llama_context_weights_resident(const struct llama_context * ctx);
     LLAMA_API  enum llama_pooling_type   llama_pooling_type(const struct llama_context * ctx); // TODO: rename to llama_get_pooling_type
 
     LLAMA_API const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model);

models/templates/deepseek-ai-DeepSeek-V4.jinja

@@ -8,12 +8,15 @@
     {%- set thinking = false -%}
   {%- endif -%}
 {%- endif -%}
+{%- if not drop_thinking is defined -%}
+  {%- set drop_thinking = false -%}
+{%- endif -%}
 {%- set dsml_token = '｜DSML｜' -%}
 {%- set thinking_start_token = '<think>' -%}
 {%- set thinking_end_token = '</think>' -%}
 {%- set tools_header = '## Tools\n\nYou have access to a set of tools to help answer the user\'s question. You can invoke tools by writing a "<' + dsml_token + 'tool_calls>" block like the following:\n\n<' + dsml_token + 'tool_calls>\n<' + dsml_token + 'invoke name="$TOOL_NAME">\n<' + dsml_token + 'parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</' + dsml_token + 'parameter>\n...\n</' + dsml_token + 'invoke>\n<' + dsml_token + 'invoke name="$TOOL_NAME2">\n...\n</' + dsml_token + 'invoke>\n</' + dsml_token + 'tool_calls>\n\nString parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`.\n\nIf thinking_mode is enabled (triggered by ' + thinking_start_token + '), you MUST output your complete reasoning inside ' + thinking_start_token + '...' + thinking_end_token + ' BEFORE any tool calls or final response.\n\nOtherwise, output directly after ' + thinking_end_token + ' with tool calls or final response.\n\n### Available Tool Schemas\n\n' -%}
 {%- set tools_footer = '\nYou MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls.\n' -%}
-{%- set ns = namespace(system_prompt='', is_first_sp=true) -%}
+{%- set ns = namespace(system_prompt='', is_first_sp=true, has_tool_calls=false) -%}
 {%- for message in messages -%}
   {%- if message['role'] == 'system' -%}
     {%- if ns.is_first_sp -%}
@@ -46,6 +49,11 @@
   {%- endif -%}
 {%- endfor -%}
 {%- set state = namespace(in_user=false) -%}
+{%- for message in messages -%}
+  {%- if message['role'] == 'tool' -%}
+    {%- set ns.has_tool_calls = true -%}
+  {%- endif -%}
+{%- endfor -%}
 {%- for message in messages -%}
   {%- if message['role'] == 'user' or message['role'] == 'developer' -%}
     {%- if state.in_user -%}
@@ -67,7 +75,8 @@
     {%- set state.in_user = false -%}
     {{- '<｜Assistant｜>' -}}
     {%- set is_after_last_user = loop.index0 > last_user_idx.value -%}
-    {%- if is_after_last_user and thinking -%}
+    {%- set retain_reasoning = (not drop_thinking) or (is_after_last_user or ns.has_tool_calls) -%}
+    {%- if retain_reasoning and thinking -%}
       {{- thinking_start_token -}}
       {%- if message['reasoning_content'] is defined and message['reasoning_content'] -%}
         {{- message['reasoning_content'] -}}

models/templates/poolside-Laguna-S-2.1.jinja

@@ -0,0 +1,93 @@
+{#- Iteration on laguna_glm_thinking_v8/chat_template.jinja -#}
+{#- No formatting instructions -#}
+{{- "〈|EOS|〉" -}}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set add_generation_prompt = add_generation_prompt | default(false) -%}
+
+{#- ───── header (system message) ───── -#}
+{#- A caller-supplied system message with empty content opts out of the default below, producing no <system> block — used to train without a system message. -#}
+{%- set system_message = "You are a helpful, conversationally-fluent assistant made by Poolside. You are here to be helpful to users through natural language conversations." -%}
+{%- if messages and messages[0].role == "system" -%}
+  {%- set system_message = messages[0].content -%}
+  {%- set messages = messages[1:] -%}
+{%- endif -%}
+
+{%- set has_sys = system_message and system_message.strip() -%}
+{%- if has_sys or tools or enable_thinking -%}
+  {{- "<system>" -}}
+
+  {%- if has_sys -%}
+    {{- system_message.rstrip() -}}
+    {%- if tools -%}{{- "\n\n" -}}{%- endif -%}
+  {%- endif -%}
+
+  {%- if tools -%}
+    {{- "### Tools\n\n" -}}
+    {{- "You may call functions to assist with the user query.\n" -}}
+    {{- "All available function signatures are listed below:\n" -}}
+    {{- "<available_tools>\n" -}}
+    {%- for tool in tools -%}
+      {{- (tool | tojson) ~ "\n" -}}
+    {%- endfor -%}
+    {{- "</available_tools>" -}}
+  {%- endif -%}
+
+  {{- "</system>\n" -}}
+{%- endif -%}
+
+{#- ───── main loop ───── -#}
+{%- for message in messages -%}
+  {%- set content = message.content if message.content is string else "" -%}
+  {%- if message.role == "user" -%}
+    {{- "<user>" + content + "</user>\n" -}}
+  {%- elif message.role == "assistant" -%}
+    {%- generation -%}
+      {{- "<assistant>" -}}
+      {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content -#}
+      {%- set reasoning_content = '' -%}
+      {%- if message.reasoning is string -%}
+        {%- set reasoning_content = message.reasoning -%}
+      {%- elif message.reasoning_content is string -%}
+        {%- set reasoning_content = message.reasoning_content -%}
+      {%- endif -%}
+      {#- Display reasoning content for all messages if enable_thinking -#}
+      {%- if enable_thinking -%}
+        {{- '<think>' + reasoning_content + '</think>' -}}
+      {%- else -%}
+        {{- '</think>' -}}
+      {%- endif -%}
+      {#- Display main content (trailing newline only when no tool_calls follow) -#}
+      {%- if content -%}
+        {{- content -}}
+      {%- endif -%}
+      {%- if message.tool_calls -%}
+        {%- for tool_call in message.tool_calls -%}
+          {%- set function_data = tool_call.function -%}
+          {{- '<tool_call>' + function_data.name -}}
+          {%- set _args = function_data.arguments -%}
+          {%- for k, v in _args.items() -%}
+            {{- "<arg_key>" ~ k ~ "</arg_key>" -}}
+            {{- "<arg_value>" -}}{{- v | tojson(ensure_ascii=False) if v is not string else v -}}{{- "</arg_value>" -}}
+          {%- endfor -%}
+          {{- "</tool_call>" -}}
+        {%- endfor -%}
+      {%- endif -%}
+      {{- "</assistant>\n" -}}
+    {%- endgeneration -%}
+  {%- elif message.role == "tool" -%}
+    {{- "<tool_response>" + content + "</tool_response>\n" -}}
+  {%- elif message.role == "system" -%}
+    {#- Render additional system messages (the first one, if any, is handled separately in the header and was sliced off above) -#}
+    {{- "<system>" + content + "</system>\n" -}}
+  {%- endif -%}
+{%- endfor -%}
+{#- ───── generation prompt ───── -#}
+{%- if add_generation_prompt -%}
+  {{- "<assistant>" -}}
+  {#- ───── Include reasoning mode directive ───── -#}
+  {%- if enable_thinking -%}
+    {{- '<think>' -}}
+  {%- else -%}
+    {{- '</think>' -}}
+  {%- endif -%}
+{%- endif -%}

models/templates/poolside-Laguna-XS-2.1.jinja

@@ -0,0 +1,132 @@
+{#- Copied from laguna_glm_thinking_v4/chat_template.jinja -#}
+{#- Removes prefix that references <think> token, and replaces message.reasoning_content reference with message.reasoning -#}
+{{- "〈|EOS|〉" -}}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set render_assistant_messages_raw = render_assistant_messages_raw | default(false) -%}
+{%- set add_generation_prompt = add_generation_prompt | default(false) -%}
+
+{#- ───── header (system message) ───── -#}
+{%- set system_message = "" -%}
+{%- if messages and messages[0].role == "system" -%}
+  {%- set system_message = messages[0].content -%}
+{%- endif -%}
+
+{%- if (system_message and system_message.strip()) or tools -%}
+  {{- "<system>\n" -}}
+
+  {%- if system_message and system_message.strip() -%}
+    {{- "\n" -}}
+    {{- system_message.rstrip() -}}
+  {%- endif -%}
+
+  {%- if tools -%}
+    {{- "\n\n### Tools\n\n" -}}
+    {%- set ns = namespace(tool_string="You may call functions to assist with the user query.\n"
+          ~ "All available function signatures are listed below:\n"
+          ~ "<available_tools>\n") -%}
+    {%- for tool in tools -%}
+      {%- set ns.tool_string = ns.tool_string ~ (tool | tojson) ~ "\n" -%}
+    {%- endfor -%}
+    {%- if enable_thinking -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "Wrap your thinking in '<think>', '</think>' tags, followed by a function call. For each function call, return an unescaped XML-like object with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<think> your thoughts here </think>\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- else -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "For each function call, return an unescaped XML-like object " ~
+            "with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- endif -%}
+    {{- tool_string -}}
+  {%- endif -%}
+
+  {{- "\n</system>\n" -}}
+{%- endif -%}
+
+{#- ───── main loop ───── -#}
+{%- for message in messages -%}
+  {%- set content = message.content if message.content is string else "" -%}
+  {%- if message.role == "user" -%}
+    {{- "<user>\n" + content + "\n</user>\n" -}}
+  {%- elif message.role == "assistant" -%}
+    {%- generation -%}
+      {{- "<assistant>\n" -}}
+      {%- if render_assistant_messages_raw -%}
+        {#- Raw mode: prepend the generation prompt token, then dump content verbatim. -#}
+        {#- The generation prompt is <think> when enable_thinking, </think> otherwise. -#}
+        {#- Only prepend if content doesn't already start with it. -#}
+        {%- if enable_thinking -%}
+          {%- if not content.startswith('<think>') -%}
+            {{- '<think>' -}}
+          {%- endif -%}
+        {%- else -%}
+          {%- if not content.startswith('</think>') -%}
+            {{- '</think>' -}}
+          {%- endif -%}
+        {%- endif -%}
+        {{- content -}}
+        {#- Append closing tag if content doesn't already end with it. -#}
+        {%- if not content.endswith('</assistant>\n') and not content.endswith('</assistant>') -%}
+          {{- '\n</assistant>' -}}
+        {%- endif -%}
+        {{- "\n" -}}
+      {%- else -%}
+        {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content, or from <think> tags -#}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning is string %}
+          {%- set reasoning_content = message.reasoning %}
+        {%- elif message.reasoning_content is string %}
+          {%- set reasoning_content = message.reasoning_content %}
+        {%- endif %}
+        {#- Always strip <think> tags from content if present to avoid duplication -#}
+        {%- if '</think>' in content %}
+          {%- if not reasoning_content %}
+            {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+          {%- endif %}
+          {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+        {%- endif %}
+        {#- Display reasoning content for all messages -#}
+        {%- if reasoning_content -%}
+          {{- '<think>\n' + reasoning_content.strip() + '\n</think>\n' -}}
+        {%- else -%}
+          {{- '</think>\n' -}}
+        {%- endif -%}
+        {#- Display main content -#}
+        {%- if content.strip() -%}
+          {{- content.strip() ~ "\n" -}}
+        {%- endif -%}
+        {%- if message.tool_calls -%}
+          {%- for tool_call in message.tool_calls -%}
+            {%- set function_data = tool_call.function -%}
+            {{- '<tool_call>' + function_data.name }}
+            {% set _args = function_data.arguments %}
+            {%- for k, v in _args.items() -%}
+              {{- "<arg_key>" ~ k ~ "</arg_key>\n" -}}
+              {{- "<arg_value>"}}{{ v | tojson(ensure_ascii=False) if v is not string else v }}{{ "</arg_value>\n" -}}
+            {%- endfor -%}
+            {{- "</tool_call>\n" -}}
+          {%- endfor -%}
+        {%- endif -%}
+        {{- "</assistant>\n" -}}
+      {%- endif -%}
+    {%- endgeneration -%}
+  {%- elif message.role == "tool" -%}
+    {{- "<tool_response>\n" + content + "\n</tool_response>\n" -}}
+  {%- elif message.role == "system" and loop.index0 != 0 -%}
+    {#- Render additional system messages (skip the first one which is handled separately in the header) -#}
+    {{- "<system>\n" + content + "\n</system>\n" -}}
+  {%- endif -%}
+{%- endfor -%}
+{#- ───── generation prompt ───── -#}
+{%- if add_generation_prompt -%}
+  {{- "<assistant>\n" -}}
+  {#- ───── Include reasoning mode directive ───── -#}
+  {%- if not enable_thinking %}
+    {{- '</think>' -}}
+  {%- else %}
+    {{- '<think>' -}}
+  {%- endif %}
+{%- endif -%}

models/templates/poolside-Laguna-XS.2.jinja

@@ -0,0 +1,132 @@
+{#- Iteration on laguna_glm_thinking_v5/chat_template.jinja -#}
+{#- Adds a default system message (used when no system message is provided in `messages`). -#}
+{{- "〈|EOS|〉" -}}
+{%- set enable_thinking = enable_thinking | default(false) -%}
+{%- set render_assistant_messages_raw = render_assistant_messages_raw | default(false) -%}
+{%- set add_generation_prompt = add_generation_prompt | default(false) -%}
+
+{#- ───── header (system message) ───── -#}
+{%- set system_message = "You are a helpful, conversationally-fluent assistant made by Poolside. You are here to be helpful to users through natural language conversations." -%}
+{%- if messages and messages[0].role == "system" -%}
+  {%- set system_message = messages[0].content -%}
+{%- endif -%}
+
+{%- if (system_message and system_message.strip()) or tools -%}
+  {{- "<system>\n" -}}
+
+  {%- if system_message and system_message.strip() -%}
+    {{- "\n" -}}
+    {{- system_message.rstrip() -}}
+  {%- endif -%}
+
+  {%- if tools -%}
+    {{- "\n\n### Tools\n\n" -}}
+    {%- set ns = namespace(tool_string="You may call functions to assist with the user query.\n"
+          ~ "All available function signatures are listed below:\n"
+          ~ "<available_tools>\n") -%}
+    {%- for tool in tools -%}
+      {%- set ns.tool_string = ns.tool_string ~ (tool | tojson) ~ "\n" -%}
+    {%- endfor -%}
+    {%- if enable_thinking -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "Wrap your thinking in '<think>', '</think>' tags, followed by a function call. For each function call, return an unescaped XML-like object with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<think> your thoughts here </think>\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- else -%}
+      {%- set tool_string = ns.tool_string + "</available_tools>\n\n" ~
+            "For each function call, return an unescaped XML-like object " ~
+            "with function name and arguments within '<tool_call>' and '</tool_call>' tags, like here:\n" ~
+            "<tool_call>function-name\n<arg_key>argument-key</arg_key>\n<arg_value>value-of-argument-key</arg_value>\n" ~
+            "</tool_call>" -%}
+    {%- endif -%}
+    {{- tool_string -}}
+  {%- endif -%}
+
+  {{- "\n</system>\n" -}}
+{%- endif -%}
+
+{#- ───── main loop ───── -#}
+{%- for message in messages -%}
+  {%- set content = message.content if message.content is string else "" -%}
+  {%- if message.role == "user" -%}
+    {{- "<user>\n" + content + "\n</user>\n" -}}
+  {%- elif message.role == "assistant" -%}
+    {%- generation -%}
+      {{- "<assistant>\n" -}}
+      {%- if render_assistant_messages_raw -%}
+        {#- Raw mode: prepend the generation prompt token, then dump content verbatim. -#}
+        {#- The generation prompt is <think> when enable_thinking, </think> otherwise. -#}
+        {#- Only prepend if content doesn't already start with it. -#}
+        {%- if enable_thinking -%}
+          {%- if not content.startswith('<think>') -%}
+            {{- '<think>' -}}
+          {%- endif -%}
+        {%- else -%}
+          {%- if not content.startswith('</think>') -%}
+            {{- '</think>' -}}
+          {%- endif -%}
+        {%- endif -%}
+        {{- content -}}
+        {#- Append closing tag if content doesn't already end with it. -#}
+        {%- if not content.endswith('</assistant>\n') and not content.endswith('</assistant>') -%}
+          {{- '\n</assistant>' -}}
+        {%- endif -%}
+        {{- "\n" -}}
+      {%- else -%}
+        {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content, or from <think> tags -#}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning is string %}
+          {%- set reasoning_content = message.reasoning %}
+        {%- elif message.reasoning_content is string %}
+          {%- set reasoning_content = message.reasoning_content %}
+        {%- endif %}
+        {#- Always strip <think> tags from content if present to avoid duplication -#}
+        {%- if '</think>' in content %}
+          {%- if not reasoning_content %}
+            {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+          {%- endif %}
+          {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+        {%- endif %}
+        {#- Display reasoning content for all messages -#}
+        {%- if reasoning_content -%}
+          {{- '<think>\n' + reasoning_content.strip() + '\n</think>\n' -}}
+        {%- else -%}
+          {{- '</think>\n' -}}
+        {%- endif -%}
+        {#- Display main content -#}
+        {%- if content.strip() -%}
+          {{- content.strip() ~ "\n" -}}
+        {%- endif -%}
+        {%- if message.tool_calls -%}
+          {%- for tool_call in message.tool_calls -%}
+            {%- set function_data = tool_call.function -%}
+            {{- '<tool_call>' + function_data.name }}
+            {% set _args = function_data.arguments %}
+            {%- for k, v in _args.items() -%}
+              {{- "<arg_key>" ~ k ~ "</arg_key>\n" -}}
+              {{- "<arg_value>"}}{{ v | tojson(ensure_ascii=False) if v is not string else v }}{{ "</arg_value>\n" -}}
+            {%- endfor -%}
+            {{- "</tool_call>\n" -}}
+          {%- endfor -%}
+        {%- endif -%}
+        {{- "</assistant>\n" -}}
+      {%- endif -%}
+    {%- endgeneration -%}
+  {%- elif message.role == "tool" -%}
+    {{- "<tool_response>\n" + content + "\n</tool_response>\n" -}}
+  {%- elif message.role == "system" and loop.index0 != 0 -%}
+    {#- Render additional system messages (skip the first one which is handled separately in the header) -#}
+    {{- "<system>\n" + content + "\n</system>\n" -}}
+  {%- endif -%}
+{%- endfor -%}
+{#- ───── generation prompt ───── -#}
+{%- if add_generation_prompt -%}
+  {{- "<assistant>\n" -}}
+  {#- ───── Include reasoning mode directive ───── -#}
+  {%- if not enable_thinking %}
+    {{- '</think>' -}}
+  {%- else %}
+    {{- '<think>' -}}
+  {%- endif %}
+{%- endif -%}

src/llama-arch.cpp

@@ -108,6 +108,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
     { LLM_ARCH_DOTS1,            "dots1"            },
     { LLM_ARCH_ARCEE,            "arcee"            },
     { LLM_ARCH_AFMOE,            "afmoe"            },
+    { LLM_ARCH_LAGUNA,           "laguna"           },
     { LLM_ARCH_ERNIE4_5,         "ernie4_5"         },
     { LLM_ARCH_ERNIE4_5_MOE,     "ernie4_5-moe"     },
     { LLM_ARCH_HUNYUAN_MOE,      "hunyuan-moe"      },
@@ -665,7 +666,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_HC_FFN_SCALE,               {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
     {LLM_TENSOR_ATTN_COMPRESSOR_WKV,        {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_ATTN_COMPRESSOR_WGATE,      {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
-    {LLM_TENSOR_ATTN_COMPRESSOR_APE,        {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}},
+    {LLM_TENSOR_ATTN_COMPRESSOR_APE,        {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
     {LLM_TENSOR_ATTN_COMPRESSOR_NORM,       {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
     {LLM_TENSOR_ATTN_K_B,                   {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_ATTN_V_B,                   {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
@@ -832,7 +833,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
     {LLM_TENSOR_INDEXER_ATTN_Q_B,           {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_INDEXER_COMPRESSOR_WKV,     {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
     {LLM_TENSOR_INDEXER_COMPRESSOR_WGATE,   {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
-    {LLM_TENSOR_INDEXER_COMPRESSOR_APE,     {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}},
+    {LLM_TENSOR_INDEXER_COMPRESSOR_APE,     {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
     {LLM_TENSOR_INDEXER_COMPRESSOR_NORM,    {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
     {LLM_TENSOR_FFN_GATE_TID2EID,           {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
     {LLM_TENSOR_NEXTN_PROJ_PRE,             {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},

src/llama-arch.h

@@ -113,6 +113,7 @@ enum llm_arch {
     LLM_ARCH_DOTS1,
     LLM_ARCH_ARCEE,
     LLM_ARCH_AFMOE,
+    LLM_ARCH_LAGUNA,
     LLM_ARCH_ERNIE4_5,
     LLM_ARCH_ERNIE4_5_MOE,
     LLM_ARCH_HUNYUAN_MOE,

src/llama-context.cpp

@@ -728,6 +728,47 @@ void llama_context::synchronize() {
     t_compute_start_us = 0;
 }
 
+void llama_context::release_device(bool evict_kv) {
+    if (!model.weights_resident() && !(evict_kv && !kv_device_evicted)) {
+        return;
+    }
+    // ensure no compute is in flight before freeing the device buffers
+    synchronize();
+    model.release_device_weights();
+    if (evict_kv && memory && !kv_device_evicted) {
+        memory->release_device_buffers();
+        kv_device_evicted = true;
+        // also free the scheduler and its worst-case compute buffer (hundreds of MiB) so a cold
+        // model holds essentially no VRAM. Rebuilt lazily by sched_reserve() on restore.
+        sched.reset();
+        sched_need_reserve = true;
+        // finally, free each backend's own compute-scratch (Vulkan prealloc/staging buffers, which
+        // are owned by the backend and survive sched.reset()). Reallocated lazily on next compute.
+        for (auto & backend : backends) {
+            ggml_backend_free_scratch(backend.get());
+        }
+    }
+}
+
+void llama_context::restore_device() {
+    if (model.weights_resident() && !kv_device_evicted && sched) {
+        return;
+    }
+    model.restore_device_weights();
+    if (kv_device_evicted && memory) {
+        memory->restore_device_buffers();
+        kv_device_evicted = false;
+    }
+    // rebuild the scheduler + compute buffer if they were freed on release (evict_kv mode)
+    if (!sched) {
+        sched_reserve();
+    }
+    // make sure all weight/KV uploads have completed before any compute reads them
+    if (sched) {
+        ggml_backend_sched_synchronize(sched.get());
+    }
+}
+
 const llama_model & llama_context::get_model() const {
     return model;
 }
@@ -1705,6 +1746,12 @@ int llama_context::decode(const llama_batch & batch_inp) {
         return -1;
     }
 
+    // on-demand VRAM: if the weights were released while this model was idle, bring them back
+    // to the device before building the compute graph
+    if (!model.weights_resident()) {
+        restore_device();
+    }
+
     const auto & vocab   = model.vocab;
     const auto & hparams = model.hparams;
 
@@ -3223,6 +3270,9 @@ llama_memory_breakdown llama_context::memory_breakdown() const {
             ret[buft].context += size;
         }
     }
+    // the scheduler (and its compute buffers) may have been freed while the model is cold
+    // (on-demand VRAM eviction, see release_device); it contributes no compute VRAM then.
+    if (sched) {
         if (model.hparams.no_alloc) {
             for (size_t i = 0; i < backends.size(); ++i) {
                 ggml_backend_t             backend = backends[i].get();
@@ -3236,6 +3286,7 @@ llama_memory_breakdown llama_context::memory_breakdown() const {
                 ret[buft].compute += ggml_backend_sched_get_buffer_size(sched.get(), backend);
             }
         }
+    }
     return ret;
 }
 
@@ -3674,6 +3725,18 @@ void llama_synchronize(llama_context * ctx) {
     ctx->synchronize();
 }
 
+void llama_context_release_device(llama_context * ctx, bool evict_kv) {
+    ctx->release_device(evict_kv);
+}
+
+void llama_context_restore_device(llama_context * ctx) {
+    ctx->restore_device();
+}
+
+bool llama_context_weights_resident(const llama_context * ctx) {
+    return ctx->get_model().weights_resident();
+}
+
 float * llama_get_logits(llama_context * ctx) {
     ctx->synchronize();
 

src/llama-context.h

@@ -57,6 +57,14 @@ struct llama_context {
 
     void synchronize();
 
+    // on-demand device (VRAM) residency: free / rebuild the model's GPU weight buffers so that
+    // VRAM can be time-shared with another model. release_device() synchronizes first; decode()
+    // auto-restores when it finds the weights have been released. If evict_kv is set, the KV cache
+    // device buffers are also evicted to a host shadow (for when the KV would not fit alongside the
+    // other model) - this adds a D2H/H2D copy of the live KV on each cycle.
+    void release_device(bool evict_kv = false);
+    void restore_device();
+
     const llama_model   & get_model()   const;
     const llama_cparams & get_cparams() const;
 
@@ -345,6 +353,9 @@ private:
 
     bool sched_need_reserve = true;
 
+    // on-demand device residency: true while the KV cache device buffers have been evicted to host
+    bool kv_device_evicted = false;
+
     ggml_backend_t backend_cpu = nullptr;
     std::vector<ggml_backend_ptr> backends;
 

src/llama-kv-cache-iswa.cpp

@@ -272,6 +272,17 @@ void llama_kv_cache_iswa::state_read(llama_io_read_i & io, llama_seq_id seq_id,
     kv_swa->state_read(io, seq_id, flags);
 }
 
+void llama_kv_cache_iswa::release_device_buffers() {
+    kv_base->release_device_buffers();
+    kv_swa->release_device_buffers();
+}
+
+bool llama_kv_cache_iswa::restore_device_buffers() {
+    bool ok = kv_base->restore_device_buffers();
+    ok = kv_swa->restore_device_buffers() && ok;
+    return ok;
+}
+
 llama_kv_cache * llama_kv_cache_iswa::get_base() const {
     return kv_base.get();
 }

src/llama-kv-cache-iswa.h

@@ -83,6 +83,9 @@ public:
     void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
     void state_read (llama_io_read_i  & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
 
+    void release_device_buffers() override;
+    bool restore_device_buffers() override;
+
     //
     // llama_kv_cache_iswa specific API
     //

src/llama-kv-cache.cpp

@@ -371,10 +371,12 @@ void llama_kv_cache::clear(bool data) {
 
     if (data) {
         for (auto & [_, buf] : ctxs_bufs) {
+            if (buf) { // may be null if evicted for on-demand VRAM sharing
                 ggml_backend_buffer_clear(buf.get(), 0);
             }
         }
     }
+}
 
 bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
@@ -681,6 +683,9 @@ llama_pos llama_kv_cache::seq_pos_max(llama_seq_id seq_id) const {
 std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache::memory_breakdown() const {
     std::map<ggml_backend_buffer_type_t, size_t> ret;
     for (const auto & [ctx, buf] : ctxs_bufs) {
+        if (!buf) { // may be null if evicted for on-demand VRAM sharing
+            continue;
+        }
         ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());
 
         if (hparams.no_alloc) {
@@ -1801,12 +1806,88 @@ size_t llama_kv_cache::total_size() const {
     size_t size = 0;
 
     for (const auto & [_, buf] : ctxs_bufs) {
+        if (buf) { // may be null if evicted for on-demand VRAM sharing
             size += ggml_backend_buffer_get_size(buf.get());
         }
+    }
 
     return size;
 }
 
+void llama_kv_cache::release_device_buffers() {
+    // NOTE: the caller must have synchronized the backend so nothing references these buffers.
+    // The KV cache is read-write, so its host shadow is (re)captured fresh on every release.
+    if (dev_released) {
+        return;
+    }
+    dev_shadows.assign(ctxs_bufs.size(), device_buffer_shadow{});
+    size_t freed = 0;
+    for (size_t i = 0; i < ctxs_bufs.size(); ++i) {
+        ggml_context * ctx = ctxs_bufs[i].first.get();
+        ggml_backend_buffer_t buf = ctxs_bufs[i].second.get();
+        if (buf == nullptr || ggml_backend_buffer_is_host(buf) || ggml_backend_buffer_get_size(buf) == 0) {
+            continue; // only real device (VRAM) buffers are evictable
+        }
+        auto & sh = dev_shadows[i];
+        sh.releasable = true;
+        sh.buft = ggml_backend_buffer_get_type(buf);
+
+        // capture live contents compactly in stable iteration order (skip views, which alias a base)
+        size_t total = 0;
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->view_src == nullptr) { total += ggml_nbytes(t); }
+        }
+        sh.data.resize(total);
+        size_t off = 0;
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->view_src != nullptr) { continue; }
+            const size_t n = ggml_nbytes(t);
+            ggml_backend_tensor_get(t, sh.data.data() + off, 0, n);
+            off += n;
+        }
+
+        freed += ggml_backend_buffer_get_size(buf);
+        ctxs_bufs[i].second.reset(); // free the device buffer
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            t->buffer = nullptr;
+            t->data   = nullptr;
+        }
+    }
+    dev_released = true;
+    if (freed > 0) {
+        LLAMA_LOG_INFO("%s: released %.2f MiB of KV cache from device\n", __func__, freed / 1024.0 / 1024.0);
+    }
+}
+
+bool llama_kv_cache::restore_device_buffers() {
+    if (!dev_released) {
+        return true;
+    }
+    for (size_t i = 0; i < ctxs_bufs.size(); ++i) {
+        auto & sh = dev_shadows[i];
+        if (!sh.releasable) {
+            continue;
+        }
+        ggml_context * ctx = ctxs_bufs[i].first.get();
+        ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, sh.buft);
+        if (buf == nullptr) {
+            LLAMA_LOG_ERROR("%s: failed to reallocate KV cache device buffer (out of VRAM?)\n", __func__);
+            return false;
+        }
+        size_t off = 0;
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->view_src != nullptr) { continue; }
+            const size_t n = ggml_nbytes(t);
+            ggml_backend_tensor_set(t, sh.data.data() + off, 0, n);
+            off += n;
+        }
+        ctxs_bufs[i].second.reset(buf);
+    }
+    dev_released = false;
+    dev_shadows.clear(); // recaptured on next release
+    return true;
+}
+
 size_t llama_kv_cache::size_k_bytes() const {
     size_t size_k_bytes = 0;
 

src/llama-kv-cache.h

@@ -149,6 +149,10 @@ public:
     void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
     void state_read (llama_io_read_i  & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
 
+    // on-demand device (VRAM) residency (see llama_memory_i)
+    void release_device_buffers() override;
+    bool restore_device_buffers() override;
+
     //
     // llama_kv_cache specific API
     //
@@ -265,6 +269,16 @@ private:
     // ggml contexts for the KV cache along with the allocated backend buffers:
     std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
 
+    // on-demand device eviction: host shadow of each device buffer's live contents (re-captured on
+    // every release since the KV is read-write), plus the buffer type needed to reallocate it.
+    struct device_buffer_shadow {
+        ggml_backend_buffer_type_t buft = nullptr;
+        bool releasable = false;      // device (VRAM) buffer that was freed
+        std::vector<uint8_t> data;    // captured tensor bytes (compact, iteration order)
+    };
+    std::vector<device_buffer_shadow> dev_shadows; // parallel to ctxs_bufs
+    bool dev_released = false;
+
     // the current index from where we start searching for a free slot in the ring buffer of KV cells (see find_slot())
     // note: this is not part of the KV state and it's only used to speed-up the find_slot() method
     std::vector<uint32_t> v_heads;

src/llama-memory-hybrid.cpp

@@ -201,6 +201,18 @@ void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id,
     mem_recr->state_read(io, seq_id, flags);
 }
 
+void llama_memory_hybrid::release_device_buffers() {
+    // evict both the attention KV (grows with context) and the recurrent/SSM state
+    mem_attn->release_device_buffers();
+    mem_recr->release_device_buffers();
+}
+
+bool llama_memory_hybrid::restore_device_buffers() {
+    bool ok = mem_attn->restore_device_buffers();
+    ok = mem_recr->restore_device_buffers() && ok;
+    return ok;
+}
+
 llama_kv_cache * llama_memory_hybrid::get_mem_attn() const {
     return mem_attn.get();
 }

src/llama-memory-hybrid.h

@@ -76,6 +76,9 @@ public:
     void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
     void state_read (llama_io_read_i  & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0)       override;
 
+    void release_device_buffers() override;
+    bool restore_device_buffers() override;
+
     //
     // llama_memory_hybrid specific API
     //

src/llama-memory-recurrent.cpp

@@ -140,9 +140,11 @@ void llama_memory_recurrent::clear(bool data) {
 
     if (data) {
         for (auto & [_, buf] : ctxs_bufs) {
+            if (buf) { // may be null if evicted for on-demand VRAM sharing
                 ggml_backend_buffer_clear(buf.get(), 0);
             }
         }
+    }
 
     std::fill(rs_idx.begin(), rs_idx.end(), 0);
 }
@@ -399,6 +401,7 @@ void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) {
 std::map<ggml_backend_buffer_type_t, size_t> llama_memory_recurrent::memory_breakdown() const {
     std::map<ggml_backend_buffer_type_t, size_t> ret;
     for (const auto & [_, buf] : ctxs_bufs) {
+        if (!buf) { continue; } // may be null if evicted for on-demand VRAM sharing
         ret[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get());
     }
     return ret;
@@ -700,12 +703,87 @@ bool llama_memory_recurrent::get_can_shift() const {
 size_t llama_memory_recurrent::total_size() const {
     size_t size = 0;
     for (const auto & [_, buf] : ctxs_bufs) {
+        if (buf) { // may be null if evicted for on-demand VRAM sharing
             size += ggml_backend_buffer_get_size(buf.get());
         }
+    }
 
     return size;
 }
 
+void llama_memory_recurrent::release_device_buffers() {
+    // Same mechanism as llama_kv_cache: the recurrent (SSM/conv) state is read-write, so its host
+    // shadow is (re)captured on every release. The caller must have synchronized the backend.
+    if (dev_released) {
+        return;
+    }
+    dev_shadows.assign(ctxs_bufs.size(), device_buffer_shadow{});
+    size_t freed = 0;
+    for (size_t i = 0; i < ctxs_bufs.size(); ++i) {
+        ggml_context * ctx = ctxs_bufs[i].first.get();
+        ggml_backend_buffer_t buf = ctxs_bufs[i].second.get();
+        if (buf == nullptr || ggml_backend_buffer_is_host(buf) || ggml_backend_buffer_get_size(buf) == 0) {
+            continue;
+        }
+        auto & sh = dev_shadows[i];
+        sh.releasable = true;
+        sh.buft = ggml_backend_buffer_get_type(buf);
+
+        size_t total = 0;
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->view_src == nullptr) { total += ggml_nbytes(t); }
+        }
+        sh.data.resize(total);
+        size_t off = 0;
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->view_src != nullptr) { continue; }
+            const size_t n = ggml_nbytes(t);
+            ggml_backend_tensor_get(t, sh.data.data() + off, 0, n);
+            off += n;
+        }
+
+        freed += ggml_backend_buffer_get_size(buf);
+        ctxs_bufs[i].second.reset();
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            t->buffer = nullptr;
+            t->data   = nullptr;
+        }
+    }
+    dev_released = true;
+    if (freed > 0) {
+        LLAMA_LOG_INFO("%s: released %.2f MiB of recurrent state from device\n", __func__, freed / 1024.0 / 1024.0);
+    }
+}
+
+bool llama_memory_recurrent::restore_device_buffers() {
+    if (!dev_released) {
+        return true;
+    }
+    for (size_t i = 0; i < ctxs_bufs.size(); ++i) {
+        auto & sh = dev_shadows[i];
+        if (!sh.releasable) {
+            continue;
+        }
+        ggml_context * ctx = ctxs_bufs[i].first.get();
+        ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, sh.buft);
+        if (buf == nullptr) {
+            LLAMA_LOG_ERROR("%s: failed to reallocate recurrent device buffer (out of VRAM?)\n", __func__);
+            return false;
+        }
+        size_t off = 0;
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->view_src != nullptr) { continue; }
+            const size_t n = ggml_nbytes(t);
+            ggml_backend_tensor_set(t, sh.data.data() + off, 0, n);
+            off += n;
+        }
+        ctxs_bufs[i].second.reset(buf);
+    }
+    dev_released = false;
+    dev_shadows.clear();
+    return true;
+}
+
 size_t llama_memory_recurrent::size_r_bytes() const {
     size_t size_r_bytes = 0;
 

src/llama-memory-recurrent.h

@@ -66,6 +66,10 @@ public:
     void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
     void state_read (llama_io_read_i  & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
 
+    // on-demand device (VRAM) residency (see llama_memory_i)
+    void release_device_buffers() override;
+    bool restore_device_buffers() override;
+
     uint32_t head = 0; // the location where the batch will be placed in the cache (see find_slot())
     uint32_t size = 0; // total number of cells, shared across all sequences
     uint32_t used = 0; // used cells (i.e. at least one seq_id)
@@ -121,6 +125,16 @@ private:
     // ggml contexts for the KV cache along with the allocated backend buffers:
     std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
 
+    // on-demand device eviction (see release_device_buffers): host shadow of each device buffer's
+    // live contents (recaptured on every release since the recurrent state is read-write)
+    struct device_buffer_shadow {
+        ggml_backend_buffer_type_t buft = nullptr;
+        bool releasable = false;
+        std::vector<uint8_t> data;
+    };
+    std::vector<device_buffer_shadow> dev_shadows; // parallel to ctxs_bufs
+    bool dev_released = false;
+
     size_t total_size() const;
 
     size_t size_r_bytes() const;

src/llama-memory.h

@@ -124,6 +124,16 @@ struct llama_memory_i {
 
     virtual void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const = 0;
     virtual void state_read (llama_io_read_i  & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) = 0;
+
+    //
+    // on-demand device (VRAM) residency
+    //
+
+    // Free this memory's device (VRAM) buffers to a host shadow and rebuild them on demand, to
+    // hand VRAM to another model while keeping the cached data (no re-prefill). The caller must
+    // have synchronized the backend first. Default: no-op (the memory stays resident).
+    virtual void release_device_buffers() {}
+    virtual bool restore_device_buffers() { return true; }
 };
 
 using llama_memory_ptr = std::unique_ptr<llama_memory_i>;

src/llama-mmap.cpp

@@ -618,13 +618,47 @@ struct llama_mmap::impl {
 };
 
 llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa) : pimpl(std::make_unique<impl>(file, prefetch, numa)) {}
-llama_mmap::~llama_mmap() = default;
+
+llama_mmap::~llama_mmap() {
+    // unpin before the pages are unmapped by the impl destructor
+    if (host_reg_addr && host_unreg_fn) {
+        host_unreg_fn(host_reg_addr);
+    }
+}
 
 size_t llama_mmap::size() const { return pimpl->size; }
 void * llama_mmap::addr() const { return pimpl->addr; }
 
 void llama_mmap::unmap_fragment(size_t first, size_t last) { pimpl->unmap_fragment(first, last); }
 
+size_t llama_mmap::register_host(size_t first, size_t last, bool (*reg_fn)(void *, size_t), void (*unreg_fn)(void *)) {
+#ifdef _POSIX_MAPPED_FILES
+    if (host_reg_addr || !reg_fn || !unreg_fn || last <= first) {
+        return 0;
+    }
+
+    // expand outward to the page boundaries retained by unmap_fragment
+    const size_t page_size = sysconf(_SC_PAGESIZE);
+    first = first & ~(page_size - 1);
+    last  = (last + page_size - 1) & ~(page_size - 1);
+
+    void * reg_addr = (uint8_t *) pimpl->addr + first;
+    if (!reg_fn(reg_addr, last - first)) {
+        return 0;
+    }
+
+    host_reg_addr  = reg_addr;
+    host_unreg_fn  = unreg_fn;
+    return last - first;
+#else
+    GGML_UNUSED(first);
+    GGML_UNUSED(last);
+    GGML_UNUSED(reg_fn);
+    GGML_UNUSED(unreg_fn);
+    return 0;
+#endif
+}
+
 #if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32)
 const bool llama_mmap::SUPPORTED  = true;
 #else

src/llama-mmap.h

@@ -50,11 +50,19 @@ struct llama_mmap {
 
     void unmap_fragment(size_t first, size_t last);
 
+    // pin the pages backing [first, last) with a backend allocator for faster H2D copies,
+    // unpinned in the destructor before the pages are unmapped
+    // returns the number of bytes registered, 0 on failure
+    size_t register_host(size_t first, size_t last, bool (*reg_fn)(void *, size_t), void (*unreg_fn)(void *));
+
     static const bool SUPPORTED;
 
 private:
     struct impl;
     std::unique_ptr<impl> pimpl;
+
+    void * host_reg_addr = nullptr;
+    void (*host_unreg_fn)(void *) = nullptr;
 };
 
 struct llama_mlock {

src/llama-model-loader.cpp

@@ -1677,6 +1677,15 @@ bool llama_model_loader::load_all_data(
     if (size_done >= size_data) {
         // unmap offloaded tensors and metadata
         if (use_mmap) {
+            // pin the pages backing the weights kept in system memory for faster H2D copies
+            bool (*reg_fn)(void *, size_t) = nullptr;
+            void (*unreg_fn)(void *) = nullptr;
+            for (size_t i = 0; i < ggml_backend_dev_count() && !reg_fn; i++) {
+                ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_dev_get(i));
+                reg_fn   = (bool (*)(void *, size_t)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_register_host_buffer");
+                unreg_fn = (void (*)(void *))         ggml_backend_reg_get_proc_address(reg, "ggml_backend_unregister_host_buffer");
+            }
+
             for (uint32_t idx = 0; idx < mappings.size(); idx++) {
                 const auto & mmap_used = mmaps_used.at(idx);
                 auto & mapping = mappings.at(idx);
@@ -1684,6 +1693,13 @@ bool llama_model_loader::load_all_data(
                 if (mmap_used.second != 0) {
                     mapping->unmap_fragment(mmap_used.second, mapping->size());
                 }
+                if (mmap_used.second > mmap_used.first) {
+                    size_t n_registered = mapping->register_host(mmap_used.first, mmap_used.second, reg_fn, unreg_fn);
+                    if (n_registered > 0) {
+                        LLAMA_LOG_INFO("%s: pinned %.2f MiB of mapped model memory for faster H2D transfers\n",
+                                __func__, n_registered / 1024.0 / 1024.0);
+                    }
+                }
             }
         }
         if (progress_callback) {

src/llama-model-saver.cpp

@@ -28,6 +28,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) {
         case LLM_ARCH_MIMO2:
         case LLM_ARCH_STEP35:
         case LLM_ARCH_MELLUM:
+        case LLM_ARCH_LAGUNA:
             return false;
         default:
             return true;

src/llama-model.cpp

@@ -250,6 +250,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
             return new llama_model_arcee(params);
         case LLM_ARCH_AFMOE:
             return new llama_model_afmoe(params);
+        case LLM_ARCH_LAGUNA:
+            return new llama_model_laguna(params);
         case LLM_ARCH_ERNIE4_5:
             return new llama_model_ernie4_5(params);
         case LLM_ARCH_ERNIE4_5_MOE:
@@ -1013,6 +1015,16 @@ struct llama_model::impl {
     // contexts where the model tensors metadata is stored as well as the corresponding buffers:
     std::vector<std::pair<ggml_context_ptr, std::vector<ggml_backend_buffer_ptr>>> ctxs_bufs;
 
+    // on-demand device (VRAM) weight residency: per-ctxs_bufs slot describing whether the
+    // device weight buffer can be freed and rebuilt from a host shadow (see release/restore_device_weights)
+    struct weight_release_slot {
+        ggml_backend_buffer_type_t buft = nullptr; // buffer type to reallocate on restore
+        bool releasable = false;                   // single-buffer alloc-path device (VRAM) buffer
+        bool resident   = true;                    // currently allocated on device
+        std::vector<uint8_t> shadow;               // host copy, captured lazily on first release
+    };
+    std::vector<weight_release_slot> weight_release; // parallel to ctxs_bufs
+
     buft_list_t cpu_buft_list;
     std::map<ggml_backend_dev_t, buft_list_t> gpu_buft_list;
 
@@ -1606,6 +1618,23 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
 
         pimpl->ctxs_bufs.emplace_back(std::move(ctx_ptr), std::move(bufs));
 
+        // record on-demand release metadata (parallel to ctxs_bufs). Only the single-buffer
+        // alloc path on a device (non-host) buffer can be freed and rebuilt from a host shadow;
+        // the mmap buffer_from_host_ptr path and host (CPU) buffers are left resident.
+        {
+            llama_model::impl::weight_release_slot slot;
+            const bool mmap_host_ptr_path = ml.use_mmap && use_mmap_buffer && buffer_from_host_ptr_supported && is_default_buft;
+            auto & last_bufs = pimpl->ctxs_bufs.back().second;
+            if (!mmap_host_ptr_path && last_bufs.size() == 1) {
+                ggml_backend_buffer_t b = last_bufs[0].get();
+                if (b != nullptr && !ggml_backend_buffer_is_host(b)) {
+                    slot.releasable = true;
+                    slot.buft = buft;
+                }
+            }
+            pimpl->weight_release.push_back(std::move(slot));
+        }
+
         ctx_buf_maps.emplace_back(ctx, buf_map);
     }
 
@@ -1660,6 +1689,97 @@ ggml_tensor * llama_model_base::create_tensor(llama_model_loader & ml, const LLM
         tn, ne, flags);
 }
 
+void llama_model::release_device_weights() const {
+    // NOTE: the caller is responsible for synchronizing the backend scheduler first, so that no
+    // compute is in flight referencing these buffers when they are freed.
+    size_t freed = 0;
+    for (size_t i = 0; i < pimpl->ctxs_bufs.size(); ++i) {
+        auto & slot = pimpl->weight_release[i];
+        if (!slot.releasable || !slot.resident) {
+            continue;
+        }
+        ggml_context * ctx = pimpl->ctxs_bufs[i].first.get();
+        ggml_backend_buffer_t buf = pimpl->ctxs_bufs[i].second[0].get();
+        const size_t sz = ggml_backend_buffer_get_size(buf);
+
+        // lazily capture a host shadow of the weights (read-only, so captured only once).
+        // Store tensor bytes compactly in stable iteration order (independent of buffer layout /
+        // alignment padding) so restore does not depend on the re-allocated offsets matching.
+        // View tensors alias their base, so they are skipped (restored implicitly via their base).
+        if (slot.shadow.empty()) {
+            size_t total = 0;
+            for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+                if (t->view_src != nullptr) continue;
+                total += ggml_nbytes(t);
+            }
+            slot.shadow.resize(total);
+            size_t off = 0;
+            for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+                if (t->view_src != nullptr) continue;
+                const size_t n = ggml_nbytes(t);
+                ggml_backend_tensor_get(t, slot.shadow.data() + off, 0, n);
+                off += n;
+            }
+        }
+
+        // free the device (VRAM) buffer and clear the now-dangling tensor pointers so that
+        // ggml_backend_alloc_ctx_tensors_from_buft reallocates them cleanly on restore
+        pimpl->ctxs_bufs[i].second[0].reset();
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            t->buffer = nullptr;
+            t->data   = nullptr;
+        }
+        slot.resident = false;
+        freed += sz;
+    }
+    if (freed > 0) {
+        LLAMA_LOG_INFO("%s: released %.2f MiB of device weights\n", __func__, freed / 1024.0 / 1024.0);
+    }
+}
+
+bool llama_model::restore_device_weights() const {
+    size_t restored = 0;
+    for (size_t i = 0; i < pimpl->ctxs_bufs.size(); ++i) {
+        auto & slot = pimpl->weight_release[i];
+        if (!slot.releasable || slot.resident) {
+            continue;
+        }
+        ggml_context * ctx = pimpl->ctxs_bufs[i].first.get();
+        ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, slot.buft);
+        if (buf == nullptr) {
+            LLAMA_LOG_ERROR("%s: failed to reallocate device weight buffer (out of VRAM?)\n", __func__);
+            return false;
+        }
+        ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+
+        // re-upload weights from the compact host shadow, using the same stable iteration order
+        // and view-skipping as the capture above
+        size_t off = 0;
+        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
+            if (t->view_src != nullptr) continue;
+            const size_t n = ggml_nbytes(t);
+            ggml_backend_tensor_set(t, slot.shadow.data() + off, 0, n);
+            off += n;
+        }
+        pimpl->ctxs_bufs[i].second[0].reset(buf);
+        slot.resident = true;
+        restored += ggml_backend_buffer_get_size(buf);
+    }
+    if (restored > 0) {
+        LLAMA_LOG_INFO("%s: restored %.2f MiB of device weights\n", __func__, restored / 1024.0 / 1024.0);
+    }
+    return true;
+}
+
+bool llama_model::weights_resident() const {
+    for (const auto & slot : pimpl->weight_release) {
+        if (slot.releasable && !slot.resident) {
+            return false;
+        }
+    }
+    return true;
+}
+
 std::string llama_model::arch_name() const {
     return llm_arch_name(arch);
 }
@@ -1712,6 +1832,11 @@ std::map<ggml_backend_buffer_type_t, size_t> llama_model::memory_breakdown() con
             ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);
         } else {
             for (const auto & buf : bufs) {
+                // buf may be null if the device weights were released for on-demand VRAM sharing
+                // (see release_device_weights); skip it in the breakdown
+                if (!buf) {
+                    continue;
+                }
                 // GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) != nullptr); // multi_buffer does not have a defined base
                 ret[ggml_backend_buffer_get_type(buf.get())] += ggml_backend_buffer_get_size(buf.get());
             }
@@ -2549,6 +2674,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
         case LLM_ARCH_COGVLM:
         case LLM_ARCH_PANGU_EMBED:
         case LLM_ARCH_AFMOE:
+        case LLM_ARCH_LAGUNA:
         case LLM_ARCH_QWEN3NEXT:
         case LLM_ARCH_MIMO2:
         case LLM_ARCH_STEP35:

src/llama-model.h

@@ -663,6 +663,14 @@ struct llama_model {
 
     const struct ggml_tensor * get_tensor(const char * name) const;
 
+    // on-demand device (VRAM) weight residency: free the GPU weight buffers (keeping a host
+    // shadow) and rebuild them on demand. Used to time-share VRAM between models without
+    // reloading or losing the KV cache. release_device_weights() requires the caller to have
+    // synchronized the backend scheduler first.
+    void release_device_weights() const;
+    bool restore_device_weights() const;
+    bool weights_resident() const;
+
     float get_rope_freq_base (const llama_cparams & cparams, int il) const;
     float get_rope_freq_scale(const llama_cparams & cparams, int il) const;
 

src/llama-vocab.cpp

@@ -496,6 +496,12 @@ struct llm_tokenizer_bpe : llm_tokenizer {
                     "[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\\r\\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
                 };
                 break;
+            case LLAMA_VOCAB_PRE_TYPE_LAGUNA:
+                regex_exprs = {
+                    "[^\\n]+|[\\n]+",
+                    "(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+                };
+                break;
             case LLAMA_VOCAB_PRE_TYPE_EXAONE_MOE:
                 regex_exprs = {
                     // original regex from tokenizer.json
@@ -2342,6 +2348,10 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
                 tokenizer_pre == "afmoe") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_AFMOE;
                 clean_spaces = false;
+            } else if (
+                tokenizer_pre == "laguna") {
+                pre_type = LLAMA_VOCAB_PRE_TYPE_LAGUNA;
+                clean_spaces = false;
             } else if (
                 tokenizer_pre == "minimax-m2") {
                 pre_type = LLAMA_VOCAB_PRE_TYPE_MINIMAX_M2;

src/llama-vocab.h

@@ -64,6 +64,7 @@ enum llama_vocab_pre_type {
     LLAMA_VOCAB_PRE_TYPE_WHITESPACE        = 53,
     LLAMA_VOCAB_PRE_TYPE_GRANITE_EMB_MULTI = 54,
     LLAMA_VOCAB_PRE_TYPE_MELLUM2           = 55,
+    LLAMA_VOCAB_PRE_TYPE_LAGUNA            = 56,
 };
 
 struct LLM_KV;

src/models/laguna.cpp

@@ -0,0 +1,332 @@
+// Laguna (poolside): sigmoid-routed MoE with a score-correction bias, one shared
+// expert, a softplus attention output gate, QK-norm, and per-layer-type RoPE
+// (YaRN on full-attention layers, plain RoPE on sliding-window layers). XS.2 is
+// hybrid full/SWA with a per-head gate; M.1 is full-attention with a per-element
+// gate. Shares the MoE/gate structure with afmoe.
+
+#include "models.h"
+
+void llama_model_laguna::load_arch_hparams(llama_model_loader & ml) {
+    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+    ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT,   hparams.n_layer_dense_lead);
+    ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH,  hparams.n_ff_exp);
+    ml.get_key(LLM_KV_EXPERT_GATING_FUNC,          hparams.expert_gating_func, false);
+    ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,        hparams.expert_weights_scale, false);
+    ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM,         hparams.expert_weights_norm, false);
+
+    // Laguna ships one shared expert and stores its size directly (routed and
+    // shared experts may differ), so read the size from expert_shared_feed_forward_length.
+    // The count is not in the config; default to 1 but read the key if present.
+    hparams.n_expert_shared = 1;
+    ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,               hparams.n_expert_shared, false);
+    ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
+    if (hparams.n_ff_shexp == 0) {
+        // Weightless fixtures (test-llama-archs) omit this key; derive a nonzero
+        // size so the shared expert is still built. Real GGUFs always carry the
+        // exact value (routed and shared FF lengths may differ).
+        hparams.n_ff_shexp = hparams.n_ff_exp * hparams.n_expert_shared;
+    }
+
+    // Sliding-window attention is OPTIONAL. XS.2 is hybrid (full / SWA / SWA /
+    // SWA repeating, period 4 starting with full); M.1 has no sliding window
+    // (all layers full attention). When sliding_window is absent or zero we
+    // leave swa_type = NONE and skip the SWA-specific per-layer-type RoPE.
+    hparams.n_swa = 0;
+    ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
+    if (hparams.n_swa > 0) {
+        hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
+
+        uint32_t swa_period = 4;
+        ml.get_key_or_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, swa_period, false);
+        hparams.set_swa_pattern(swa_period, /*dense_first=*/true);  // XS.2: FULL at il%4==0
+
+        // Per-layer-type RoPE: full layers use YaRN θ=500000 over 64 dims;
+        // SWA layers use default RoPE θ=10000 over 128 dims. Base load_hparams
+        // already reads ROPE_FREQ_BASE and ROPE_DIMENSION_COUNT into the
+        // non-SWA fields; we explicitly pull the SWA mirrors here.
+        hparams.rope_freq_base_train_swa  = hparams.rope_freq_base_train;
+        hparams.rope_freq_scale_train_swa = 1.0f;  // SWA uses plain RoPE (no YaRN scaling); do NOT inherit full layers 1/factor
+        ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);
+        ml.get_key(LLM_KV_ROPE_DIMENSION_COUNT_SWA, hparams.n_rot_swa, false);
+    }
+
+    // Default the expert gating function to SIGMOID when the key is absent
+    // (matches the HF reference).
+    if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {
+        hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;
+    }
+
+    switch (hparams.n_layer()) {
+        case 40: type = LLM_TYPE_30B_A3B;   break;  // Laguna-XS.2
+        case 70: type = LLM_TYPE_230B_A10B; break;  // Laguna-M.1
+        default: type = LLM_TYPE_UNKNOWN;
+    }
+}
+
+void llama_model_laguna::load_arch_tensors(llama_model_loader & ml) {
+    LLAMA_LOAD_LOCALS;
+
+    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+    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) {
+        // tied embeddings fallback
+        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
+    }
+
+    const int64_t n_ff_exp   = hparams.n_ff_exp;
+    const int64_t n_ff_shexp = hparams.n_ff_shexp;
+
+    for (int i = 0; i < n_layer; ++i) {
+        auto & layer = layers[i];
+
+        // Per-layer head count — Laguna varies n_head between full and SWA
+        // layers (48 vs 64 in XS.2). KV head count is uniform.
+        const int64_t n_head_il    = hparams.n_head(i);
+        const int64_t n_head_kv_il = hparams.n_head_kv(i);
+        const int64_t n_embd_q_il  = n_embd_head_k * n_head_il;
+        const int64_t n_embd_k_il  = n_embd_head_k * n_head_kv_il;
+        const int64_t n_embd_v_il  = n_embd_head_v * n_head_kv_il;
+
+        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+
+        create_tensor_qkv(layer, i, n_embd, n_embd_q_il, n_embd_k_il, n_embd_v_il, 0);
+        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_q_il, n_embd}, 0);
+
+        layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0);
+        layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0);
+
+        // Attention output gate. XS.2 is per-head (g_proj -> n_head, one scalar
+        // per head broadcast over head_dim at multiply time); M.1 is per-element
+        // (g_proj -> n_head*head_dim, like afmoe). Detect from the stored tensor
+        // shape so a single arch handles both; the graph mirrors this check.
+        // Gate width selects per-head vs per-element. Real GGUFs always carry the
+        // gate tensor, so read the width from it and require EXACTLY one of the two
+        // valid widths -- never guess between them. Weightless fixtures
+        // (test-llama-archs) have no gate tensor; fall back to the per-head layout so
+        // the per-head reshape path is still exercised.
+        const int64_t n_gate_per_head = n_head_il;
+        const int64_t n_gate_per_elem = n_embd_head_k * n_head_il;
+        const ggml_tensor * gate_meta = ml.get_tensor_meta(tn(LLM_TENSOR_ATTN_GATE, "weight", i).str().c_str());
+        int64_t n_gate_out;
+        if (gate_meta != nullptr) {
+            n_gate_out = gate_meta->ne[1];
+            if (n_gate_out != n_gate_per_head && n_gate_out != n_gate_per_elem) {
+                GGML_ABORT("Laguna: unexpected attention gate width %lld at layer %d "
+                           "(expected %lld per-head or %lld per-element)",
+                           (long long) n_gate_out, i, (long long) n_gate_per_head, (long long) n_gate_per_elem);
+            }
+        } else {
+            n_gate_out = n_gate_per_head;
+        }
+        layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, n_gate_out}, 0);
+
+        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+        if ((uint32_t)i >= hparams.n_layer_dense_lead) {
+            // MoE layer
+            layer.ffn_gate_inp    = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,    "weight", i), {n_embd, n_expert}, 0);
+            layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias",   i), {n_expert}, 0);
+
+            layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd,   n_ff_exp, n_expert}, 0);
+            layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {n_embd,   n_ff_exp, n_expert}, 0);
+            layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd,   n_expert}, 0);
+
+            // Always-on shared expert.
+            layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd,    n_ff_shexp}, 0);
+            layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd,    n_ff_shexp}, 0);
+            layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd},    0);
+        } else {
+            // Dense layer (the leading n_layer_dense_lead layers)
+            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
+            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, 0);
+            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff,   n_embd}, 0);
+        }
+    }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_laguna::build_arch_graph(const llm_graph_params & params) const {
+    return std::make_unique<graph>(*this, params);
+}
+
+llama_model_laguna::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
+    const int64_t n_embd_head = hparams.n_embd_head_v();
+    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+
+    ggml_tensor * cur;
+    ggml_tensor * inpL;
+
+    inpL = build_inp_embd(model.tok_embd);
+    // No MuP embedding scale (laguna omits this; afmoe scales by sqrt(hidden)).
+
+    ggml_tensor * inp_pos = build_inp_pos();
+    // XS.2 is hybrid SWA -> interleaved-SWA KV input; M.1 is all-full -> plain
+    // KV input. Pick the matching input (and build_attn overload) per swa_type.
+    const bool has_swa = hparams.swa_type != LLAMA_SWA_TYPE_NONE;
+    llm_graph_input_attn_kv      * inp_attn_kv   = has_swa ? nullptr : build_attn_inp_kv();
+    llm_graph_input_attn_kv_iswa * inp_attn_iswa = has_swa ? build_attn_inp_kv_iswa() : nullptr;
+    ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+    const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
+
+    for (int il = 0; il < n_layer; ++il) {
+        const bool    is_swa_il   = hparams.is_swa(il);
+        const int64_t n_head_il   = hparams.n_head(il);
+        const int64_t n_head_kv_il = hparams.n_head_kv(il);
+
+        // Per-layer-type RoPE config. SWA layers run plain rope (no YaRN),
+        // achieved by zeroing the YaRN ext/beta params for those layers.
+        const int   n_rot_l       = is_swa_il ? hparams.n_rot_swa : n_rot;
+        const float freq_base_l   = is_swa_il ? hparams.rope_freq_base_train_swa : freq_base;
+        const float freq_scale_l  = is_swa_il ? hparams.rope_freq_scale_train_swa : freq_scale;
+        const float ext_factor_l  = is_swa_il ? 0.0f : ext_factor;
+        // YaRN magnitude scaling (mscale) is already handled by the framework:
+        // llama_context pre-divides cparams.yarn_attn_factor by (1 + 0.1*ln(factor))
+        // to cancel ggml rope_yarn's internal mscale *= 1 + 0.1*ln(1/freq_scale).
+        // Pass attn_factor straight through (like every other arch); SWA layers run
+        // plain RoPE (ext_factor 0, no mscale) so force 1.0 there.
+        const float attn_factor_l = is_swa_il ? 1.0f : attn_factor;
+        const float beta_fast_l   = is_swa_il ? 0.0f : beta_fast;
+        const float beta_slow_l   = is_swa_il ? 0.0f : beta_slow;
+        const int   n_ctx_orig_l  = is_swa_il ? hparams.n_ctx_train : n_ctx_orig;
+
+        ggml_tensor * inpSA = inpL;
+
+        // Pre-norm
+        cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);
+        cb(cur, "attn_norm", il);
+
+        // Self-attention
+        {
+            ggml_tensor * attn_inp = cur;  // saved for the gate projection
+
+            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
+                    n_embd_head, n_head_il, n_head_kv_il, il);
+
+            // g_proj on the *pre-attention* hidden state (matches HF
+            // reference: gate is computed from the same `hidden_states`
+            // input as q/k/v, not from the attn output).
+            ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp);
+            cb(gate, "attn_gate_proj", il);
+
+            // QK RMSNorm at head_dim level (Qwen3 style)
+            Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);
+            Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);
+            cb(Qcur, "Qcur_normed", il);
+            cb(Kcur, "Kcur_normed", il);
+
+            Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr,
+                    n_rot_l, rope_type, n_ctx_orig_l, freq_base_l, freq_scale_l,
+                    ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);
+            Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr,
+                    n_rot_l, rope_type, n_ctx_orig_l, freq_base_l, freq_scale_l,
+                    ext_factor_l, attn_factor_l, beta_fast_l, beta_slow_l);
+            cb(Qcur, "Qcur_rope", il);
+            cb(Kcur, "Kcur_rope", il);
+
+            cur = has_swa
+                ? build_attn(inp_attn_iswa,
+                        NULL, NULL, NULL,    // o_proj deferred until after gating
+                        Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il)
+                : build_attn(inp_attn_kv,
+                        NULL, NULL, NULL,
+                        Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+            cb(cur, "attn_out", il);
+
+            // Softplus output gate (the unary kernel computes softplus in fp32
+            // and casts back). Two shapes, distinguished by the g_proj output
+            // dim (matching the load-time detection):
+            //   XS.2 per-head     : gate [n_head_il, n_tokens] -> reshape to
+            //                       [1, n_head_il, n_tokens] and broadcast over
+            //                       head_dim against cur [head_dim, n_head, T].
+            //   M.1  per-element  : gate [n_head_il*head_dim, n_tokens] spans the
+            //                       full attention output -> direct ggml_mul.
+            gate = ggml_softplus(ctx0, gate);
+            cb(gate, "attn_gate_softplus", il);
+
+            const int64_t n_tokens = cur->ne[1];
+            if (model.layers[il].wqkv_gate->ne[1] == n_head_il) {
+                cur  = ggml_reshape_3d(ctx0, cur,  n_embd_head, n_head_il, n_tokens);
+                gate = ggml_reshape_3d(ctx0, gate, 1,           n_head_il, n_tokens);
+                cur  = ggml_mul(ctx0, cur, gate);
+                cur  = ggml_reshape_2d(ctx0, cur, n_embd_head * n_head_il, n_tokens);
+            } else {
+                cur = ggml_mul(ctx0, cur, gate);
+            }
+            cb(cur, "attn_gated", il);
+
+            cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);
+            cb(cur, "attn_o_proj", il);
+        }
+
+        if (il == n_layer - 1 && inp_out_ids) {
+            cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);
+            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+        }
+
+        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+        cb(ffn_inp, "ffn_inp", il);
+
+        // Pre-norm only (no post-attn norm)
+        cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);
+        cb(cur, "ffn_norm", il);
+
+        if ((uint32_t)il >= hparams.n_layer_dense_lead) {
+            // MoE: sigmoid routing + score-correction bias + sum-norm +
+            // routed_scaling_factor (all handled by build_moe_ffn).
+            ggml_tensor * moe_out = build_moe_ffn(cur,
+                    model.layers[il].ffn_gate_inp,
+                    model.layers[il].ffn_up_exps,
+                    model.layers[il].ffn_gate_exps,
+                    model.layers[il].ffn_down_exps,
+                    model.layers[il].ffn_exp_probs_b,
+                    n_expert, n_expert_used,
+                    LLM_FFN_SILU,
+                    hparams.expert_weights_norm,
+                    hparams.expert_weights_scale,
+                    (llama_expert_gating_func_type) hparams.expert_gating_func,
+                    il);
+            cb(moe_out, "ffn_moe_out", il);
+
+            // Always-on shared expert, summed in parallel.
+            ggml_tensor * ffn_shexp = build_ffn(cur,
+                    model.layers[il].ffn_up_shexp,   NULL, NULL,
+                    model.layers[il].ffn_gate_shexp, NULL, NULL,
+                    model.layers[il].ffn_down_shexp, NULL, NULL,
+                    NULL,
+                    LLM_FFN_SILU, LLM_FFN_PAR, il);
+            cb(ffn_shexp, "ffn_shexp", il);
+
+            cur = ggml_add(ctx0, moe_out, ffn_shexp);
+            cb(cur, "ffn_out", il);
+        } else {
+            // Dense FFN for the leading n_layer_dense_lead layers (XS.2: 1, M.1: 3)
+            cur = build_ffn(cur,
+                    model.layers[il].ffn_up,   NULL, NULL,
+                    model.layers[il].ffn_gate, NULL, NULL,
+                    model.layers[il].ffn_down, NULL, NULL,
+                    NULL,
+                    LLM_FFN_SILU, LLM_FFN_PAR, il);
+            cb(cur, "ffn_out", il);
+        }
+
+        // No post-ffn norm
+        cur = ggml_add(ctx0, cur, ffn_inp);
+        cur = build_cvec(cur, il);
+        cb(cur, "l_out", il);
+
+        inpL = cur;
+    }
+
+    cur = inpL;
+    cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);
+    cb(cur, "result_norm", -1);
+    res->t_embd = cur;
+
+    cur = build_lora_mm(model.output, cur);
+    cb(cur, "result_output", -1);
+    res->t_logits = cur;
+
+    ggml_build_forward_expand(gf, cur);
+}

src/models/models.h

@@ -1681,6 +1681,19 @@ struct llama_model_afmoe : public llama_model_base {
 };
 
 
+struct llama_model_laguna : public llama_model_base {
+    llama_model_laguna(const struct llama_model_params & params) : llama_model_base(params) {}
+    void load_arch_hparams(llama_model_loader & ml) override;
+    void load_arch_tensors(llama_model_loader & ml) override;
+
+    struct graph : public llm_graph_context {
+        graph(const llama_model & model, const llm_graph_params & params);
+    };
+
+    std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
+
 struct llama_model_ernie4_5 : public llama_model_base {
     llama_model_ernie4_5(const struct llama_model_params & params) : llama_model_base(params) {}
     void load_arch_hparams(llama_model_loader & ml) override;

tests/test-backend-ops.cpp

@@ -5960,6 +5960,7 @@ enum MoeGatingFunc {
     GATING_FUNC_SOFTMAX,
     GATING_FUNC_SIGMOID,
     GATING_FUNC_SOFTMAX_WEIGHT,
+    GATING_FUNC_SQRT_SOFTPLUS,
 };
 
 struct test_topk_moe : public test_case {
@@ -6003,7 +6004,8 @@ struct test_topk_moe : public test_case {
         ggml_tensor * logits = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data());
         ggml_tensor * probs            =
             (gating_func == GATING_FUNC_SOFTMAX) ? ggml_soft_max(ctx, logits) :
-            (gating_func == GATING_FUNC_SIGMOID) ? ggml_sigmoid(ctx, logits) : logits;
+            (gating_func == GATING_FUNC_SIGMOID) ? ggml_sigmoid(ctx, logits) :
+            (gating_func == GATING_FUNC_SQRT_SOFTPLUS) ? ggml_sqrt(ctx, ggml_softplus(ctx, logits)) : logits;
         ggml_set_name(probs, "probs");
 
         ggml_tensor * selection_probs = probs;
@@ -9584,7 +9586,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
 
-    for (auto gate : {GATING_FUNC_SOFTMAX, GATING_FUNC_SIGMOID, GATING_FUNC_SOFTMAX_WEIGHT}) {
+    for (auto gate : {GATING_FUNC_SOFTMAX, GATING_FUNC_SIGMOID, GATING_FUNC_SOFTMAX_WEIGHT, GATING_FUNC_SQRT_SOFTPLUS}) {
         for (bool with_norm : {false, true}) {
             for (bool bias_probs : {false, true}) {
                 for (float scale_w : {0.0f, 2.0f}) {
@@ -9596,6 +9598,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                     test_cases.emplace_back(new test_topk_moe({128, 1, 1, 1}, 128, with_norm, bias_probs, gate, scale_w));
                     test_cases.emplace_back(new test_topk_moe({129, 1, 1, 1}, 128, with_norm, bias_probs, gate, scale_w));
                     test_cases.emplace_back(new test_topk_moe({160, 4, 1, 1}, 160, with_norm, bias_probs, gate, scale_w));
+                    test_cases.emplace_back(new test_topk_moe({256, 22, 1, 1}, 6, with_norm, bias_probs, gate, scale_w)); // Used by DeepSeek-V4
                     test_cases.emplace_back(new test_topk_moe({288, 22, 1, 1}, 8, with_norm, bias_probs, gate, scale_w)); // Used by StepFun 3.7
                 }
             }

tests/test-chat-auto-parser.cpp

@@ -57,6 +57,15 @@ static void test_seed_oss_tool_with_reasoning(testing & t);
 static void test_nemotron_analysis(testing & t);
 static void test_nemotron_reasoning_detection(testing & t);
 static void test_nemotron_tool_format(testing & t);
+static void test_laguna_analysis(testing & t);
+static void test_laguna_reasoning_detection(testing & t);
+static void test_laguna_tool_format(testing & t);
+static void test_laguna_s_analysis(testing & t);
+static void test_laguna_s_reasoning_detection(testing & t);
+static void test_laguna_s_tool_format(testing & t);
+static void test_laguna_xs2_analysis(testing & t);
+static void test_laguna_xs2_reasoning_detection(testing & t);
+static void test_laguna_xs2_tool_format(testing & t);
 
 // CohereForAI template analysis tests
 static void test_cohere_reasoning_detection(testing & t);
@@ -101,6 +110,9 @@ int main(int argc, char * argv[]) {
     t.test("seed_oss_diffs", test_seed_oss_tool_analysis);
     t.test("cohere", test_cohere_analysis);
     t.test("nemotron", test_nemotron_analysis);
+    t.test("laguna", test_laguna_analysis);
+    t.test("laguna-s", test_laguna_s_analysis);
+    t.test("laguna-xs2", test_laguna_xs2_analysis);
     t.test("smollm3", test_smollm3_analysis);
     t.test("standard_json_tools", test_standard_json_tools_formats);
     t.test("normalize_quotes_to_json", test_normalize_quotes_to_json);
@@ -1378,6 +1390,94 @@ static void test_nemotron_tool_format(testing & t) {
     t.assert_true("should support tools", analysis.jinja_caps.supports_tools);
 }
 
+// ============================================================================
+// Laguna Template Analysis Tests
+// ============================================================================
+static common_chat_template load_laguna_template(testing & t) {
+    return load_template(t, "models/templates/poolside-Laguna-XS-2.1.jinja");
+}
+
+static void test_laguna_reasoning_detection(testing & t) {
+    common_chat_template tmpl = load_laguna_template(t);
+    struct autoparser analysis;
+    analysis.analyze_template(tmpl);
+    // Laguna's template renders reasoning delimiters with formatting whitespace
+    // ("<think>\n") that the model does not emit; the Laguna patch trims them.
+    t.assert_equal("reasoning_start should be '<think>'", "<think>", analysis.reasoning.start);
+    t.assert_equal("reasoning_end should be '</think>'", "</think>", analysis.reasoning.end);
+    t.assert_equal("reasoning should be TAG_BASED", reasoning_mode::TAG_BASED, analysis.reasoning.mode);
+}
+
+static void test_laguna_tool_format(testing & t) {
+    common_chat_template tmpl = load_laguna_template(t);
+    struct autoparser analysis;
+    analysis.analyze_template(tmpl);
+    t.assert_equal("arg_value_suffix should be '</arg_value>'", "</arg_value>", analysis.tools.arguments.value_suffix);
+}
+
+static void test_laguna_stop_string(testing & t) {
+    // The </assistant> turn terminator can be emitted as ordinary text tokens
+    // (not the single eot token), so it must also be a literal stop string.
+    common_chat_template tmpl = load_laguna_template(t);
+    struct autoparser analysis;
+    analysis.analyze_template(tmpl);
+    bool has_stop = false;
+    for (const auto & stop : analysis.additional_stops) {
+        if (stop == "</assistant>") { has_stop = true; break; }
+    }
+    t.assert_true("Laguna additional_stops contains </assistant>", has_stop);
+}
+
+static void test_laguna_analysis(testing & t) {
+    t.test("Laguna reasoning detection", test_laguna_reasoning_detection);
+    t.test("Laguna tool format", test_laguna_tool_format);
+    t.test("Laguna stop string", test_laguna_stop_string);
+}
+
+static common_chat_template load_laguna_s_template(testing & t) {
+    return load_template(t, "models/templates/poolside-Laguna-S-2.1.jinja");
+}
+static void test_laguna_s_reasoning_detection(testing & t) {
+    common_chat_template tmpl = load_laguna_s_template(t);
+    struct autoparser analysis;
+    analysis.analyze_template(tmpl);
+    t.assert_equal("Laguna-S(v8) reasoning_start should be '<think>'", "<think>", analysis.reasoning.start);
+    t.assert_equal("Laguna-S(v8) reasoning_end should be '</think>'", "</think>", analysis.reasoning.end);
+    t.assert_equal("Laguna-S(v8) reasoning should be TAG_BASED", reasoning_mode::TAG_BASED, analysis.reasoning.mode);
+}
+static void test_laguna_s_tool_format(testing & t) {
+    common_chat_template tmpl = load_laguna_s_template(t);
+    struct autoparser analysis;
+    analysis.analyze_template(tmpl);
+    t.assert_equal("Laguna-S(v8) arg_value_suffix should be '</arg_value>'", "</arg_value>", analysis.tools.arguments.value_suffix);
+}
+static void test_laguna_s_analysis(testing & t) {
+    t.test("Laguna-S(v8) reasoning detection", test_laguna_s_reasoning_detection);
+    t.test("Laguna-S(v8) tool format", test_laguna_s_tool_format);
+}
+
+static common_chat_template load_laguna_xs2_template(testing & t) {
+    return load_template(t, "models/templates/poolside-Laguna-XS.2.jinja");
+}
+static void test_laguna_xs2_reasoning_detection(testing & t) {
+    common_chat_template tmpl = load_laguna_xs2_template(t);
+    struct autoparser analysis;
+    analysis.analyze_template(tmpl);
+    t.assert_equal("Laguna-XS.2(v5) reasoning_start should be '<think>'", "<think>", analysis.reasoning.start);
+    t.assert_equal("Laguna-XS.2(v5) reasoning_end should be '</think>'", "</think>", analysis.reasoning.end);
+    t.assert_equal("Laguna-XS.2(v5) reasoning should be TAG_BASED", reasoning_mode::TAG_BASED, analysis.reasoning.mode);
+}
+static void test_laguna_xs2_tool_format(testing & t) {
+    common_chat_template tmpl = load_laguna_xs2_template(t);
+    struct autoparser analysis;
+    analysis.analyze_template(tmpl);
+    t.assert_equal("Laguna-XS.2(v5) arg_value_suffix should be '</arg_value>'", "</arg_value>", analysis.tools.arguments.value_suffix);
+}
+static void test_laguna_xs2_analysis(testing & t) {
+    t.test("Laguna-XS.2(v5) reasoning detection", test_laguna_xs2_reasoning_detection);
+    t.test("Laguna-XS.2(v5) tool format", test_laguna_xs2_tool_format);
+}
+
 static common_chat_template load_cohere_template(testing & t) {
     return load_template(t, "models/templates/CohereForAI-c4ai-command-r7b-12-2024-tool_use.jinja");
 }

tests/test-chat.cpp

@@ -109,6 +109,15 @@ static void assert_contains(const std::string & haystack, const std::string & ne
     }
 }
 
+static void assert_not_contains(const std::string & haystack, const std::string & needle) {
+    if (haystack.find(needle) != std::string::npos) {
+        LOG_ERR("Expected NOT to contain: %s\n", needle.c_str());
+        LOG_ERR("Actual: %s\n", haystack.c_str());
+        common_log_flush(common_log_main());
+        throw std::runtime_error("Test failed");
+    }
+}
+
 static void assert_ends_with(const std::string & str, const std::string & suffix) {
     if (str.size() < suffix.size() ||
         str.compare(str.size() - suffix.size(), suffix.size(), suffix) != 0) {
@@ -4016,6 +4025,132 @@ static void test_template_output_peg_parsers(bool detailed_debug) {
             .run();
     }
 
+    // DeepSeek V4 tests - same DSML markup as V3.2, but the tool call block is named
+    // "tool_calls" and the non-thinking generation prompt ends in a bare </think>
+    // instead of an empty <think></think> pair.
+    {
+        auto tst = peg_tester("models/templates/deepseek-ai-DeepSeek-V4.jinja", detailed_debug);
+
+        // Pure content (non-thinking mode; generation prompt ends with </think>)
+        tst.test("Hello, world!\nWhat's up?")
+            .enable_thinking(false)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .expect(message_assist)
+            .run();
+
+        // Thinking + content
+        tst.test("I'm\nthinking</think>Hello, world!\nWhat's up?")
+            .enable_thinking(true)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .expect(message_assist_thoughts)
+            .run();
+
+        // Thinking + tool call (single, string param)
+        tst.test(
+               "Let me check the time</think>\n\n"
+               "<｜DSML｜tool_calls>\n"
+               "<｜DSML｜invoke name=\"get_time\">\n"
+               "<｜DSML｜parameter name=\"city\" string=\"true\">Tokyo</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "</｜DSML｜tool_calls>")
+            .enable_thinking(true)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .tools({ get_time_tool })
+            .expect(message_with_tool_calls_and_reasoning("get_time", R"({"city": "Tokyo"})", "Let me check the time"))
+            .run();
+
+        // Tool call without reasoning (non-thinking mode), integer param (string="false")
+        tst.test(
+               "<｜DSML｜tool_calls>\n"
+               "<｜DSML｜invoke name=\"special_function\">\n"
+               "<｜DSML｜parameter name=\"arg1\" string=\"false\">1</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "</｜DSML｜tool_calls>")
+            .enable_thinking(false)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .tools({ special_function_tool })
+            .expect(message_assist_call)
+            .run();
+
+        // Multiple parallel tool calls with reasoning
+        tst.test(
+               "Calling both</think>\n\n"
+               "<｜DSML｜tool_calls>\n"
+               "<｜DSML｜invoke name=\"get_time\">\n"
+               "<｜DSML｜parameter name=\"city\" string=\"true\">Paris</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "<｜DSML｜invoke name=\"get_weather\">\n"
+               "<｜DSML｜parameter name=\"city\" string=\"true\">Paris</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "</｜DSML｜tool_calls>")
+            .enable_thinking(true)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .parallel_tool_calls(true)
+            .tools({ get_time_tool, get_weather_tool })
+            .expect(message_with_reasoning_content_and_multiple_tool_calls(
+                "Calling both", "",
+                { { "get_time", R"({"city": "Paris"})" }, { "get_weather", R"({"city": "Paris"})" } }))
+            .run();
+
+        // Tool call with content before tool calls
+        tst.test(
+               "Thinking about it</think>"
+               "Let me call the function.\n\n"
+               "<｜DSML｜tool_calls>\n"
+               "<｜DSML｜invoke name=\"special_function\">\n"
+               "<｜DSML｜parameter name=\"arg1\" string=\"false\">1</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "</｜DSML｜tool_calls>")
+            .enable_thinking(true)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .tools({ special_function_tool })
+            .expect_reasoning("Thinking about it")
+            .expect_content("Let me call the function.")
+            .expect_tool_calls({
+                { "special_function", R"({"arg1": 1})", {} },
+            })
+            .run();
+
+        // Tool call with multiple params (mixed types)
+        tst.test(
+               "Multi-arg call</think>\n\n"
+               "<｜DSML｜tool_calls>\n"
+               "<｜DSML｜invoke name=\"magic_int\">\n"
+               "<｜DSML｜parameter name=\"ref\" string=\"false\">42</｜DSML｜parameter>\n"
+               "<｜DSML｜parameter name=\"name\" string=\"true\">foo bar</｜DSML｜parameter>\n"
+               "</｜DSML｜invoke>\n"
+               "</｜DSML｜tool_calls>")
+            .enable_thinking(true)
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .tools({ magic_int_tool })
+            .expect_reasoning("Multi-arg call")
+            .expect_tool_calls({
+                { "magic_int", R"({"ref": 42, "name": "foo bar"})", {} },
+            })
+            .run();
+
+        // Continuation tests
+        tst.test("world!\nWhat's up?")
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .enable_thinking(true)
+            .messages({ message_user, message_assist_prefill_content })
+            .add_generation_prompt(false)
+            .continue_final_message(COMMON_CHAT_CONTINUATION_CONTENT)
+            .expect_reasoning("I'm thinking")
+            .expect_content("Hello, world!\nWhat's up?")
+            .run();
+
+        tst.test(" thinking</think>Hello, world!\nWhat's up?")
+            .reasoning_format(COMMON_REASONING_FORMAT_DEEPSEEK)
+            .enable_thinking(true)
+            .messages({ message_user, message_assist_prefill_reasoning })
+            .add_generation_prompt(false)
+            .continue_final_message(COMMON_CHAT_CONTINUATION_REASONING)
+            .expect_reasoning("I'm thinking")
+            .expect_content("Hello, world!\nWhat's up?")
+            .run();
+    }
+
     // GLM-4.6 tests - format: <tool_call>function_name\n<arg_key>...</arg_key>\n<arg_value>...</arg_value>\n</tool_call>
     {
         auto tst = peg_tester("models/templates/GLM-4.6.jinja", detailed_debug);
@@ -5918,6 +6053,144 @@ static void test_developer_role_to_system_workaround() {
     }
 }
 
+// Verify reasoning-trace retention rules in the DeepSeek-V4 template:
+// all traces are retained unless drop_thinking is true AND the conversation
+// has no tool calls, in which case only the last (after-final-user) trace is
+// kept and earlier ones are dropped.
+static void test_deepseek_v4_thinking_retention() {
+    LOG_DBG("%s\n", __func__);
+
+    auto tmpls = read_templates("models/templates/deepseek-ai-DeepSeek-V4.jinja");
+
+    common_chat_msg user_q1; user_q1.role = "user"; user_q1.content = "Question 1";
+    common_chat_msg user_q2; user_q2.role = "user"; user_q2.content = "Question 2";
+    common_chat_msg asst_a1 = simple_assist_msg("Answer 1", "thinking A1");
+    common_chat_msg asst_a2 = simple_assist_msg("Answer 2", "thinking A2");
+
+    common_chat_msg tool_assist = message_with_tool_calls("special_function", "{\"arg1\": 1}");
+    common_chat_msg tool_result; tool_result.role = "tool";
+    tool_result.tool_name = "special_function"; tool_result.tool_call_id = "0"; tool_result.content = "result";
+
+    // The template uses U+FF5C as the role separator and literal think tags
+    // for the reasoning block.
+    const std::string asst_marker   = "<\xef\xbd\x9c" "Assistant" "\xef\xbd\x9c>";
+    // Built via concatenation so the thinking tokens are not interpreted by
+    // tooling processing this source file.
+    const std::string think_start   = "<" "think" ">";
+    const std::string think_end     = "</" "think" ">";
+
+    const std::string think_a1      = asst_marker + think_start + "thinking A1" + think_end;
+    const std::string think_a2      = asst_marker + think_start + "thinking A2" + think_end;
+    const std::string asst_no_think = asst_marker + think_end;
+
+    auto render = [&](const std::vector<common_chat_msg> & messages, bool drop_thinking) {
+        common_chat_templates_inputs inputs;
+        inputs.messages = messages;
+        inputs.add_generation_prompt = false;
+        inputs.chat_template_kwargs["thinking"]     = "true";
+        inputs.chat_template_kwargs["drop_thinking"] = drop_thinking ? "true" : "false";
+        return common_chat_templates_apply(tmpls.get(), inputs).prompt;
+    };
+
+    // No tools, drop_thinking=false: all reasoning is retained.
+    {
+        auto prompt = render({ user_q1, asst_a1, user_q2, asst_a2 }, /* drop_thinking = */ false);
+        assert_contains(prompt, think_a1);
+        assert_contains(prompt, think_a2);
+    }
+
+    // No tools, drop_thinking=true: only the last reasoning trace is kept,
+    // earlier ones are dropped (the assistant block emits just the end token).
+    {
+        auto prompt = render({ user_q1, asst_a1, user_q2, asst_a2 }, /* drop_thinking = */ true);
+        assert_not_contains(prompt, think_a1);
+        assert_contains(prompt, think_a2);
+        // The dropped assistant turn still opens with the marker + bare end token.
+        assert_contains(prompt, asst_no_think + "Answer 1");
+    }
+
+    // Single assistant turn, drop_thinking=true: the only trace is the last
+    // one, so it must be retained even with drop_thinking set.
+    {
+        auto prompt = render({ user_q1, asst_a1 }, /* drop_thinking = */ true);
+        assert_contains(prompt, think_a1);
+    }
+
+    // Single assistant turn, drop_thinking=false: reasoning is retained.
+    {
+        auto prompt = render({ user_q1, asst_a1 }, /* drop_thinking = */ false);
+        assert_contains(prompt, think_a1);
+    }
+
+    // With tool calls, drop_thinking=true: tool presence forces all reasoning
+    // to be retained, including the pre-tool-call trace.
+    {
+        auto prompt = render({ user_q1, asst_a1, user_q2, tool_assist, tool_result, asst_a2 },
+                             /* drop_thinking = */ true);
+        assert_contains(prompt, think_a1);
+        assert_contains(prompt, think_a2);
+    }
+
+    // With tool calls, drop_thinking=false: all reasoning retained.
+    {
+        auto prompt = render({ user_q1, asst_a1, user_q2, tool_assist, tool_result, asst_a2 },
+                             /* drop_thinking = */ false);
+        assert_contains(prompt, think_a1);
+        assert_contains(prompt, think_a2);
+    }
+}
+
+// Verify that consecutive tool results are rendered in the tool call order of the
+// preceding assistant message (matched by tool call id), as required by the reference
+// DeepSeek-V4 implementation.
+static void test_deepseek_v4_tool_result_ordering() {
+    LOG_DBG("%s\n", __func__);
+
+    auto tmpls = read_templates("models/templates/deepseek-ai-DeepSeek-V4.jinja");
+
+    common_chat_msg user_q; user_q.role = "user"; user_q.content = "Question";
+
+    common_chat_msg assist_calls;
+    assist_calls.role = "assistant";
+    assist_calls.tool_calls.push_back({ "get_time",    "{\"city\": \"Paris\"}", "call_1" });
+    assist_calls.tool_calls.push_back({ "get_weather", "{\"city\": \"Paris\"}", "call_2" });
+
+    common_chat_msg time_result; time_result.role = "tool";
+    time_result.tool_name = "get_time"; time_result.tool_call_id = "call_1"; time_result.content = "12:00";
+    common_chat_msg weather_result; weather_result.role = "tool";
+    weather_result.tool_name = "get_weather"; weather_result.tool_call_id = "call_2"; weather_result.content = "sunny";
+
+    auto render = [&](const std::vector<common_chat_msg> & messages) {
+        common_chat_templates_inputs inputs;
+        inputs.messages = messages;
+        inputs.add_generation_prompt = false;
+        return common_chat_templates_apply(tmpls.get(), inputs).prompt;
+    };
+
+    // Results sent out of order are reordered to match the tool call order.
+    {
+        auto prompt = render({ user_q, assist_calls, weather_result, time_result });
+        assert_contains(prompt, "<tool_result>12:00</tool_result>\n\n<tool_result>sunny</tool_result>");
+    }
+
+    // Results already in call order stay put.
+    {
+        auto prompt = render({ user_q, assist_calls, time_result, weather_result });
+        assert_contains(prompt, "<tool_result>12:00</tool_result>\n\n<tool_result>sunny</tool_result>");
+    }
+
+    // Without tool call ids there is nothing to match against; order is preserved.
+    {
+        auto no_id_calls = assist_calls;
+        no_id_calls.tool_calls[0].id = "";
+        no_id_calls.tool_calls[1].id = "";
+        auto no_id_weather = weather_result; no_id_weather.tool_call_id = "";
+        auto no_id_time    = time_result;    no_id_time.tool_call_id = "";
+        auto prompt = render({ user_q, no_id_calls, no_id_weather, no_id_time });
+        assert_contains(prompt, "<tool_result>sunny</tool_result>\n\n<tool_result>12:00</tool_result>");
+    }
+}
+
 static void test_reasoning_budget_tokens_per_request() {
     LOG_DBG("%s\n", __func__);
     // Use Qwen3 template which has <think>...</think> reasoning markers.
@@ -6139,6 +6412,8 @@ int main(int argc, char ** argv) {
         test_tools_oaicompat_json_conversion();
         test_convert_responses_to_chatcmpl();
         test_developer_role_to_system_workaround();
+        test_deepseek_v4_thinking_retention();
+        test_deepseek_v4_tool_result_ordering();
         test_template_generation_prompt();
         test_reasoning_budget_tokens_per_request();
         test_reasoning_budget_message_per_request();

tests/test-llama-archs.cpp

@@ -362,6 +362,7 @@ static bool moe_mandatory(const llm_arch arch) {
         case LLM_ARCH_STEP35:
         case LLM_ARCH_MISTRAL4:
         case LLM_ARCH_MELLUM:
+        case LLM_ARCH_LAGUNA:
             return true;
         default:
             return false;

tools/mtmd/clip.cpp

@@ -158,6 +158,13 @@ struct clip_ctx {
     ggml_backend_t backend_cpu = nullptr;
     ggml_backend_buffer_ptr buf;
 
+    // on-demand device (VRAM) residency: the vision/audio encoder weights are read-only, so a host
+    // shadow is captured once and the device buffer can be freed while the model is cold, then
+    // rebuilt on wake (mirrors llama_model::release_device_weights). See clip_release_device().
+    ggml_backend_buffer_type_t dev_buft = nullptr;
+    std::vector<uint8_t>       dev_shadow;
+    bool                       dev_released = false;
+
 
     int max_nodes = 8192;
     ggml_backend_sched_ptr sched;
@@ -3241,6 +3248,74 @@ void clip_free(clip_ctx * ctx) {
     delete ctx;
 }
 
+void clip_release_device(struct clip_ctx * ctx) {
+    if (ctx == nullptr || ctx->dev_released) {
+        return;
+    }
+    ggml_backend_buffer_t buf = ctx->buf.get();
+    if (buf == nullptr || ggml_backend_buffer_is_host(buf) || ggml_backend_buffer_get_size(buf) == 0) {
+        return; // CPU-backed encoder: nothing in VRAM to free
+    }
+    // ensure no encode is in flight before freeing the weights
+    if (ctx->sched) {
+        ggml_backend_sched_synchronize(ctx->sched.get());
+    }
+    ctx->dev_buft = ggml_backend_buffer_get_type(buf);
+
+    ggml_context * cd = ctx->ctx_data.get();
+    // capture the host shadow once (weights are read-only): compact, stable iteration order,
+    // skipping view tensors (which alias a base and are restored implicitly)
+    if (ctx->dev_shadow.empty()) {
+        size_t total = 0;
+        for (ggml_tensor * t = ggml_get_first_tensor(cd); t != nullptr; t = ggml_get_next_tensor(cd, t)) {
+            if (t->view_src == nullptr) { total += ggml_nbytes(t); }
+        }
+        ctx->dev_shadow.resize(total);
+        size_t off = 0;
+        for (ggml_tensor * t = ggml_get_first_tensor(cd); t != nullptr; t = ggml_get_next_tensor(cd, t)) {
+            if (t->view_src != nullptr) { continue; }
+            const size_t n = ggml_nbytes(t);
+            ggml_backend_tensor_get(t, ctx->dev_shadow.data() + off, 0, n);
+            off += n;
+        }
+    }
+
+    // free the device buffer and clear the now-dangling tensor pointers so restore reallocates cleanly
+    ctx->buf.reset();
+    for (ggml_tensor * t = ggml_get_first_tensor(cd); t != nullptr; t = ggml_get_next_tensor(cd, t)) {
+        t->buffer = nullptr;
+        t->data   = nullptr;
+    }
+    ctx->dev_released = true;
+}
+
+bool clip_restore_device(struct clip_ctx * ctx) {
+    if (ctx == nullptr || !ctx->dev_released) {
+        return true;
+    }
+    ggml_context * cd = ctx->ctx_data.get();
+    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(cd, ctx->dev_buft);
+    if (buf == nullptr) {
+        LOG_ERR("%s: failed to reallocate encoder device buffer (out of VRAM?)\n", __func__);
+        return false;
+    }
+    ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+    size_t off = 0;
+    for (ggml_tensor * t = ggml_get_first_tensor(cd); t != nullptr; t = ggml_get_next_tensor(cd, t)) {
+        if (t->view_src != nullptr) { continue; }
+        const size_t n = ggml_nbytes(t);
+        ggml_backend_tensor_set(t, ctx->dev_shadow.data() + off, 0, n);
+        off += n;
+    }
+    ctx->buf.reset(buf);
+    ctx->dev_released = false;
+    return true;
+}
+
+bool clip_weights_resident(const struct clip_ctx * ctx) {
+    return ctx == nullptr || !ctx->dev_released;
+}
+
 const char * clip_patch_merge_type(const struct clip_ctx * ctx) {
     return ctx->model.hparams.mm_patch_merge_type == PATCH_MERGE_SPATIAL_UNPAD ? "spatial_unpad" : "flat";
 }

tools/mtmd/clip.h

@@ -67,6 +67,12 @@ struct clip_init_result clip_init(const char * fname, struct clip_context_params
 
 void clip_free(struct clip_ctx * ctx);
 
+// on-demand device (VRAM) residency: free/rebuild the encoder weight buffer to shrink an idle
+// (cold) multimodal model's VRAM footprint. No-op for a CPU-backed encoder. See clip.cpp.
+void clip_release_device(struct clip_ctx * ctx);
+bool clip_restore_device(struct clip_ctx * ctx);
+bool clip_weights_resident(const struct clip_ctx * ctx);
+
 // TODO: should be enum, not string
 const char * clip_patch_merge_type(const struct clip_ctx * ctx);
 

tools/mtmd/models/qwen3vl.cpp

@@ -37,7 +37,7 @@ ggml_cgraph * clip_graph_qwen3vl::build() {
     }
 
     // calculate absolute position embedding and apply
-    ggml_tensor * learned_pos_embd = resize_position_embeddings();
+    ggml_tensor * learned_pos_embd = resize_position_embeddings(GGML_SCALE_MODE_BILINEAR | GGML_SCALE_FLAG_ALIGN_CORNERS);
     learned_pos_embd = ggml_cont_4d(
         ctx0, learned_pos_embd,
         n_embd * 2, n_patches_x / 2, n_patches_y, batch_size);

tools/mtmd/mtmd-helper.cpp

@@ -238,6 +238,29 @@ struct decode_embd_batch {
     }
 };
 
+// Helper class to set non-causal attention via RAII
+class scope_non_causal {
+public:
+    scope_non_causal(llama_context * context, bool enabled) : context_(context), enabled_(enabled) {
+        if (enabled_) {
+            // TODO @ngxson : need to make sure only one image is processed at a time, and n_ubatch must be enough to hold the image
+            llama_set_causal_attn(context_, false);
+        }
+    }
+    ~scope_non_causal() {
+        if (enabled_) {
+            llama_set_causal_attn(context_, true);
+        }
+    }
+
+    scope_non_causal(const scope_non_causal &) = delete;
+    scope_non_causal & operator=(const scope_non_causal &) = delete;
+
+private:
+    llama_context * context_;
+    bool enabled_;
+};
+
 // Helper function for decoding an image whose embeddings have already been calculated
 int32_t mtmd_helper_decode_image_chunk(
         mtmd_context * ctx,
@@ -288,10 +311,7 @@ int32_t mtmd_helper_decode_image_chunk(
     }
 
     const bool use_non_causal = mtmd_decode_use_non_causal(ctx, chunk);
-    if (use_non_causal) {
+    const scope_non_causal non_causal(lctx, use_non_causal);
-        llama_set_causal_attn(lctx, false);
-        // TODO @ngxson : need to make sure only one image is processed at a time, and n_ubatch must be enough to hold the image
-    }
 
     while (i_batch < n_img_batches) { // split into batches
         int pos_offset = i_batch*n_batch;
@@ -304,9 +324,6 @@ int32_t mtmd_helper_decode_image_chunk(
         int32_t ret = llama_decode(lctx, batch_embd_view);
         if (ret != 0) {
             LOG_ERR("failed to decode %s\n", name);
-            if (use_non_causal) {
-                llama_set_causal_attn(lctx, true);
-            }
             return ret;
         }
 
@@ -314,9 +331,6 @@ int32_t mtmd_helper_decode_image_chunk(
             ret = callback(batch_embd_view, user_data);
             if (ret != 0) {
                 LOG_ERR("post-decode callback failed\n");
-                if (use_non_causal) {
-                    llama_set_causal_attn(lctx, true);
-                }
                 return ret;
             }
         }
@@ -329,9 +343,6 @@ int32_t mtmd_helper_decode_image_chunk(
     n_past += mtmd_input_chunk_get_n_pos(chunk);
     *new_n_past = n_past;
 
-    if (use_non_causal) {
-        llama_set_causal_attn(lctx, true);
-    }
     return 0;
 }
 

tools/mtmd/mtmd.cpp

@@ -806,6 +806,24 @@ void mtmd_free(mtmd_context * ctx) {
     delete ctx;
 }
 
+void mtmd_release_device(mtmd_context * ctx) {
+    if (ctx == nullptr) {
+        return;
+    }
+    if (ctx->ctx_v) { clip_release_device(ctx->ctx_v); }
+    if (ctx->ctx_a) { clip_release_device(ctx->ctx_a); }
+}
+
+bool mtmd_restore_device(mtmd_context * ctx) {
+    if (ctx == nullptr) {
+        return true;
+    }
+    bool ok = true;
+    if (ctx->ctx_v) { ok = clip_restore_device(ctx->ctx_v) && ok; }
+    if (ctx->ctx_a) { ok = clip_restore_device(ctx->ctx_a) && ok; }
+    return ok;
+}
+
 struct mtmd_tokenizer {
     mtmd_context * ctx;
 

tools/mtmd/mtmd.h

@@ -127,6 +127,12 @@ MTMD_API mtmd_context * mtmd_init_from_file(const char * mmproj_fname,
 
 MTMD_API void mtmd_free(mtmd_context * ctx);
 
+// on-demand device (VRAM) residency: free / rebuild the vision+audio encoder weight buffers so an
+// idle (cold) multimodal model releases its encoder VRAM and reclaims it on wake. No-op for a
+// CPU-backed encoder. Restore returns false if reallocation failed (out of VRAM).
+MTMD_API void mtmd_release_device(mtmd_context * ctx);
+MTMD_API bool mtmd_restore_device(mtmd_context * ctx);
+
 // whether we need to set non-causal mask before llama_decode
 // if chunk is nullptr, we assume the default case where chunk is an image chunk
 MTMD_API bool mtmd_decode_use_non_causal(const mtmd_context * ctx, const mtmd_input_chunk * chunk);

tools/server/server-context.cpp

@@ -25,6 +25,8 @@
 #include <filesystem>
 #include <utility>
 #include <fstream>
+#include <thread>
+#include <atomic>
 
 // fix problem with std::min and std::max
 #if defined(_WIN32)
@@ -35,6 +37,16 @@
 #include <windows.h>
 #endif
 
+// POSIX file locking + inotify doorbell for the cross-process VRAM arbiter (see vram_share_* below)
+#if !defined(_WIN32)
+#include <sys/file.h>
+#include <sys/stat.h>
+#include <sys/inotify.h>
+#include <fcntl.h>
+#include <unistd.h>
+#include <poll.h>
+#endif
+
 using json = nlohmann::ordered_json;
 
 constexpr int HTTP_POLLING_SECONDS = 1;
@@ -768,7 +780,36 @@ struct server_slot {
         // TODO @ngxson : move this log line to debug when it become more stable
         SLT_TRC(*this, "encoding mtmd batch from idx = %zu, n_chunks = %d\n", idx, n_added);
 
+        // Bring the vision/audio encoder into VRAM just for this encode. In on-demand mode the
+        // encoder is normally kept in RAM (its VRAM freed for KV / expert cache). If it does not
+        // fit, evict the LLM backbone weights first: they are NOT needed while the encoder runs (the
+        // decode that uses them happens afterwards) and their host shadow is read-only, so this is a
+        // cheap free (no D2H) and leaves the KV cache / prompt cache untouched.
+        static const bool mmproj_ondemand = getenv("LLAMA_MMPROJ_ONDEMAND") != nullptr;
+        bool weights_evicted = false;
+        if (mctx && !mtmd_restore_device(mctx)) {
+            if (mmproj_ondemand) {
+                llama_context_release_device(ctx_tgt, /* evict_kv = */ false); // free backbone, keep KV
+                weights_evicted = true;
+            }
+            if (!mtmd_restore_device(mctx)) {
+                if (weights_evicted) { llama_context_restore_device(ctx_tgt); }
+                SLT_ERR(*this, "%s", "failed to bring the multimodal encoder into VRAM for encoding\n");
+                return -1;
+            }
+        }
+
         res = mtmd_batch_encode(mbatch.get());
+
+        // release the encoder VRAM again (on-demand), then restore the backbone weights for decode.
+        // Order matters: free the encoder BEFORE re-uploading the weights so the peak stays within VRAM.
+        if (mctx && mmproj_ondemand) {
+            mtmd_release_device(mctx);
+        }
+        if (weights_evicted) {
+            llama_context_restore_device(ctx_tgt);
+        }
+
         if (res != 0) {
             SLT_ERR(*this, "failed to encode mtmd batch for chunk idx = %zu, res = %d\n", idx, res);
             return -1;
@@ -871,6 +912,8 @@ public:
     }
 
     ~server_context_impl() {
+        // stop the VRAM warden thread (and release the token) before tearing anything down
+        vram_share_shutdown();
         if (!sleeping) {
             // destroy() is already called when entering sleeping state
             // we don't call it again here to avoid double free
@@ -894,6 +937,198 @@ private:
     llama_model   * model_dft = nullptr;
     llama_context * ctx_dft   = nullptr;
 
+    // Cross-process VRAM arbiter: when LLAMA_SLEEP_VRAM_ONLY is set, several always-loaded
+    // llama-server processes time-share one GPU. A single flock() on <arena>/token.lock is the
+    // baton: "resident (weights in VRAM) iff I hold the token". A model warms up (locks + restores
+    // weights) right before it decodes, and goes cold (releases weights + unlocks) when it idles,
+    // so only one model holds VRAM at a time and the KV cache is never evicted (no re-prefill).
+    bool vram_only    = false;   // release-mode sleep active (LLAMA_SLEEP_VRAM_ONLY)
+    bool vram_evict_kv = false;  // also evict the KV cache to host on release (LLAMA_SLEEP_EVICT_KV)
+    bool vram_flock   = false;   // cross-process flock coordination available
+    std::atomic<bool> vram_cold{true}; // weights currently released (read by warden thread)
+    int  vram_lock_fd = -1;      // fd for <arena>/token.lock
+
+    // inotify "doorbell": a waiter that wants the VRAM token touches <arena>/doorbell/<pid>, which
+    // wakes the current holder's warden thread so it releases immediately instead of only on idle.
+    std::string       vram_doorbell_dir;
+    std::string       vram_pid_str;
+    int               vram_inotify_fd = -1;
+    std::thread       vram_warden;
+    std::atomic<bool> vram_warden_run{false};
+
+    // open the shared arena (flock token + doorbell dir). Idempotent; sets vram_only/vram_flock.
+    void vram_arena_open() {
+        if (getenv("LLAMA_SLEEP_VRAM_ONLY") == nullptr) {
+            return;
+        }
+        vram_only = true;
+        vram_evict_kv = getenv("LLAMA_SLEEP_EVICT_KV") != nullptr;
+#if !defined(_WIN32)
+        if (vram_lock_fd >= 0) {
+            return; // already open
+        }
+        const char * arena_env = getenv("LLAMA_VRAM_ARENA");
+        const std::string arena = arena_env ? arena_env : "/dev/shm/llama-vram";
+        mkdir(arena.c_str(), 0777);
+        const std::string lock_path = arena + "/token.lock";
+        vram_lock_fd = open(lock_path.c_str(), O_RDWR | O_CREAT, 0666);
+        if (vram_lock_fd >= 0) {
+            vram_flock = true;
+            vram_pid_str = std::to_string(getpid());
+            vram_doorbell_dir = arena + "/doorbell";
+            mkdir(vram_doorbell_dir.c_str(), 0777);
+            SRV_INF("VRAM arbiter: cross-process GPU sharing via %s (doorbell %s)\n",
+                    lock_path.c_str(), vram_doorbell_dir.c_str());
+        } else {
+            SRV_WRN("VRAM arbiter: cannot open %s, running VRAM-only sleep without cross-process lock\n", lock_path.c_str());
+        }
+#endif
+    }
+
+    // Acquire the VRAM token BEFORE uploading this model's weights, so any model currently resident
+    // on the GPU releases first and the load uploads into free VRAM instead of racing it (which
+    // could OOM, e.g. the 4B task model holding VRAM while a large model loads). Blocks until free.
+    // Called from load_model() right before common_init_from_params().
+    void vram_acquire_for_load() {
+        vram_arena_open();
+        if (!vram_only) {
+            return;
+        }
+#if !defined(_WIN32)
+        if (vram_flock) {
+            vram_ring_doorbell();         // nudge whoever is resident to release
+            flock(vram_lock_fd, LOCK_EX); // block until the GPU is free
+        }
+#endif
+        vram_cold = false; // we hold the token; weights will be resident after the upload
+    }
+
+    // Start the doorbell warden. Called from init() after the (coordinated) load. The model stays
+    // warm (holding the token acquired in vram_acquire_for_load) and serves its first request
+    // without a re-warm; the warden releases it when another model rings the doorbell.
+    void vram_share_init() {
+        vram_arena_open();
+        if (!vram_only) {
+            return;
+        }
+        vram_cold = false; // resident and holding the token after the coordinated load
+#if !defined(_WIN32)
+        if (vram_flock && vram_inotify_fd < 0) {
+            vram_inotify_fd = inotify_init1(IN_NONBLOCK);
+            if (vram_inotify_fd >= 0) {
+                inotify_add_watch(vram_inotify_fd, vram_doorbell_dir.c_str(), IN_CLOSE_WRITE);
+                vram_warden_run = true;
+                vram_warden = std::thread([this]{ vram_warden_loop(); });
+            }
+        }
+#endif
+    }
+
+    // warden thread: block on the doorbell; when another process rings (wants the token) and we
+    // currently hold it, ask the loop to yield (release) at its next idle point. Only touches the
+    // thread-safe queue - never the GPU - so it cannot race a decode.
+    void vram_warden_loop() {
+#if !defined(_WIN32)
+        char buf[4096];
+        while (vram_warden_run.load()) {
+            struct pollfd pfd { vram_inotify_fd, POLLIN, 0 };
+            int pr = poll(&pfd, 1, 500); // 500ms so we periodically re-check the run flag
+            if (pr <= 0) {
+                continue;
+            }
+            ssize_t n = read(vram_inotify_fd, buf, sizeof(buf));
+            if (n <= 0) {
+                continue;
+            }
+            bool foreign_ring = false;
+            for (char * p = buf; p < buf + n; ) {
+                struct inotify_event * ev = (struct inotify_event *) p;
+                if (ev->len > 0 && vram_pid_str != ev->name) {
+                    foreign_ring = true; // someone else wants the GPU
+                }
+                p += sizeof(struct inotify_event) + ev->len;
+            }
+            if (foreign_ring && !vram_cold) {
+                queue_tasks.request_yield();
+            }
+        }
+#endif
+    }
+
+    // ring the doorbell so the current token holder releases promptly
+    void vram_ring_doorbell() {
+#if !defined(_WIN32)
+        if (!vram_flock || vram_doorbell_dir.empty()) {
+            return;
+        }
+        const std::string f = vram_doorbell_dir + "/" + vram_pid_str;
+        int fd = open(f.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0666);
+        if (fd >= 0) {
+            ssize_t w = write(fd, "1", 1);
+            (void) w;
+            close(fd);
+        }
+#endif
+    }
+
+    // acquire the VRAM token (blocking) and bring weights back to the device. Called on the loop
+    // thread right before a decode, so it never races compute.
+    void vram_ensure_warm() {
+        if (!vram_only || !vram_cold) {
+            return;
+        }
+#if !defined(_WIN32)
+        if (vram_flock) {
+            vram_ring_doorbell();         // nudge the current holder to release
+            flock(vram_lock_fd, LOCK_EX); // blocks until the current holder goes cold
+        }
+#endif
+        llama_context_restore_device(ctx_tgt);
+        if (ctx_dft != nullptr) {
+            llama_context_restore_device(ctx_dft);
+        }
+        // NOTE: the multimodal (vision/audio) encoder is intentionally NOT restored here. It is
+        // brought into VRAM just-in-time before an image/audio encode (process_mtmd_chunk) and, in
+        // on-demand mode, released again right after - so a warm text-only model holds no encoder
+        // VRAM (that space is free for KV / expert cache). A cold->warm wake for a *text* request
+        // therefore leaves the encoder in RAM; an image request restores it at encode time.
+        vram_cold = false;
+    }
+
+    void vram_share_shutdown() {
+#if !defined(_WIN32)
+        vram_warden_run = false;
+        if (vram_warden.joinable()) {
+            vram_warden.join();
+        }
+        if (vram_inotify_fd >= 0) { close(vram_inotify_fd); vram_inotify_fd = -1; }
+        if (vram_lock_fd    >= 0) { flock(vram_lock_fd, LOCK_UN); close(vram_lock_fd); vram_lock_fd = -1; }
+#endif
+    }
+
+    // release the device weights (keeping KV + host shadow) and drop the VRAM token so another
+    // model can warm up. Called on the loop thread when the server goes idle.
+    void vram_go_cold() {
+        if (!vram_only || vram_cold) {
+            return;
+        }
+        llama_context_release_device(ctx_tgt, vram_evict_kv);
+        if (ctx_dft != nullptr) {
+            llama_context_release_device(ctx_dft, vram_evict_kv);
+        }
+        // also release the multimodal (vision/audio) encoder weights (dead weight while cold);
+        // its read-only host shadow is captured once and rebuilt on wake by vram_ensure_warm()
+        if (mctx != nullptr) {
+            mtmd_release_device(mctx);
+        }
+#if !defined(_WIN32)
+        if (vram_flock) {
+            flock(vram_lock_fd, LOCK_UN);
+        }
+#endif
+        vram_cold = true;
+    }
+
     common_speculative_init_result_ptr spec_init;
 
     common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO;
@@ -951,6 +1186,29 @@ private:
 
     void handle_sleeping_state(bool new_state) {
         GGML_ASSERT(sleeping != new_state);
+
+        // Lightweight VRAM-only sleep: instead of a full unload/reload (which frees RAM and the
+        // KV cache and pays a full reload on wake), just release the model's device (VRAM) weight
+        // buffers to a host shadow (keeping the context, KV cache and host weights) and drop the
+        // shared VRAM token. Waking is handled lazily by vram_ensure_warm() right before the next
+        // decode (which re-acquires the token first), so a single GPU is time-shared between models
+        // with no re-prefill.
+        if (vram_only && ctx_tgt != nullptr) {
+            if (new_state) {
+                SRV_INF("%s", "server entering sleeping state (VRAM-only: releasing device weights, keeping KV cache)\n");
+                vram_go_cold();
+            } else {
+                SRV_INF("%s", "server exiting sleeping state (VRAM-only: restoring device weights/KV)\n");
+                // Restore NOW, on wake, before update_slots runs. update_slots touches the KV cache
+                // (e.g. SWA checkpoint creation reads it via ggml_backend_tensor_get) before the
+                // decode-time vram_ensure_warm(), so with KV eviction the KV must already be resident
+                // here or those reads hit a freed (null) buffer.
+                vram_ensure_warm();
+            }
+            sleeping = new_state;
+            return;
+        }
+
         if (new_state) {
             SRV_INF("%s", "server is entering sleeping state\n");
             destroy();
@@ -1142,6 +1400,11 @@ private:
             params_base.load_progress_callback_user_data = &load_progress_text;
         }
 
+        // VRAM arbiter: acquire the GPU token before uploading weights, so any resident model (e.g.
+        // the warm 4B task model) releases first and this load uploads into free VRAM instead of
+        // racing it and OOM-ing. No-op unless LLAMA_SLEEP_VRAM_ONLY is set.
+        vram_acquire_for_load();
+
         llama_init = common_init_from_params(params_base);
 
         model_tgt = llama_init->model();
@@ -1152,6 +1415,11 @@ private:
             return false;
         }
 
+        if (ctx_tgt == nullptr) {
+            SRV_ERR("failed to create_context with model '%s'\n", params_base.model.path.c_str());
+            return false;
+        }
+
         vocab = llama_model_get_vocab(model_tgt);
 
         n_ctx = llama_n_ctx(ctx_tgt);
@@ -1207,6 +1475,13 @@ private:
             }
             SRV_INF("loaded multimodal model, '%s'\n", mmproj_path.c_str());
 
+            // on-demand encoder: keep the vision/audio encoder weights in RAM (out of VRAM) until an
+            // image/audio actually needs encoding, freeing that VRAM for KV / expert cache on the
+            // common text-only path. process_mtmd_chunk() brings the encoder in just-in-time.
+            if (getenv("LLAMA_MMPROJ_ONDEMAND") != nullptr) {
+                mtmd_release_device(mctx);
+            }
+
             if (params_base.ctx_shift) {
                 params_base.ctx_shift = false;
                 SRV_WRN("%s\n", "ctx_shift is not supported by multimodal, it will be disabled");
@@ -1404,6 +1679,11 @@ private:
             handle_sleeping_state(sleeping);
         });
 
+        // VRAM arbiter: start the doorbell warden. The load was coordinated (vram_acquire_for_load
+        // grabbed the token before uploading), so we stay warm holding the token and serve the first
+        // request without a re-warm; the warden releases us when another model rings the doorbell.
+        vram_share_init();
+
         metrics.init();
 
         if (params_base.cache_idle_slots) {
@@ -3584,6 +3864,10 @@ private:
             n_empty_consecutive = 0;
         }
 
+        // VRAM arbiter: make sure our weights are resident (acquiring the shared GPU token first)
+        // before we decode. Runs on the loop thread, so it never races an in-flight decode.
+        vram_ensure_warm();
+
         const int ret = llama_decode(ctx_tgt, batch_view);
 
         metrics.on_decoded(slots);
@@ -4005,8 +4289,12 @@ struct server_res_generator : server_res_spipe {
     server_response_reader rd;
     server_res_generator(server_queue & queue_tasks, server_response & queue_results, int sleep_idle_seconds, bool bypass_sleep = false)
             : rd(queue_tasks, queue_results, HTTP_POLLING_SECONDS) {
-        // fast path in case sleeping is disabled
+        // fast path in case sleeping is disabled. Note: the VRAM arbiter (LLAMA_SLEEP_VRAM_ONLY)
-        bypass_sleep |= sleep_idle_seconds < 0;
+        // can put the server to sleep via the cross-process doorbell even when idle-sleep is
+        // disabled (sleep_idle_seconds < 0), so in that case requests must still wake it - otherwise
+        // a doorbell-slept server would hang, never returning from a request.
+        static const bool vram_arbiter = getenv("LLAMA_SLEEP_VRAM_ONLY") != nullptr;
+        bypass_sleep |= (sleep_idle_seconds < 0 && !vram_arbiter);
         if (!bypass_sleep) {
             queue_tasks.wait_until_no_sleep();
         }

tools/server/server-queue.cpp

@@ -116,6 +116,12 @@ void server_queue::wait_until_no_sleep() {
     }
 }
 
+void server_queue::request_yield() {
+    std::unique_lock<std::mutex> lock(mutex_tasks);
+    yield_requested = true;
+    condition_tasks.notify_all();
+}
+
 void server_queue::terminate() {
     std::unique_lock<std::mutex> lock(mutex_tasks);
     running = false;
@@ -129,6 +135,9 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
     constexpr auto max_wait_time = std::chrono::seconds(1);
     auto should_sleep = [&]() -> bool {
         // caller must hold mutex_tasks
+        if (yield_requested) {
+            return true; // another process rang the VRAM doorbell - release now
+        }
         if (idle_sleep_ms < 0) {
             return false;
         }
@@ -178,6 +187,7 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
             if (should_sleep()) {
                 QUE_INF("%s", "entering sleeping state\n");
                 sleeping = true;
+                yield_requested = false; // consumed
                 callback_sleeping_state(true);
                 req_stop_sleeping = false;
                 // wait until we are requested to exit sleeping state
@@ -195,13 +205,17 @@ void server_queue::start_loop(int64_t idle_sleep_ms) {
                 condition_tasks.notify_all(); // notify wait_until_no_sleep()
                 break; // process new tasks
             } else {
-                // wait for new tasks or timeout for checking sleeping condition
+                // wait for new tasks, a VRAM yield request, or timeout for checking sleeping condition
                 bool res = condition_tasks.wait_for(lock, max_wait_time, [&]{
-                    return (!queue_tasks.empty() || !running);
+                    return (!queue_tasks.empty() || !running || yield_requested);
                 });
-                if (res) {
+                if (res && !queue_tasks.empty()) {
                     break; // new task arrived or terminate
                 }
+                if (!running) {
+                    break;
+                }
+                // otherwise (timeout or yield request), loop again to re-check should_sleep
                 // otherwise, loop again to check sleeping condition
             }
         }

tools/server/server-queue.h

@@ -16,6 +16,7 @@ private:
     bool running  = false;
     bool sleeping = false;
     bool req_stop_sleeping = false;
+    bool yield_requested = false; // set by request_yield() when another process wants the VRAM token
     int64_t time_last_task = 0;
 
     // queues
@@ -51,6 +52,11 @@ public:
     // returns immediately if not sleeping
     void wait_until_no_sleep();
 
+    // request that the loop go to sleep (release VRAM) as soon as it is idle - called from the
+    // VRAM-arbiter warden thread when another process rings the doorbell for the GPU token.
+    // Thread-safe; wakes the loop so it releases promptly instead of at the next idle poll.
+    void request_yield();
+
     bool is_sleeping() {
         std::unique_lock<std::mutex> lock(mutex_tasks);
         return sleeping;

tools/server/server-stream.cpp

@@ -632,6 +632,13 @@ void server_res_spipe::on_complete() {
     if (!spipe || next_finished) {
         return;
     }
+    // an empty next_orig means set_next() never ran: the request failed before streaming
+    // started, typically a params validation throw. evict the session installed by set_req()
+    // so the failed request leaves nothing behind for discovery or replay
+    if (!next_orig) {
+        g_stream_sessions.evict(server_stream_conv_id_from_headers(req->headers));
+        return;
+    }
     std::string chunk;
     while (!spipe->is_cancelled()) {
         chunk.clear();

tools/ui/.gitignore

@@ -36,3 +36,7 @@ static/favicon*
 *storybook.log
 storybook-static
 *.code-workspace
+
+# Vitest browser mode failure artifacts
+.vitest-attachments/
+tests/**/__screenshots__/

tools/ui/.prettierignore

@@ -16,3 +16,6 @@ build/
 /build/
 /.svelte-kit/
 test-results
+
+# Vendored third party sources, kept byte identical to upstream
+src/lib/vendors/

tools/ui/eslint.config.js

@@ -59,7 +59,8 @@ export default ts.config(
 			'.svelte-kit/**',
 			'test-results/**',
 			'.storybook/**/*',
-			'src/lib/services/sandbox-worker.js'
+			'src/lib/services/sandbox-worker.js',
+			'src/lib/vendors/**'
 		]
 	},
 	storybook.configs['flat/recommended']

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

@@ -0,0 +1,49 @@
+import { build } from 'esbuild';
+import { dirname, resolve } from 'path';
+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;
+
+/**
+ * Bundle the vendored nerdamer-prime source into a minified IIFE string,
+ * exposed as the `virtual:nerdamer` module. Flags mirror the upstream
+ * build (esbuild --bundle --minify --format=iife --global-name=nerdamer),
+ * so only human readable source lives in the repo and minification is a
+ * build artifact. Vendored under src/lib/vendors/, upstream snapshot:
+ * https://github.com/together-science/nerdamer-prime/commit/1936145f8af306ec0d883b9bfd7730aedd175c24
+ */
+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'
+				});
+				bundled = result.outputFiles[0].text;
+			}
+			return `export default ${JSON.stringify(bundled)};`;
+		}
+	};
+}

tools/ui/src/lib/constants/sandbox.ts

@@ -13,14 +13,27 @@ export const SANDBOX_EMPTY_OUTPUT = '(no output)';
 
 export const SANDBOX_TRUNCATION_NOTICE = '[output truncated]';
 
-export const SANDBOX_TOOL_DEFINITION: OpenAIToolDefinition = {
+const NERDAMER_DESCRIPTION = `
+Symbolic/numeric math via \`nerdamer\` (pre-loaded, do not require, use it directly).
+nerdamer('diff(sin(x)/x,x)') or nerdamer.diff('sin(x)/x','x') → Expression; convert with .toString()/.text()/.toTeX(), or .evaluate() (→ still Expression, then .toString()).
+nerdamer(expr,{x:2}) substitutes only; chain .evaluate() or pass 'numer' for numeric result.
+solve(expr,var)→Symbol[]; solveEquations([eq1,..])→[[var,val],..] pairs.
+Functions: simplify/expand/factor(expr), diff(expr,var[,n]), integrate(expr,var), defint(expr,from,to,var), limit(expr,var,to), laplace(expr,t,s), ilt(expr,s,t), gcd/lcm(a,b), roots/coeffs/partfrac(expr,var), pfactor(n), numer/decimals/erf(expr), product/sum(expr,var,from,to), mean/median/stdev/variance(...vals).
+Object.keys(nerdamer).filter(k=>typeof nerdamer[k]==='function') lists all available functions. If you need a function not documented above, list them first — do not guess function names.`;
+
+/**
+ * Build the sandbox tool definition. When `includeSymbolicMath` is true,
+ * the description includes nerdamer API documentation; otherwise it
+ * describes a plain JavaScript sandbox.
+ */
+export function buildSandboxToolDefinition(includeSymbolicMath: boolean): OpenAIToolDefinition {
+	return {
 		type: ToolCallType.FUNCTION,
 		function: {
 			name: SANDBOX_TOOL_NAME,
-		description:
+			description: includeSymbolicMath
-			'Execute JavaScript in a sandboxed browser worker (no DOM, no page access). ' +
+				? `Execute JS in a sandboxed browser worker (no DOM/page access). Top-level await ok; console.log for intermediates; top-level return is captured as result.${NERDAMER_DESCRIPTION}`
-			'Top level await is supported. Use console.log to print intermediate values; ' +
+				: 'Execute JS in a sandboxed browser worker (no DOM/page access). Top-level await ok; console.log for intermediates; top-level return is captured as result.',
-			'a top level return statement is captured as the result.',
 			parameters: {
 				type: JsonSchemaType.OBJECT,
 				properties: {
@@ -37,3 +50,7 @@ export const SANDBOX_TOOL_DEFINITION: OpenAIToolDefinition = {
 			}
 		}
 	};
+}
+
+/** @deprecated Use {@link buildSandboxToolDefinition} instead. Kept for backward compatibility. */
+export const SANDBOX_TOOL_DEFINITION = buildSandboxToolDefinition(true);

tools/ui/src/lib/constants/settings-keys.ts

@@ -67,6 +67,7 @@ export const SETTINGS_KEYS = {
 	EXCLUDE_REASONING_FROM_CONTEXT: 'excludeReasoningFromContext',
 	SHOW_RAW_OUTPUT_SWITCH: 'showRawOutputSwitch',
 	JS_SANDBOX_ENABLED: 'jsSandboxEnabled',
+	SYMBOLIC_MATH_ENABLED: 'symbolicMathEnabled',
 	// PY_INTERPRETER_ENABLED: 'pyInterpreterEnabled',
 	CUSTOM_JSON: 'customJson',
 	CUSTOM_CSS: 'customCss'

tools/ui/src/lib/constants/settings-registry.ts

@@ -724,6 +724,15 @@ const SETTINGS_REGISTRY: Record<string, SettingsSectionEntry> = {
 					paramType: SyncableParameterType.BOOLEAN
 				}
 			},
+			{
+				key: SETTINGS_KEYS.SYMBOLIC_MATH_ENABLED,
+				label: 'Symbolic math (nerdamer)',
+				help: 'Pre-load nerdamer in the sandbox for symbolic computation: simplify, diff, integrate, solve, and more. Requires "JavaScript sandbox tool" to be enabled.',
+				defaultValue: false,
+				type: SettingsFieldType.CHECKBOX,
+				section: SETTINGS_SECTION_SLUGS.DEVELOPER,
+				dependsOn: SETTINGS_KEYS.JS_SANDBOX_ENABLED
+			},
 			{
 				key: SETTINGS_KEYS.CUSTOM_JSON,
 				label: 'Custom JSON',

tools/ui/src/lib/services/index.ts

@@ -276,7 +276,7 @@ export { MCPService } from './mcp.service';
  * - **toolsStore**: Exposes the tool definition when the sandbox is enabled
  * - **agenticStore**: Dispatches ToolSource.FRONTEND calls here
  *
- * @see SANDBOX_TOOL_DEFINITION in constants/sandbox.ts - tool schema sent to the LLM
+ * @see buildSandboxToolDefinition in constants/sandbox.ts - tool schema sent to the LLM
  * @see agenticStore in stores/agentic.svelte.ts - tool dispatch
  */
 export { SandboxService } from './sandbox.service';

tools/ui/src/lib/services/sandbox-harness.ts

@@ -1,14 +1,25 @@
+import { NEWLINE } from '$lib/constants';
 import WORKER_SHIM from './sandbox-worker.js?raw';
 
+/**
+ * CSP for the harness document, inherited by the blob worker. connect-src
+ * falls back to default-src, removing network egress for model and vendored
+ * code. 'unsafe-eval' is required by the worker's AsyncFunction constructor,
+ * 'unsafe-inline' by the inline script below, worker-src by the blob worker.
+ */
+const HARNESS_CSP = `default-src 'none'; script-src 'unsafe-inline' 'unsafe-eval'; worker-src blob:`;
+
 /**
  * Harness loaded as srcdoc into a sandboxed iframe (allow-scripts only).
  * The opaque origin is the security boundary: no access to the app origin,
  * its storage or its API. The harness spawns a worker so model code never
  * runs on a main thread, which makes the parent timeout enforceable by
- * removing the iframe.
+ * removing the iframe. The prelude runs in the worker before the shim,
+ * exposing globals such as `nerdamer` to model code.
  */
-export const SANDBOX_HARNESS_HTML = `<!doctype html><script>
+export function buildSandboxHarness(preludeJs: string): string {
-const SHIM = ${JSON.stringify(WORKER_SHIM)};
+	return `<!doctype html><meta http-equiv="Content-Security-Policy" content="${HARNESS_CSP}"><script>
+const SHIM = ${JSON.stringify(preludeJs + NEWLINE + WORKER_SHIM)};
 addEventListener('message', (event) => {
 	const respond = (payload) => parent.postMessage(payload, '*');
 	let worker;
@@ -23,3 +34,4 @@ addEventListener('message', (event) => {
 	worker.postMessage({ code: event.data.code });
 });
 </script>`;
+}

tools/ui/src/lib/services/sandbox-worker.js

@@ -21,7 +21,9 @@ self.onmessage = async (event) => {
 	const reply = { logs, result: null, error: null };
 	try {
 		const AsyncFunction = Object.getPrototypeOf(async function () {}).constructor;
-		const value = await new AsyncFunction(event.data.code)();
+		// The prelude bundled ahead of this shim defines self.nerdamer,
+		// passed into the execution scope as the `nerdamer` parameter.
+		const value = await new AsyncFunction('nerdamer', event.data.code)(self.nerdamer);
 		if (value !== undefined) reply.result = fmt(value);
 	} catch (err) {
 		reply.error = err instanceof Error ? err.stack || err.message : String(err);

tools/ui/src/lib/services/sandbox.service.ts

@@ -7,9 +7,32 @@ import {
 	SANDBOX_TOOL_NAME,
 	SANDBOX_TRUNCATION_NOTICE
 } from '$lib/constants';
-import { SANDBOX_HARNESS_HTML } from './sandbox-harness';
+import { buildSandboxHarness } from './sandbox-harness';
+import { config } from '$lib/stores/settings.svelte';
 import type { ToolExecutionResult } from '$lib/types';
 
+/** Cached harnesses keyed by whether nerdamer is included. */
+const harnessCache: Record<string, string> = {};
+
+/**
+ * Build the sandbox harness. When symbolic math is enabled, loads the
+ * nerdamer prelude lazily; otherwise builds a plain harness with an empty
+ * prelude. Cached per variant so toggling the setting is instant.
+ */
+async function getHarness(): Promise<string> {
+	const enabled = !!config().symbolicMathEnabled;
+	const key = enabled ? 'nerdamer' : 'plain';
+	if (!harnessCache[key]) {
+		if (enabled) {
+			const { default: nerdamerJs } = await import('virtual:nerdamer');
+			harnessCache[key] = buildSandboxHarness(nerdamerJs);
+		} else {
+			harnessCache[key] = buildSandboxHarness('');
+		}
+	}
+	return harnessCache[key];
+}
+
 interface SandboxReply {
 	logs?: unknown;
 	result?: unknown;
@@ -45,20 +68,22 @@ export class SandboxService {
 	 * timeout or abort. Removing the iframe terminates the worker
 	 * at the browser level, so runaway code cannot outlive it.
 	 */
-	static executeTool(
+	static async executeTool(
 		toolName: string,
 		params: Record<string, unknown>,
 		signal?: AbortSignal
 	): Promise<ToolExecutionResult> {
 		if (toolName !== SANDBOX_TOOL_NAME) {
-			return Promise.resolve({ content: `Unknown frontend tool: ${toolName}`, isError: true });
+			return { content: `Unknown frontend tool: ${toolName}`, isError: true };
 		}
 
 		const code = typeof params.code === 'string' ? params.code : '';
 		if (!code) {
-			return Promise.resolve({ content: 'Missing required parameter: code', isError: true });
+			return { content: 'Missing required parameter: code', isError: true };
 		}
 
+		const harness = await getHarness();
+
 		const requested = Number(params.timeout_ms);
 		const timeoutMs =
 			Number.isFinite(requested) && requested > 0
@@ -69,7 +94,7 @@ export class SandboxService {
 			const iframe = document.createElement('iframe');
 			iframe.setAttribute('sandbox', 'allow-scripts');
 			iframe.style.display = 'none';
-			iframe.srcdoc = SANDBOX_HARNESS_HTML;
+			iframe.srcdoc = harness;
 
 			let settled = false;
 

tools/ui/src/lib/stores/chat.svelte.ts

@@ -2428,6 +2428,7 @@ class ChatStore {
 
 		if (currentConfig.samplers) apiOptions.samplers = currentConfig.samplers;
 
+		if (hasValue(currentConfig.backend_sampling))
 			apiOptions.backend_sampling = currentConfig.backend_sampling;
 
 		if (currentConfig.customJson) apiOptions.custom = currentConfig.customJson;

tools/ui/src/lib/stores/tools.svelte.ts

@@ -5,7 +5,7 @@ import { HealthCheckStatus, JsonSchemaType, ToolCallType, ToolSource } from '$li
 import { config } from '$lib/stores/settings.svelte';
 import {
 	DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY,
-	SANDBOX_TOOL_DEFINITION,
+	buildSandboxToolDefinition,
 	TOOL_GROUP_LABELS,
 	TOOL_SERVER_LABELS
 } from '$lib/constants';
@@ -143,7 +143,9 @@ class ToolsStore {
 	}
 
 	get frontendTools(): OpenAIToolDefinition[] {
-		return config().jsSandboxEnabled ? [SANDBOX_TOOL_DEFINITION] : [];
+		return config().jsSandboxEnabled
+			? [buildSandboxToolDefinition(!!config().symbolicMathEnabled)]
+			: [];
 	}
 
 	get customTools(): OpenAIToolDefinition[] {